feat(agent): typed budget errors, per-step cost tracking, and structred output mode

[AseemPrasad]

Three inter-related improvements to the agentic tool-call runtime in backend/agent/engine.py and
supporting modules. All changes are fully opt-in via the new AGENT_STRUCTURED_OUTPUT environment
variable — existing deployments are completely unaffected.

Feature	Backend change	Frontend change
Typed budget errors	✅	none (new WS message type)
Per-step cost tracking	✅	none
Structured output mode	✅	none

---

Feature 1: Typed BudgetExceededError

Problem

BudgetExceededError was raised with a bare string and re-stringified by the caller, losing the error type.
Both "global budget exceeded" and "per-step budget exceeded" looked identical to the frontend.

Solution

BudgetExceededError and its new subclass PerStepBudgetExceededError carry two class attributes:

class BudgetExceededError(Exception):                                                                        
    is_per_step: bool = False                                                                                
    typed_message: str = "Generation stopped: this variant exceeded its resource limit."                     
                                                                                                             
class PerStepBudgetExceededError(BudgetExceededError):                                                       
    is_per_step = True                                                                                       
    typed_message = "Generation stopped: this variant exceeded its per-step spend budget."                   

Both are caught in PipelineContext._run_variant() and sent as a new budgetExceeded WebSocket message:

  type: "budgetExceeded"                                                                                       
  value: "Generation stopped: this variant exceeded its per-step spend budget."                                
  data: {"is_per_step": true}                                                                                  

Key design decision: unlike variantError, budgetExceeded does not close the session. The user sees a clear
banner and can retry with a different model or settings immediately.

────────────────────────────────────────────────────────────────────────────────

Feature 2: Per-Step Cost Tracking

Problem

Cost was only visible in aggregate at the end of a run (and only if PROMPT_REPORTS_ENABLED=true). There was no
way to observe which step caused a cost spike.

Solution

  [STEP] variant=0 step=3 step_cost=$0.0042 cum_cost=$0.0913                                                   

• pre_step_spend is captured before session.stream_turn() every loop iteration
• post_step_spend is captured after session.append_tool_results()
• The delta (LLM turn + all tool executions for that step) is logged and recorded

AgentRunRecorder.record_step_cost() writes a step_cost event to events.jsonl live, and appends a step_costs[]
array to run.json at finalisation:

  "step_costs": [                                                                                              
    {"step": 1, "step_cost_usd": 0.0182, "cumulative_cost_usd": 0.0182},                                       
    {"step": 2, "step_cost_usd": 0.0041, "cumulative_cost_usd": 0.0223},                                       
    {"step": 3, "step_cost_usd": 0.0008, "cumulative_cost_usd": 0.0231}                                        
  ]                                                                                                            

────────────────────────────────────────────────────────────────────────────────

Feature 3: Structured Output Mode (Opt-In)

Problem

Models with weaker JSON-mode support (especially Gemini family members) occasionally emit malformed JSON tool
calls, causing InvalidJsonToolCallError and retry loops. There was no way to force tool-call mode or to switch
on structured output without a full config refactor.

Solution

New file: backend/agent/modes.py

  class StructuredOutputMode(Enum):                                                                            
      FREE = "free"           # model may return text or tool call (default)                                   
      PREFER_JSON = "prefer_json"  # JSON-mode guarantee without forcing tool                                  
      FORCE_TOOL = "force_tool"   # model must call a tool every turn                                          
                                                                                                               
  class ToolCallPolicy(Enum):                                                                                  
      FREE = "free"     # no constraint (default)                                                              
      WARN = "warn"     # log warning if model emits text instead of tool call                                 
      REQUIRED = "required"  # raise error if no tool call on a step                                           

Config (all in backend/.env, opt-in):

  # Master switch — must be true for the next three to have any effect                                         
  AGENT_STRUCTURED_OUTPUT=true                                                                                 
                                                                                                               
  # StructuredOutputMode: free | prefer_json | force_tool (default: force_tool)                                
  AGENT_STRUCTURED_OUTPUT_MODE=force_tool                                                                      
                                                                                                               
  # ToolCallPolicy: free | warn | required (default: free)                                                     
  AGENT_TOOL_CALL_POLICY=required                                                                              
                                                                                                               
  # Optional per-step spend ceiling in USD (default: none)                                                     
  AGENT_STEP_SPEND_BUDGET_USD=0.05                                                                             

Provider wiring:

┌───────────┬────────────────────┬────────────────────────────────┬──────────────────────────────────────────┐
│ Provider │ FREE │ PREFER_JSON │ FORCE_TOOL │
├───────────┼────────────────────┼────────────────────────────────┼──────────────────────────────────────────┤
│ OpenAI │ tool_choice="auto" │ tool_choice="auto" + JSON mode │ tool_choice="required" │
├───────────┼────────────────────┼────────────────────────────────┼──────────────────────────────────────────┤
│ Anthropic │ no system nudge │ no system nudge │ prepends invisible nudge to system │
│ │ │ │ prompt │
├───────────┼────────────────────┼────────────────────────────────┼──────────────────────────────────────────┤
│ Gemini │ no change │ no change │ no change (see below) │
└───────────┴────────────────────┴────────────────────────────────┴──────────────────────────────────────────┘

│ Note on Gemini: Gemini tool-calling uses forced_function_calling which is all-or-nothing at the API level.
│ Full FORCE_TOOL support requires deeper API changes and is deferred to a future PR.

────────────────────────────────────────────────────────────────────────────────

Files Changed

┌─────────────────────────────────────────┬──────────────────────────────────────────────────────────────────┐
│ File │ Change │
├─────────────────────────────────────────┼──────────────────────────────────────────────────────────────────┤
│ backend/agent/modes.py │ New — StructuredOutputMode, ToolCallPolicy, mapping dicts │
├─────────────────────────────────────────┼──────────────────────────────────────────────────────────────────┤
│ backend/config.py │ 4 new env vars with validation + typed annotations │
├─────────────────────────────────────────┼──────────────────────────────────────────────────────────────────┤
│ backend/ws/constants.py │ BUDGET_EXCEEDED_CODE = 4333 added │
├─────────────────────────────────────────┼──────────────────────────────────────────────────────────────────┤
│ backend/agent/engine.py │ BudgetExceededError refactored; per-step budget gate; step cost │
│ │ tracking; structured output wiring │
├─────────────────────────────────────────┼──────────────────────────────────────────────────────────────────┤
│ backend/routes/generate_code.py │ budgetExceeded message type; typed exception handler │
├─────────────────────────────────────────┼──────────────────────────────────────────────────────────────────┤
│ backend/fs_logging/agent_runs.py │ record_step_cost() method; step_costs[] in run.json │
├─────────────────────────────────────────┼──────────────────────────────────────────────────────────────────┤
│ backend/agent/providers/factory.py │ structured_output_mode param → provider-native params │
├─────────────────────────────────────────┼──────────────────────────────────────────────────────────────────┤
│ backend/agent/providers/openai.py │ tool_choice in init + stream_turn │
├─────────────────────────────────────────┼──────────────────────────────────────────────────────────────────┤
│ backend/agent/providers/anthropic/provi │ tool_nudge in init + prepended to system prompt │
│ der.py │

Testing Notes

• Zero new errors in pyright on all changed files (checked with python -m pyright agent/engine.py
  agent/modes.py routes/generate_code.py fs_logging/agent_runs.py)
• All modified files compile successfully with python -m py_compile
• No changes to the frontend are required — the new budgetExceeded WebSocket message is additive
• The budgetExceeded message type should be documented in the frontend's WebSocket message handler
  (frontend/src/lib/ws-messages.ts or equivalent) to surface a distinct UI banner

[AseemPrasad]

@abi
would love to get this contribution reviewed and incorporated..
Thank you..