mirror of
https://github.com/coder/coder.git
synced 2026-09-22 05:05:20 +08:00
Michael Suchacz
ad100452d4
## Problem CODAGT-906: chats using OpenAI-compatible backends (e.g. poolside) through the AI Bridge persist token usage inflated 105x-640x, which falsely triggers automatic chat compaction on every turn. The chat-completions streaming interceptor summed usage across every SSE chunk of one upstream stream and rewrote each relayed usage-bearing chunk with that running sum. Spec-compliant OpenAI emits usage once (final chunk with `stream_options.include_usage`), so the sum equals the final value. vLLM-style backends emit cumulative usage snapshots on every chunk, so the relayed final usage becomes roughly `N_chunks x prompt_tokens` (e.g. 417,012 persisted for a ~6,000-token context). chatd persists that value per assistant message and its compaction trigger reads it as context occupancy. ## Fix Track the latest usage-bearing chunk's raw usage (last-wins) in the stream processor, updating only when a chunk actually carries usage so a trailing usage-less chunk cannot zero it. `marshalChunk` relays that value and `recordTokenUsage` records the same value, unifying relayed and recorded usage. Last-wins is correct for both shapes: a single final usage chunk, and cumulative snapshots where each snapshot already includes all prior tokens. Per-iteration semantics are unchanged: each tool-loop iteration has its own processor, and the final iteration's usage is what the client sees. ## Tests - `TestStreamProcessorUsage` (internal): cumulative snapshots with a trailing usage-less chunk, and the spec-compliant final-only shape; asserts relayed and recorded usage equal the last snapshot. - New txtar fixture `streaming_cumulative_usage_injected_tool.txtar` with per-chunk cumulative usage plus an injected tool call; asserts client-visible final usage through the full interceptor. - Red-green verified: with the fix reverted, the internal test reports zeroed usage (trailing chunk overwrite) and the fixture test reports 18000 summed prompt tokens instead of 6000. The blocking (non-streaming) path deliberately keeps its cross-iteration summation for external clients and is untouched. Remote dogfood UAT validated chat streaming, tool calls, plausible usage numbers, and zero spurious compactions. > Mux acted on Mike's behalf to author this change.
aibridge
aibridge provides an HTTP handler that intercepts AI client requests bound for upstream AI providers (Anthropic, OpenAI, Copilot). It records token usage, prompts, and tool invocations per user. Optionally supports centralized MCP tool injection with allowlist/denylist filtering.
The handler is mounted by a host process. Today that host is coderd, which mounts the handler at /api/v2/ai-gateway/<provider>/*. Running aibridge as a separate process is planned for the future.
Architecture
┌─────────────────┐ ┌───────────────────────────────────────────┐
│ AI Client │ │ aibridge │
│ (Claude Code, │────▶│ ┌─────────────────┐ ┌─────────────┐ │
│ Cursor, etc.) │ │ │ RequestBridge │───▶│ Providers │ │
└─────────────────┘ │ │ (http.Handler) │ │ (Anthropic │ │
│ └─────────────────┘ │ OpenAI) │ │
│ └──────┬──────┘ │
│ │ │
│ ▼ │ ┌─────────────┐
│ ┌─────────────────┐ ┌─────────────┐ │ │ Upstream │
│ │ Recorder │◀───│ Interceptor │─── ───▶│ API │
│ │ (tokens, tools, │ │ (streaming/ │ │ │ (Anthropic │
│ │ prompts) │ │ blocking) │ │ │ OpenAI) │
│ └────────┬────────┘ └──────┬──────┘ │ └─────────────┘
│ │ │ │
│ ▼ ┌──────▼──────┐ │
│ ┌ ─ ─ ─ ─ ─ ─ ─ ┐ │ MCP Proxy │ │
│ │ Database │ │ (tools) │ │
│ └ ─ ─ ─ ─ ─ ─ ─ ┘ └─────────────┘ │
└───────────────────────────────────────────┘
Components
- RequestBridge: The main
http.Handlerthat routes requests to providers - Provider: Defines bridged routes (intercepted) and passthrough routes (proxied)
- Interceptor: Handles request/response processing and streaming
- Recorder: Interface for capturing usage data (tokens, prompts, tools)
- MCP Proxy (optional): Connects to MCP servers to list tool, inject them into requests, and invoke them in an inner agentic loop
Request Flow
- Client sends request to
/anthropic/v1/messagesor/openai/v1/chat/completions - Actor extraction: Request must have an actor in context (via
AsActor()). The host is responsible for authenticating the caller before invoking the handler. - Upstream call: Request forwarded to the AI provider
- Response relay: Response streamed/sent to client
- Recording: Token usage, prompts, and tool invocations recorded
With MCP enabled: Tools from configured MCP servers are centrally defined and injected into requests (prefixed bmcp_). Allowlist/denylist regex patterns control which tools are available. When the model selects an injected tool, the gateway invokes it in an inner agentic loop, and continues the conversation loop until complete.
Passthrough routes (/v1/models, /v1/messages/count_tokens) are reverse-proxied directly.
Observability
Prometheus Metrics
Create metrics with NewMetrics(prometheus.Registerer):
| Metric | Type | Description |
|---|---|---|
interceptions_total |
Counter | Intercepted request count |
interceptions_inflight |
Gauge | Currently processing requests |
interceptions_duration_seconds |
Histogram | Request duration |
passthrough_total |
Counter | Non-intercepted requests forwarded to the upstream |
prompts_total |
Counter | User prompt count |
tokens_total |
Counter | Token usage (input, output, cache read/write, provider extras) |
injected_tool_invocations_total |
Counter | Injected MCP tool invocations performed by the handler |
non_injected_tool_selections_total |
Counter | Client-defined tool selections returned by the model |
circuit_breaker_state |
Gauge | Circuit breaker state per provider/endpoint (0=closed, 0.5=half, 1=open) |
circuit_breaker_trips_total |
Counter | Times the circuit breaker transitioned to open |
circuit_breaker_rejects_total |
Counter | Requests rejected due to an open circuit breaker |
Recorder Interface
Implement Recorder to persist usage data to your database:
aibridge_interceptions- request metadata (provider, model, initiator, timestamps)aibridge_token_usages- input/output and cache read/write token counts per responseaibridge_user_prompts- user promptsaibridge_tool_usages- tool invocations (injected and client-defined)aibridge_model_thoughts- model reasoning content (thinking, reasoning summaries, commentary)
type Recorder interface {
RecordInterception(ctx context.Context, req *InterceptionRecord) error
RecordInterceptionEnded(ctx context.Context, req *InterceptionRecordEnded) error
RecordTokenUsage(ctx context.Context, req *TokenUsageRecord) error
RecordPromptUsage(ctx context.Context, req *PromptUsageRecord) error
RecordToolUsage(ctx context.Context, req *ToolUsageRecord) error
RecordModelThought(ctx context.Context, req *ModelThoughtRecord) error
}
Supported Routes
Each provider instance is mounted under /api/v2/ai-gateway/<name>, where <name> is the provider's configured name. For example, with an Anthropic provider named my-anthropic, its /messages endpoint would be reachable at /api/v2/ai-gateway/my-anthropic/v1/messages.
If a name is not set, the route path defaults to the provider's type: anthropic, openai, or copilot. The table below uses the default names.
(/*) denotes a route that handles both the exact path and any subpaths. A trailing /* denotes subpaths only.
| Provider | Route | Type |
|---|---|---|
| Anthropic | /anthropic/v1/messages |
Bridged (intercepted) |
| Anthropic | /anthropic/v1/messages/count_tokens |
Passthrough |
| Anthropic | /anthropic/v1/models(/*) |
Passthrough |
| Anthropic | /anthropic/api/event_logging/* |
Passthrough |
| OpenAI | /openai/v1/chat/completions |
Bridged (intercepted) |
| OpenAI | /openai/v1/responses |
Bridged (intercepted) |
| OpenAI | /openai/v1/responses/* |
Passthrough |
| OpenAI | /openai/v1/conversations(/*) |
Passthrough |
| OpenAI | /openai/v1/models(/*) |
Passthrough |
| Copilot | /copilot/chat/completions |
Bridged (intercepted) |
| Copilot | /copilot/responses |
Bridged (intercepted) |
| Copilot | /copilot/models(/*) |
Passthrough |
| Copilot | /copilot/agents/* |
Passthrough |
| Copilot | /copilot/mcp/* |
Passthrough |
| Copilot | /copilot/.well-known/* |
Passthrough |