Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/workflows/ci.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@ jobs:
# Main repo directories for completeness in case other files are
# touched:
- "agent/**"
- "aibridge/**"
- "cli/**"
- "cmd/**"
- "coderd/**"
Expand Down
2 changes: 2 additions & 0 deletions .github/workflows/typos.toml
Original file line number Diff line number Diff line change
Expand Up @@ -57,4 +57,6 @@ extend-exclude = [
# notifications' golden files confuse the detector because of quoted-printable encoding
"coderd/notifications/testdata/**",
"agent/agentcontainers/testdata/devcontainercli/**",
# aibridge fixtures contain truncated streaming chunks that look like typos
"aibridge/fixtures/**",
]
16 changes: 16 additions & 0 deletions .golangci.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,21 @@ linters-settings:
# goal: 100
threshold: 412

depguard:
rules:
aibridge_import_isolation:
list-mode: lax
files:
- "aibridge/*.go"
- "aibridge/**/*.go"
allow:
- $gostd
- github.com/coder/coder/v2/aibridge
- github.com/coder/coder/v2/buildinfo
deny:
- pkg: github.com/coder/coder/v2
desc: aibridge code must not import coder packages outside aibridge; buildinfo is the only exception

exhaustruct:
include:
# Gradually extend to cover more of the codebase.
Expand Down Expand Up @@ -227,6 +242,7 @@ linters:
- asciicheck
- bidichk
- bodyclose
- depguard
- dogsled
- errcheck
- errname
Expand Down
56 changes: 28 additions & 28 deletions aibridge/AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -7,46 +7,46 @@ over clever ones. Readability is a primary concern.

## Tone & Relationship

We're colleagues — push back on bad ideas and speak up when something
We're colleagues. Push back on bad ideas and speak up when something
doesn't make sense. Honesty over agreeableness.

- Disagree when I'm wrong — act as a critical peer reviewer.
- Disagree when I'm wrong. Act as a critical peer reviewer.
- Call out bad ideas, unreasonable expectations, and mistakes.
- **Ask for clarification** instead of assuming. Say when you don't know something.
- Architectural decisions require discussion; routine fixes do not.

## Foundational Rules

- Doing it right is better than doing it fast.
- YAGNI — don't add features we don't need right now.
- YAGNI. Don't add features we don't need right now.
- Make the smallest reasonable changes to achieve the goal.
- Reduce code duplication, even if it takes extra effort.
- Match the style of surrounding code — consistency within a file matters.
- Match the style of surrounding code. Consistency within a file matters.
- Fix bugs immediately when you find them.

## Essential Commands

| Task | Command | Notes |
| ----------- | ---------------------- | --------------------------------- |
| Test | `make test` | All tests, no race detector |
| Test (race) | `make test-race` | CGO_ENABLED=1, use for CI |
| Coverage | `make coverage` | Prints summary to stdout |
| Format | `make fmt` | gofumpt; single file: `make fmt FILE=path` |
| Mocks | `make mocks` | Regenerate from `mcp/api.go` |
| Task | Command | Notes |
|-------------|------------------|--------------------------------------------|
| Test | `make test` | All tests, no race detector |
| Test (race) | `make test-race` | CGO_ENABLED=1, use for CI |
| Coverage | `make coverage` | Prints summary to stdout |
| Format | `make fmt` | gofumpt; single file: `make fmt FILE=path` |
| Mocks | `make mocks` | Regenerate from `mcp/api.go` |

**Always use these commands** instead of running `go test` or `gofumpt` directly.

## Code Navigation

Use LSP tools (go to definition, find references, hover) **before** resorting to grep.
This codebase has 90+ Go files across multiple packages LSP is faster and more accurate.
This codebase has 90+ Go files across multiple packages, so LSP is faster and more accurate.

## Architecture Overview

AI Bridge is a smart gateway that sits between AI clients (Claude Code, Cursor,
etc.) and upstream providers (Anthropic, OpenAI). It intercepts all AI traffic
to provide centralized authn/z, auditing, token attribution, and MCP tool
administration. It runs as part of `coderd` (the Coder control plane) — users
administration. It runs as part of `coderd` (the Coder control plane). Users
authenticate with their Coder session tokens.

```
Expand All @@ -68,28 +68,28 @@ authenticate with their Coder session tokens.
```

Key packages:
- `intercept/` request/response interception, per-provider subdirs (`messages/`, `responses/`, `chatcompletions/`)
- `provider/` upstream provider definitions (Anthropic, OpenAI, Copilot)
- `mcp/` MCP protocol integration
- `circuitbreaker/` circuit breaker for upstream calls
- `context/` request-scoped context helpers
- `internal/integrationtest/` integration tests with mock upstreams
- `intercept/`: request/response interception, per-provider subdirs (`messages/`, `responses/`, `chatcompletions/`)
- `provider/`: upstream provider definitions (Anthropic, OpenAI, Copilot)
- `mcp/`: MCP protocol integration
- `circuitbreaker/`: circuit breaker for upstream calls
- `context/`: request-scoped context helpers
- `internal/integrationtest/`: integration tests with mock upstreams

## Go Patterns

- Follow the [Uber Go Style Guide](https://github.com/uber-go/guide/blob/master/style.md).
- Use `gofumpt` for formatting (enforced by `make fmt`).
- Prefer table-driven tests.
- **Never use `time.Sleep` in tests** — use `github.com/coder/quartz` or channels/contexts for synchronization.
- **Never use `time.Sleep` in tests**. Use `github.com/coder/quartz` or channels/contexts for synchronization.
- Use unique identifiers in tests: `fmt.Sprintf("test-%s-%d", t.Name(), time.Now().UnixNano())`.
- Test observable behavior, not implementation details.

## Streaming Code

This codebase heavily uses SSE streaming. When modifying interceptors:
- Always handle both blocking and streaming paths.
- Test with `*_test.go` files in the same package — they cover edge cases for chunked responses.
- Be careful with goroutine lifecycle — ensure proper cleanup on context cancellation.
- Test with `*_test.go` files in the same package. They cover edge cases for chunked responses.
- Be careful with goroutine lifecycle. Ensure proper cleanup on context cancellation.

## Commit Style

Expand All @@ -110,9 +110,9 @@ type(scope): message

## Common Pitfalls

| Problem | Fix |
| ------- | --- |
| Race in streaming tests | Use `t.Cleanup()` and proper synchronization, never `time.Sleep` |
| Mock not updated | Run `make mocks` after changing `mcp/api.go` |
| Formatting failures | Run `make fmt` before committing |
| `retract` directive in go.mod | Don't remove — it's intentional (v1.0.8 conflict marker) |
| Problem | Fix |
|-------------------------------|------------------------------------------------------------------|
| Race in streaming tests | Use `t.Cleanup()` and proper synchronization, never `time.Sleep` |
| Mock not updated | Run `make mocks` after changing `mcp/api.go` |
| Formatting failures | Run `make fmt` before committing |
| `retract` directive in go.mod | Don't remove. It's intentional (v1.0.8 conflict marker) |
93 changes: 93 additions & 0 deletions aibridge/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
# aibridge

aibridge is an HTTP gateway that sits between AI clients and upstream AI providers (Anthropic, OpenAI). It intercepts requests to record token usage, prompts, and tool invocations per user. Optionally supports centralized [MCP](https://modelcontextprotocol.io/) tool injection with allowlist/denylist filtering.

## 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.Handler` that 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

1. Client sends request to `/anthropic/v1/messages` or `/openai/v1/chat/completions`
2. **Actor extraction**: Request must have an actor in context (via `AsActor()`).
3. **Upstream call**: Request forwarded to the AI provider
4. **Response relay**: Response streamed/sent to client
5. **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 |
| `tokens_total` | Counter | Token usage (input/output) |
| `prompts_total` | Counter | User prompt count |
| `injected_tool_invocations_total` | Counter | MCP tool invocations |
| `passthrough_total` | Counter | Non-intercepted requests |

### Recorder Interface

Implement `Recorder` to persist usage data to your database:

- `aibridge_interceptions` - request metadata (provider, model, initiator, timestamps)
- `aibridge_token_usages` - input/output token counts per response
- `aibridge_user_prompts` - user prompts
- `aibridge_tool_usages` - tool invocations (injected and client-defined)

```go
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
}
```

## Supported Routes

| Provider | Route | Type |
|----------|-------|------|
| Anthropic | `/anthropic/v1/messages` | Bridged (intercepted) |
| Anthropic | `/anthropic/v1/models` | Passthrough |
| Anthropic | `/anthropic/v1/messages/count_tokens` | Passthrough |
| OpenAI | `/openai/v1/chat/completions` | Bridged (intercepted) |
| OpenAI | `/openai/v1/models` | Passthrough |
10 changes: 5 additions & 5 deletions aibridge/api.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,11 @@ import (
"go.opentelemetry.io/otel/trace"

"cdr.dev/slog/v3"
"github.com/coder/aibridge/config"
aibcontext "github.com/coder/aibridge/context"
"github.com/coder/aibridge/metrics"
"github.com/coder/aibridge/provider"
"github.com/coder/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/config"
aibcontext "github.com/coder/coder/v2/aibridge/context"
"github.com/coder/coder/v2/aibridge/metrics"
"github.com/coder/coder/v2/aibridge/provider"
"github.com/coder/coder/v2/aibridge/recorder"
)

// Const + Type + function aliases for backwards compatibility.
Expand Down
14 changes: 7 additions & 7 deletions aibridge/bridge.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,13 @@ import (
"golang.org/x/xerrors"

"cdr.dev/slog/v3"
"github.com/coder/aibridge/circuitbreaker"
aibcontext "github.com/coder/aibridge/context"
"github.com/coder/aibridge/mcp"
"github.com/coder/aibridge/metrics"
"github.com/coder/aibridge/provider"
"github.com/coder/aibridge/recorder"
"github.com/coder/aibridge/tracing"
"github.com/coder/coder/v2/aibridge/circuitbreaker"
aibcontext "github.com/coder/coder/v2/aibridge/context"
"github.com/coder/coder/v2/aibridge/mcp"
"github.com/coder/coder/v2/aibridge/metrics"
"github.com/coder/coder/v2/aibridge/provider"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/tracing"
)

const (
Expand Down
8 changes: 4 additions & 4 deletions aibridge/bridge_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,10 @@ import (
"go.opentelemetry.io/otel"

"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/aibridge"
"github.com/coder/aibridge/config"
"github.com/coder/aibridge/internal/testutil"
"github.com/coder/aibridge/provider"
"github.com/coder/coder/v2/aibridge"
"github.com/coder/coder/v2/aibridge/config"
"github.com/coder/coder/v2/aibridge/internal/testutil"
"github.com/coder/coder/v2/aibridge/provider"
)

var bridgeTestTracer = otel.Tracer("bridge_test")
Expand Down
4 changes: 2 additions & 2 deletions aibridge/circuitbreaker/circuitbreaker.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,8 @@ import (
"github.com/sony/gobreaker/v2"
"golang.org/x/xerrors"

"github.com/coder/aibridge/config"
"github.com/coder/aibridge/metrics"
"github.com/coder/coder/v2/aibridge/config"
"github.com/coder/coder/v2/aibridge/metrics"
)

// ErrCircuitOpen is returned by Execute when the circuit breaker is open
Expand Down
4 changes: 2 additions & 2 deletions aibridge/circuitbreaker/circuitbreaker_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,8 @@ import (
"github.com/sony/gobreaker/v2"
"github.com/stretchr/testify/assert"

"github.com/coder/aibridge/circuitbreaker"
"github.com/coder/aibridge/config"
"github.com/coder/coder/v2/aibridge/circuitbreaker"
"github.com/coder/coder/v2/aibridge/config"
)

func TestExecute_PerModelIsolation(t *testing.T) {
Expand Down
2 changes: 1 addition & 1 deletion aibridge/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import (

"github.com/stretchr/testify/require"

"github.com/coder/aibridge"
"github.com/coder/coder/v2/aibridge"
)

func TestGuessClient(t *testing.T) {
Expand Down
2 changes: 1 addition & 1 deletion aibridge/context/context.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ package context
import (
"context"

"github.com/coder/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/recorder"
)

type (
Expand Down
4 changes: 2 additions & 2 deletions aibridge/context/context_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"

aibcontext "github.com/coder/aibridge/context"
"github.com/coder/aibridge/recorder"
aibcontext "github.com/coder/coder/v2/aibridge/context"
"github.com/coder/coder/v2/aibridge/recorder"
)

func TestAsActor(t *testing.T) {
Expand Down
2 changes: 1 addition & 1 deletion aibridge/intercept/actor_headers.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ import (
ant_option "github.com/anthropics/anthropic-sdk-go/option"
oai_option "github.com/openai/openai-go/v3/option"

"github.com/coder/aibridge/context"
"github.com/coder/coder/v2/aibridge/context"
)

const (
Expand Down
6 changes: 3 additions & 3 deletions aibridge/intercept/actor_headers_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@ import (
"github.com/google/uuid"
"github.com/stretchr/testify/require"

"github.com/coder/aibridge/context"
"github.com/coder/aibridge/intercept"
"github.com/coder/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/context"
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/recorder"
)

func TestNilActor(t *testing.T) {
Expand Down
2 changes: 1 addition & 1 deletion aibridge/intercept/apidump/apidump.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ import (
"golang.org/x/xerrors"

"cdr.dev/slog/v3"
"github.com/coder/aibridge/utils"
"github.com/coder/coder/v2/aibridge/utils"
"github.com/coder/quartz"
)

Expand Down
Loading
Loading