Status: Layer 1 implemented, Layers 2-3 spec Date: 2026-02-25 (spec), 2026-03-03 (Layer 1 implemented)

Problem¶

Token counts and costs are extracted from LLM API responses by the LiteLLM provider but discarded in the agent loop. The infrastructure exists at both ends — LLMResponse carries prompt_tokens, completion_tokens, total_tokens, cost_usd, and ToolCallLogger has record_turn_usage() ready to persist them — but nothing connects the two. process_direct() returns a bare str, so callers (CLI, eval runner, daemon) have no visibility into resource consumption.

There's also a Tracer class at picobot/agent/trace.py with per-turn TurnTrace dataclass (prompt_tokens, completion_tokens, cost_usd, llm_latency_ms, per-tool-call timing) and JSONL output to ~/.picobot/data/traces/traces.jsonl — but it's imported nowhere. This is a third piece of disconnected plumbing that should be wired alongside the ToolCallLogger.

This blocks three things: 1. Eval framework — can't measure tokens/cost per task without metrics from the loop 2. Session cost monitoring — users can't see what a session cost 3. Agent/sub-agent budgets — can't enforce token limits on tasks or sub-agents

Current State¶

What works¶

What's fixed (Layer 1 — implemented 2026-03-03)¶

What's still open¶

Design¶

Three layers, each building on the previous.

Layer 1: Wire Token Tracking (prerequisite for everything)¶

Connect the existing plumbing so token data flows from LLM responses through to the log and callers.

1a. Record turn usage in the loop — DONE¶

Implemented at agent/loop.py:493. After each _call_llm() call, record_turn_usage() feeds real token counts to the logger:

response = await self._call_llm(messages, _log_event if on_event else None)
tc_logger.record_turn_usage(
    response.prompt_tokens, response.completion_tokens,
    response.total_tokens, response.cost_usd,
)

Streaming gap: _call_llm_stream() still returns zero tokens. Options remain: - Option B (recommended): Parse the final stream chunk for usage data (LiteLLM stream_options={"include_usage": True}). Zero overhead, provider-dependent. - Option C (fallback): Fall back to chars/4 estimate for streamed responses.

1b. Enrich process_direct() return type — DONE¶

Implemented at agent/result.py. process_direct() returns AgentResult:

@dataclass
class AgentResult:
    response: str
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    cost_usd: float = 0.0
    tool_calls: int = 0
    failed_tool_calls: int = 0
    iterations: int = 0
    elapsed_seconds: float = 0.0
    tool_call_log: list[dict] = field(default_factory=list)
    subagent_tokens: int = 0

__str__ returns self.response for backwards compatibility. Populated from tc_logger.session_totals() + wall-clock timer + subagents._subagent_tokens at the end of process_direct() (loop.py:630-643).

1c. Feed real counts into TokenBudget¶

After each LLM call, update a running total of actual tokens consumed. Use this to inform compression decisions instead of (or alongside) the chars/4 heuristic.

# In the loop, after each LLM call:
self.budget.record_actual(response.total_tokens)

# In TokenBudget:
def record_actual(self, tokens: int) -> None:
    self._actual_total += tokens

def should_compress(self, messages):
    # Use actual count if available, fall back to estimate
    estimated = self._actual_total or self.estimate_tokens(messages)
    return estimated > int(self.budget * 0.8)

Note: _actual_total tracks cumulative tokens sent to the API across all turns, not the current message list size. The chars/4 estimate of the current message list is still useful for deciding when to compress — actual API token counts tell you what you've spent, but the message list estimate tells you what the next call will cost. Keep both. Use the estimate for compression decisions, actual counts for reporting.

Layer 2: Session Cost Monitoring¶

Surface token/cost data to users during and after sessions.

2a. Session summary on exit¶

When a session ends (CLI exit, daemon session close), print a summary:

Session: 14 turns, 23 tool calls
Tokens: 45,231 in / 12,847 out (58,078 total)
Cost: $0.087

Data source: tc_logger.session_totals() (already accumulated if Layer 1 is wired).

2b. JSONL log enrichment¶

write_session_end() already writes token totals — once Layer 1 feeds real data in, the JSONL logs become a complete cost record with no additional work. Structure already defined:

{
  "type": "session_end",
  "tokens": {"prompt": N, "completion": N, "total": N, "cost_usd": X},
  "turns": N,
  "tool_calls": N
}

2c. Cost dashboard (optional, lower priority)¶

Extend the TUI dashboard or add a CLI command (picobot cost [--since 7d]) that reads JSONL logs and aggregates: - Cost per session - Cost per day/week - Cost by model - Running total

This is a read-only view over existing log files — no new data collection needed once Layer 1 is in place.

Layer 3: Budget Enforcement¶

Allow setting token/cost limits on agents and sub-agents, with configurable enforcement (warn, soft-stop, hard-stop).

3a. Config schema additions¶

class BudgetConfig(BaseModel):
    """Token/cost budget for an agent."""
    max_tokens: int | None = None       # Total tokens (in+out) before action
    max_cost_usd: float | None = None   # Dollar cost before action
    action: str = "warn"                # "warn" | "stop" | "none"

class AgentDefaults(BaseModel):
    # ... existing fields ...
    budget: BudgetConfig = Field(default_factory=BudgetConfig)

class SubagentConfig(BaseModel):
    """Sub-agent specific limits."""
    max_iterations: int = 15            # Already exists (hardcoded)
    budget: BudgetConfig = Field(default_factory=BudgetConfig)

3b. Budget checking in the loop¶

After each LLM call in process_direct(), check against the budget:

totals = tc_logger.session_totals()
if self.budget_config.max_tokens and totals["total_tokens"] > self.budget_config.max_tokens:
    if self.budget_config.action == "stop":
        final_content = f"Token budget exceeded ({totals['total_tokens']} / {self.budget_config.max_tokens}). Stopping."
        break
    elif self.budget_config.action == "warn":
        logger.warning(f"Token budget exceeded: {totals['total_tokens']} / {self.budget_config.max_tokens}")
        # Continue but warn — agent sees this in next system message

Same pattern for max_cost_usd.

3c. Sub-agent budgets — PARTIALLY DONE¶

Token tracking (done): Sub-agent provider.chat() responses are accumulated in SubagentManager._subagent_tokens. This total is added to AgentResult.total_tokens and exposed separately via AgentResult.subagent_tokens. The eval runner and CLI get honest totals including sub-agent work.

Budget enforcement (not yet done):

1. Pass a BudgetConfig to SubagentManager.spawn() — caller specifies per-task limits
2. Inside _run_subagent(), create a local ToolCallLogger and check budget after each LLM call
3. Default sub-agent budget: inherit from parent config's subagent.budget, overridable per-spawn
4. Surface sub-agent cost in the registry entry so parent sees what each sub-agent spent

# Future: per-sub-agent cost in registry entry
{
    "id": "abc123",
    "status": "completed",
    "result": "...",
    "tokens": {"prompt": N, "completion": N, "total": N, "cost_usd": X}
}

3d. Eval framework integration¶

The eval runner uses budgets differently from interactive sessions:

• Observe, don't enforce: Default for evals is action: "none" — record tokens/cost but don't stop the agent. The eval framework compares actual usage against expected budgets in post-processing.
• Hard limits for runaway prevention: Set a generous max (e.g., 500k tokens) as a safety net to kill tasks that enter infinite loops, not as a capability test.
• Budget as a metric, not a gate: The eval results JSON includes token/cost data per task. Whether the agent was "efficient" is a separate analysis from whether it succeeded.

Implementation Order¶

Layer 1 (1a + 1b) is done. Sub-agent token tracking (3c partial) is done. Remaining items are incremental.

Callers Affected by AgentResult Change¶

The __str__ fallback means no caller broke during the migration. The eval runner is the primary consumer of metrics.