diff --git a/aibridge/intercept/responses/streaming.go b/aibridge/intercept/responses/streaming.go index 492783f4de72b..572de64b432b2 100644 --- a/aibridge/intercept/responses/streaming.go +++ b/aibridge/intercept/responses/streaming.go @@ -241,6 +241,14 @@ func (i *StreamingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r return err } + // Record token usage for every iteration, whether or not tools are + // injected. Usage is reported by upstream independently of the MCP + // proxy, so gating this on the proxy drops usage entirely for + // deployments that run without one. + if completedResponse != nil { + i.recordTokenUsage(ctx, completedResponse) + } + if i.mcpProxy != nil && completedResponse != nil { pending := i.getPendingInjectedToolCalls(completedResponse) shouldLoop, innerLoopErr = i.handleInnerAgenticLoop(ctx, pending, completedResponse) @@ -248,9 +256,6 @@ func (i *StreamingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r i.sendCustomErr(ctx, w, http.StatusInternalServerError, innerLoopErr) shouldLoop = false } - - // Record token usage for each inner loop iteration - i.recordTokenUsage(ctx, completedResponse) } i.recordModelThoughts(ctx, completedResponse) diff --git a/aibridge/internal/integrationtest/bridge_internal_test.go b/aibridge/internal/integrationtest/bridge_internal_test.go index 77a7566f05e09..60756131c9bc6 100644 --- a/aibridge/internal/integrationtest/bridge_internal_test.go +++ b/aibridge/internal/integrationtest/bridge_internal_test.go @@ -2486,3 +2486,65 @@ func extractSigV4Field(authHeader, prefix string) string { } return strings.TrimSpace(val) } + +// TestTokenUsageRecordedWithoutMCPProxier asserts that an interception records +// token usage when no MCP server proxier is configured. Upstream reports usage +// independently of tool injection, and coderd/aibridged tolerates a nil +// proxier when proxier construction fails, so usage must not depend on one. +func TestTokenUsageRecordedWithoutMCPProxier(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + fixture []byte + path string + expectedInputTokens, expectedOutputTokens int64 + }{ + { + name: "openai responses", + fixture: fixtures.OaiResponsesStreamingSimple, + path: pathOpenAIResponses, + expectedInputTokens: 11, + expectedOutputTokens: 18, + }, + { + name: "anthropic messages", + fixture: fixtures.AntSimple, + path: pathAnthropicMessages, + expectedInputTokens: 18, + expectedOutputTokens: 241, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong) + t.Cleanup(cancel) + + fix := fixtures.Parse(t, tc.fixture) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) + + bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, func(c *bridgeConfig) { + c.noMCPProxy = true + }) + + inputBefore := bridgeServer.Recorder.TotalInputTokens() + outputBefore := bridgeServer.Recorder.TotalOutputTokens() + + reqBody, err := sjson.SetBytes(fix.Request(), "stream", true) + require.NoError(t, err) + resp, err := bridgeServer.makeRequest(t, http.MethodPost, tc.path, reqBody) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + + require.NotEmpty(t, bridgeServer.Recorder.RecordedTokenUsages(), "token usage must be recorded without an MCP proxier") + assert.EqualValues(t, tc.expectedInputTokens, bridgeServer.Recorder.TotalInputTokens()-inputBefore, "input tokens miscalculated") + assert.EqualValues(t, tc.expectedOutputTokens, bridgeServer.Recorder.TotalOutputTokens()-outputBefore, "output tokens miscalculated") + }) + } +} diff --git a/aibridge/internal/integrationtest/setupbridge.go b/aibridge/internal/integrationtest/setupbridge.go index efa94074295be..fdad3610e7a58 100644 --- a/aibridge/internal/integrationtest/setupbridge.go +++ b/aibridge/internal/integrationtest/setupbridge.go @@ -50,9 +50,12 @@ type bridgeConfig struct { metrics *metrics.Metrics tracer trace.Tracer mcpProxy mcp.ServerProxier - userID string - metadata recorder.Metadata - logger slog.Logger + // noMCPProxy leaves the proxier nil instead of falling back to + // NoopMCPManager, which is non-nil and reports zero tools. + noMCPProxy bool + userID string + metadata recorder.Metadata + logger slog.Logger } // bridgeTestServer wraps an httptest.Server running a RequestBridge. @@ -150,7 +153,7 @@ func newBridgeTestServer( cfg.tracer = defaultTracer } cfg.logger = newLogger(t) - if cfg.mcpProxy == nil { + if cfg.mcpProxy == nil && !cfg.noMCPProxy { cfg.mcpProxy = newNoopMCPManager() }