feat(llc): parse CLI token usage and report to LLC budget (#10220)

autobot-backend/llc/adapters/base.py

@@ -25,6 +25,10 @@ class AdapterRunStatus:
     status: LLCRunStatus
     exit_code: Optional[int] = None
     error: Optional[str] = None
+    # GH#10220: token usage parsed from the CLI output on a completed run, so the
+    # scheduler can forward it to BudgetService.ingest_cost_event (None = unknown).
+    tokens_in: Optional[int] = None
+    tokens_out: Optional[int] = None
 
 
 @runtime_checkable

autobot-backend/llc/adapters/claude_code_adapter.py

@@ -34,6 +34,7 @@
 import shutil
 import time
 import uuid
+from dataclasses import replace
 from typing import Optional
 
 from autobot_shared.logging_manager import get_logger
@@ -46,6 +47,7 @@
 from .subprocess_base import SubprocessLifecycleAdapter
 from .subprocess_base import resolve_timeout as _resolve_timeout
 from .subprocess_support import (
+    extract_usage,
     final_result_event,
     inject_agent_credentials,
     is_rate_limit_output,
@@ -260,8 +262,10 @@ async def _status(self, agent_config: dict, run_id: str) -> AdapterRunStatus:
         # Clean success: skip keyword scan entirely to avoid false-positive
         # reclassification of completed runs whose transcript mentions rate-limit
         # terms incidentally (e.g. in tool output, issue numbers, SHAs).
+        # GH#10220: attach parsed token usage so the scheduler can bill it.
         if result_event is not None and not result_event.get("is_error") and result_event.get("subtype") == "success":
-            return base
+            tokens_in, tokens_out = extract_usage(result_event)
+            return replace(base, tokens_in=tokens_in, tokens_out=tokens_out)
 
         # Either no result event (mid-stream kill) or an error/non-success result:
         # scan the stdout tail AND the stderr sidecar for rate-limit markers —

autobot-backend/llc/adapters/subprocess_support.py

@@ -125,6 +125,31 @@ def final_result_event(tail: str) -> dict | None:
     return None
 
 
+def extract_usage(result_event: dict | None) -> tuple[int | None, int | None]:
+    """Return ``(tokens_in, tokens_out)`` from a stream-json result event (GH#10220).
+
+    The Claude CLI result event carries a ``usage`` object. Input tokens are the
+    sum of fresh + cache-read + cache-creation input tokens (all billed input);
+    output is ``output_tokens``. Returns ``(None, None)`` when no usable usage is
+    present so callers can skip budget ingestion rather than record zeros.
+    """
+    if not isinstance(result_event, dict):
+        return (None, None)
+    usage = result_event.get("usage")
+    if not isinstance(usage, dict):
+        return (None, None)
+
+    def _int(key: str) -> int:
+        val = usage.get(key)
+        return val if isinstance(val, int) and val >= 0 else 0
+
+    tokens_in = _int("input_tokens") + _int("cache_read_input_tokens") + _int("cache_creation_input_tokens")
+    tokens_out = _int("output_tokens")
+    if tokens_in == 0 and tokens_out == 0:
+        return (None, None)
+    return (tokens_in, tokens_out)
+
+
 # Context keys rendered by dedicated prompt sections or consumed as env vars —
 # excluded from the generic "Additional Context" catch-all.
 _RENDERED_CONTEXT_KEYS = frozenset(

autobot-backend/llc/adapters/tests/test_subprocess_support.py

@@ -12,6 +12,7 @@
 
 from llc.adapters.subprocess_support import (
     AGENT_API_KEY_PLACEHOLDER,
+    extract_usage,
     inject_agent_credentials,
     probe_pid,
     render_context_markdown,
@@ -191,3 +192,34 @@ def kill_fn(pid, sig):
 
         assert result is False
         assert sigkill_count[0] == 1  # SIGKILL was attempted once
+
+
+class TestExtractUsage:
+    """Token-usage parsing from the stream-json result event (GH#10220)."""
+
+    def test_sums_input_plus_cache(self) -> None:
+        ev = {
+            "type": "result",
+            "usage": {
+                "input_tokens": 100,
+                "output_tokens": 50,
+                "cache_read_input_tokens": 20,
+                "cache_creation_input_tokens": 5,
+            },
+        }
+        assert extract_usage(ev) == (125, 50)
+
+    def test_basic_input_output(self) -> None:
+        assert extract_usage({"usage": {"input_tokens": 10, "output_tokens": 3}}) == (10, 3)
+
+    def test_missing_usage_returns_none(self) -> None:
+        assert extract_usage({"type": "result"}) == (None, None)
+        assert extract_usage({}) == (None, None)
+        assert extract_usage(None) == (None, None)
+
+    def test_all_zero_returns_none(self) -> None:
+        assert extract_usage({"usage": {"input_tokens": 0, "output_tokens": 0}}) == (None, None)
+
+    def test_non_int_values_ignored(self) -> None:
+        # input_tokens "x" is non-int → counted as 0; output 7 still recorded.
+        assert extract_usage({"usage": {"input_tokens": "x", "output_tokens": 7}}) == (0, 7)

autobot-backend/llc/scheduler/heartbeat_scheduler.py

@@ -46,13 +46,14 @@
 from autobot_shared.singleton_factory import lazy_singleton
 from user_management.database import get_async_session_factory
 
-from ..adapters import AutoBotAgentAdapter, get_adapter
+from ..adapters import AdapterRunStatus, AutoBotAgentAdapter, get_adapter
 from ..adapters.subprocess_base import is_subprocess_adapter
 from ..config import AGENT_API_BASE_URL
-from ..exceptions import AdapterRunFailed, HeartbeatDispatchSkipped, ProviderRateLimited
+from ..exceptions import AdapterRunFailed, BudgetExhausted, HeartbeatDispatchSkipped, ProviderRateLimited
 from ..models.enums import HeartbeatInvocationSource, LLCRunStatus
 from ..models.heartbeat_run import LLCHeartbeatRun
 from ..services.api_key import ApiKeyService
+from ..services.budget import BudgetService
 from ..services.replay_service import RunReplayService, parse_jsonl_events
 
 logger = logging.getLogger(__name__)
@@ -823,7 +824,8 @@ async def _dispatch_registry_adapter(adapter: Any, agent: Dict[str, Any], contex
     external_run_id: Optional[str] = None
     try:
         external_run_id = await adapter.invoke(agent_config, enriched)
-        final_status = await _await_adapter_completion(adapter, agent_config, external_run_id)
+        result = await _await_adapter_completion(adapter, agent_config, external_run_id)
+        final_status = result.status
     except asyncio.CancelledError:
         # Shutdown / task cancellation — kill the external run so it does not
         # outlive its about-to-be-revoked key, then propagate.
@@ -851,9 +853,37 @@ async def _dispatch_registry_adapter(adapter: Any, agent: Dict[str, Any], contex
     if final_status != LLCRunStatus.COMPLETED:
         raise AdapterRunFailed(agent.get("adapter_type") or "", str(external_run_id), final_status)
 
+    # GH#10220: forward parsed token usage to the agent's budget. BudgetExhausted
+    # is re-raised (GH#8215 hard-stop → run marked FAILED); other cost errors are
+    # best-effort and never fail an otherwise-successful run.
+    await _ingest_adapter_usage(agent, result)
+
     return external_run_id
 
 
+async def _ingest_adapter_usage(agent: Dict[str, Any], result: AdapterRunStatus) -> None:
+    """Record a completed registry run's token usage against the agent budget."""
+    if result.tokens_in is None and result.tokens_out is None:
+        return  # adapter did not report usage (e.g. CLI without stream-json result)
+    agent_id: str = agent["agent_id"]
+    model = (agent.get("adapter_config") or {}).get("model") or agent.get("model") or ""
+    try:
+        factory = get_async_session_factory()
+        async with factory() as session:
+            await BudgetService().ingest_cost_event(
+                session,
+                agent_id,
+                int(result.tokens_in or 0),
+                int(result.tokens_out or 0),
+                model,
+            )
+            await session.commit()
+    except BudgetExhausted:
+        raise  # GH#8215 hard-stop — propagate so the run is recorded FAILED
+    except Exception:
+        logger.warning("Budget ingest failed for agent %s (best-effort)", agent_id)
+
+
 async def _safe_cancel(adapter: Any, agent_config: Dict[str, Any], run_id: str) -> None:
     """Cancel an external run, swallowing adapter errors (best-effort)."""
     try:
@@ -894,22 +924,23 @@ async def _revoke_run_key(agent_id: str, key_id: uuid.UUID) -> None:
         logger.exception("Failed to revoke ephemeral heartbeat key %s for agent %s", key_id, agent_id)
 
 
-async def _await_adapter_completion(adapter: Any, agent_config: Dict[str, Any], run_id: str) -> LLCRunStatus:
+async def _await_adapter_completion(adapter: Any, agent_config: Dict[str, Any], run_id: str) -> AdapterRunStatus:
     """Poll ``adapter.status`` until the run is terminal or the max wait elapses.
 
-    Cancels the run if it overruns ``_ADAPTER_MAX_WAIT_SECONDS`` so the
-    ephemeral key is never left live indefinitely.
+    Returns the full terminal :class:`AdapterRunStatus` (carrying token usage for
+    billing, GH#10220). Cancels the run if it overruns
+    ``_ADAPTER_MAX_WAIT_SECONDS`` so the ephemeral key is never left live forever.
     """
     started = time.monotonic()
     while time.monotonic() - started < _ADAPTER_MAX_WAIT_SECONDS:
         result = await adapter.status(agent_config, run_id)
         if result.status.is_terminal():
-            return result.status
+            return result
         await asyncio.sleep(_ADAPTER_POLL_INTERVAL)
 
     logger.warning("Adapter run %s exceeded max wait — cancelling", run_id)
     await adapter.cancel(agent_config, run_id)
-    return LLCRunStatus.TIMEOUT
+    return AdapterRunStatus(status=LLCRunStatus.TIMEOUT)
 
 
 async def _fetch_recent_decisions(company_id: str, n: int = 5) -> list[Dict[str, Any]]:

autobot-backend/llc/tests/test_heartbeat_scheduler.py

@@ -30,6 +30,7 @@
     _dispatch_adapter,
     _dispatch_autobot_agent,
     _dispatch_registry_adapter,
+    _ingest_adapter_usage,
     _next_fire,
 )
 
@@ -355,7 +356,10 @@ async def fake_invoke(agent_config, context):
         with (
             patch(f"{_HBS}._issue_run_key", new=AsyncMock(return_value=(key_record, "llc_rawkey"))),
             patch(f"{_HBS}._revoke_run_key", new=AsyncMock()) as mock_revoke,
-            patch(f"{_HBS}._await_adapter_completion", new=AsyncMock(return_value=LLCRunStatus.COMPLETED)),
+            patch(
+                f"{_HBS}._await_adapter_completion",
+                new=AsyncMock(return_value=AdapterRunStatus(status=LLCRunStatus.COMPLETED)),
+            ),
         ):
             await _dispatch_registry_adapter(fake_adapter, agent, {"recent_decisions": []})
 
@@ -397,7 +401,10 @@ async def fake_invoke(agent_config, context):
         with (
             patch(f"{_HBS}._issue_run_key", new=AsyncMock()) as mock_issue,
             patch(f"{_HBS}._revoke_run_key", new=AsyncMock()) as mock_revoke,
-            patch(f"{_HBS}._await_adapter_completion", new=AsyncMock(return_value=LLCRunStatus.COMPLETED)),
+            patch(
+                f"{_HBS}._await_adapter_completion",
+                new=AsyncMock(return_value=AdapterRunStatus(status=LLCRunStatus.COMPLETED)),
+            ),
         ):
             await _dispatch_registry_adapter(fake_adapter, agent, {})
 
@@ -418,7 +425,7 @@ async def test_returns_when_terminal(self):
         )
         with patch(f"{_HBS}.asyncio.sleep", new=AsyncMock()):
             result = await _await_adapter_completion(fake_adapter, {"agent_id": "a"}, "1/s")
-        assert result == LLCRunStatus.COMPLETED
+        assert result.status == LLCRunStatus.COMPLETED
         assert fake_adapter.status.await_count == 2
 
     async def test_cancels_on_max_wait(self):
@@ -427,7 +434,7 @@ async def test_cancels_on_max_wait(self):
         fake_adapter.cancel = AsyncMock()
         with patch(f"{_HBS}._ADAPTER_MAX_WAIT_SECONDS", 0):
             result = await _await_adapter_completion(fake_adapter, {"agent_id": "a"}, "1/s")
-        assert result == LLCRunStatus.TIMEOUT
+        assert result.status == LLCRunStatus.TIMEOUT
         fake_adapter.cancel.assert_awaited_once()
 
 
@@ -443,7 +450,10 @@ async def test_failed_terminal_status_raises_and_revokes(self):
         with (
             patch(f"{_HBS}._issue_run_key", new=AsyncMock(return_value=(key_record, "llc_k"))),
             patch(f"{_HBS}._revoke_run_key", new=AsyncMock()) as mock_revoke,
-            patch(f"{_HBS}._await_adapter_completion", new=AsyncMock(return_value=LLCRunStatus.FAILED)),
+            patch(
+                f"{_HBS}._await_adapter_completion",
+                new=AsyncMock(return_value=AdapterRunStatus(status=LLCRunStatus.FAILED)),
+            ),
         ):
             with pytest.raises(AdapterRunFailed):
                 await _dispatch_registry_adapter(fake_adapter, agent, {})
@@ -457,7 +467,10 @@ async def test_completed_terminal_status_does_not_raise(self):
         with (
             patch(f"{_HBS}._issue_run_key", new=AsyncMock(return_value=(None, None))),
             patch(f"{_HBS}._revoke_run_key", new=AsyncMock()),
-            patch(f"{_HBS}._await_adapter_completion", new=AsyncMock(return_value=LLCRunStatus.COMPLETED)),
+            patch(
+                f"{_HBS}._await_adapter_completion",
+                new=AsyncMock(return_value=AdapterRunStatus(status=LLCRunStatus.COMPLETED)),
+            ),
         ):
             await _dispatch_registry_adapter(fake_adapter, agent, {})  # no raise
 
@@ -780,3 +793,70 @@ async def fake_get_resumable(agent_id: str):
                 pass
 
         assert len(resume_called) == 1, "resume SHOULD be called for normal runs"
+
+
+@pytest.mark.asyncio
+class TestIngestAdapterUsage:
+    """GH#10220: completed registry runs forward token usage to the budget."""
+
+    async def test_ingests_tokens_with_model(self):
+        agent = _make_agent(adapter_type="claude_code", adapter_config={"model": "claude-x"})
+        result = AdapterRunStatus(status=LLCRunStatus.COMPLETED, tokens_in=120, tokens_out=40)
+
+        session = AsyncMock()
+        session.commit = AsyncMock()
+        factory = MagicMock()
+        factory.return_value.__aenter__ = AsyncMock(return_value=session)
+        factory.return_value.__aexit__ = AsyncMock(return_value=False)
+        mock_ingest = AsyncMock()
+
+        with (
+            patch(f"{_HBS}.get_async_session_factory", return_value=factory),
+            patch(f"{_HBS}.BudgetService.ingest_cost_event", new=mock_ingest),
+        ):
+            await _ingest_adapter_usage(agent, result)
+
+        mock_ingest.assert_awaited_once()
+        args = mock_ingest.await_args.args
+        assert args[1] == agent["agent_id"] and args[2] == 120 and args[3] == 40 and args[4] == "claude-x"
+
+    async def test_noop_when_usage_unknown(self):
+        agent = _make_agent(adapter_type="claude_code")
+        result = AdapterRunStatus(status=LLCRunStatus.COMPLETED)  # tokens None
+        mock_ingest = AsyncMock()
+        with patch(f"{_HBS}.BudgetService.ingest_cost_event", new=mock_ingest):
+            await _ingest_adapter_usage(agent, result)
+        mock_ingest.assert_not_awaited()
+
+    async def test_budget_exhausted_propagates(self):
+        from llc.exceptions import BudgetExhausted
+
+        agent = _make_agent(adapter_type="claude_code", adapter_config={"model": "m"})
+        result = AdapterRunStatus(status=LLCRunStatus.COMPLETED, tokens_in=10, tokens_out=5)
+        session = AsyncMock()
+        factory = MagicMock()
+        factory.return_value.__aenter__ = AsyncMock(return_value=session)
+        factory.return_value.__aexit__ = AsyncMock(return_value=False)
+
+        with (
+            patch(f"{_HBS}.get_async_session_factory", return_value=factory),
+            patch(
+                f"{_HBS}.BudgetService.ingest_cost_event",
+                new=AsyncMock(side_effect=BudgetExhausted("agent", 100.0, 50.0)),
+            ),
+            pytest.raises(BudgetExhausted),
+        ):
+            await _ingest_adapter_usage(agent, result)
+
+    async def test_other_errors_best_effort(self):
+        agent = _make_agent(adapter_type="claude_code", adapter_config={"model": "m"})
+        result = AdapterRunStatus(status=LLCRunStatus.COMPLETED, tokens_in=10, tokens_out=5)
+        session = AsyncMock()
+        factory = MagicMock()
+        factory.return_value.__aenter__ = AsyncMock(return_value=session)
+        factory.return_value.__aexit__ = AsyncMock(return_value=False)
+        with (
+            patch(f"{_HBS}.get_async_session_factory", return_value=factory),
+            patch(f"{_HBS}.BudgetService.ingest_cost_event", new=AsyncMock(side_effect=RuntimeError("db down"))),
+        ):
+            await _ingest_adapter_usage(agent, result)  # must not raise