Merge pull request #630 from Molecule-AI/fix/issue-615-cap-token-counts

fix(platform): cap token counts before upsert to prevent NUMERIC overflow (#615)

platform/internal/handlers/workspace_metrics.go

@@ -98,10 +98,32 @@ func todayUTC() time.Time {
 	return time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC)
 }
 
+// maxTokensPerCall is the per-call sanity cap applied before upsert (#615).
+// An adversarial or buggy agent reporting INT64_MAX would otherwise cause a
+// NUMERIC(12,6) overflow in Postgres (silent failure, no cross-workspace
+// impact, but corrupts the workspace's cost accounting). 10 M tokens/call is
+// generous for any real LLM API response; anything above is clamped.
+const maxTokensPerCall = int64(10_000_000)
+
 // upsertTokenUsage accumulates input/output token counts for workspaceID's
 // current UTC day. Cost is estimated using the default per-token pricing
 // constants. Always call in a detached goroutine — never block the A2A path.
 func upsertTokenUsage(ctx context.Context, workspaceID string, inputTokens, outputTokens int64) {
+	// Clamp to safe range before any arithmetic — prevents NUMERIC overflow
+	// from adversarial or buggy agent responses (#615).
+	if inputTokens < 0 {
+		inputTokens = 0
+	}
+	if outputTokens < 0 {
+		outputTokens = 0
+	}
+	if inputTokens > maxTokensPerCall {
+		inputTokens = maxTokensPerCall
+	}
+	if outputTokens > maxTokensPerCall {
+		outputTokens = maxTokensPerCall
+	}
+
 	if inputTokens == 0 && outputTokens == 0 {
 		return
 	}

platform/internal/handlers/workspace_metrics_test.go

@@ -176,6 +176,79 @@ func TestGetMetrics_CostFormat(t *testing.T) {
 	}
 }
 
+// ---- upsertTokenUsage cap tests (#615) ----
+
+// TestUpsertTokenUsage_615_CapsInt64Max verifies that an adversarial
+// INT64_MAX token count is clamped to maxTokensPerCall before the upsert,
+// preventing NUMERIC(12,6) overflow in Postgres.
+func TestUpsertTokenUsage_615_CapsInt64Max(t *testing.T) {
+	mock := setupTestDB(t)
+
+	// We expect the INSERT to be called with maxTokensPerCall, not math.MaxInt64.
+	mock.ExpectExec(`INSERT INTO workspace_token_usage`).
+		WithArgs("ws-1", sqlmock.AnyArg(),
+			maxTokensPerCall,  // input clamped
+			maxTokensPerCall,  // output clamped
+			sqlmock.AnyArg()). // cost
+		WillReturnResult(sqlmock.NewResult(0, 1))
+
+	// INT64_MAX overflows — must be clamped.
+	const int64Max = int64(^uint64(0) >> 1)
+	upsertTokenUsage(t.Context(), "ws-1", int64Max, int64Max)
+
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("expected clamped values in upsert: %v", err)
+	}
+}
+
+// TestUpsertTokenUsage_615_CapsNegative verifies negative token counts are
+// clamped to 0 before upsert (no negative accumulation in cost rows).
+func TestUpsertTokenUsage_615_CapsNegative(t *testing.T) {
+	// Negative input + negative output → both become 0 → early return, no DB call.
+	setupTestDB(t) // no expectations
+
+	upsertTokenUsage(t.Context(), "ws-1", -100, -200)
+	// If any DB call were made the mock would error — passing here is the assertion.
+}
+
+// TestUpsertTokenUsage_615_NormalValuesUnchanged verifies that token counts
+// within the valid range pass through to the DB unchanged.
+func TestUpsertTokenUsage_615_NormalValuesUnchanged(t *testing.T) {
+	mock := setupTestDB(t)
+
+	mock.ExpectExec(`INSERT INTO workspace_token_usage`).
+		WithArgs("ws-1", sqlmock.AnyArg(),
+			int64(1500),      // input unchanged
+			int64(300),       // output unchanged
+			sqlmock.AnyArg()). // cost
+		WillReturnResult(sqlmock.NewResult(0, 1))
+
+	upsertTokenUsage(t.Context(), "ws-1", 1500, 300)
+
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("normal values altered unexpectedly: %v", err)
+	}
+}
+
+// TestUpsertTokenUsage_615_ExactlyAtCap verifies that a count exactly equal
+// to maxTokensPerCall is accepted without clamping.
+func TestUpsertTokenUsage_615_ExactlyAtCap(t *testing.T) {
+	mock := setupTestDB(t)
+
+	mock.ExpectExec(`INSERT INTO workspace_token_usage`).
+		WithArgs("ws-1", sqlmock.AnyArg(),
+			maxTokensPerCall,
+			maxTokensPerCall,
+			sqlmock.AnyArg()).
+		WillReturnResult(sqlmock.NewResult(0, 1))
+
+	upsertTokenUsage(t.Context(), "ws-1", maxTokensPerCall, maxTokensPerCall)
+
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("at-cap values should not be altered: %v", err)
+	}
+}
+
 // ---- parseUsageFromA2AResponse tests ----
 
 func TestParseUsage_JSONRPCResultEnvelope(t *testing.T) {