opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-10-16 14:48:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: extract recorder construction (#30224)

Browse Source 
Extract `newRecorder` from the interception pool into helper.
This lets the proxy use the same recorder construction code.

Generated by Coder Agents on behalf of @pawbana.
This commit is contained in:
Paweł Banaszewski
2026-10-02 16:18:52 +02:00
committed by GitHub
parent ff158c0589
commit f96d2696a4
3 changed files with 117 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/aibridged/aibridged.go
+24
View File
@@ -16,6 +16,7 @@ import (
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/aibridge"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/x/proxy"
"github.com/coder/coder/v2/coderd/aibridged/proto"
"github.com/coder/coder/v2/codersdk"
@@ -380,6 +381,29 @@ func (s *Server) ReplaceProviders(ctx context.Context, providers []aibridge.Prov
return nil
}
// newRecorder builds the recorder for one API key.
//
// revive:disable-next-line:flag-parameter // Constructor configuration flags.
func newRecorder(logger slog.Logger, tracer trace.Tracer, apiKeyID string, structuredLogging bool, disableContentRecording bool, clientFn ClientFunc) recorder.Recorder {
var middleware []recorder.Middleware
if disableContentRecording {
middleware = append(middleware, recorder.WithoutRecords(recorder.DisabledRecords{
PromptUsage: true,
ToolUsage: true,
ModelThought: true,
}))
}
return aibridge.NewRecorder(
logger.Named("recorder"), tracer, apiKeyID, structuredLogging,
recorder.NewDRPCRecorder(apiKeyID, func(ctx context.Context) (proto.DRPCRecorderClient, error) {
// The recorder outlives its caller, so acquire the client with each
// record call's context.
return clientFn(ctx)
}),
middleware...,
)
}
// PoolOptions reports the options the published interception pool was built
// with, so a caller can assert that its configuration reached the pool serving
// requests. It reports the zero value in proxy mode, which has no pool.
coderd/aibridged/aibridged_internal_test.go
+90
View File
@@ -2,22 +2,112 @@ package aibridged
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"go.opentelemetry.io/otel/trace/noop"
"golang.org/x/xerrors"
"storj.io/drpc"
"cdr.dev/slog/v3"
"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/coder/v2/aibridge"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/coder/v2/coderd/aibridged/proto"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/testutil"
)
func TestNewRecorder(t *testing.T) {
t.Parallel()
drpcErr := xerrors.New("DRPC recording was attempted")
for _, tc := range []struct {
structuredLogging bool
disableContentRecording bool
wantClientCalls int
wantLogs int
wantContentErr error // nil means no content recording is expected
}{
{
wantClientCalls: 6,
wantContentErr: drpcErr,
},
{
structuredLogging: true,
wantClientCalls: 6,
wantLogs: 6,
wantContentErr: drpcErr,
},
{
disableContentRecording: true,
wantClientCalls: 3,
},
{
structuredLogging: true,
disableContentRecording: true,
wantClientCalls: 3,
wantLogs: 6,
},
} {
t.Run(fmt.Sprintf("Structured=%t/DisableContent=%t", tc.structuredLogging, tc.disableContentRecording), func(t *testing.T) {
t.Parallel()
apiKeyID, interceptionID := uuid.NewString(), uuid.NewString()
now := time.Now().UTC()
type contextKey struct{}
sink := testutil.NewFakeSink(t)
clientCalls := 0
rec := newRecorder(sink.Logger(), noop.NewTracerProvider().Tracer(t.Name()), apiKeyID, tc.structuredLogging, tc.disableContentRecording, func(ctx context.Context) (DRPCClient, error) {
clientCalls++
require.Equal(t, "record context", ctx.Value(contextKey{}))
return nil, drpcErr
})
require.Zero(t, clientCalls, "creating a recorder must not acquire a client")
ctx := context.WithValue(t.Context(), contextKey{}, "record context")
require.ErrorIs(t, rec.RecordInterception(ctx, &recorder.InterceptionRecord{ID: interceptionID, StartedAt: now}), drpcErr)
require.ErrorIs(t, rec.RecordInterceptionEnded(ctx, &recorder.InterceptionRecordEnded{ID: interceptionID, EndedAt: now}), drpcErr)
require.ErrorIs(t, rec.RecordTokenUsage(ctx, &recorder.TokenUsageRecord{InterceptionID: interceptionID, CreatedAt: now, Input: 1}), drpcErr)
require.ErrorIs(t, rec.RecordPromptUsage(ctx, &recorder.PromptUsageRecord{InterceptionID: interceptionID, CreatedAt: now, Prompt: "prompt"}), tc.wantContentErr)
require.ErrorIs(t, rec.RecordToolUsage(ctx, &recorder.ToolUsageRecord{InterceptionID: interceptionID, CreatedAt: now, Tool: "tool"}), tc.wantContentErr)
require.ErrorIs(t, rec.RecordModelThought(ctx, &recorder.ModelThoughtRecord{InterceptionID: interceptionID, CreatedAt: now, Content: "thought"}), tc.wantContentErr)
require.Equal(t, tc.wantClientCalls, clientCalls, "each forwarded record acquires its own client")
logs := sink.Entries(func(entry slog.SinkEntry) bool { return entry.Message == recorder.InterceptionLogMarker })
require.Len(t, logs, tc.wantLogs, "logging must precede content filtering")
require.ErrorIs(t, rec.RecordPromptUsage(ctx, &recorder.PromptUsageRecord{}), recorder.ErrInvalidRecord, "validation must precede content filtering")
require.Equal(t, tc.wantClientCalls, clientCalls)
})
}
t.Run("APIKeyID", func(t *testing.T) {
t.Parallel()
apiKeyID, interceptionID := uuid.NewString(), uuid.NewString()
client := &interceptionRecordingClient{}
rec := newRecorder(testutil.NewFakeSink(t).Logger(), noop.NewTracerProvider().Tracer(t.Name()), apiKeyID, false, false, func(context.Context) (DRPCClient, error) {
return client, nil
})
require.NoError(t, rec.RecordInterception(t.Context(), &recorder.InterceptionRecord{ID: interceptionID, StartedAt: time.Now().UTC()}))
require.Len(t, client.requests, 1)
require.Equal(t, apiKeyID, client.requests[0].GetApiKeyId())
require.Equal(t, interceptionID, client.requests[0].GetId())
})
}
// interceptionRecordingClient captures interception records. Its nil embedded
// client makes any other call fail loudly.
type interceptionRecordingClient struct {
DRPCClient
requests []*proto.RecordInterceptionRequest
}
func (c *interceptionRecordingClient) RecordInterception(_ context.Context, in *proto.RecordInterceptionRequest) (*proto.RecordInterceptionResponse, error) {
c.requests = append(c.requests, in)
return &proto.RecordInterceptionResponse{}, nil
}
// blockingHandler returns a handler that signals started, then blocks until
// release is closed or the request context is canceled.
func blockingHandler() (h http.Handler, started <-chan struct{}, release func()) {
coderd/aibridged/pool.go
+3 -30
View File
@@ -20,9 +20,7 @@ import (
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/mcp"
"github.com/coder/coder/v2/aibridge/provider"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/tracing"
"github.com/coder/coder/v2/coderd/aibridged/proto"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/quartz"
)
@@ -57,7 +55,8 @@ type PoolOptions struct {
Clock quartz.Clock
// StructuredLogging makes each bridge emit AI Gateway interception
// records in the format described by [recorder.InterceptionLogMarker].
// records in the format described by
// [github.com/coder/coder/v2/aibridge/recorder.InterceptionLogMarker].
StructuredLogging bool
// DisableContentRecording stops prompts, tool call arguments and model
// thoughts from being recorded. Interceptions and token usage are still
@@ -99,10 +98,6 @@ type CachedBridgePool struct {
logger slog.Logger
options PoolOptions
// recorderMiddleware is the record policy derived from options, resolved
// once here rather than on every cache miss.
recorderMiddleware []recorder.Middleware
singleflight *singleflight.Group[string, *aibridge.RequestBridge]
metrics *aibridge.Metrics
@@ -155,15 +150,6 @@ func NewCachedBridgePool(options PoolOptions, providers []aibridge.Provider, log
clk = quartz.NewReal()
}
var recorderMiddleware []recorder.Middleware
if options.DisableContentRecording {
recorderMiddleware = append(recorderMiddleware, recorder.WithoutRecords(recorder.DisabledRecords{
PromptUsage: true,
ToolUsage: true,
ModelThought: true,
}))
}
pool := &CachedBridgePool{
cache: cache,
clock: clk,
@@ -172,8 +158,6 @@ func NewCachedBridgePool(options PoolOptions, providers []aibridge.Provider, log
tracer: tracer,
logger: logger,
recorderMiddleware: recorderMiddleware,
singleflight: &singleflight.Group[string, *aibridge.RequestBridge]{},
shuttingDownCh: make(chan struct{}),
@@ -274,18 +258,7 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
span.AddEvent("cache_miss")
providerVersion := p.providerVersion.Load()
rec := aibridge.NewRecorder(
p.logger.Named("recorder"),
p.tracer,
req.APIKeyID,
p.options.StructuredLogging,
recorder.NewDRPCRecorder(req.APIKeyID, func(clientCtx context.Context) (proto.DRPCRecorderClient, error) {
// The recorder outlives this Acquire call, so the client is acquired
// against the context of the record call being served.
return clientFn(clientCtx)
}),
p.recorderMiddleware...,
)
rec := newRecorder(p.logger, p.tracer, req.APIKeyID, p.options.StructuredLogging, p.options.DisableContentRecording, clientFn)
// Slow path.
// Creating an *aibridge.RequestBridge may take some time, so gate all subsequent callers behind the initial request and return the resulting value.