fix(agent): support flat_rate cost handling (#2116)

huaiyuWangh created 3 days ago

Change summary

internal/agent/agent.go       | 17 ++++++++++++++---
internal/agent/coordinator.go |  2 ++
internal/config/config.go     |  3 +++
schema.json                   |  4 ++++
4 files changed, 23 insertions(+), 3 deletions(-)

Detailed changes

internal/agent/agent.go

@@ -103,6 +103,7 @@ type Model struct {
 	Model      fantasy.LanguageModel
 	CatwalkCfg catwalk.Model
 	ModelCfg   config.SelectedModel
+	FlatRate   bool
 }
 
 type sessionAgent struct {
@@ -1079,6 +1080,11 @@ func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, user
 		cost = *openrouterCost
 	}
 
+	// Skip cost accumulation
+	if model.FlatRate {
+		cost = 0
+	}
+
 	promptTokens := resp.TotalUsage.InputTokens + resp.TotalUsage.CacheCreationTokens
 	completionTokens := resp.TotalUsage.OutputTokens
 
@@ -1113,12 +1119,17 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
 
 	a.eventTokensUsed(session.ID, model, usage, cost)
 
+	// Use override cost if available (e.g., from OpenRouter).
 	if overrideCost != nil {
-		session.Cost += *overrideCost
-	} else {
-		session.Cost += cost
+		cost = *overrideCost
+	}
+
+	// Skip cost accumulation
+	if model.FlatRate {
+		cost = 0
 	}
 
+	session.Cost += cost
 	session.CompletionTokens = usage.OutputTokens
 	session.PromptTokens = usage.InputTokens + usage.CacheReadTokens
 }

internal/agent/coordinator.go

@@ -626,10 +626,12 @@ func (c *coordinator) buildAgentModels(ctx context.Context, isSubAgent bool) (Mo
 			Model:      largeModel,
 			CatwalkCfg: *largeCatwalkModel,
 			ModelCfg:   largeModelCfg,
+			FlatRate:   largeProviderCfg.FlatRate,
 		}, Model{
 			Model:      smallModel,
 			CatwalkCfg: *smallCatwalkModel,
 			ModelCfg:   smallModelCfg,
+			FlatRate:   smallProviderCfg.FlatRate,
 		}, nil
 }
 

internal/config/config.go

@@ -120,6 +120,9 @@ type ProviderConfig struct {
 	// Used to pass extra parameters to the provider.
 	ExtraParams map[string]string `json:"-"`
 
+	// Skip cost accumulation for this provider when using subscription or flat rate billing.
+	FlatRate bool `json:"flat_rate,omitempty" jsonschema:"description=Flat-rate mode for this provider"`
+
 	// The provider models
 	Models []catwalk.Model `json:"models,omitempty" jsonschema:"description=List of models available from this provider"`
 }

schema.json

@@ -588,6 +588,10 @@
           "type": "object",
           "description": "Additional provider-specific options for this provider"
         },
+        "flat_rate": {
+          "type": "boolean",
+          "description": "Flat-rate mode for this provider"
+        },
         "models": {
           "items": {
             "$ref": "#/$defs/Model"