From 5816fada9363c15bd9bc211c70f97e0e38af2f34 Mon Sep 17 00:00:00 2001
From: huaiyuWangh <34158348+huaiyuWangh@users.noreply.github.com>
Date: Fri, 8 May 2026 21:55:28 +0800
Subject: [PATCH] fix(agent): support flat_rate cost handling (#2116)

---
 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(-)

diff --git a/internal/agent/agent.go b/internal/agent/agent.go
index 028911f686ed9a7f4800e0d04f98c96b7cd14038..3f21fa7f504ebb3fc6212354127fc5f0d15ae62a 100644
--- a/internal/agent/agent.go
+++ b/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
 }
diff --git a/internal/agent/coordinator.go b/internal/agent/coordinator.go
index 2eef9fe377c1d1cd3ad22ccd5a26a5f2a5e1cc62..468397c07336fc1b52288462d53a71772ea566d7 100644
--- a/internal/agent/coordinator.go
+++ b/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
 }
 
diff --git a/internal/config/config.go b/internal/config/config.go
index 33251615d44252aaeb2b5db58577759b5dfdff51..d4870e7e063c87a070e8643c8c327f135ded7125 100644
--- a/internal/config/config.go
+++ b/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"`
 }
diff --git a/schema.json b/schema.json
index 9d610579b90dedb636522532b3648d441fc0057f..173d9fdebca8c67e9c25c4d3c4612c901663ad10 100644
--- a/schema.json
+++ b/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"