From be0268cf03ed382535bb7abd177f4a6052314610 Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:17:23 +0530 Subject: [PATCH 1/2] fix(aibridge): record token usage without an MCP proxier Streaming Responses interceptions called recordTokenUsage from inside the `i.mcpProxy != nil` branch, so a bridge built with a nil proxier served requests normally but recorded no token usage. Upstream reports usage on the response.completed event independently of tool injection. coderd/aibridged treats proxier construction failure as non-fatal and caches the resulting bridge, so a transient config-retrieval error suppressed usage recording for every streaming Responses request served by that bridge until its cache TTL expired. Record usage for every completed response, guarded only on completedResponse, matching responses/blocking.go and both chatcompletions implementations. Per-iteration semantics are preserved for the inner agentic loop. The integration harness substituted a non-nil noop manager whenever no proxier was supplied, so the nil path was never exercised. Add withoutMCP() to cover it. Fixes #27885 --- aibridge/intercept/responses/streaming.go | 11 ++- .../internal/integrationtest/setupbridge.go | 10 ++- .../tokenusage_internal_test.go | 74 +++++++++++++++++++ 3 files changed, 91 insertions(+), 4 deletions(-) create mode 100644 aibridge/internal/integrationtest/tokenusage_internal_test.go 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/setupbridge.go b/aibridge/internal/integrationtest/setupbridge.go index efa94074295be..60f04b6019dbc 100644 --- a/aibridge/internal/integrationtest/setupbridge.go +++ b/aibridge/internal/integrationtest/setupbridge.go @@ -50,6 +50,7 @@ type bridgeConfig struct { metrics *metrics.Metrics tracer trace.Tracer mcpProxy mcp.ServerProxier + noMCPProxy bool userID string metadata recorder.Metadata logger slog.Logger @@ -120,6 +121,13 @@ func withMCP(p mcp.ServerProxier) bridgeOption { return func(c *bridgeConfig) { c.mcpProxy = p } } +// withoutMCP runs the bridge with a nil MCP server proxier, matching a +// deployment where proxier construction failed and injection degraded. This +// differs from NoopMCPManager, which is non-nil and reports zero tools. +func withoutMCP() bridgeOption { + return func(c *bridgeConfig) { c.noMCPProxy = true } +} + // withActor sets the actor ID and metadata for the BaseContext. func withActor(id string, md recorder.Metadata) bridgeOption { return func(c *bridgeConfig) { c.userID = id; c.metadata = md } @@ -150,7 +158,7 @@ func newBridgeTestServer( cfg.tracer = defaultTracer } cfg.logger = newLogger(t) - if cfg.mcpProxy == nil { + if cfg.mcpProxy == nil && !cfg.noMCPProxy { cfg.mcpProxy = newNoopMCPManager() } diff --git a/aibridge/internal/integrationtest/tokenusage_internal_test.go b/aibridge/internal/integrationtest/tokenusage_internal_test.go new file mode 100644 index 0000000000000..a27870a7cd2fc --- /dev/null +++ b/aibridge/internal/integrationtest/tokenusage_internal_test.go @@ -0,0 +1,74 @@ +package integrationtest + +import ( + "context" + "io" + "net/http" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/aibridge/fixtures" + "github.com/coder/coder/v2/aibridge/internal/testutil" +) + +// TestResponsesStreamingRecordsTokenUsageWithoutMCP asserts that a streaming +// Responses interception records token usage regardless of whether an MCP +// server proxier is configured. Upstream reports usage independently of tool +// injection, and coderd/aibridged tolerates a nil proxier when construction +// fails, so usage must not depend on it. +func TestResponsesStreamingRecordsTokenUsageWithoutMCP(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + name string + opts []bridgeOption + }{ + {name: "without_mcp_proxy", opts: []bridgeOption{withoutMCP()}}, + {name: "with_noop_mcp_proxy", opts: nil}, + } { + 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, fixtures.OaiResponsesStreamingSimple) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) + bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, tc.opts...) + + resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIResponses, fix.Request()) + require.NoError(t, err) + defer resp.Body.Close() + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + + usages := bridgeServer.Recorder.RecordedTokenUsages() + require.Len(t, usages, 1, "exactly one token usage record per completed response") + require.Positive(t, bridgeServer.Recorder.TotalInputTokens()) + require.Positive(t, bridgeServer.Recorder.TotalOutputTokens()) + }) + } +} + +// TestResponsesStreamingRecordsTokenUsagePerAgenticIteration asserts that +// decoupling token recording from the MCP proxier does not double-count usage +// when the inner agentic loop iterates. The injected-tool fixture drives two +// upstream calls, so two records are expected. +func TestResponsesStreamingRecordsTokenUsagePerAgenticIteration(t *testing.T) { + t.Parallel() + + bridgeServer, _, resp := setupInjectedToolTest( + t, + fixtures.OaiResponsesStreamingSingleInjectedTool, + true, + defaultTracer, + pathOpenAIResponses, + nil, + ) + defer resp.Body.Close() + + usages := bridgeServer.Recorder.RecordedTokenUsages() + require.Len(t, usages, 2, "one token usage record per agentic iteration") +} From 0e8127795d9cc1447388532894f6d2361f66b338 Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Wed, 12 Aug 2026 14:57:09 +0530 Subject: [PATCH 2/2] test(aibridge): address review on token usage coverage Inline the nil-proxier bridge option at its only call site and document the noMCPProxy field, which newBridgeTestServer still needs to tell "unset" apart from "explicitly nil". Compare token counts against the values recorded before the request, and assert the exact fixture totals rather than only that they are positive, matching TestOpenAIChatCompletions. Fold the standalone test file into bridge_internal_test.go and cover the Anthropic Messages route alongside OpenAI Responses. Drop the per-agentic-iteration test: TestResponsesInjectedTool/streaming_success already drives the same fixture and asserts an exact record count, so it covers the double-counting case. --- .../integrationtest/bridge_internal_test.go | 62 ++++++++++++++++ .../internal/integrationtest/setupbridge.go | 17 ++--- .../tokenusage_internal_test.go | 74 ------------------- 3 files changed, 68 insertions(+), 85 deletions(-) delete mode 100644 aibridge/internal/integrationtest/tokenusage_internal_test.go diff --git a/aibridge/internal/integrationtest/bridge_internal_test.go b/aibridge/internal/integrationtest/bridge_internal_test.go index 0b2389d3390a9..552fe0e6b724b 100644 --- a/aibridge/internal/integrationtest/bridge_internal_test.go +++ b/aibridge/internal/integrationtest/bridge_internal_test.go @@ -2451,3 +2451,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 60f04b6019dbc..fdad3610e7a58 100644 --- a/aibridge/internal/integrationtest/setupbridge.go +++ b/aibridge/internal/integrationtest/setupbridge.go @@ -50,10 +50,12 @@ type bridgeConfig struct { metrics *metrics.Metrics tracer trace.Tracer mcpProxy mcp.ServerProxier - noMCPProxy bool - 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. @@ -121,13 +123,6 @@ func withMCP(p mcp.ServerProxier) bridgeOption { return func(c *bridgeConfig) { c.mcpProxy = p } } -// withoutMCP runs the bridge with a nil MCP server proxier, matching a -// deployment where proxier construction failed and injection degraded. This -// differs from NoopMCPManager, which is non-nil and reports zero tools. -func withoutMCP() bridgeOption { - return func(c *bridgeConfig) { c.noMCPProxy = true } -} - // withActor sets the actor ID and metadata for the BaseContext. func withActor(id string, md recorder.Metadata) bridgeOption { return func(c *bridgeConfig) { c.userID = id; c.metadata = md } diff --git a/aibridge/internal/integrationtest/tokenusage_internal_test.go b/aibridge/internal/integrationtest/tokenusage_internal_test.go deleted file mode 100644 index a27870a7cd2fc..0000000000000 --- a/aibridge/internal/integrationtest/tokenusage_internal_test.go +++ /dev/null @@ -1,74 +0,0 @@ -package integrationtest - -import ( - "context" - "io" - "net/http" - "testing" - - "github.com/stretchr/testify/require" - - "github.com/coder/coder/v2/aibridge/fixtures" - "github.com/coder/coder/v2/aibridge/internal/testutil" -) - -// TestResponsesStreamingRecordsTokenUsageWithoutMCP asserts that a streaming -// Responses interception records token usage regardless of whether an MCP -// server proxier is configured. Upstream reports usage independently of tool -// injection, and coderd/aibridged tolerates a nil proxier when construction -// fails, so usage must not depend on it. -func TestResponsesStreamingRecordsTokenUsageWithoutMCP(t *testing.T) { - t.Parallel() - - for _, tc := range []struct { - name string - opts []bridgeOption - }{ - {name: "without_mcp_proxy", opts: []bridgeOption{withoutMCP()}}, - {name: "with_noop_mcp_proxy", opts: nil}, - } { - 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, fixtures.OaiResponsesStreamingSimple) - upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) - bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, tc.opts...) - - resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIResponses, fix.Request()) - require.NoError(t, err) - defer resp.Body.Close() - _, err = io.ReadAll(resp.Body) - require.NoError(t, err) - require.Equal(t, http.StatusOK, resp.StatusCode) - - usages := bridgeServer.Recorder.RecordedTokenUsages() - require.Len(t, usages, 1, "exactly one token usage record per completed response") - require.Positive(t, bridgeServer.Recorder.TotalInputTokens()) - require.Positive(t, bridgeServer.Recorder.TotalOutputTokens()) - }) - } -} - -// TestResponsesStreamingRecordsTokenUsagePerAgenticIteration asserts that -// decoupling token recording from the MCP proxier does not double-count usage -// when the inner agentic loop iterates. The injected-tool fixture drives two -// upstream calls, so two records are expected. -func TestResponsesStreamingRecordsTokenUsagePerAgenticIteration(t *testing.T) { - t.Parallel() - - bridgeServer, _, resp := setupInjectedToolTest( - t, - fixtures.OaiResponsesStreamingSingleInjectedTool, - true, - defaultTracer, - pathOpenAIResponses, - nil, - ) - defer resp.Body.Close() - - usages := bridgeServer.Recorder.RecordedTokenUsages() - require.Len(t, usages, 2, "one token usage record per agentic iteration") -}