diff --git a/backend/agent/engine.py b/backend/agent/engine.py index 11ead964f..765d5c268 100644 --- a/backend/agent/engine.py +++ b/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: diff --git a/backend/agent/modes.py b/backend/agent/modes.py new file mode 100644 index 000000000..8cb16d23f --- /dev/null +++ b/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." + ), +} diff --git a/backend/agent/providers/anthropic/provider.py b/backend/agent/providers/anthropic/provider.py index 49d880da4..9c3d2ebc1 100644 --- a/backend/agent/providers/anthropic/provider.py +++ b/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"}, diff --git a/backend/agent/providers/factory.py b/backend/agent/providers/factory.py index 767bfb473..5adc33e6b 100644 --- a/backend/agent/providers/factory.py +++ b/backend/agent/providers/factory.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import Optional from anthropic import AsyncAnthropic @@ -5,6 +7,7 @@ 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,12 +47,18 @@ 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: @@ -56,12 +66,19 @@ def create_provider_session( 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: @@ -69,6 +86,8 @@ def create_provider_session( 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, diff --git a/backend/agent/providers/openai.py b/backend/agent/providers/openai.py index d1e2cdbea..096059e0f 100644 --- a/backend/agent/providers/openai.py +++ b/backend/agent/providers/openai.py @@ -424,12 +424,14 @@ def __init__( prompt_messages: List[ChatCompletionMessageParam], tools: List[Dict[str, Any]], recorder: Optional[AgentRunRecorder] = None, + tool_choice: str | None = None, ): self._client = client self._model = model self._tools = tools self._total_usage = TokenUsage() self._recorder = recorder + self._tool_choice = tool_choice or "auto" self._prompt_report_logger = PromptReportLogger( provider="openai", model=model, @@ -447,7 +449,7 @@ async def stream_turn(self, on_event: EventSink) -> ProviderTurn: "model": model_name, "input": self._input_items, "tools": self._tools, - "tool_choice": "auto", + "tool_choice": self._tool_choice, "stream": True, "max_output_tokens": 50000, } diff --git a/backend/config.py b/backend/config.py index 6631629a9..0d406f704 100644 --- a/backend/config.py +++ b/backend/config.py @@ -35,3 +35,47 @@ # Set to True when running in production (on the hosted version) # Used as a feature flag to enable or disable certain features IS_PROD = os.environ.get("IS_PROD", False) + +# -------------------------------------------------------------------------- # +# Agent / tool-runtime settings +# -------------------------------------------------------------------------- # + +# Opt-in switch for structured-output / forced-tool-call behaviour. +# When True the factory switches to AGENT_STRUCTURED_OUTPUT_MODE and +# AGENT_TOOL_CALL_POLICY below, overriding the provider defaults. +# Backward-compatible: defaults to False so existing deployments are unaffected. +AGENT_STRUCTURED_OUTPUT = ( + os.environ.get("AGENT_STRUCTURED_OUTPUT", "").strip().lower() in {"1", "true", "yes", "on"} +) + +# Which StructuredOutputMode to use when AGENT_STRUCTURED_OUTPUT is True. +# Options: "free" | "prefer_json" | "force_tool" +# Default: "force_tool" — the safest setting for unreliable JSON-emitting models. +_AGENT_STRUCTURED_OUTPUT_MODE = os.environ.get( + "AGENT_STRUCTURED_OUTPUT_MODE", "force_tool" +).strip().lower() +if _AGENT_STRUCTURED_OUTPUT_MODE not in {"free", "prefer_json", "force_tool"}: + _AGENT_STRUCTURED_OUTPUT_MODE = "force_tool" +AGENT_STRUCTURED_OUTPUT_MODE: str = _AGENT_STRUCTURED_OUTPUT_MODE # consumed by factory + +# Tool-call policy applied on every step of the agent loop. +# Options: "free" | "warn" | "required" +# "required" causes UnexpectedTextStepError (AgentStepError subclass) when the +# model emits plain text instead of a tool call on a step. +_AGENT_TOOL_CALL_POLICY = os.environ.get("AGENT_TOOL_CALL_POLICY", "free").strip().lower() +if _AGENT_TOOL_CALL_POLICY not in {"free", "warn", "required"}: + _AGENT_TOOL_CALL_POLICY = "free" +AGENT_TOOL_CALL_POLICY: str = _AGENT_TOOL_CALL_POLICY # consumed by factory + +# Optional per-step spend ceiling in USD. +# If set, any step that would bring cumulative spend above this threshold is +# aborted with BudgetExceededError (sent to the frontend as budgetExceeded). +# Default: None (no per-step cap; only the global GENERATION_MAX_COST_USD applies). +_AGENT_STEP_SPEND_BUDGET_USD = os.environ.get("AGENT_STEP_SPEND_BUDGET_USD", "").strip() +AGENT_STEP_SPEND_BUDGET_USD: float | None = ( + float(_AGENT_STEP_SPEND_BUDGET_USD) if _AGENT_STEP_SPEND_BUDGET_USD else None +) + +# -------------------------------------------------------------------------- # +# End agent settings +# -------------------------------------------------------------------------- # diff --git a/backend/fs_logging/agent_runs.py b/backend/fs_logging/agent_runs.py index abb144460..0b5be920e 100644 --- a/backend/fs_logging/agent_runs.py +++ b/backend/fs_logging/agent_runs.py @@ -300,6 +300,7 @@ def __init__( self._tool_asset_urls: set[str] = set() self._tool_asset_tasks: list["asyncio.Task[None]"] = [] self._tool_asset_manifest: list[dict[str, Any]] = [] + self._step_cost_log: list[dict[str, Any]] = [] # ------------------------------------------------------------------ paths @@ -721,6 +722,28 @@ def record_set_code(self, content_len: int, source: str) -> None: except Exception as exc: print(f"[AGENT RUN] Failed to record set_code: {exc}") + def record_step_cost( + self, step: int, step_cost_usd: float, cumulative_cost_usd: float + ) -> None: + """Record the cost incurred by a single tool-call step. + + Called from ``AgentEngine._run_with_session`` after ``append_tool_results`` + so the full cost (LLM turn + all tool executions) is included. + Written to JSONL live and aggregated into run.json at finalisation. + """ + if not self.enabled: + return + try: + entry = { + "step": step, + "step_cost_usd": step_cost_usd, + "cumulative_cost_usd": cumulative_cost_usd, + } + self._step_cost_log.append(entry) + self._append_event("step_cost", entry) + except Exception as exc: + print(f"[AGENT RUN] Failed to record step cost: {exc}") + # --------------------------------------------------------------- finalize async def record_run_end( @@ -790,6 +813,7 @@ async def record_run_end( "has_unpriced_calls": self._has_unpriced_calls, "llm_calls": self._llm_call_summaries, "tool_calls": self._tool_call_summaries, + "step_costs": self._step_cost_log, "tool_assets": self._tool_asset_manifest, "final_html": final_html, } diff --git a/backend/routes/generate_code.py b/backend/routes/generate_code.py index 091874ba6..56a8c1d37 100644 --- a/backend/routes/generate_code.py +++ b/backend/routes/generate_code.py @@ -52,6 +52,8 @@ "assistant", "toolStart", "toolResult", + "budgetExceeded", + "variantCost", ] from prompts.pipeline import build_prompt_messages from prompts.request_parsing import parse_prompt_content, parse_prompt_history @@ -62,6 +64,7 @@ infer_local_asset_base_url, ) from agent.runner import Agent +from agent.engine import BudgetExceededError, PerStepBudgetExceededError from fs_logging.agent_runs import AgentRunRecorder from routes.model_choice_sets import ( ALL_KEYS_MODELS_DEFAULT, @@ -656,6 +659,18 @@ async def send_runner_message( recorder=recorder, ) completion = await runner.run(model, prompt_messages) + # Emit per-variant cost attribution once the session is finalised. + # Only sent when the backend can compute a dollar figure (i.e. when + # the provider session had token-usage data and a pricing entry). + final_cost = runner.last_cost_usd + if final_cost is not None: + await self.send_message( + "variantCost", + None, + index, + {"costUsd": final_cost}, + None, + ) if completion: await self.send_message("setCode", completion, index, None, None) await self.send_message( @@ -706,6 +721,23 @@ async def send_runner_message( ) await self.send_message("variantError", error_message, index, None, None) return "" + except (BudgetExceededError, PerStepBudgetExceededError) as exc: + # ``BudgetExceededError`` carries a typed human-readable message + # delivered as a distinct ``budgetExceeded`` message so the frontend + # can render a distinct UI banner. The WebSocket session stays alive + # (unlike ``variantError``) so the user can try again immediately. + print( + f"[VARIANT {index + 1}] Budget exceeded " + f"(per_step={exc.is_per_step}): {exc.typed_message}" + ) + await self.send_message( + "budgetExceeded", + exc.typed_message, + index, + {"is_per_step": exc.is_per_step}, + None, + ) + return "" except Exception as e: print(f"Error in variant {index + 1}: {e}") traceback.print_exception(type(e), e, e.__traceback__) diff --git a/backend/ws/constants.py b/backend/ws/constants.py index e992ecbd1..ac7bdbdeb 100644 --- a/backend/ws/constants.py +++ b/backend/ws/constants.py @@ -1,2 +1,10 @@ # WebSocket protocol (RFC 6455) allows for the use of custom close codes in the range 4000-4999 +# RFC 6455 custom-code range: 4000-4999 +# Used when the backend encounters an application-level error that requires +# the client to reconnect rather than continue the session. APP_ERROR_WEB_SOCKET_CODE = 4332 + +# Used when a generation is aborted because a per-step or cumulative spend +# budget was exceeded. Unlike APP_ERROR_WEB_SOCKET_CODE this does not invalidate +# the session — the frontend may surface the message and stay connected. +BUDGET_EXCEEDED_CODE = 4333 diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index 1cbab73a1..e134ff56b 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -64,6 +64,7 @@ function App() { updateVariantStatus, resizeVariants, setVariantModels, + setVariantCost, appendVariantHistoryMessage, startAgentEvent, appendAgentEventContent, @@ -477,6 +478,9 @@ function App() { onVariantModels: (models) => { setVariantModels(commit.hash, models); }, + onVariantCost: (variantIndex, costUsd, inputTokens, outputTokens) => { + setVariantCost(commit.hash, variantIndex, costUsd, inputTokens, outputTokens); + }, onThinking: (content, variantIndex, eventId) => { if (!eventId) return; lastThinkingEventIdRef.current[variantIndex] = eventId; diff --git a/frontend/src/components/commits/types.ts b/frontend/src/components/commits/types.ts index f95450f84..e508fb6d3 100644 --- a/frontend/src/components/commits/types.ts +++ b/frontend/src/components/commits/types.ts @@ -38,6 +38,13 @@ export type Variant = { thinkingDuration?: number; agentEvents?: AgentEvent[]; model?: string; + /** USD cost for this variant, populated after generation completes. */ + costUsd?: number; + /** Token usage for this variant, populated after generation completes. */ + tokens?: { + input: number; + output: number; + }; }; export type BaseCommit = { diff --git a/frontend/src/generateCode.ts b/frontend/src/generateCode.ts index 50033fd1f..35b5e682a 100644 --- a/frontend/src/generateCode.ts +++ b/frontend/src/generateCode.ts @@ -24,7 +24,8 @@ type WebSocketResponse = { | "thinking" | "assistant" | "toolStart" - | "toolResult"; + | "toolResult" + | "variantCost"; value?: string; data?: any; eventId?: string; @@ -43,6 +44,13 @@ interface CodeGenerationCallbacks { onAssistant: (content: string, variantIndex: number, eventId?: string) => void; onToolStart: (data: any, variantIndex: number, eventId?: string) => void; onToolResult: (data: any, variantIndex: number, eventId?: string) => void; + /** Called with the final cost for a variant after its agent run completes. */ + onVariantCost: ( + variantIndex: number, + costUsd: number, + inputTokens: number, + outputTokens: number + ) => void; onCancel: ( reason: "user_cancelled" | "request_failed" | "connection_error", errorMessage?: string @@ -92,6 +100,13 @@ export function generateCode( } else if (response.type === "error") { console.error("Error generating code", response.value); toast.error(response.value || ERROR_MESSAGE); + } else if (response.type === "variantCost") { + callbacks.onVariantCost( + response.variantIndex, + response.data?.costUsd ?? 0, + response.data?.inputTokens ?? 0, + response.data?.outputTokens ?? 0 + ); } }); diff --git a/frontend/src/store/project-store.ts b/frontend/src/store/project-store.ts index 03286dc5d..186292a17 100644 --- a/frontend/src/store/project-store.ts +++ b/frontend/src/store/project-store.ts @@ -48,6 +48,14 @@ interface ProjectStore { message: VariantHistoryMessage ) => void; updateSelectedVariantIndex: (hash: CommitHash, index: number) => void; + /** Records cost and token usage for a specific variant after generation completes. */ + setVariantCost: ( + hash: CommitHash, + numVariant: number, + costUsd: number, + inputTokens: number, + outputTokens: number + ) => void; updateVariantStatus: ( hash: CommitHash, numVariant: number, @@ -276,6 +284,28 @@ export const useProjectStore = create((set, get) => ({ }, }; }), + setVariantCost: (hash, numVariant, costUsd, inputTokens, outputTokens) => + set((state) => { + const commit = state.commits[hash]; + if (!commit) return state; + return { + commits: { + ...state.commits, + [hash]: { + ...commit, + variants: commit.variants.map((variant, index) => + index === numVariant + ? { + ...variant, + costUsd, + tokens: { input: inputTokens, output: outputTokens }, + } + : variant + ), + }, + }, + }; + }), updateVariantStatus: ( hash: CommitHash, numVariant: number,