diff --git a/platform/internal/handlers/workspace_metrics.go b/platform/internal/handlers/workspace_metrics.go index db6400a3..92d65a2e 100644 --- a/platform/internal/handlers/workspace_metrics.go +++ b/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 } diff --git a/platform/internal/handlers/workspace_metrics_test.go b/platform/internal/handlers/workspace_metrics_test.go index 63e64d49..5741a6fb 100644 --- a/platform/internal/handlers/workspace_metrics_test.go +++ b/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) {