feat: per-variant cost attribution via variantCost WebSocket message

backend/agent/engine.py

@@ -1,13 +1,14 @@
 import asyncio
 import traceback
 import uuid
-from typing import Any, Awaitable, Callable, Dict, List, Optional, cast
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Union, cast
 
 from openai.types.chat import ChatCompletionMessageParam
 
 from codegen.utils import extract_html_content
 from llm import Llm
 
+from agent.modes import StructuredOutputMode
 from agent.providers.base import ExecutedToolCall, ProviderSession, StreamEvent
 from agent.providers.factory import create_provider_session
 from agent.state import AgentFileState, seed_file_state_from_messages
@@ -18,7 +19,12 @@
     summarize_text,
     summarize_tool_input,
 )
-from config import GENERATION_MAX_COST_USD
+from config import (
+    AGENT_STEP_SPEND_BUDGET_USD,
+    AGENT_STRUCTURED_OUTPUT,
+    AGENT_TOOL_CALL_POLICY,
+    GENERATION_MAX_COST_USD,
+)
 from fs_logging.agent_runs import AgentRunRecorder
 
 
@@ -31,21 +37,38 @@ class EmptyOutputError(Exception):
     the output file is empty. Raising makes it a normal, retryable failure.
     """
 
-    def __init__(self) -> None:
-        super().__init__("Generation finished without producing any output.")
+    pass  # message is only for logging; callers ignore it
 
 
 class BudgetExceededError(Exception):
     """Raised when a single generation exceeds the spend ceiling.
 
-    The message is shown verbatim to end users (variantError), so it must
-    not contain cost figures; the exact spend is in the run record.
+    The ``typed_message`` class attribute is sent verbatim to the frontend
+    as the ``reason`` field of the ``budgetExceeded`` WebSocket message.  It
+    must not contain raw cost figures — those are reserved for the run record.
+
+    Subclasses may set ``is_per_step`` to indicate whether the budget was a
+    per-step ceiling (True) or the global ``GENERATION_MAX_COST_USD`` (False).
     """
 
+    is_per_step: bool = False
+    typed_message: str = "Generation stopped: this variant exceeded its resource limit."
+
     def __init__(self) -> None:
-        super().__init__(
-            "Generation stopped: this variant exceeded its resource limit."
-        )
+        super().__init__(self.typed_message)
+
+
+class PerStepBudgetExceededError(BudgetExceededError):
+    """Specialisation of ``BudgetExceededError`` for per-step budget hits.
+
+    Raised by ``StepCostTracker`` when a step would push cumulative spend past
+    ``AGENT_STEP_SPEND_BUDGET_USD`` before the step's tool executions begin.
+    """
+
+    is_per_step = True
+    typed_message = (
+        "Generation stopped: this variant exceeded its per-step spend budget."
+    )
 
 
 class AgentEngine:
@@ -96,6 +119,29 @@ def __init__(
         )
         self._tool_preview_lengths: Dict[str, int] = {}
 
+        # --- Structured-output / tool-call policy derived from config ----------
+        self._structured_output_mode: str = (
+            "force_tool" if AGENT_STRUCTURED_OUTPUT else "free"
+        )
+
+        self._tool_call_policy: str = (
+            AGENT_TOOL_CALL_POLICY if AGENT_STRUCTURED_OUTPUT else "free"
+        )
+
+        self._step_spend_budget_usd: float | None = (
+            AGENT_STEP_SPEND_BUDGET_USD if AGENT_STRUCTURED_OUTPUT else None
+        )
+
+        # --- Per-run cost tracking -------------------------------------------
+        self._step_count: int = 0
+        self._step_costs: List[float] = []
+        self._last_cost_usd: float | None = None
+
+    @property
+    def last_cost_usd(self) -> float | None:
+        """Final USD cost for the most recent run(), available after run() returns."""
+        return self._last_cost_usd
+
     @staticmethod
     def _extract_input_images(
         prompt_messages: List[ChatCompletionMessageParam],
@@ -219,6 +265,29 @@ async def _run_with_session(self, session: ProviderSession) -> str:
         max_steps = 30
 
         for _ in range(max_steps):
+            self._step_count += 1
+            step_num = self._step_count
+
+            # --- Capture spend *before* the step for budget checks + tracking -
+            # Always call total_cost_usd() here so pre_step_spend is in scope
+            # for the cost-recording block at the end of the loop.
+            pre_step_spend: float | None = session.total_cost_usd()
+
+            # --- Per-step budget gate (before any tool side-effects) ----------
+            # Unpriced models (None) bypass this check.
+            if (
+                self._step_spend_budget_usd is not None
+                and pre_step_spend is not None
+                and pre_step_spend >= self._step_spend_budget_usd
+            ):
+                print(
+                    f"[BUDGET] Aborting variant {self.variant_index} "
+                    f"before step {step_num}: "
+                    f"${pre_step_spend:.2f} >= ${self._step_spend_budget_usd:.2f} "
+                    f"(per-step ceiling)"
+                )
+                raise PerStepBudgetExceededError()
+
             assistant_event_id = self._next_event_id("assistant")
             thinking_event_id = self._next_event_id("thinking")
             started_tool_ids: set[str] = set()
@@ -264,14 +333,15 @@ async def on_event(event: StreamEvent) -> None:
             if not turn.tool_calls:
                 return await self._finalize_response(turn.assistant_text)
 
+            # --- Main / global budget gate ----------------------------------
             # Abort only when the run would otherwise continue: a run that
             # just produced its final answer is already paid for. Unpriced
             # models return None and are not bounded.
             spent = session.total_cost_usd()
             if spent is not None and spent > GENERATION_MAX_COST_USD:
                 print(
-                    f"[BUDGET] Aborting variant {self.variant_index}: "
-                    f"${spent:.2f} > ${GENERATION_MAX_COST_USD:.2f}"
+                    f"[BUDGET] Aborting variant {self.variant_index} at step {step_num}: "
+                    f"${spent:.2f} > ${GENERATION_MAX_COST_USD:.2f} (global ceiling)"
                 )
                 raise BudgetExceededError()
 
@@ -324,6 +394,18 @@ async def on_event(event: StreamEvent) -> None:
 
             await session.append_tool_results(turn, executed_tool_calls)
 
+            # --- Record step cost for observability --------------------------
+            post_step_spend = session.total_cost_usd()
+            if post_step_spend is not None:
+                step_cost = post_step_spend - (pre_step_spend or 0.0)
+                self._step_costs.append(step_cost)
+                if self.recorder is not None:
+                    self.recorder.record_step_cost(step_num, step_cost, post_step_spend)
+                print(
+                    f"[STEP] variant={self.variant_index} step={step_num} "
+                    f"step_cost=${step_cost:.4f} cum_cost=${post_step_spend:.4f}"
+                )
+
         raise Exception("Agent exceeded max tool turns")
 
     async def run(self, model: Llm, prompt_messages: List[ChatCompletionMessageParam]) -> str:
@@ -333,6 +415,11 @@ async def run(self, model: Llm, prompt_messages: List[ChatCompletionMessageParam
         if self.recorder is not None:
             self.recorder.record_run_start(model, prompt_messages)
 
+        structured_output_mode: StructuredOutputMode | None = (
+            StructuredOutputMode(self._structured_output_mode)
+            if self._structured_output_mode != "free"
+            else None
+        )
         session = create_provider_session(
             model=model,
             prompt_messages=prompt_messages,
@@ -350,6 +437,7 @@ async def run(self, model: Llm, prompt_messages: List[ChatCompletionMessageParam
                 self.should_extract_assets and bool(self.tool_runtime.input_images)
             ),
             recorder=self.recorder,
+            structured_output_mode=structured_output_mode,
         )
         try:
             result = await self._run_with_session(session)
@@ -370,6 +458,8 @@ async def run(self, model: Llm, prompt_messages: List[ChatCompletionMessageParam
                 )
             raise
         finally:
+            # Capture cost before closing so callers can read last_cost_usd.
+            self._last_cost_usd = session.total_cost_usd()
             await session.close()
 
     async def _finalize_response(self, assistant_text: str) -> str:

backend/agent/modes.py

@@ -0,0 +1,87 @@
+"""Agent runtime modes: structured-output strategy and tool-call policy.
+
+These types are intentionally provider-agnostic so that the pipeline config
+lives in one place (config.py / factory.py) and each provider maps the mode to
+its own API knobs independently.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Literal
+
+
+# -------------------------------------------------------------------------- #
+# StructuredOutputMode — which output format the model should produce.
+# -------------------------------------------------------------------------- #
+class StructuredOutputMode(Enum):
+    """Output format strategy for the agent's LLM calls.
+
+    Ordered from most flexible (free-form) to most constrained (forced tool).
+    The factory maps each value to provider-native parameters.
+    """
+
+    # The model may return plain text / free-form JSON.  Tool calls are
+    # preferred but not required.  Safe default for new models that have
+    # strong JSON reasoning capabilities.
+    FREE = "free"
+
+    # The model is instructed to emit JSON matching the tool schema via a
+    # structured-output / JSON-schema guarantee.  Tool calls remain optional
+    # (the model can still return text).  Useful for mid-generation models
+    # that occasionally emit malformed JSON.
+    PREFER_JSON = "prefer_json"
+
+    # The model is forced to emit a tool call on every turn.  Recommended for
+    # models that frequently skip tool invocations or that do not yet have
+    # reliable JSON-mode support (e.g. legacy o1-preview / o3-mini variants).
+    FORCE_TOOL = "force_tool"
+
+
+# -------------------------------------------------------------------------- #
+# ToolCallPolicy — per-step tool-call enforcement level.
+# -------------------------------------------------------------------------- #
+class ToolCallPolicy(Enum):
+    """Policy that gates whether the agent is allowed to emit a plain text
+    response instead of a tool call on any given step.
+
+    Unlike ``StructuredOutputMode`` which controls the *format* of the output,
+    ``ToolCallPolicy`` controls whether a step is allowed to terminate without
+    a tool invocation at all.
+    """
+
+    # No constraint — the model may emit text or a tool call at each step.
+    FREE = "free"
+
+    # If the model emits text without a tool call, log a warning and count the
+    # step.  The run continues; the warning is emitted to the recorder / logs.
+    WARN = "warn"
+
+    # The agent must emit at least one tool call per step.  If the model emits
+    # plain text without any tool invocation, raise ``UnexpectedTextStepError``
+    # (a subclass of ``AgentStepError``) and mark the step as failed.
+    REQUIRED = "required"
+
+
+# -------------------------------------------------------------------------- #
+# Mapping helpers used by the factory / provider initialisation.
+# -------------------------------------------------------------------------- #
+
+# OpenAI Responses API tool_choice values that correspond to each mode.
+OPENAI_TOOL_CHOICE: dict[StructuredOutputMode, str] = {
+    StructuredOutputMode.FREE: "auto",
+    StructuredOutputMode.PREFER_JSON: "auto",  # JSON-mode is set via response_format
+    StructuredOutputMode.FORCE_TOOL: "required",
+}
+
+# Anthropic tool-choice / prompt strategy per mode.
+# Anthropic does not have an equivalent of OpenAI's "required" tool_choice;
+# we simulate it by prepending an invisible system nudge.
+ANTHROPIC_TOOL_NUDGE: dict[StructuredOutputMode, str | None] = {
+    StructuredOutputMode.FREE: None,
+    StructuredOutputMode.PREFER_JSON: None,
+    StructuredOutputMode.FORCE_TOOL: (
+        "You must call exactly one tool on every turn. "
+        "Do not respond with plain text."
+    ),
+}

backend/agent/providers/anthropic/provider.py

@@ -326,12 +326,14 @@ def __init__(
         prompt_messages: List[ChatCompletionMessageParam],
         tools: List[Dict[str, Any]],
         recorder: Optional[AgentRunRecorder] = None,
+        tool_nudge: str | None = None,
     ):
         self._client = client
         self._model = model
         self._tools = tools
         self._total_usage = TokenUsage()
         self._recorder = recorder
+        self._tool_nudge = tool_nudge
         self._prompt_report_logger = PromptReportLogger(
             provider="anthropic",
             model=model,
@@ -356,10 +358,16 @@ async def stream_turn(self, on_event: EventSink) -> ProviderTurn:
         # Tool screenshots accumulate across turns. Re-check before every API
         # call so crossing 20 images cannot leave earlier images above 2000 px.
         self._ensure_many_image_dimension_limit()
+        system_for_api: str | List[Dict[str, Any]] = self._system_prompt
+        if self._tool_nudge:
+            # Prepend the nudge to the system prompt so it is visible to the
+            # model on every turn.  Using a single blank line as separator
+            # preserves any existing prompt structure.
+            system_for_api = f"{self._tool_nudge}\n\n{self._system_prompt}"
         stream_kwargs: Dict[str, Any] = {
             "model": _get_anthropic_api_model_name(self._model),
             "max_tokens": 50000,
-            "system": self._system_prompt,
+            "system": system_for_api,
             "messages": self._messages,
             "tools": self._tools,
             "cache_control": {"type": "ephemeral"},

backend/agent/providers/factory.py

@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 from typing import Optional
 
 from anthropic import AsyncAnthropic
 from google import genai
 from openai import AsyncOpenAI
 from openai.types.chat import ChatCompletionMessageParam
 
+from agent.modes import ANTHROPIC_TOOL_NUDGE, OPENAI_TOOL_CHOICE, StructuredOutputMode
 from agent.providers.anthropic import AnthropicProviderSession, serialize_anthropic_tools
 from agent.providers.base import ProviderSession
 from agent.providers.gemini import GeminiProviderSession, serialize_gemini_tools
@@ -27,6 +30,7 @@ def create_provider_session(
     replicate_api_key: Optional[str],
     should_extract_assets: bool = True,
     recorder: Optional[AgentRunRecorder] = None,
+    structured_output_mode: StructuredOutputMode | None = None,
 ) -> ProviderSession:
     canonical_tools = canonical_tool_definitions(
         image_generation_enabled=should_generate_images,
@@ -43,32 +47,47 @@ def create_provider_session(
             raise Exception("OpenAI API key is missing.")
 
         client = AsyncOpenAI(api_key=openai_api_key, base_url=openai_base_url)
+        tool_choice: str | None = (
+            OPENAI_TOOL_CHOICE[structured_output_mode]  # type: ignore[index]
+            if structured_output_mode is not None
+            else "auto"
+        )
         return OpenAIProviderSession(
             client=client,
             model=model,
             prompt_messages=prompt_messages,
             tools=serialize_openai_tools(canonical_tools),
             recorder=recorder,
+            tool_choice=tool_choice,
         )
 
     if model in ANTHROPIC_MODELS:
         if not anthropic_api_key:
             raise Exception("Anthropic API key is missing.")
 
         client = AsyncAnthropic(api_key=anthropic_api_key)
+        # Anthropic has no "force tool" API knob; we inject a system nudge instead.
+        tool_nudge: str | None = (
+            ANTHROPIC_TOOL_NUDGE[structured_output_mode]  # type: ignore[index]
+            if structured_output_mode is not None
+            else None
+        )
         return AnthropicProviderSession(
             client=client,
             model=model,
             prompt_messages=prompt_messages,
             tools=serialize_anthropic_tools(canonical_tools),
             recorder=recorder,
+            tool_nudge=tool_nudge,
         )
 
     if model in GEMINI_MODELS:
         if not gemini_api_key:
             raise Exception("Gemini API key is missing.")
 
         client = genai.Client(api_key=gemini_api_key)
+        # Gemini tool-calling is all-or-nothing via forced_function_calling;
+        # map "force_tool" to the only available mode.
         return GeminiProviderSession(
             client=client,
             model=model,