mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
fix(usage): capture tokens by reading the agent transcript at finalize
Spawn sessions recorded 0 tokens in production despite the full capture chain being deployed (hook in the image, settings registering it, /usage/sync, the finalize fetch). The chain is inherently racy: token counts live in the SDK server's in-memory state inside a short-lived agent container, and the orchestrator races to fetch /usage/status before the container is torn down — so for cold-respawned agents the fetch returns nothing (verified: parser and transcripts are correct, every spawn-session row was still 0). Read the durable source of truth instead. The host ~/.claude is already mounted into the orchestrator, so at finalize, when the SDK fetch yields nothing, sum the agent's newest Claude Code transcript directly (_usage_from_transcript + the existing _sum_transcript_usage). Removes the race entirely — usage is captured regardless of container timing.
This commit is contained in:
@@ -3207,6 +3207,34 @@ class AgentOrchestrator:
|
|||||||
)
|
)
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _usage_from_transcript(agent_id: str) -> tuple[int, int, int, int]:
|
||||||
|
"""Sum token usage from the agent's newest Claude Code transcript.
|
||||||
|
|
||||||
|
The host ``~/.claude`` is mounted into the orchestrator, so each agent's
|
||||||
|
transcripts are readable here under ``projects/*-{slug}/`` (Claude Code
|
||||||
|
encodes the agent's workspace cwd into the dir name, which ends in the
|
||||||
|
agent slug). The durable fallback for the live SDK ``/usage/status``
|
||||||
|
fetch, which misses whenever the agent container is short-lived or
|
||||||
|
already torn down. Returns zeros when no transcript is found.
|
||||||
|
"""
|
||||||
|
from roboco.agent_sdk.server import _sum_transcript_usage
|
||||||
|
|
||||||
|
projects = Path.home() / ".claude" / "projects"
|
||||||
|
try:
|
||||||
|
jsonl = [
|
||||||
|
f
|
||||||
|
for d in projects.glob(f"*-{agent_id}")
|
||||||
|
if d.is_dir()
|
||||||
|
for f in d.glob("*.jsonl")
|
||||||
|
]
|
||||||
|
except OSError:
|
||||||
|
return (0, 0, 0, 0)
|
||||||
|
if not jsonl:
|
||||||
|
return (0, 0, 0, 0)
|
||||||
|
newest = max(jsonl, key=lambda f: f.stat().st_mtime)
|
||||||
|
return _sum_transcript_usage(newest)
|
||||||
|
|
||||||
async def _finalize_spawn_session(
|
async def _finalize_spawn_session(
|
||||||
self,
|
self,
|
||||||
agent_id: str,
|
agent_id: str,
|
||||||
@@ -3248,6 +3276,19 @@ class AgentOrchestrator:
|
|||||||
error=str(sdk_exc),
|
error=str(sdk_exc),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# The live SDK fetch above races the container teardown and misses
|
||||||
|
# for short-lived agents (counts live in the SDK server's memory,
|
||||||
|
# which dies with the container). Fall back to the durable source of
|
||||||
|
# truth: the agent's Claude Code transcript, mounted into this
|
||||||
|
# container — so usage is captured regardless of container timing.
|
||||||
|
if not tokens_input and not tokens_output:
|
||||||
|
tin, tout, cr, cw = self._usage_from_transcript(agent_id)
|
||||||
|
if tin or tout:
|
||||||
|
tokens_input = tin
|
||||||
|
tokens_output = tout
|
||||||
|
tokens_cache_read = cr
|
||||||
|
tokens_cache_write = cw
|
||||||
|
|
||||||
# Look up the model and usage_session_id from the running instance config.
|
# Look up the model and usage_session_id from the running instance config.
|
||||||
instance = self._instances.get(agent_id)
|
instance = self._instances.get(agent_id)
|
||||||
if instance and instance.config:
|
if instance and instance.config:
|
||||||
|
|||||||
Reference in New Issue
Block a user