diff --git a/aibridge/api.go b/aibridge/api.go index 587c2c38f7..99c58fc4f9 100644 --- a/aibridge/api.go +++ b/aibridge/api.go @@ -70,6 +70,8 @@ func NewMetrics(reg prometheus.Registerer) *metrics.Metrics { return metrics.NewMetrics(reg) } -func NewRecorder(logger slog.Logger, tracer trace.Tracer, clientFn func() (Recorder, error)) Recorder { +// NewRecorder creates a [Recorder] which acquires a client per call. +// clientFn receives the context of the call it serves. +func NewRecorder(logger slog.Logger, tracer trace.Tracer, clientFn func(context.Context) (Recorder, error)) Recorder { return recorder.NewWrappedRecorder(logger, tracer, clientFn) } diff --git a/aibridge/internal/integrationtest/setupbridge.go b/aibridge/internal/integrationtest/setupbridge.go index d2e8c0929b..efa9407429 100644 --- a/aibridge/internal/integrationtest/setupbridge.go +++ b/aibridge/internal/integrationtest/setupbridge.go @@ -169,7 +169,7 @@ func newBridgeTestServer( } mockRec := &testutil.MockRecorder{} - rec := aibridge.NewRecorder(cfg.logger, cfg.tracer, func() (aibridge.Recorder, error) { + rec := aibridge.NewRecorder(cfg.logger, cfg.tracer, func(context.Context) (aibridge.Recorder, error) { return mockRec, nil }) diff --git a/aibridge/recorder/recorder.go b/aibridge/recorder/recorder.go index 3f2435db35..29c1ccdafd 100644 --- a/aibridge/recorder/recorder.go +++ b/aibridge/recorder/recorder.go @@ -23,14 +23,14 @@ var ( type WrappedRecorder struct { logger slog.Logger tracer trace.Tracer - clientFn func() (Recorder, error) + clientFn func(context.Context) (Recorder, error) } func (r *WrappedRecorder) RecordInterception(ctx context.Context, req *InterceptionRecord) (outErr error) { ctx, span := r.tracer.Start(ctx, "Intercept.RecordInterception", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) defer tracing.EndSpanErr(span, &outErr) - client, err := r.clientFn() + client, err := r.clientFn(ctx) if err != nil { return xerrors.Errorf("acquire client: %w", err) } @@ -48,7 +48,7 @@ func (r *WrappedRecorder) RecordInterceptionEnded(ctx context.Context, req *Inte ctx, span := r.tracer.Start(ctx, "Intercept.RecordInterceptionEnded", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) defer tracing.EndSpanErr(span, &outErr) - client, err := r.clientFn() + client, err := r.clientFn(ctx) if err != nil { return xerrors.Errorf("acquire client: %w", err) } @@ -66,7 +66,7 @@ func (r *WrappedRecorder) RecordPromptUsage(ctx context.Context, req *PromptUsag ctx, span := r.tracer.Start(ctx, "Intercept.RecordPromptUsage", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) defer tracing.EndSpanErr(span, &outErr) - client, err := r.clientFn() + client, err := r.clientFn(ctx) if err != nil { return xerrors.Errorf("acquire client: %w", err) } @@ -84,7 +84,7 @@ func (r *WrappedRecorder) RecordTokenUsage(ctx context.Context, req *TokenUsageR ctx, span := r.tracer.Start(ctx, "Intercept.RecordTokenUsage", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) defer tracing.EndSpanErr(span, &outErr) - client, err := r.clientFn() + client, err := r.clientFn(ctx) if err != nil { return xerrors.Errorf("acquire client: %w", err) } @@ -102,7 +102,7 @@ func (r *WrappedRecorder) RecordToolUsage(ctx context.Context, req *ToolUsageRec ctx, span := r.tracer.Start(ctx, "Intercept.RecordToolUsage", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) defer tracing.EndSpanErr(span, &outErr) - client, err := r.clientFn() + client, err := r.clientFn(ctx) if err != nil { return xerrors.Errorf("acquire client: %w", err) } @@ -120,7 +120,7 @@ func (r *WrappedRecorder) RecordModelThought(ctx context.Context, req *ModelThou ctx, span := r.tracer.Start(ctx, "Intercept.RecordModelThought", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) defer tracing.EndSpanErr(span, &outErr) - client, err := r.clientFn() + client, err := r.clientFn(ctx) if err != nil { return xerrors.Errorf("acquire client: %w", err) } @@ -134,7 +134,9 @@ func (r *WrappedRecorder) RecordModelThought(ctx context.Context, req *ModelThou return err } -func NewWrappedRecorder(logger slog.Logger, tracer trace.Tracer, clientFn func() (Recorder, error)) *WrappedRecorder { +// NewWrappedRecorder creates a [WrappedRecorder]. clientFn receives the +// context of the call it serves. +func NewWrappedRecorder(logger slog.Logger, tracer trace.Tracer, clientFn func(context.Context) (Recorder, error)) *WrappedRecorder { return &WrappedRecorder{ logger: logger, tracer: tracer, diff --git a/cli/aibridged.go b/cli/aibridged.go index a3316609e2..c7383151ed 100644 --- a/cli/aibridged.go +++ b/cli/aibridged.go @@ -36,10 +36,11 @@ import ( // // SubscribeProviderReload performs a best-effort initial reload synchronously, // so the pool is populated before this returns whenever the fetch succeeds. -// That reload blocks on srv.Client(), but the embedded daemon's connection is -// an in-memory pipe that comes up immediately, and the env seed (which holds -// the seed lock) has already completed earlier in startup, so the wait is -// negligible. +// That reload blocks while acquiring a client, and it passes a background +// context, so only the daemon lifecycle bounds the wait. That is acceptable +// here: the embedded daemon's connection is an in-memory pipe that comes up +// immediately, and the env seed (which holds the seed lock) has already +// completed earlier in startup, so the wait is negligible. func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg prometheus.Registerer, metrics *aibridge.Metrics) (*aibridged.Server, func(), error) { ctx := context.Background() coderAPI.Logger.Debug(ctx, "starting in-memory aibridge daemon") @@ -61,7 +62,7 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr reg.MustRegister(keypool.NewStateCollector(pool.KeyPools)) // Create daemon. Construct it before subscribing so the reloader can use - // srv.ClientContext to fetch providers over the in-memory RPC. + // srv.Client to fetch providers over the in-memory RPC. srv, err := aibridged.New(ctx, pool, func(dialCtx context.Context) (aibridged.DRPCClient, error) { return coderAPI.CreateInMemoryAIBridgeServer(dialCtx) }, logger, tracer) @@ -72,7 +73,7 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr // Subscribe to ai_providers change events so the pool tracks the database // without a restart, and perform the initial reload. The reload data path // is the in-memory RPC. - reloader := NewPoolRPCReloader(pool, srv.ClientContext, cfg, logger.Named("provider-loader"), metrics, providerMetrics) + reloader := NewPoolRPCReloader(pool, srv.Client, cfg, logger.Named("provider-loader"), metrics, providerMetrics) unsubscribe, err := aibridged.SubscribeProviderReload(ctx, coderAPI.Pubsub, reloader, logger.Named("provider-reload")) if err != nil { // Without the subscription the pool can never track provider changes, @@ -91,7 +92,7 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr // build, replace, and reload-metric accounting live in one place. type poolRPCReloader struct { pool aibridged.Pooler - client aibridged.ClientFuncWithContext + client aibridged.ClientFunc cfg codersdk.AIBridgeConfig logger slog.Logger aibridgeMetrics *aibridge.Metrics @@ -105,7 +106,7 @@ type poolRPCReloader struct { // canceled. func NewPoolRPCReloader( pool aibridged.Pooler, - client aibridged.ClientFuncWithContext, + client aibridged.ClientFunc, cfg codersdk.AIBridgeConfig, logger slog.Logger, aibridgeMetrics *aibridge.Metrics, diff --git a/coderd/aibridged/aibridged.go b/coderd/aibridged/aibridged.go index 52682c5f1d..4bc433bf47 100644 --- a/coderd/aibridged/aibridged.go +++ b/coderd/aibridged/aibridged.go @@ -167,11 +167,9 @@ func (s *Server) Err() error { return s.lifecycleCtx.Err() } -func (s *Server) Client() (DRPCClient, error) { - return s.ClientContext(context.Background()) -} - -func (s *Server) ClientContext(ctx context.Context) (DRPCClient, error) { +// Client acquires a [DRPCClient], blocking until the daemon is connected to +// coderd, the server lifecycle ends, or ctx is canceled. +func (s *Server) Client(ctx context.Context) (DRPCClient, error) { select { case <-ctx.Done(): return nil, ctx.Err() diff --git a/coderd/aibridged/aibridged_test.go b/coderd/aibridged/aibridged_test.go index f954311c9d..217d1851e2 100644 --- a/coderd/aibridged/aibridged_test.go +++ b/coderd/aibridged/aibridged_test.go @@ -123,7 +123,7 @@ func TestClient_TransientDialErrorRetries(t *testing.T) { require.NoError(t, err) t.Cleanup(func() { _ = srv.Shutdown(context.Background()) }) - _, err = srv.ClientContext(testutil.Context(t, testutil.WaitShort)) + _, err = srv.Client(testutil.Context(t, testutil.WaitShort)) require.NoError(t, err) require.Equal(t, int32(2), calls.Load()) } diff --git a/coderd/aibridged/client.go b/coderd/aibridged/client.go index 3144b1a874..693f54eb0f 100644 --- a/coderd/aibridged/client.go +++ b/coderd/aibridged/client.go @@ -10,12 +10,10 @@ import ( type Dialer func(ctx context.Context) (DRPCClient, error) -type ClientFunc func() (DRPCClient, error) - -// ClientFuncWithContext acquires a DRPCClient, honoring the passed context so a -// blocking acquisition (e.g. waiting for the daemon to connect to coderd) -// unblocks when the context is canceled. Server.ClientContext satisfies it. -type ClientFuncWithContext func(context.Context) (DRPCClient, error) +// ClientFunc acquires a DRPCClient, honoring the passed context so a blocking +// acquisition (e.g. waiting for the daemon to connect to coderd) unblocks when +// the context is canceled. Server.Client satisfies it. +type ClientFunc func(context.Context) (DRPCClient, error) // DRPCClient is the union of various service interfaces the client must support. type DRPCClient interface { diff --git a/coderd/aibridged/http.go b/coderd/aibridged/http.go index 0f8c099eaa..50e74c19fe 100644 --- a/coderd/aibridged/http.go +++ b/coderd/aibridged/http.go @@ -113,7 +113,7 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) { r.Header.Del("X-Api-Key") } - client, err := s.ClientContext(ctx) + client, err := s.Client(ctx) if err != nil { logger.Warn(ctx, "failed to connect to coderd", slog.Error(err)) http.Error(rw, ErrConnect.Error(), http.StatusServiceUnavailable) diff --git a/coderd/aibridged/mcp.go b/coderd/aibridged/mcp.go index 72e1ed0f5e..ef6c8a2628 100644 --- a/coderd/aibridged/mcp.go +++ b/coderd/aibridged/mcp.go @@ -60,7 +60,7 @@ func (m *MCPProxyFactory) Build(ctx context.Context, req Request, tracer trace.T } func (m *MCPProxyFactory) retrieveMCPServerConfigs(ctx context.Context, req Request) (map[string]mcp.ServerProxier, error) { - client, err := m.clientFn() + client, err := m.clientFn(ctx) if err != nil { return nil, xerrors.Errorf("acquire client: %w", err) } diff --git a/coderd/aibridged/pool.go b/coderd/aibridged/pool.go index cee88a98d0..3eb3630a5d 100644 --- a/coderd/aibridged/pool.go +++ b/coderd/aibridged/pool.go @@ -228,8 +228,10 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl span.AddEvent("cache_miss") providerVersion := p.providerVersion.Load() - recorder := aibridge.NewRecorder(p.logger.Named("recorder"), p.tracer, func() (aibridge.Recorder, error) { - client, err := clientFn() + recorder := aibridge.NewRecorder(p.logger.Named("recorder"), p.tracer, func(clientCtx context.Context) (aibridge.Recorder, error) { + // The recorder outlives this Acquire call, so the client is acquired + // against the context of the record call being served. + client, err := clientFn(clientCtx) if err != nil { return nil, xerrors.Errorf("acquire client: %w", err) } diff --git a/coderd/aibridged/pool_test.go b/coderd/aibridged/pool_test.go index b5a6b97ad1..4ddc9abfa6 100644 --- a/coderd/aibridged/pool_test.go +++ b/coderd/aibridged/pool_test.go @@ -47,7 +47,7 @@ func TestPool(t *testing.T) { t.Cleanup(func() { pool.Shutdown(context.Background()) }) id, id2, apiKeyID1, apiKeyID2 := uuid.New(), uuid.New(), uuid.New(), uuid.New() - clientFn := func() (aibridged.DRPCClient, error) { + clientFn := func(context.Context) (aibridged.DRPCClient, error) { return client, nil } @@ -149,7 +149,7 @@ func TestPoolReplaceProvidersClearsCacheAndUsesNewProviders(t *testing.T) { InitiatorID: uuid.New(), APIKeyID: uuid.New().String(), } - clientFn := func() (aibridged.DRPCClient, error) { + clientFn := func(context.Context) (aibridged.DRPCClient, error) { return client, nil } @@ -195,7 +195,7 @@ func TestPoolReplaceProvidersDoesNotJoinStaleSingleflight(t *testing.T) { InitiatorID: uuid.New(), APIKeyID: uuid.New().String(), } - clientFn := func() (aibridged.DRPCClient, error) { + clientFn := func(context.Context) (aibridged.DRPCClient, error) { return client, nil } @@ -266,7 +266,7 @@ func TestPoolReplaceProvidersAfterShutdownIsNoop(t *testing.T) { SessionKey: "key", InitiatorID: uuid.New(), APIKeyID: uuid.New().String(), - }, func() (aibridged.DRPCClient, error) { + }, func(context.Context) (aibridged.DRPCClient, error) { return nil, context.Canceled }, newMockMCPFactory(nil)) require.ErrorContains(t, err, "pool shutting down") @@ -294,7 +294,7 @@ func TestPool_Expiry(t *testing.T) { InitiatorID: uuid.New(), APIKeyID: uuid.New().String(), } - clientFn := func() (aibridged.DRPCClient, error) { + clientFn := func(context.Context) (aibridged.DRPCClient, error) { return client, nil } @@ -384,7 +384,7 @@ func TestPoolShutdownReplaceProviders(t *testing.T) { }, logger, nil, testTracer) require.NoError(t, err) - clientFn := func() (aibridged.DRPCClient, error) { return client, nil } + clientFn := func(context.Context) (aibridged.DRPCClient, error) { return client, nil } // Populate the cache so ReplaceProviders' Clear has an entry to evict. _, err = pool.Acquire(ctx, aibridged.Request{ diff --git a/coderd/aibridged/reload.go b/coderd/aibridged/reload.go index 305606b47f..71ede0e994 100644 --- a/coderd/aibridged/reload.go +++ b/coderd/aibridged/reload.go @@ -70,12 +70,12 @@ func SubscribeProviderReload( // before serving. // // It runs until ctx is canceled, then returns ctx.Err(). clientFn receives ctx, -// so a client acquisition that blocks (e.g. Server.ClientContext waiting for +// so a client acquisition that blocks (e.g. Server.Client waiting for // the daemon to connect to coderd) unblocks when ctx is canceled, leaving no // goroutine behind. func WatchProviderReload( ctx context.Context, - clientFn ClientFuncWithContext, + clientFn ClientFunc, reloader ProviderReloader, logger slog.Logger, ) error { @@ -109,7 +109,7 @@ func WatchProviderReload( // watchProviderReloadOnce opens a single WatchAIProviders stream and reloads on // each signal until the stream fails. received reports whether at least one // signal was received before the error. -func watchProviderReloadOnce(ctx context.Context, clientFn ClientFuncWithContext, reloader ProviderReloader, logger slog.Logger) (received bool, err error) { +func watchProviderReloadOnce(ctx context.Context, clientFn ClientFunc, reloader ProviderReloader, logger slog.Logger) (received bool, err error) { // clientFn blocks until the daemon connects to coderd or ctx is canceled. c, err := clientFn(ctx) if err != nil { diff --git a/coderd/aibridged/reload_test.go b/coderd/aibridged/reload_test.go index abb15a44cc..b73ff68968 100644 --- a/coderd/aibridged/reload_test.go +++ b/coderd/aibridged/reload_test.go @@ -186,7 +186,7 @@ func TestWatchProviderReloadCancelUnblocksClient(t *testing.T) { logger := slogtest.Make(t, nil) // clientFn blocks until its context is canceled, modeling - // Server.ClientContext waiting for the daemon to connect to coderd. Only + // Server.Client waiting for the daemon to connect to coderd. Only // watchCancel is exercised (no stream activity, no daemon close), so the // loop can return only if clientFn honors the context it receives. var once sync.Once diff --git a/coderd/aibridgedtest/aibridgedtest.go b/coderd/aibridgedtest/aibridgedtest.go index 97ccbbb184..13e3acbc1d 100644 --- a/coderd/aibridgedtest/aibridgedtest.go +++ b/coderd/aibridgedtest/aibridgedtest.go @@ -78,7 +78,7 @@ func StartTestAIBridgeDaemonWithPubsub( // The reloader fetches providers from coderd over srv's DRPC client; the // subscription drives an initial load and refreshes on change events. - reloader := cli.NewPoolRPCReloader(pool, srv.ClientContext, cfg, logger.Named("reloader"), nil, metrics) + reloader := cli.NewPoolRPCReloader(pool, srv.Client, cfg, logger.Named("reloader"), nil, metrics) unsubscribe, err := aibridged.SubscribeProviderReload(ctx, ps, reloader, logger.Named("subscriber")) if err != nil { t.Fatalf("subscribe provider reload: %v", err) diff --git a/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index ad0caa3515..c32c6bdd51 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -314,7 +314,7 @@ func newStandaloneGateway(params standaloneGatewayParams) (*standaloneGateway, e providerLogger := params.logger.Named("providers") gateway := &standaloneGateway{ daemon: daemon, - reloader: agpl.NewPoolRPCReloader(params.pool, daemon.ClientContext, params.bridgeConfig, providerLogger, params.metrics, params.providerMetrics), + reloader: agpl.NewPoolRPCReloader(params.pool, daemon.Client, params.bridgeConfig, providerLogger, params.metrics, params.providerMetrics), coderURL: params.coderURL, httpAddress: params.httpAddress, @@ -384,7 +384,7 @@ func (s *standaloneGateway) serve(ctx context.Context) error { return } // WatchProviderReload reconnects internally and normally returns only when canceled. - err := aibridged.WatchProviderReload(provReloadCtx, s.daemon.ClientContext, s.reloader, s.providerLogger) + err := aibridged.WatchProviderReload(provReloadCtx, s.daemon.Client, s.reloader, s.providerLogger) if err != nil && provReloadCtx.Err() == nil { s.providerLogger.Error(provReloadCtx, "ai provider reload watch stopped", slog.Error(err)) }