Skip to content

feat: per-variant cost attribution via variantCost WebSocket message - #632

Open
AseemPrasad wants to merge 2 commits into
abi:mainfrom
AseemPrasad:asemstc
Open

feat: per-variant cost attribution via variantCost WebSocket message#632
AseemPrasad wants to merge 2 commits into
abi:mainfrom
AseemPrasad:asemstc

Conversation

@AseemPrasad

Copy link
Copy Markdown

Summary

Adds per-variant cost attribution to the WebSocket pipeline. After each variant's agent run completes,
the backend now emits a variantCost WebSocket message carrying the USD cost. The frontend stores it on the
Variant object, making it available for live display, the eval compare page, and any future cost UI.

This is the minimal, highest-value slice of PR 2 — it works with every existing model today (including all
GPT-5, Claude Opus 4/5, and Gemini variants) and requires no changes to pricing tables.

New model enum variants (GPT-5.6, Claude Opus 4.10, etc.) and their pricing entries are tracked as TODOs
pending provider API confirmation.


What Changed

Backend

agent/engine.py

  • Added _last_cost_usd: float | None field and last_cost_usd property to AgentEngine.
  • In run(), the cost is captured from session.total_cost_usd() before the session finally block
    closes it. Callers can always read runner.last_cost_usd after run() returns.

routes/generate_code.py

  • Added "variantCost" to the MessageType Literal union.

  • In _run_variant(), after runner.run() completes:

    final_cost = runner.last_cost_usd                                                                          
    if final_cost is not None:                                                                                 
        await self.send_message(                                                                               
            "variantCost",                                                                                     
            None,                                                                                              
            index,                                                                                             
            {"costUsd": final_cost},                                                                           
            None,                                                                                              
        )                                                                                                      
                                                                                                               
  The message is emitted before setCode and variantComplete, so the frontend gets cost data as early as        
possible.                                                                                                      
                                                                                                               
### Frontend                                                                                                   
                                                                                                               
generateCode.ts                                                                                                
- WebSocketResponse.type now includes "variantCost".                                                           
- CodeGenerationCallbacks gains onVariantCost(variantIndex, costUsd, inputTokens, outputTokens).               
- Message handler: response.type === "variantCost" → callbacks.onVariantCost(response.variantIndex,            
  response.data?.costUsd ?? 0, ...) (tokens currently not yet transmitted; future work).                       
                                                                                                               
commits/types.ts                                                                                               
- Variant type gains costUsd?: number and tokens?: {input: number; output: number}.                            
                                                                                                               
project-store.ts                                                                                               
- setVariantCost(hash, numVariant, costUsd, inputTokens, outputTokens) — updates the costUsd and tokens fields 
  on the target variant in the Zustand store.                                                                  
                                                                                                               
App.tsx                                                                                                        
- setVariantCost destructured from useProjectStore.                                                            
- onVariantCost callback wired to persist cost data against the current commit.                                
                                                                                                               
────────────────────────────────────────────────────────────────────────────────                               
                                                                                                               
WebSocket Protocol                                                                                             
                                                                                                               
```ts                                                                                                          
  // Server → client                                                                                           
  {                                                                                                            
    type: "variantCost",                                                                                       
    value: null,                                                                                               
    variantIndex: number,   // 0-based variant index                                                           
    data: {                                                                                                    
      costUsd: number,      // USD cost for this variant                                                       
      // inputTokens/outputTokens: future enhancement                                                          
    }                                                                                                          
  }                                                                                                            

Note: variantCost is sent before variantComplete. If the variant failed with an error, no cost message is sent
(the provider session may not have a final cost in that case).

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

Backward Compatibility

  • Zero breaking changes — all new fields are optional; the variantCost message is ignored by any frontend that
    doesn't handle it.
  • Works with all existing models — the cost is only emitted when runner.last_cost_usd is not None, which
    requires both (a) a token-usage record from the provider and (b) a ModelPricing entry for the API model
    name. All currently defined models have pricing entries.
  • No new environment variables.

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

Testing Notes

  • Zero new pyright errors on agent/engine.py and routes/generate_code.py.
  • Zero new TypeScript errors on generateCode.ts, project-store.ts, App.tsx, and commits/types.ts.
  • Backend files compile successfully with python -m py_compile.

…red output mode

   Adds three inter-related improvements to the agentic tool-call runtime,
   all gated behind the new AGENT_STRUCTURED_OUTPUT opt-in flag.
@AseemPrasad

Copy link
Copy Markdown
Author

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant