From b395fa7c0b5b2804e146ca115aeca4ec33e3776c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Thu, 16 Jul 2026 16:38:15 +0000 Subject: [PATCH 1/6] fix: remove standalone AI Gateway http listener dependency on loading providers --- enterprise/cli/aigatewaystart.go | 339 +++++++++++------- .../cli/aigatewaystart_internal_test.go | 240 ++++++++++--- .../tests/testdata/default_values.golden | 12 + .../testdata/listener_tls_with_ingress.golden | 12 + .../networking_ai-gateway-test.golden | 12 + .../ai-gateway/tests/testdata/nodeport.golden | 12 + helm/ai-gateway/values.yaml | 6 +- 7 files changed, 458 insertions(+), 175 deletions(-) diff --git a/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index 7faba2f9703..fcb97ac5250 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -11,6 +11,7 @@ import ( "os" "strings" "sync" + "sync/atomic" "time" "github.com/prometheus/client_golang/prometheus" @@ -35,11 +36,13 @@ import ( ) const ( - // helm/ai-gateway's terminationGracePeriodSeconds must exceed - // shutdownTimeout so graceful shutdown completes before Kubernetes sends - // SIGKILL. - shutdownTimeout = 5 * time.Minute - traceShutdownTimeout = 5 * time.Second + // The sum of shutdownTimeout, daemonShutdownTimeout, and + // traceShutdownTimeout must stay below terminationGracePeriodSeconds + // in helm/ai-gateway/values.yaml so the process can complete graceful + // shutdown before Kubernetes sends SIGKILL. + daemonShutdownTimeout = 5 * time.Second + shutdownTimeout = 5 * time.Minute + traceShutdownTimeout = 5 * time.Second healthzPath = "/healthz" readyzPath = "/readyz" @@ -143,9 +146,14 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { providerMetrics := aibridged.NewMetrics(registry) tracerProvider, _, closeTracing := agpl.ConfigureTraceProviderWithService(signalCtx, logger, vals, "coder-ai-gateway") + // The tracer is shared by the gateway's HTTP middleware, pool, and + // daemon, so it must be flushed only after run() returns and all + // span producers have stopped, hence the handler-level defer. defer func() { logger.Debug(signalCtx, "closing tracing") - traceCloseErr := shutdownWithTimeout(closeTracing, traceShutdownTimeout) + traceCtx, traceCancel := context.WithTimeout(context.Background(), traceShutdownTimeout) + defer traceCancel() + traceCloseErr := closeTracing(traceCtx) logger.Debug(signalCtx, "tracing closed", slog.Error(traceCloseErr)) }() tracer := tracerProvider.Tracer("ai-gateway") @@ -167,98 +175,26 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { } registry.MustRegister(keypool.NewStateCollector(pool.KeyPools)) - dialer := aibridged.NewWebsocketDialer(serverURL, transport, resolvedKey) - aibridgedCtx, aibridgedCancel := context.WithCancel(context.Background()) - defer aibridgedCancel() - srv, err := aibridged.New(aibridgedCtx, pool, dialer, gatewayLogger, tracer) - if err != nil { - return xerrors.Errorf("start AI Gateway daemon: %w", err) - } - defer srv.Close() - - // Fetch the initial provider set from coderd, retrying until - // success. Subsequent changes are delivered by the watch loop - // started below. The reloader's client acquisition honors the - // context of each Reload call, so loadProviders is bounded by - // signalCtx and the watch loop by watchCtx. - providerLogger := gatewayLogger.Named("providers") - reloader := agpl.NewPoolRPCReloader(pool, srv.ClientContext, vals.AI.BridgeConfig, providerLogger, metrics, providerMetrics) - if err := loadProviders(signalCtx, reloader, providerLogger, srv.Done()); err != nil { - if signalCtx.Err() != nil { - logger.Info(signalCtx, "shutting down standalone AI Gateway") - return nil - } - return xerrors.Errorf("initialize ai providers: %w", err) - } - - mw := gatewayMiddleware(vals.AI.BridgeConfig, tracer) - - // Watch coderd for provider changes and refresh the pool on each - // signal. - watchCtx, watchCancel := context.WithCancel(signalCtx) - var watchWG sync.WaitGroup - watchWG.Go(func() { - // srv.ClientContext observes watchCtx, so watchCancel below - // unblocks a pending client acquisition and drains this - // goroutine without relying on srv.Close. - if err := aibridged.WatchProviderReload(watchCtx, srv.ClientContext, reloader, providerLogger); err != nil && watchCtx.Err() == nil { - providerLogger.Warn(watchCtx, "ai provider watch loop exited", slog.Error(err)) - } - }) - defer func() { - watchCancel() - watchWG.Wait() - }() + gateway, err := newStandaloneGateway(standaloneGatewayParams{ + bridgeConfig: vals.AI.BridgeConfig, + coderURL: serverURL.String(), + httpAddress: httpAddress, + tlsCertFile: tlsCertFile, + tlsKeyFile: tlsKeyFile, - mux := newGatewayMux(srv, srv.Ready, mw) + dialer: aibridged.NewWebsocketDialer(serverURL, transport, resolvedKey), + pool: pool, - listener, err := net.Listen("tcp", httpAddress) + logger: gatewayLogger, + metrics: metrics, + providerMetrics: providerMetrics, + tracer: tracer, + }) if err != nil { - return xerrors.Errorf("listen on %q: %w", httpAddress, err) - } - defer listener.Close() - - logger.Info(signalCtx, "standalone AI Gateway listening", - slog.F("address", listener.Addr().String()), - slog.F("coder_url", serverURL.String()), - slog.F("tls", tlsCertFile != ""), - ) - - httpServer := &http.Server{ - Handler: mux, - ReadHeaderTimeout: time.Minute, - } - - serveErr := make(chan error, 1) - go func() { - if tlsCertFile != "" { - serveErr <- httpServer.ServeTLS(listener, tlsCertFile, tlsKeyFile) - } else { - serveErr <- httpServer.Serve(listener) - } - }() - - var aibridgedErr error - select { - case <-signalCtx.Done(): - logger.Info(signalCtx, "shutting down standalone AI Gateway") - case <-srv.Done(): - aibridgedErr = srv.Err() - case err := <-serveErr: - if err != nil && !errors.Is(err, http.ErrServerClosed) { - return xerrors.Errorf("serve: %w", err) - } + return err } - shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), shutdownTimeout) - defer shutdownCancel() - if err := httpServer.Shutdown(shutdownCtx); err != nil { - return xerrors.Errorf("shutdown http server: %w", err) - } - if aibridgedErr != nil { - return xerrors.Errorf("AI Gateway daemon exited: %w", aibridgedErr) - } - return nil + return gateway.run(signalCtx) }, } @@ -314,6 +250,190 @@ func gatewayMiddleware(cfg codersdk.AIBridgeConfig, tracer trace.Tracer) func(ht } } +type standaloneGatewayParams struct { + // Configuration. + bridgeConfig codersdk.AIBridgeConfig + coderURL string + httpAddress string + tlsCertFile string + tlsKeyFile string + + // Runtime dependencies. + dialer aibridged.Dialer + pool *aibridged.CachedBridgePool + + // Observability. + // logger is the gateway-scoped logger; derived loggers (daemon, + // providers) are named under it. + logger slog.Logger + metrics *aibridge.Metrics + providerMetrics *aibridged.Metrics + tracer trace.Tracer +} + +type standaloneGateway struct { + // Services. + daemon *aibridged.Server + httpServer *http.Server + reloader aibridged.ProviderReloader + + // Configuration. + coderURL string + httpAddress string + tlsCertFile string + tlsKeyFile string + + // State. + providersLoaded atomic.Bool + + // Observability. + logger slog.Logger + providerLogger slog.Logger +} + +func newStandaloneGateway(params standaloneGatewayParams) (*standaloneGateway, error) { + // Using context.Background() since daemon must outlive the run context + // so in-flight HTTP requests retain their DRPC connection during graceful shutdown. + srv, err := aibridged.New(context.Background(), params.pool, params.dialer, params.logger.Named("aibridged"), params.tracer) + if err != nil { + return nil, xerrors.Errorf("start AI Gateway daemon: %w", err) + } + + providerLogger := params.logger.Named("providers") + server := &standaloneGateway{ + coderURL: params.coderURL, + daemon: srv, + httpAddress: params.httpAddress, + logger: params.logger, + providerLogger: providerLogger, + reloader: agpl.NewPoolRPCReloader(params.pool, srv.ClientContext, params.bridgeConfig, providerLogger, params.metrics, params.providerMetrics), + tlsCertFile: params.tlsCertFile, + tlsKeyFile: params.tlsKeyFile, + } + server.httpServer = &http.Server{ + Handler: newGatewayMux(server.daemon, server.ready, gatewayMiddleware(params.bridgeConfig, params.tracer)), + ReadHeaderTimeout: time.Minute, + } + return server, nil +} + +func (s *standaloneGateway) run(ctx context.Context) error { + listener, err := net.Listen("tcp", s.httpAddress) + if err != nil { + _ = s.daemon.Close() + return xerrors.Errorf("listen on %q: %w", s.httpAddress, err) + } + + serveErr := make(chan error, 1) + var serveWG sync.WaitGroup + serveWG.Go(func() { + if s.tlsCertFile != "" { + serveErr <- s.httpServer.ServeTLS(listener, s.tlsCertFile, s.tlsKeyFile) + return + } + serveErr <- s.httpServer.Serve(listener) + }) + + s.logger.Info(ctx, "standalone AI Gateway listening", + slog.F("address", listener.Addr().String()), + slog.F("coder_url", s.coderURL), + slog.F("tls", s.tlsCertFile != ""), + ) + + provReloadCtx, provReloadCancel := context.WithCancel(ctx) + provReloadDone := make(chan struct{}) + provReloadErr := make(chan error, 1) + go func() { + defer close(provReloadDone) + if err := s.loadProviders(provReloadCtx); err != nil { + if provReloadCtx.Err() == nil { + provReloadErr <- xerrors.Errorf("initialize ai providers: %w", err) + } + return + } + if err := aibridged.WatchProviderReload(provReloadCtx, s.daemon.ClientContext, s.reloader, s.providerLogger); err != nil && provReloadCtx.Err() == nil { + provReloadErr <- xerrors.Errorf("watch ai providers: %w", err) + } + }() + + var runErr error + select { + case <-ctx.Done(): + case <-s.daemon.Done(): + if ctx.Err() == nil { + runErr = xerrors.Errorf("AI Gateway daemon exited: %w", s.daemon.Err()) + } + case err := <-provReloadErr: + runErr = err + case err := <-serveErr: + if err != nil && !errors.Is(err, http.ErrServerClosed) { + runErr = xerrors.Errorf("serve: %w", err) + } + } + s.logger.Info(ctx, "shutting down standalone AI Gateway") + + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), shutdownTimeout) + defer shutdownCancel() + provReloadCancel() + + var ( + shutdownWG sync.WaitGroup + provReloadJoinErr error + httpShutdownErr error + ) + shutdownWG.Go(func() { + select { + case <-provReloadDone: + case <-shutdownCtx.Done(): + provReloadJoinErr = xerrors.Errorf("provider synchronization did not stop in time") + } + }) + + // Closes the listener immediately, rejecting new requests, + // then drains in-flight requests. The daemon remains connected + // so draining requests retain their DRPC connection. + shutdownWG.Go(func() { + if err := s.httpServer.Shutdown(shutdownCtx); err != nil { + httpShutdownErr = xerrors.Errorf("shutdown http server: %w", err) + if closeErr := s.httpServer.Close(); closeErr != nil { + httpShutdownErr = errors.Join(httpShutdownErr, xerrors.Errorf("force close http server: %w", closeErr)) + } + } + serveWG.Wait() + }) + shutdownWG.Wait() + + daemonCtx, daemonCancel := context.WithTimeout(context.Background(), daemonShutdownTimeout) + defer daemonCancel() + var daemonShutdownErr error + if err := s.daemon.Shutdown(daemonCtx); err != nil { + daemonShutdownErr = xerrors.Errorf("shutdown AI Gateway daemon: %w", err) + } + return errors.Join(runErr, provReloadJoinErr, httpShutdownErr, daemonShutdownErr) +} + +func (s *standaloneGateway) loadProviders(ctx context.Context) error { + for r := retry.New(50*time.Millisecond, 10*time.Second); r.Wait(ctx); { + if err := s.reloader.Reload(ctx); err != nil { + select { + case <-s.daemon.Done(): + return err + default: + } + s.providerLogger.Warn(ctx, "failed to load ai providers, will retry", slog.Error(err)) + continue + } + s.providersLoaded.Store(true) + s.providerLogger.Info(ctx, "loaded ai providers from coderd") + return nil + } + return context.Cause(ctx) +} + +func (s *standaloneGateway) ready() bool { + return s.daemon.Ready() && s.providersLoaded.Load() +} + func newGatewayMux(aibridgedHandler http.Handler, aibridgedReady func() bool, middleware func(http.Handler) http.Handler) *http.ServeMux { mux := http.NewServeMux() mux.Handle("/api/v2/aibridge/", middleware(http.StripPrefix("/api/v2/aibridge", aibridgedHandler))) @@ -327,7 +447,8 @@ func newGatewayMux(aibridgedHandler http.Handler, aibridgedReady func() bool, mi }) mux.HandleFunc(readyzPath, func(w http.ResponseWriter, _ *http.Request) { - // readyz: returns 200 only when the DRPC connection to coderd is established. + // readyz: returns 200 after the initial provider load while the + // DRPC connection to coderd remains active. if aibridgedReady() { w.WriteHeader(http.StatusOK) return @@ -381,33 +502,3 @@ func resolveAIGatewayKey(key string, keyFile string) (string, error) { } return strings.TrimSpace(string(data)), nil } - -// loadProviders performs the standalone gateway's initial provider -// load by driving reloader until it succeeds or ctx is canceled. The reloader -// owns the actual fetch/build/replace/metrics work; the reloader's underlying -// client blocks until the daemon connects to coderd, and the fetch may still -// fail transiently (e.g. mid-seed contention or a dropped connection), so the -// reload is retried with backoff. A successful empty provider list is a valid -// result and ends the loop. -// -// Subsequent provider changes are delivered by WatchProviderReload, started -// after this initial load returns. -func loadProviders(ctx context.Context, reloader aibridged.ProviderReloader, logger slog.Logger, aibridgedDone <-chan struct{}) error { - for r := retry.New(50*time.Millisecond, 10*time.Second); r.Wait(ctx); { - if err := reloader.Reload(ctx); err != nil { - select { - case <-aibridgedDone: - return err - default: - } - logger.Warn(ctx, "failed to load ai providers, will retry", slog.Error(err)) - continue - } - logger.Info(ctx, "loaded ai providers from coderd") - return nil - } - if cause := context.Cause(ctx); cause != nil { - return cause - } - return ctx.Err() -} diff --git a/enterprise/cli/aigatewaystart_internal_test.go b/enterprise/cli/aigatewaystart_internal_test.go index cd0e97ede76..a7ac76d8b86 100644 --- a/enterprise/cli/aigatewaystart_internal_test.go +++ b/enterprise/cli/aigatewaystart_internal_test.go @@ -4,6 +4,8 @@ package cli import ( "context" + "fmt" + "net" "net/http" "net/http/httptest" "os" @@ -16,7 +18,9 @@ import ( "golang.org/x/xerrors" "cdr.dev/slog/v3" + "github.com/coder/coder/v2/cli/clitest" agplaibridge "github.com/coder/coder/v2/coderd/aibridge" + "github.com/coder/coder/v2/coderd/aibridged" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" ) @@ -25,7 +29,8 @@ import ( // returns its error. It models the standalone gateway's initial reload // waiting on a daemon connection to an unreachable coderd. type blockingReloader struct { - started chan struct{} + canceled chan struct{} + started chan struct{} } func (r *blockingReloader) Reload(ctx context.Context) error { @@ -34,6 +39,12 @@ func (r *blockingReloader) Reload(ctx context.Context) error { default: } <-ctx.Done() + if r.canceled != nil { + select { + case r.canceled <- struct{}{}: + default: + } + } return ctx.Err() } @@ -52,32 +63,9 @@ func (r *failThenSucceedReloader) Reload(_ context.Context) error { return nil } -// alwaysFailReloader returns the same error every time Reload is called. -type alwaysFailReloader struct { - calls atomic.Int32 - err error - after func() - called chan struct{} -} - -func (r *alwaysFailReloader) Reload(context.Context) error { - r.calls.Add(1) - if r.after != nil { - r.after() - } - select { - case r.called <- struct{}{}: - default: - } - return r.err -} - -// TestLoadProviders_Interruptible verifies that a stop signal, -// modeled by canceling the context, unblocks the initial provider load even -// when the reloader is stuck waiting for coderd. This guards the standalone -// "ai-gateway start" command against the regression where startup could not -// be interrupted. -func TestLoadProviders_Interruptible(t *testing.T) { +// Canceling the run context must interrupt a provider load that is waiting for +// coderd. +func TestStandaloneGatewayLoadProviders_Interruptible(t *testing.T) { t.Parallel() // testCtx bounds the test and drives the channel receives; runCtx is the @@ -91,9 +79,14 @@ func TestLoadProviders_Interruptible(t *testing.T) { reloader := &blockingReloader{started: make(chan struct{}, 1)} logger := slog.Make() + gateway := &standaloneGateway{ + daemon: newTestStandaloneDaemon(t, logger), + providerLogger: logger, + reloader: reloader, + } done := make(chan error, 1) go func() { - done <- loadProviders(runCtx, reloader, logger, nil) + done <- gateway.loadProviders(runCtx) }() // Wait for the reload to be in-flight, then cancel as a signal would. @@ -104,36 +97,191 @@ func TestLoadProviders_Interruptible(t *testing.T) { require.ErrorIs(t, err, context.Canceled) } -// TestLoadProviders_RetrySucceeds verifies loadProviders keeps retrying past -// transient failures and returns nil once a reload succeeds. This guards the -// retry contract: replacing the loop's continue with a return would fail here. -func TestLoadProviders_RetrySucceeds(t *testing.T) { +// Provider loading must retry transient failures and mark providers as loaded +// after a successful reload. +func TestStandaloneGatewayLoadProviders_RetrySucceeds(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) reloader := &failThenSucceedReloader{failUntil: 2} + logger := slog.Make() + gateway := &standaloneGateway{ + daemon: newTestStandaloneDaemon(t, logger), + providerLogger: logger, + reloader: reloader, + } - require.NoError(t, loadProviders(ctx, reloader, slog.Make(), nil)) - require.GreaterOrEqual(t, reloader.calls.Load(), int32(3)) + require.NoError(t, gateway.loadProviders(ctx)) + require.True(t, gateway.providersLoaded.Load()) + require.Equal(t, int32(3), reloader.calls.Load()) } -func TestLoadProviders_AIBridgedDoneStopsRetry(t *testing.T) { +func TestAIGatewayStart_HealthBeforeProviders(t *testing.T) { t.Parallel() - errMsg := "aibridged fatal" - ctx := testutil.Context(t, testutil.WaitShort) - aibridgedDone := make(chan struct{}) - reloader := &alwaysFailReloader{ - err: xerrors.New(errMsg), - called: make(chan struct{}, 1), - after: func() { - close(aibridgedDone) + gatewayAddress := fmt.Sprintf("127.0.0.1:%d", testutil.RandomPort(t)) + coderAddress := fmt.Sprintf("127.0.0.1:%d", testutil.RandomPort(t)) + + var root RootCmd + cmd, err := root.Command(root.enterpriseOnly()) + require.NoError(t, err) + inv, _ := clitest.NewWithCommand(t, cmd, + "--url", "http://"+coderAddress, + "ai-gateway", "start", + "--key", "test-key", + "--http-address", gatewayAddress, + ) + clitest.Start(t, inv.WithContext(testutil.Context(t, testutil.WaitLong))) + + client := &http.Client{Timeout: testutil.WaitShort} + baseURL := "http://" + gatewayAddress + require.Eventually(t, func() bool { + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, baseURL+healthzPath, nil) + if err != nil { + return false + } + resp, err := client.Do(req) + if err != nil { + return false + } + defer resp.Body.Close() + return resp.StatusCode == http.StatusOK + }, testutil.WaitShort, testutil.IntervalFast) + + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, baseURL+readyzPath, nil) + require.NoError(t, err) + resp, err := client.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) +} + +func TestStandaloneGatewayRun_ShutdownLifetimes(t *testing.T) { + t.Parallel() + + testCtx := testutil.Context(t, testutil.WaitLong) + logger := slog.Make() + pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger, nil, sdktrace.NewTracerProvider().Tracer("test")) + require.NoError(t, err) + + dialCtxCh := make(chan context.Context, 1) + daemon, err := aibridged.New(context.Background(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) { + select { + case dialCtxCh <- ctx: + default: + } + <-ctx.Done() + return nil, ctx.Err() + }, logger, sdktrace.NewTracerProvider().Tracer("test")) + require.NoError(t, err) + + httpAddress := fmt.Sprintf("127.0.0.1:%d", testutil.RandomPort(t)) + reloader := &blockingReloader{ + canceled: make(chan struct{}, 1), + started: make(chan struct{}, 1), + } + handlerStarted := make(chan struct{}, 1) + releaseHandler := make(chan struct{}) + gateway := &standaloneGateway{ + daemon: daemon, + httpServer: &http.Server{ + Handler: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + select { + case handlerStarted <- struct{}{}: + default: + } + <-releaseHandler + w.WriteHeader(http.StatusNoContent) + }), + ReadHeaderTimeout: testutil.WaitShort, }, + httpAddress: httpAddress, + logger: logger, + providerLogger: logger, + reloader: reloader, + } + + runCtx, cancelRun := context.WithCancel(testCtx) + runDone := make(chan error, 1) + go func() { + runDone <- gateway.run(runCtx) + }() + + dialCtx := testutil.RequireReceive(testCtx, t, dialCtxCh) + testutil.RequireReceive(testCtx, t, reloader.started) + require.Eventually(t, func() bool { + conn, err := net.Dial("tcp", httpAddress) + if err != nil { + return false + } + _ = conn.Close() + return true + }, testutil.WaitShort, testutil.IntervalFast) + + requestDone := make(chan error, 1) + go func() { + req, err := http.NewRequestWithContext(testCtx, http.MethodGet, "http://"+httpAddress, nil) + if err != nil { + requestDone <- err + return + } + resp, err := http.DefaultClient.Do(req) + if err == nil { + defer resp.Body.Close() + if resp.StatusCode != http.StatusNoContent { + err = xerrors.Errorf("unexpected status code: %d", resp.StatusCode) + } + } + requestDone <- err + }() + testutil.RequireReceive(testCtx, t, handlerStarted) + + cancelRun() + testutil.RequireReceive(testCtx, t, reloader.canceled) + select { + case <-dialCtx.Done(): + t.Fatal("daemon context canceled before the in-flight HTTP request drained") + case err := <-runDone: + t.Fatalf("server returned before the in-flight HTTP request drained: %v", err) + default: + } + + close(releaseHandler) + require.NoError(t, testutil.RequireReceive(testCtx, t, requestDone)) + require.NoError(t, testutil.RequireReceive(testCtx, t, runDone)) + select { + case <-dialCtx.Done(): + case <-testCtx.Done(): + t.Fatal("daemon dial context was not canceled") + } + select { + case <-daemon.Done(): + case <-testCtx.Done(): + t.Fatal("daemon did not stop") } - err := loadProviders(ctx, reloader, slog.Make(), aibridgedDone) - require.ErrorContains(t, err, errMsg) - require.Equal(t, int32(1), reloader.calls.Load()) + conn, err := net.Dial("tcp", httpAddress) + if err == nil { + _ = conn.Close() + } + require.Error(t, err, "HTTP listener must be closed before run returns") +} + +func newTestStandaloneDaemon(t *testing.T, logger slog.Logger) *aibridged.Server { + t.Helper() + + tracer := sdktrace.NewTracerProvider().Tracer("test") + pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger, nil, tracer) + require.NoError(t, err) + daemon, err := aibridged.New(context.Background(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) { + <-ctx.Done() + return nil, ctx.Err() + }, logger, tracer) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, daemon.Close()) + }) + return daemon } func TestResolveAIGatewayKey(t *testing.T) { diff --git a/helm/ai-gateway/tests/testdata/default_values.golden b/helm/ai-gateway/tests/testdata/default_values.golden index cdee78607ff..295336ffa55 100644 --- a/helm/ai-gateway/tests/testdata/default_values.golden +++ b/helm/ai-gateway/tests/testdata/default_values.golden @@ -101,6 +101,12 @@ spec: image: ghcr.io/coder/coder:v2.36.0 imagePullPolicy: IfNotPresent lifecycle: {} + livenessProbe: + httpGet: + path: /healthz + port: http + scheme: HTTP + initialDelaySeconds: 0 name: coder ports: - containerPort: 4001 @@ -126,6 +132,12 @@ spec: runAsUser: 1000 seccompProfile: type: RuntimeDefault + startupProbe: + httpGet: + path: /healthz + port: http + scheme: HTTP + initialDelaySeconds: 0 volumeMounts: [] restartPolicy: Always serviceAccountName: coder-ai-gateway diff --git a/helm/ai-gateway/tests/testdata/listener_tls_with_ingress.golden b/helm/ai-gateway/tests/testdata/listener_tls_with_ingress.golden index 75418126d3a..a10031fa8ad 100644 --- a/helm/ai-gateway/tests/testdata/listener_tls_with_ingress.golden +++ b/helm/ai-gateway/tests/testdata/listener_tls_with_ingress.golden @@ -102,6 +102,12 @@ spec: image: ghcr.io/coder/coder:v2.36.0 imagePullPolicy: IfNotPresent lifecycle: {} + livenessProbe: + httpGet: + path: /healthz + port: http + scheme: HTTPS + initialDelaySeconds: 0 name: coder ports: - containerPort: 4001 @@ -127,6 +133,12 @@ spec: runAsUser: 1000 seccompProfile: type: RuntimeDefault + startupProbe: + httpGet: + path: /healthz + port: http + scheme: HTTPS + initialDelaySeconds: 0 volumeMounts: - mountPath: /etc/coder/ai-gateway-auth name: ai-gateway-auth diff --git a/helm/ai-gateway/tests/testdata/networking_ai-gateway-test.golden b/helm/ai-gateway/tests/testdata/networking_ai-gateway-test.golden index 2374dbc0400..625aaf76e35 100644 --- a/helm/ai-gateway/tests/testdata/networking_ai-gateway-test.golden +++ b/helm/ai-gateway/tests/testdata/networking_ai-gateway-test.golden @@ -102,6 +102,12 @@ spec: image: ghcr.io/coder/coder:v2.36.0 imagePullPolicy: IfNotPresent lifecycle: {} + livenessProbe: + httpGet: + path: /healthz + port: http + scheme: HTTP + initialDelaySeconds: 0 name: coder ports: - containerPort: 4001 @@ -127,6 +133,12 @@ spec: runAsUser: 1000 seccompProfile: type: RuntimeDefault + startupProbe: + httpGet: + path: /healthz + port: http + scheme: HTTP + initialDelaySeconds: 0 volumeMounts: - mountPath: /etc/coder/ai-gateway-auth name: ai-gateway-auth diff --git a/helm/ai-gateway/tests/testdata/nodeport.golden b/helm/ai-gateway/tests/testdata/nodeport.golden index 423e3b1d21f..aedb723d14a 100644 --- a/helm/ai-gateway/tests/testdata/nodeport.golden +++ b/helm/ai-gateway/tests/testdata/nodeport.golden @@ -100,6 +100,12 @@ spec: image: ghcr.io/coder/coder:v2.36.0 imagePullPolicy: IfNotPresent lifecycle: {} + livenessProbe: + httpGet: + path: /healthz + port: http + scheme: HTTP + initialDelaySeconds: 0 name: coder ports: - containerPort: 4001 @@ -125,6 +131,12 @@ spec: runAsUser: 1000 seccompProfile: type: RuntimeDefault + startupProbe: + httpGet: + path: /healthz + port: http + scheme: HTTP + initialDelaySeconds: 0 volumeMounts: - mountPath: /etc/coder/ai-gateway-auth name: ai-gateway-auth diff --git a/helm/ai-gateway/values.yaml b/helm/ai-gateway/values.yaml index a3232d45661..6094538322e 100644 --- a/helm/ai-gateway/values.yaml +++ b/helm/ai-gateway/values.yaml @@ -117,17 +117,13 @@ coder: memory: 1Gi # coder.startupProbe -- Startup probe configuration for the AI Gateway. - # Enable this with a failure threshold long enough for initial provider load - # before enabling the liveness probe. startupProbe: enabled: false initialDelaySeconds: 0 # coder.livenessProbe -- Liveness probe configuration for the AI Gateway. - # Without this probe, Kubernetes does not restart a running but unresponsive - # Gateway. Enable and tune the startup probe before enabling this probe. livenessProbe: - enabled: false + enabled: true initialDelaySeconds: 0 # coder.readinessProbe -- Readiness probe configuration for the AI Gateway. From e76c6ae6b6115ca83be96a2eebec489ba7ce4b70 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Tue, 21 Jul 2026 14:15:49 +0000 Subject: [PATCH 2/6] agentic review 1 --- cli/aibridged.go | 4 +- enterprise/cli/aigatewaystart.go | 222 +++++++------- .../cli/aigatewaystart_internal_test.go | 270 ++++++++++++++---- helm/ai-gateway/values.yaml | 5 +- 4 files changed, 330 insertions(+), 171 deletions(-) diff --git a/cli/aibridged.go b/cli/aibridged.go index 75ded313c97..a3316609e2a 100644 --- a/cli/aibridged.go +++ b/cli/aibridged.go @@ -90,7 +90,7 @@ func newAIBridgeDaemon(coderAPI *coderd.API, cfg codersdk.AIBridgeConfig, reg pr // and the standalone gateway (WebSocket RPC, retried at startup) so the fetch, // build, replace, and reload-metric accounting live in one place. type poolRPCReloader struct { - pool *aibridged.CachedBridgePool + pool aibridged.Pooler client aibridged.ClientFuncWithContext cfg codersdk.AIBridgeConfig logger slog.Logger @@ -104,7 +104,7 @@ type poolRPCReloader struct { // Reload's context, so a blocking acquisition unblocks when that context is // canceled. func NewPoolRPCReloader( - pool *aibridged.CachedBridgePool, + pool aibridged.Pooler, client aibridged.ClientFuncWithContext, cfg codersdk.AIBridgeConfig, logger slog.Logger, diff --git a/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index fcb97ac5250..14381240fe2 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -36,12 +36,12 @@ import ( ) const ( - // The sum of shutdownTimeout, daemonShutdownTimeout, and + // The sum of daemonShutdownTimeout, httpShutdownTimeout, and // traceShutdownTimeout must stay below terminationGracePeriodSeconds // in helm/ai-gateway/values.yaml so the process can complete graceful // shutdown before Kubernetes sends SIGKILL. daemonShutdownTimeout = 5 * time.Second - shutdownTimeout = 5 * time.Minute + httpShutdownTimeout = 5 * time.Minute traceShutdownTimeout = 5 * time.Second healthzPath = "/healthz" @@ -147,13 +147,11 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { tracerProvider, _, closeTracing := agpl.ConfigureTraceProviderWithService(signalCtx, logger, vals, "coder-ai-gateway") // The tracer is shared by the gateway's HTTP middleware, pool, and - // daemon, so it must be flushed only after run() returns and all - // span producers have stopped, hence the handler-level defer. + // daemon, so it must be flushed only after runStandaloneGateway returns + // and all span producers have stopped, hence the handler-level defer. defer func() { logger.Debug(signalCtx, "closing tracing") - traceCtx, traceCancel := context.WithTimeout(context.Background(), traceShutdownTimeout) - defer traceCancel() - traceCloseErr := closeTracing(traceCtx) + traceCloseErr := shutdownWithTimeout(closeTracing, traceShutdownTimeout) logger.Debug(signalCtx, "tracing closed", slog.Error(traceCloseErr)) }() tracer := tracerProvider.Tracer("ai-gateway") @@ -175,7 +173,7 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { } registry.MustRegister(keypool.NewStateCollector(pool.KeyPools)) - gateway, err := newStandaloneGateway(standaloneGatewayParams{ + return runStandaloneGateway(signalCtx, standaloneGatewayParams{ bridgeConfig: vals.AI.BridgeConfig, coderURL: serverURL.String(), httpAddress: httpAddress, @@ -190,11 +188,6 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { providerMetrics: providerMetrics, tracer: tracer, }) - if err != nil { - return err - } - - return gateway.run(signalCtx) }, } @@ -241,13 +234,23 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { return cmd } -func gatewayMiddleware(cfg codersdk.AIBridgeConfig, tracer trace.Tracer) func(http.Handler) http.Handler { - mw := coderd.AIGatewayDataPlaneMiddleware(cfg) - // Tracing wraps outermost so rejected requests are still traced. - traced := tracingMiddleware(tracer) - return func(next http.Handler) http.Handler { - return traced(mw(next)) +// resolveAIGatewayKey resolves key from --key or --key-file flags. +// If both are set, an error is returned. If neither is set, an empty string is returned. +func resolveAIGatewayKey(key string, keyFile string) (string, error) { + if key != "" && keyFile != "" { + return "", xerrors.New(keyFlagsExclusiveErr) } + if key == "" && keyFile == "" { + return "", xerrors.New(keyFlagsMissingErr) + } + if keyFile == "" { + return key, nil + } + data, err := os.ReadFile(keyFile) + if err != nil { + return "", xerrors.Errorf("read AI Gateway key file %q: %w", keyFile, err) + } + return strings.TrimSpace(string(data)), nil } type standaloneGatewayParams struct { @@ -260,7 +263,7 @@ type standaloneGatewayParams struct { // Runtime dependencies. dialer aibridged.Dialer - pool *aibridged.CachedBridgePool + pool aibridged.Pooler // Observability. // logger is the gateway-scoped logger; derived loggers (daemon, @@ -284,6 +287,8 @@ type standaloneGateway struct { tlsKeyFile string // State. + // providersLoaded is an initial-load latch. Reconnects refresh providers + // through the watch loop without resetting readiness. providersLoaded atomic.Bool // Observability. @@ -291,42 +296,52 @@ type standaloneGateway struct { providerLogger slog.Logger } -func newStandaloneGateway(params standaloneGatewayParams) (*standaloneGateway, error) { - // Using context.Background() since daemon must outlive the run context - // so in-flight HTTP requests retain their DRPC connection during graceful shutdown. - srv, err := aibridged.New(context.Background(), params.pool, params.dialer, params.logger.Named("aibridged"), params.tracer) +// runStandaloneGateway establishes DRCP connection to coderd (aibridged daemon) +// and starts the standalone AI Gateway. It manages the daemon life cycle. +func runStandaloneGateway(ctx context.Context, params standaloneGatewayParams) error { + // The aibrideged daemon must outlive ctx so in-flight HTTP requests + // retain their DRPC connection during graceful HTTP shutdown. + daemon, err := aibridged.New(context.Background(), params.pool, params.dialer, params.logger.Named("aibridged"), params.tracer) if err != nil { - return nil, xerrors.Errorf("start AI Gateway daemon: %w", err) + return xerrors.Errorf("start AI Gateway daemon: %w", err) } providerLogger := params.logger.Named("providers") - server := &standaloneGateway{ - coderURL: params.coderURL, - daemon: srv, - httpAddress: params.httpAddress, + gateway := &standaloneGateway{ + daemon: daemon, + reloader: agpl.NewPoolRPCReloader(params.pool, daemon.ClientContext, params.bridgeConfig, providerLogger, params.metrics, params.providerMetrics), + + coderURL: params.coderURL, + httpAddress: params.httpAddress, + tlsCertFile: params.tlsCertFile, + tlsKeyFile: params.tlsKeyFile, + logger: params.logger, providerLogger: providerLogger, - reloader: agpl.NewPoolRPCReloader(params.pool, srv.ClientContext, params.bridgeConfig, providerLogger, params.metrics, params.providerMetrics), - tlsCertFile: params.tlsCertFile, - tlsKeyFile: params.tlsKeyFile, } - server.httpServer = &http.Server{ - Handler: newGatewayMux(server.daemon, server.ready, gatewayMiddleware(params.bridgeConfig, params.tracer)), + gateway.httpServer = &http.Server{ + Handler: newGatewayMux(gateway.daemon, gateway.ready, gatewayMiddleware(params.bridgeConfig, params.tracer)), ReadHeaderTimeout: time.Minute, } - return server, nil + + serveErr := gateway.serve(ctx) + var daemonShutdownErr error + if err := shutdownWithTimeout(daemon.Shutdown, daemonShutdownTimeout); err != nil { + daemonShutdownErr = xerrors.Errorf("shutdown AI Gateway daemon: %w", err) + } + return errors.Join(serveErr, daemonShutdownErr) } -func (s *standaloneGateway) run(ctx context.Context) error { +func (s *standaloneGateway) serve(ctx context.Context) error { listener, err := net.Listen("tcp", s.httpAddress) if err != nil { - _ = s.daemon.Close() return xerrors.Errorf("listen on %q: %w", s.httpAddress, err) } serveErr := make(chan error, 1) var serveWG sync.WaitGroup serveWG.Go(func() { + defer listener.Close() if s.tlsCertFile != "" { serveErr <- s.httpServer.ServeTLS(listener, s.tlsCertFile, s.tlsKeyFile) return @@ -342,18 +357,17 @@ func (s *standaloneGateway) run(ctx context.Context) error { provReloadCtx, provReloadCancel := context.WithCancel(ctx) provReloadDone := make(chan struct{}) - provReloadErr := make(chan error, 1) + initialProviderLoadErr := make(chan error, 1) go func() { defer close(provReloadDone) if err := s.loadProviders(provReloadCtx); err != nil { if provReloadCtx.Err() == nil { - provReloadErr <- xerrors.Errorf("initialize ai providers: %w", err) + initialProviderLoadErr <- xerrors.Errorf("initialize ai providers: %w", err) } return } - if err := aibridged.WatchProviderReload(provReloadCtx, s.daemon.ClientContext, s.reloader, s.providerLogger); err != nil && provReloadCtx.Err() == nil { - provReloadErr <- xerrors.Errorf("watch ai providers: %w", err) - } + // WatchProviderReload reconnects internally and returns only when canceled. + _ = aibridged.WatchProviderReload(provReloadCtx, s.daemon.ClientContext, s.reloader, s.providerLogger) }() var runErr error @@ -363,7 +377,7 @@ func (s *standaloneGateway) run(ctx context.Context) error { if ctx.Err() == nil { runErr = xerrors.Errorf("AI Gateway daemon exited: %w", s.daemon.Err()) } - case err := <-provReloadErr: + case err := <-initialProviderLoadErr: runErr = err case err := <-serveErr: if err != nil && !errors.Is(err, http.ErrServerClosed) { @@ -372,46 +386,34 @@ func (s *standaloneGateway) run(ctx context.Context) error { } s.logger.Info(ctx, "shutting down standalone AI Gateway") - shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), shutdownTimeout) + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), httpShutdownTimeout) defer shutdownCancel() provReloadCancel() - var ( - shutdownWG sync.WaitGroup - provReloadJoinErr error - httpShutdownErr error - ) - shutdownWG.Go(func() { - select { - case <-provReloadDone: - case <-shutdownCtx.Done(): - provReloadJoinErr = xerrors.Errorf("provider synchronization did not stop in time") - } - }) + var provReloadJoinErr error + select { + case <-provReloadDone: + case <-shutdownCtx.Done(): + provReloadJoinErr = xerrors.Errorf("provider synchronization did not stop within %s, continuing gateway shutdown", httpShutdownTimeout) + } - // Closes the listener immediately, rejecting new requests, - // then drains in-flight requests. The daemon remains connected - // so draining requests retain their DRPC connection. - shutdownWG.Go(func() { - if err := s.httpServer.Shutdown(shutdownCtx); err != nil { - httpShutdownErr = xerrors.Errorf("shutdown http server: %w", err) - if closeErr := s.httpServer.Close(); closeErr != nil { - httpShutdownErr = errors.Join(httpShutdownErr, xerrors.Errorf("force close http server: %w", closeErr)) - } + // Provider synchronization stops before HTTP draining so it cannot clear the + // bridge cache while requests are draining. The daemon remains connected so + // in-flight requests retain their DRPC connection. + var httpShutdownErr error + if err := s.httpServer.Shutdown(shutdownCtx); err != nil { + httpShutdownErr = xerrors.Errorf("shutdown http server: %w", err) + if closeErr := s.httpServer.Close(); closeErr != nil { + httpShutdownErr = errors.Join(httpShutdownErr, xerrors.Errorf("force close http server: %w", closeErr)) } - serveWG.Wait() - }) - shutdownWG.Wait() - - daemonCtx, daemonCancel := context.WithTimeout(context.Background(), daemonShutdownTimeout) - defer daemonCancel() - var daemonShutdownErr error - if err := s.daemon.Shutdown(daemonCtx); err != nil { - daemonShutdownErr = xerrors.Errorf("shutdown AI Gateway daemon: %w", err) } - return errors.Join(runErr, provReloadJoinErr, httpShutdownErr, daemonShutdownErr) + serveWG.Wait() + return errors.Join(runErr, provReloadJoinErr, httpShutdownErr) } +// loadProviders retries the initial provider load until it succeeds or the +// context or daemon stops. A successful empty provider list completes the +// initial load. Subsequent changes are handled by the watch loop. func (s *standaloneGateway) loadProviders(ctx context.Context) error { for r := retry.New(50*time.Millisecond, 10*time.Second); r.Wait(ctx); { if err := s.reloader.Reload(ctx); err != nil { @@ -434,29 +436,13 @@ func (s *standaloneGateway) ready() bool { return s.daemon.Ready() && s.providersLoaded.Load() } -func newGatewayMux(aibridgedHandler http.Handler, aibridgedReady func() bool, middleware func(http.Handler) http.Handler) *http.ServeMux { - mux := http.NewServeMux() - mux.Handle("/api/v2/aibridge/", middleware(http.StripPrefix("/api/v2/aibridge", aibridgedHandler))) - mux.Handle("/api/v2/ai-gateway/", middleware(http.StripPrefix("/api/v2/ai-gateway", aibridgedHandler))) - mux.Handle("/", middleware(aibridgedHandler)) - - // Health probes are registered without middleware. - mux.HandleFunc(healthzPath, func(w http.ResponseWriter, _ *http.Request) { - // healthz: returns 200 once the HTTP server is listening. - w.WriteHeader(http.StatusOK) - }) - - mux.HandleFunc(readyzPath, func(w http.ResponseWriter, _ *http.Request) { - // readyz: returns 200 after the initial provider load while the - // DRPC connection to coderd remains active. - if aibridgedReady() { - w.WriteHeader(http.StatusOK) - return - } - w.WriteHeader(http.StatusServiceUnavailable) - }) - - return mux +func gatewayMiddleware(cfg codersdk.AIBridgeConfig, tracer trace.Tracer) func(http.Handler) http.Handler { + mw := coderd.AIGatewayDataPlaneMiddleware(cfg) + // Tracing wraps outermost so rejected requests are still traced. + traced := tracingMiddleware(tracer) + return func(next http.Handler) http.Handler { + return traced(mw(next)) + } } // tracingMiddleware traces every request to the wrapped handler, unlike @@ -484,21 +470,27 @@ func tracingMiddleware(tracer trace.Tracer) func(http.Handler) http.Handler { } } -// resolveAIGatewayKey resolves key from --key or --key-file flags. -// If both are set, an error is returned. If neither is set, an empty string is returned. -func resolveAIGatewayKey(key string, keyFile string) (string, error) { - if key != "" && keyFile != "" { - return "", xerrors.New(keyFlagsExclusiveErr) - } - if key == "" && keyFile == "" { - return "", xerrors.New(keyFlagsMissingErr) - } - if keyFile == "" { - return key, nil - } - data, err := os.ReadFile(keyFile) - if err != nil { - return "", xerrors.Errorf("read AI Gateway key file %q: %w", keyFile, err) - } - return strings.TrimSpace(string(data)), nil +func newGatewayMux(aibridgedHandler http.Handler, aibridgedReady func() bool, middleware func(http.Handler) http.Handler) *http.ServeMux { + mux := http.NewServeMux() + mux.Handle("/api/v2/aibridge/", middleware(http.StripPrefix("/api/v2/aibridge", aibridgedHandler))) + mux.Handle("/api/v2/ai-gateway/", middleware(http.StripPrefix("/api/v2/ai-gateway", aibridgedHandler))) + mux.Handle("/", middleware(aibridgedHandler)) + + // Health probes are registered without middleware. + mux.HandleFunc(healthzPath, func(w http.ResponseWriter, _ *http.Request) { + // healthz: returns 200 once the HTTP server is listening. + w.WriteHeader(http.StatusOK) + }) + + mux.HandleFunc(readyzPath, func(w http.ResponseWriter, _ *http.Request) { + // readyz: returns 200 after the initial provider load while the + // DRPC connection to coderd remains active. + if aibridgedReady() { + w.WriteHeader(http.StatusOK) + return + } + w.WriteHeader(http.StatusServiceUnavailable) + }) + + return mux } diff --git a/enterprise/cli/aigatewaystart_internal_test.go b/enterprise/cli/aigatewaystart_internal_test.go index a7ac76d8b86..3bc8b2bee95 100644 --- a/enterprise/cli/aigatewaystart_internal_test.go +++ b/enterprise/cli/aigatewaystart_internal_test.go @@ -4,6 +4,7 @@ package cli import ( "context" + "errors" "fmt" "net" "net/http" @@ -63,38 +64,73 @@ func (r *failThenSucceedReloader) Reload(_ context.Context) error { return nil } -// Canceling the run context must interrupt a provider load that is waiting for -// coderd. -func TestStandaloneGatewayLoadProviders_Interruptible(t *testing.T) { - t.Parallel() +type failingReloader struct { + after func() + calls atomic.Int32 + err error +} - // testCtx bounds the test and drives the channel receives; runCtx is the - // context handed to loadProviders and is canceled to model a - // stop signal. They are distinct so the receives still work after the - // signal context is canceled. - testCtx := testutil.Context(t, testutil.WaitShort) - runCtx, cancel := context.WithCancel(testCtx) - defer cancel() +func (r *failingReloader) Reload(context.Context) error { + r.calls.Add(1) + if r.after != nil { + r.after() + } + return r.err +} - reloader := &blockingReloader{started: make(chan struct{}, 1)} - logger := slog.Make() +type controlledShutdownPool struct { + *aibridged.CachedBridgePool + err error + release <-chan struct{} + started chan<- struct{} +} - gateway := &standaloneGateway{ - daemon: newTestStandaloneDaemon(t, logger), - providerLogger: logger, - reloader: reloader, +func (p *controlledShutdownPool) Shutdown(ctx context.Context) error { + if p.started != nil { + p.started <- struct{}{} } - done := make(chan error, 1) - go func() { - done <- gateway.loadProviders(runCtx) - }() + if p.release != nil { + select { + case <-p.release: + case <-ctx.Done(): + return errors.Join(ctx.Err(), p.err) + } + } + return errors.Join(p.CachedBridgePool.Shutdown(ctx), p.err) +} - // Wait for the reload to be in-flight, then cancel as a signal would. - testutil.RequireReceive(testCtx, t, reloader.started) - cancel() +type standaloneGatewayTestParams struct { + address string + params standaloneGatewayParams + pool *controlledShutdownPool +} - err := testutil.RequireReceive(testCtx, t, done) - require.ErrorIs(t, err, context.Canceled) +func newStandaloneGatewayTestParams(t *testing.T) *standaloneGatewayTestParams { + t.Helper() + + logger := slog.Make() + tracer := sdktrace.NewTracerProvider().Tracer("test") + cachedPool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger, nil, tracer) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, shutdownWithTimeout(cachedPool.Shutdown, testutil.WaitShort)) + }) + + pool := &controlledShutdownPool{CachedBridgePool: cachedPool} + address := fmt.Sprintf("127.0.0.1:%d", testutil.RandomPort(t)) + return &standaloneGatewayTestParams{ + address: address, + params: standaloneGatewayParams{ + httpAddress: address, + + dialer: blockingStandaloneDaemonDialer, + pool: pool, + + logger: logger, + tracer: tracer, + }, + pool: pool, + } } // Provider loading must retry transient failures and mark providers as loaded @@ -116,6 +152,29 @@ func TestStandaloneGatewayLoadProviders_RetrySucceeds(t *testing.T) { require.Equal(t, int32(3), reloader.calls.Load()) } +func TestStandaloneGatewayLoadProviders_DaemonDoneStopsRetry(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + logger := slog.Make() + daemon := newTestStandaloneDaemon(t, logger) + reloadErr := xerrors.New("reload failed") + reloader := &failingReloader{ + after: func() { + require.NoError(t, daemon.Close()) + }, + err: reloadErr, + } + gateway := &standaloneGateway{ + daemon: daemon, + providerLogger: logger, + reloader: reloader, + } + + require.ErrorIs(t, gateway.loadProviders(ctx), reloadErr) + require.Equal(t, int32(1), reloader.calls.Load()) +} + func TestAIGatewayStart_HealthBeforeProviders(t *testing.T) { t.Parallel() @@ -131,7 +190,7 @@ func TestAIGatewayStart_HealthBeforeProviders(t *testing.T) { "--key", "test-key", "--http-address", gatewayAddress, ) - clitest.Start(t, inv.WithContext(testutil.Context(t, testutil.WaitLong))) + clitest.Start(t, inv.WithContext(testutil.Context(t, testutil.WaitShort))) client := &http.Client{Timeout: testutil.WaitShort} baseURL := "http://" + gatewayAddress @@ -156,23 +215,104 @@ func TestAIGatewayStart_HealthBeforeProviders(t *testing.T) { require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) } -func TestStandaloneGatewayRun_ShutdownLifetimes(t *testing.T) { +func TestRunStandaloneGateway_ContextCanceled(t *testing.T) { + t.Parallel() + + testCtx := testutil.Context(t, testutil.WaitShort) + runCtx, cancelRun := context.WithCancel(testCtx) + defer cancelRun() + test := newStandaloneGatewayTestParams(t) + + runDone := make(chan error, 1) + go func() { + runDone <- runStandaloneGateway(runCtx, test.params) + }() + requireListenerReady(t, test.address) + cancelRun() + + require.NoError(t, testutil.RequireReceive(testCtx, t, runDone)) + requireListenerAvailable(t, test.address, "HTTP listener must be closed before run returns") +} + +func TestRunStandaloneGateway_DaemonExited(t *testing.T) { t.Parallel() - testCtx := testutil.Context(t, testutil.WaitLong) + test := newStandaloneGatewayTestParams(t) + test.params.dialer = func(context.Context) (aibridged.DRPCClient, error) { + return nil, codersdk.NewError(http.StatusUnauthorized, codersdk.Response{Message: "invalid gateway key"}) + } + + err := runStandaloneGateway(testutil.Context(t, testutil.WaitShort), test.params) + require.ErrorContains(t, err, "AI Gateway daemon exited") + requireListenerAvailable(t, test.address, "HTTP listener must be closed before run returns") +} + +func TestRunStandaloneGateway_HTTPStopsBeforeDaemonShutdown(t *testing.T) { + t.Parallel() + + testCtx := testutil.Context(t, testutil.WaitShort) + test := newStandaloneGatewayTestParams(t) + shutdownErr := xerrors.New("pool shutdown failed") + shutdownStarted := make(chan struct{}, 1) + shutdownRelease := make(chan struct{}) + test.pool.err = shutdownErr + test.pool.started = shutdownStarted + test.pool.release = shutdownRelease + test.params.tlsCertFile = filepath.Join(t.TempDir(), "missing.crt") + test.params.tlsKeyFile = filepath.Join(t.TempDir(), "missing.key") + + runDone := make(chan error, 1) + go func() { + runDone <- runStandaloneGateway(testCtx, test.params) + }() + testutil.RequireReceive(testCtx, t, shutdownStarted) + requireListenerAvailable(t, test.address, "HTTP listener must close before daemon shutdown") + close(shutdownRelease) + + err := testutil.RequireReceive(testCtx, t, runDone) + require.ErrorContains(t, err, "serve:") + require.ErrorContains(t, err, "shutdown AI Gateway daemon:") + require.ErrorContains(t, err, shutdownErr.Error()) +} + +func TestRunStandaloneGateway_ListenAndShutdownErrors(t *testing.T) { + t.Parallel() + + test := newStandaloneGatewayTestParams(t) + shutdownErr := xerrors.New("pool shutdown failed") + test.pool.err = shutdownErr + listener, err := net.Listen("tcp", test.address) + require.NoError(t, err) + t.Cleanup(func() { + _ = listener.Close() + }) + + err = runStandaloneGateway(testutil.Context(t, testutil.WaitShort), test.params) + require.NoError(t, listener.Close()) + require.ErrorContains(t, err, "listen on") + require.ErrorContains(t, err, "shutdown AI Gateway daemon:") + require.ErrorContains(t, err, shutdownErr.Error()) +} + +func TestStandaloneGatewayServe_ShutdownOrder(t *testing.T) { + t.Parallel() + + // Set up a running daemon, provider reloader, and blocked HTTP request. + testCtx := testutil.Context(t, testutil.WaitShort) logger := slog.Make() pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger, nil, sdktrace.NewTracerProvider().Tracer("test")) require.NoError(t, err) dialCtxCh := make(chan context.Context, 1) - daemon, err := aibridged.New(context.Background(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) { + dialer := func(ctx context.Context) (aibridged.DRPCClient, error) { select { case dialCtxCh <- ctx: default: } <-ctx.Done() return nil, ctx.Err() - }, logger, sdktrace.NewTracerProvider().Tracer("test")) + } + daemon, err := aibridged.New(context.Background(), pool, dialer, logger, sdktrace.NewTracerProvider().Tracer("test")) require.NoError(t, err) httpAddress := fmt.Sprintf("127.0.0.1:%d", testutil.RandomPort(t)) @@ -201,22 +341,15 @@ func TestStandaloneGatewayRun_ShutdownLifetimes(t *testing.T) { reloader: reloader, } - runCtx, cancelRun := context.WithCancel(testCtx) - runDone := make(chan error, 1) + serveCtx, cancelServe := context.WithCancel(testCtx) + serveDone := make(chan error, 1) go func() { - runDone <- gateway.run(runCtx) + serveDone <- gateway.serve(serveCtx) }() dialCtx := testutil.RequireReceive(testCtx, t, dialCtxCh) testutil.RequireReceive(testCtx, t, reloader.started) - require.Eventually(t, func() bool { - conn, err := net.Dial("tcp", httpAddress) - if err != nil { - return false - } - _ = conn.Close() - return true - }, testutil.WaitShort, testutil.IntervalFast) + requireListenerReady(t, httpAddress) requestDone := make(chan error, 1) go func() { @@ -236,19 +369,31 @@ func TestStandaloneGatewayRun_ShutdownLifetimes(t *testing.T) { }() testutil.RequireReceive(testCtx, t, handlerStarted) - cancelRun() + // Trigger shutdown while the HTTP request is still in flight. + cancelServe() testutil.RequireReceive(testCtx, t, reloader.canceled) select { case <-dialCtx.Done(): t.Fatal("daemon context canceled before the in-flight HTTP request drained") - case err := <-runDone: + case err := <-serveDone: t.Fatalf("server returned before the in-flight HTTP request drained: %v", err) default: } + // Expect provider synchronization to stop while HTTP draining keeps the daemon alive. close(releaseHandler) require.NoError(t, testutil.RequireReceive(testCtx, t, requestDone)) - require.NoError(t, testutil.RequireReceive(testCtx, t, runDone)) + require.NoError(t, testutil.RequireReceive(testCtx, t, serveDone)) + select { + case <-dialCtx.Done(): + t.Fatal("daemon context canceled by serve") + case <-daemon.Done(): + t.Fatal("daemon stopped before its runtime owner shut it down") + default: + } + + // Expect the runtime owner to shut down the daemon after HTTP serving stops. + require.NoError(t, shutdownWithTimeout(daemon.Shutdown, daemonShutdownTimeout)) select { case <-dialCtx.Done(): case <-testCtx.Done(): @@ -260,11 +405,28 @@ func TestStandaloneGatewayRun_ShutdownLifetimes(t *testing.T) { t.Fatal("daemon did not stop") } - conn, err := net.Dial("tcp", httpAddress) - if err == nil { + requireListenerAvailable(t, httpAddress, "HTTP listener must be closed before serve returns") +} + +func requireListenerReady(t *testing.T, address string) { + t.Helper() + + require.Eventually(t, func() bool { + conn, err := net.Dial("tcp", address) + if err != nil { + return false + } _ = conn.Close() - } - require.Error(t, err, "HTTP listener must be closed before run returns") + return true + }, testutil.WaitShort, testutil.IntervalFast) +} + +func requireListenerAvailable(t *testing.T, address, message string) { + t.Helper() + + listener, err := net.Listen("tcp", address) + require.NoError(t, err, message) + require.NoError(t, listener.Close()) } func newTestStandaloneDaemon(t *testing.T, logger slog.Logger) *aibridged.Server { @@ -273,10 +435,7 @@ func newTestStandaloneDaemon(t *testing.T, logger slog.Logger) *aibridged.Server tracer := sdktrace.NewTracerProvider().Tracer("test") pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger, nil, tracer) require.NoError(t, err) - daemon, err := aibridged.New(context.Background(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) { - <-ctx.Done() - return nil, ctx.Err() - }, logger, tracer) + daemon, err := aibridged.New(context.Background(), pool, blockingStandaloneDaemonDialer, logger, tracer) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, daemon.Close()) @@ -284,6 +443,11 @@ func newTestStandaloneDaemon(t *testing.T, logger slog.Logger) *aibridged.Server return daemon } +func blockingStandaloneDaemonDialer(ctx context.Context) (aibridged.DRPCClient, error) { + <-ctx.Done() + return nil, ctx.Err() +} + func TestResolveAIGatewayKey(t *testing.T) { t.Parallel() diff --git a/helm/ai-gateway/values.yaml b/helm/ai-gateway/values.yaml index 6094538322e..62bfbe6da2a 100644 --- a/helm/ai-gateway/values.yaml +++ b/helm/ai-gateway/values.yaml @@ -122,6 +122,8 @@ coder: initialDelaySeconds: 0 # coder.livenessProbe -- Liveness probe configuration for the AI Gateway. + # This checks HTTP server responsiveness. Readiness reflects the established + # connection to coderd and initial provider configuration loaded. livenessProbe: enabled: true initialDelaySeconds: 0 @@ -186,7 +188,8 @@ aigateway: name: "" certKey: tls.crt keyKey: tls.key - # This must exceed the application's 300-second shutdown timeout. + # Allow 300 seconds for HTTP draining, 5 seconds for daemon shutdown, 5 + # seconds for tracing shutdown, and 20 seconds of termination headroom. terminationGracePeriodSeconds: 330 # Stable data-plane Service fronting the container's port 4001 listener. From 4f60862da4737dddb41a4a4b9152ab675feaf757 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Wed, 22 Jul 2026 11:06:49 +0000 Subject: [PATCH 3/6] agentic review 2 --- enterprise/cli/aigatewaystart.go | 45 +++++---- .../cli/aigatewaystart_internal_test.go | 98 +++++++++++++++++++ .../tests/testdata/default_values.golden | 6 -- .../testdata/listener_tls_with_ingress.golden | 6 -- .../networking_ai-gateway-test.golden | 6 -- .../ai-gateway/tests/testdata/nodeport.golden | 6 -- helm/ai-gateway/values.yaml | 5 +- 7 files changed, 127 insertions(+), 45 deletions(-) diff --git a/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index 14381240fe2..bc00bcb2314 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -36,13 +36,14 @@ import ( ) const ( - // The sum of daemonShutdownTimeout, httpShutdownTimeout, and - // traceShutdownTimeout must stay below terminationGracePeriodSeconds - // in helm/ai-gateway/values.yaml so the process can complete graceful - // shutdown before Kubernetes sends SIGKILL. - daemonShutdownTimeout = 5 * time.Second - httpShutdownTimeout = 5 * time.Minute - traceShutdownTimeout = 5 * time.Second + // The sum of daemonShutdownTimeout, httpShutdownTimeout, + // providerSyncShutdownTimeout, and traceShutdownTimeout must stay below + // terminationGracePeriodSeconds in helm/ai-gateway/values.yaml so the + // process can complete graceful shutdown before Kubernetes sends SIGKILL. + daemonShutdownTimeout = 5 * time.Second + httpShutdownTimeout = 5 * time.Minute + providerSyncShutdownTimeout = 5 * time.Second + traceShutdownTimeout = 5 * time.Second healthzPath = "/healthz" readyzPath = "/readyz" @@ -296,10 +297,10 @@ type standaloneGateway struct { providerLogger slog.Logger } -// runStandaloneGateway establishes DRCP connection to coderd (aibridged daemon) +// runStandaloneGateway establishes a DRPC connection to coderd (aibridged daemon) // and starts the standalone AI Gateway. It manages the daemon life cycle. func runStandaloneGateway(ctx context.Context, params standaloneGatewayParams) error { - // The aibrideged daemon must outlive ctx so in-flight HTTP requests + // The aibridged daemon must outlive ctx so in-flight HTTP requests // retain their DRPC connection during graceful HTTP shutdown. daemon, err := aibridged.New(context.Background(), params.pool, params.dialer, params.logger.Named("aibridged"), params.tracer) if err != nil { @@ -361,7 +362,10 @@ func (s *standaloneGateway) serve(ctx context.Context) error { go func() { defer close(provReloadDone) if err := s.loadProviders(provReloadCtx); err != nil { - if provReloadCtx.Err() == nil { + select { + case <-provReloadCtx.Done(): + case <-s.daemon.Done(): + default: initialProviderLoadErr <- xerrors.Errorf("initialize ai providers: %w", err) } return @@ -386,20 +390,23 @@ func (s *standaloneGateway) serve(ctx context.Context) error { } s.logger.Info(ctx, "shutting down standalone AI Gateway") - shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), httpShutdownTimeout) - defer shutdownCancel() provReloadCancel() + providerSyncCtx, providerSyncCancel := context.WithTimeout(context.Background(), providerSyncShutdownTimeout) + defer providerSyncCancel() - var provReloadJoinErr error + var providerSyncStopErr error select { case <-provReloadDone: - case <-shutdownCtx.Done(): - provReloadJoinErr = xerrors.Errorf("provider synchronization did not stop within %s, continuing gateway shutdown", httpShutdownTimeout) + case <-providerSyncCtx.Done(): + providerSyncStopErr = xerrors.Errorf("provider synchronization did not stop within %s, continuing gateway shutdown", providerSyncShutdownTimeout) } - // Provider synchronization stops before HTTP draining so it cannot clear the - // bridge cache while requests are draining. The daemon remains connected so - // in-flight requests retain their DRPC connection. + // Provider synchronization normally stops before HTTP draining so it cannot + // clear the bridge cache while requests are draining. If it does not stop + // within its timeout, continue with best-effort graceful HTTP shutdown. + // The daemon remains connected so in-flight requests retain their DRPC connection. + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), httpShutdownTimeout) + defer shutdownCancel() var httpShutdownErr error if err := s.httpServer.Shutdown(shutdownCtx); err != nil { httpShutdownErr = xerrors.Errorf("shutdown http server: %w", err) @@ -408,7 +415,7 @@ func (s *standaloneGateway) serve(ctx context.Context) error { } } serveWG.Wait() - return errors.Join(runErr, provReloadJoinErr, httpShutdownErr) + return errors.Join(runErr, providerSyncStopErr, httpShutdownErr) } // loadProviders retries the initial provider load until it succeeds or the diff --git a/enterprise/cli/aigatewaystart_internal_test.go b/enterprise/cli/aigatewaystart_internal_test.go index 3bc8b2bee95..d7af88db1af 100644 --- a/enterprise/cli/aigatewaystart_internal_test.go +++ b/enterprise/cli/aigatewaystart_internal_test.go @@ -11,12 +11,14 @@ import ( "net/http/httptest" "os" "path/filepath" + "sync" "sync/atomic" "testing" "github.com/stretchr/testify/require" sdktrace "go.opentelemetry.io/otel/sdk/trace" "golang.org/x/xerrors" + "storj.io/drpc" "cdr.dev/slog/v3" "github.com/coder/coder/v2/cli/clitest" @@ -78,6 +80,23 @@ func (r *failingReloader) Reload(context.Context) error { return r.err } +type connectedDRPCConn struct { + drpc.Conn + closed chan struct{} + once sync.Once +} + +func (c *connectedDRPCConn) Close() error { + c.once.Do(func() { + close(c.closed) + }) + return nil +} + +func (c *connectedDRPCConn) Closed() <-chan struct{} { + return c.closed +} + type controlledShutdownPool struct { *aibridged.CachedBridgePool err error @@ -175,6 +194,85 @@ func TestStandaloneGatewayLoadProviders_DaemonDoneStopsRetry(t *testing.T) { require.Equal(t, int32(1), reloader.calls.Load()) } +func TestStandaloneGatewayHealthAndReadiness(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + logger := slog.Make() + tracer := sdktrace.NewTracerProvider().Tracer("test") + pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, nil, logger, nil, tracer) + require.NoError(t, err) + connections := make(chan drpc.Conn, 2) + dialer := func(ctx context.Context) (aibridged.DRPCClient, error) { + select { + case conn := <-connections: + return &aibridged.Client{Conn: conn}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } + } + daemon, err := aibridged.New(ctx, pool, dialer, logger, tracer) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, shutdownWithTimeout(daemon.Shutdown, testutil.WaitShort)) + }) + + gateway := &standaloneGateway{ + daemon: daemon, + providerLogger: logger, + reloader: &failThenSucceedReloader{}, + } + gateway.httpServer = &http.Server{ + Handler: newGatewayMux(daemon, gateway.ready, func(next http.Handler) http.Handler { return next }), + ReadHeaderTimeout: testutil.WaitShort, + } + + // The HTTP server is healthy before the daemon connects or providers load. + require.Equal(t, http.StatusOK, healthzStatus(t, gateway)) + require.Equal(t, http.StatusServiceUnavailable, readyzStatus(t, gateway)) + + // A daemon connection alone does not make the gateway ready. + firstConn := &connectedDRPCConn{closed: make(chan struct{})} + connections <- firstConn + require.Eventually(t, daemon.Ready, testutil.WaitShort, testutil.IntervalFast) + require.Equal(t, http.StatusOK, healthzStatus(t, gateway)) + require.Equal(t, http.StatusServiceUnavailable, readyzStatus(t, gateway)) + + // The gateway becomes ready after the initial provider load completes. + require.NoError(t, gateway.loadProviders(ctx)) + require.Equal(t, http.StatusOK, healthzStatus(t, gateway)) + require.Equal(t, http.StatusOK, readyzStatus(t, gateway)) + + // Losing the daemon connection affects readiness but not HTTP health. + require.NoError(t, firstConn.Close()) + require.Eventually(t, func() bool { return !daemon.Ready() }, testutil.WaitShort, testutil.IntervalFast) + require.Equal(t, http.StatusOK, healthzStatus(t, gateway)) + require.Equal(t, http.StatusServiceUnavailable, readyzStatus(t, gateway)) + + // Readiness recovers when the daemon reconnects; providers remain loaded. + connections <- &connectedDRPCConn{closed: make(chan struct{})} + require.Eventually(t, daemon.Ready, testutil.WaitShort, testutil.IntervalFast) + require.Equal(t, http.StatusOK, healthzStatus(t, gateway)) + require.Equal(t, http.StatusOK, readyzStatus(t, gateway)) +} + +func healthzStatus(t *testing.T, gateway *standaloneGateway) int { + t.Helper() + return probeStatus(t, gateway, healthzPath) +} + +func readyzStatus(t *testing.T, gateway *standaloneGateway) int { + t.Helper() + return probeStatus(t, gateway, readyzPath) +} + +func probeStatus(t *testing.T, gateway *standaloneGateway, path string) int { + t.Helper() + rec := httptest.NewRecorder() + gateway.httpServer.Handler.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, path, nil)) + return rec.Code +} + func TestAIGatewayStart_HealthBeforeProviders(t *testing.T) { t.Parallel() diff --git a/helm/ai-gateway/tests/testdata/default_values.golden b/helm/ai-gateway/tests/testdata/default_values.golden index 295336ffa55..66be5cd0566 100644 --- a/helm/ai-gateway/tests/testdata/default_values.golden +++ b/helm/ai-gateway/tests/testdata/default_values.golden @@ -132,12 +132,6 @@ spec: runAsUser: 1000 seccompProfile: type: RuntimeDefault - startupProbe: - httpGet: - path: /healthz - port: http - scheme: HTTP - initialDelaySeconds: 0 volumeMounts: [] restartPolicy: Always serviceAccountName: coder-ai-gateway diff --git a/helm/ai-gateway/tests/testdata/listener_tls_with_ingress.golden b/helm/ai-gateway/tests/testdata/listener_tls_with_ingress.golden index a10031fa8ad..8b901086024 100644 --- a/helm/ai-gateway/tests/testdata/listener_tls_with_ingress.golden +++ b/helm/ai-gateway/tests/testdata/listener_tls_with_ingress.golden @@ -133,12 +133,6 @@ spec: runAsUser: 1000 seccompProfile: type: RuntimeDefault - startupProbe: - httpGet: - path: /healthz - port: http - scheme: HTTPS - initialDelaySeconds: 0 volumeMounts: - mountPath: /etc/coder/ai-gateway-auth name: ai-gateway-auth diff --git a/helm/ai-gateway/tests/testdata/networking_ai-gateway-test.golden b/helm/ai-gateway/tests/testdata/networking_ai-gateway-test.golden index 625aaf76e35..501ed02ad72 100644 --- a/helm/ai-gateway/tests/testdata/networking_ai-gateway-test.golden +++ b/helm/ai-gateway/tests/testdata/networking_ai-gateway-test.golden @@ -133,12 +133,6 @@ spec: runAsUser: 1000 seccompProfile: type: RuntimeDefault - startupProbe: - httpGet: - path: /healthz - port: http - scheme: HTTP - initialDelaySeconds: 0 volumeMounts: - mountPath: /etc/coder/ai-gateway-auth name: ai-gateway-auth diff --git a/helm/ai-gateway/tests/testdata/nodeport.golden b/helm/ai-gateway/tests/testdata/nodeport.golden index aedb723d14a..e8163590016 100644 --- a/helm/ai-gateway/tests/testdata/nodeport.golden +++ b/helm/ai-gateway/tests/testdata/nodeport.golden @@ -131,12 +131,6 @@ spec: runAsUser: 1000 seccompProfile: type: RuntimeDefault - startupProbe: - httpGet: - path: /healthz - port: http - scheme: HTTP - initialDelaySeconds: 0 volumeMounts: - mountPath: /etc/coder/ai-gateway-auth name: ai-gateway-auth diff --git a/helm/ai-gateway/values.yaml b/helm/ai-gateway/values.yaml index 62bfbe6da2a..151a9459d0c 100644 --- a/helm/ai-gateway/values.yaml +++ b/helm/ai-gateway/values.yaml @@ -188,8 +188,9 @@ aigateway: name: "" certKey: tls.crt keyKey: tls.key - # Allow 300 seconds for HTTP draining, 5 seconds for daemon shutdown, 5 - # seconds for tracing shutdown, and 20 seconds of termination headroom. + # Allow 5 seconds for provider synchronization, 300 seconds for HTTP + # draining, 5 seconds for daemon shutdown, 5 seconds for tracing shutdown, + # and 15 seconds of termination headroom. terminationGracePeriodSeconds: 330 # Stable data-plane Service fronting the container's port 4001 listener. From b989fe9531960087e8bd62f53bc6831d2afaac91 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Wed, 22 Jul 2026 16:56:04 +0000 Subject: [PATCH 4/6] agentic review 3 --- enterprise/cli/aigatewaystart.go | 62 ++++---- .../cli/aigatewaystart_internal_test.go | 136 +++++++++--------- helm/ai-gateway/values.yaml | 2 +- 3 files changed, 108 insertions(+), 92 deletions(-) diff --git a/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index bc00bcb2314..31e0812d865 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -37,13 +37,13 @@ import ( const ( // The sum of daemonShutdownTimeout, httpShutdownTimeout, - // providerSyncShutdownTimeout, and traceShutdownTimeout must stay below + // providerReloadShutdownTimeout, and traceShutdownTimeout must stay below // terminationGracePeriodSeconds in helm/ai-gateway/values.yaml so the // process can complete graceful shutdown before Kubernetes sends SIGKILL. - daemonShutdownTimeout = 5 * time.Second - httpShutdownTimeout = 5 * time.Minute - providerSyncShutdownTimeout = 5 * time.Second - traceShutdownTimeout = 5 * time.Second + daemonShutdownTimeout = 5 * time.Second + httpShutdownTimeout = 5 * time.Minute + providerReloadShutdownTimeout = 5 * time.Second + traceShutdownTimeout = 5 * time.Second healthzPath = "/healthz" readyzPath = "/readyz" @@ -297,8 +297,9 @@ type standaloneGateway struct { providerLogger slog.Logger } -// runStandaloneGateway establishes a DRPC connection to coderd (aibridged daemon) -// and starts the standalone AI Gateway. It manages the daemon life cycle. +// runStandaloneGateway starts the aibridged daemon and serves the standalone +// AI Gateway. The daemon dials coderd asynchronously, so HTTP serving does not +// wait for the DRPC connection. It manages the daemon life cycle. func runStandaloneGateway(ctx context.Context, params standaloneGatewayParams) error { // The aibridged daemon must outlive ctx so in-flight HTTP requests // retain their DRPC connection during graceful HTTP shutdown. @@ -358,31 +359,38 @@ func (s *standaloneGateway) serve(ctx context.Context) error { provReloadCtx, provReloadCancel := context.WithCancel(ctx) provReloadDone := make(chan struct{}) - initialProviderLoadErr := make(chan error, 1) go func() { defer close(provReloadDone) if err := s.loadProviders(provReloadCtx); err != nil { - select { - case <-provReloadCtx.Done(): - case <-s.daemon.Done(): - default: - initialProviderLoadErr <- xerrors.Errorf("initialize ai providers: %w", err) + if provReloadCtx.Err() == nil { + s.providerLogger.Error(provReloadCtx, "initial ai provider load stopped", slog.Error(err)) } return } - // WatchProviderReload reconnects internally and returns only when canceled. - _ = aibridged.WatchProviderReload(provReloadCtx, s.daemon.ClientContext, s.reloader, s.providerLogger) + // WatchProviderReload reconnects internally and normally returns only when canceled. + err := aibridged.WatchProviderReload(provReloadCtx, s.daemon.ClientContext, s.reloader, s.providerLogger) + if err != nil && provReloadCtx.Err() == nil { + s.providerLogger.Error(provReloadCtx, "ai provider reload watch stopped", slog.Error(err)) + } }() var runErr error select { case <-ctx.Done(): case <-s.daemon.Done(): + // daemon uses context.Background() so no race with ctx.Done() is possible, ctx.Err() check is not needed. + runErr = xerrors.Errorf("AI Gateway daemon exited: %w", s.daemon.Err()) + case <-provReloadDone: if ctx.Err() == nil { - runErr = xerrors.Errorf("AI Gateway daemon exited: %w", s.daemon.Err()) + select { + // reload can exit due to daemon failure + // covering race with previous daemon.Done() case. + case <-s.daemon.Done(): + runErr = xerrors.Errorf("AI Gateway daemon exited: %w", s.daemon.Err()) + default: + runErr = xerrors.New("provider reload stopped unexpectedly") + } } - case err := <-initialProviderLoadErr: - runErr = err case err := <-serveErr: if err != nil && !errors.Is(err, http.ErrServerClosed) { runErr = xerrors.Errorf("serve: %w", err) @@ -391,19 +399,19 @@ func (s *standaloneGateway) serve(ctx context.Context) error { s.logger.Info(ctx, "shutting down standalone AI Gateway") provReloadCancel() - providerSyncCtx, providerSyncCancel := context.WithTimeout(context.Background(), providerSyncShutdownTimeout) - defer providerSyncCancel() + provReloadShutdownCtx, provReloadShutdownCancel := context.WithTimeout(context.Background(), providerReloadShutdownTimeout) + defer provReloadShutdownCancel() - var providerSyncStopErr error + var provReloadStopErr error select { case <-provReloadDone: - case <-providerSyncCtx.Done(): - providerSyncStopErr = xerrors.Errorf("provider synchronization did not stop within %s, continuing gateway shutdown", providerSyncShutdownTimeout) + case <-provReloadShutdownCtx.Done(): + provReloadStopErr = xerrors.Errorf("provider reload did not stop within %s, continuing gateway shutdown", providerReloadShutdownTimeout) } - // Provider synchronization normally stops before HTTP draining so it cannot - // clear the bridge cache while requests are draining. If it does not stop - // within its timeout, continue with best-effort graceful HTTP shutdown. + // Provider reload normally stops before HTTP draining so it cannot clear the + // bridge cache while requests are draining. If it does not stop within its + // timeout, continue with best-effort graceful HTTP shutdown. // The daemon remains connected so in-flight requests retain their DRPC connection. shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), httpShutdownTimeout) defer shutdownCancel() @@ -415,7 +423,7 @@ func (s *standaloneGateway) serve(ctx context.Context) error { } } serveWG.Wait() - return errors.Join(runErr, providerSyncStopErr, httpShutdownErr) + return errors.Join(runErr, provReloadStopErr, httpShutdownErr) } // loadProviders retries the initial provider load until it succeeds or the diff --git a/enterprise/cli/aigatewaystart_internal_test.go b/enterprise/cli/aigatewaystart_internal_test.go index d7af88db1af..2d55b7d7ae7 100644 --- a/enterprise/cli/aigatewaystart_internal_test.go +++ b/enterprise/cli/aigatewaystart_internal_test.go @@ -28,29 +28,6 @@ import ( "github.com/coder/coder/v2/testutil" ) -// blockingReloader blocks in Reload until the context is canceled, then -// returns its error. It models the standalone gateway's initial reload -// waiting on a daemon connection to an unreachable coderd. -type blockingReloader struct { - canceled chan struct{} - started chan struct{} -} - -func (r *blockingReloader) Reload(ctx context.Context) error { - select { - case r.started <- struct{}{}: - default: - } - <-ctx.Done() - if r.canceled != nil { - select { - case r.canceled <- struct{}{}: - default: - } - } - return ctx.Err() -} - // failThenSucceedReloader fails the first failUntil reloads, then succeeds, // modeling a coderd connection or provider fetch that recovers after a few // transient failures. @@ -152,46 +129,76 @@ func newStandaloneGatewayTestParams(t *testing.T) *standaloneGatewayTestParams { } } -// Provider loading must retry transient failures and mark providers as loaded -// after a successful reload. -func TestStandaloneGatewayLoadProviders_RetrySucceeds(t *testing.T) { +func TestStandaloneGatewayLoadProviders(t *testing.T) { t.Parallel() - ctx := testutil.Context(t, testutil.WaitShort) - reloader := &failThenSucceedReloader{failUntil: 2} - logger := slog.Make() - gateway := &standaloneGateway{ - daemon: newTestStandaloneDaemon(t, logger), - providerLogger: logger, - reloader: reloader, + reloadErr := xerrors.New("reload failed") + tests := []struct { + name string + setup func(*testing.T, *aibridged.Server, context.CancelFunc) (aibridged.ProviderReloader, *atomic.Int32) + wantErr error + wantCalls int32 + wantLoaded bool + }{ + { + name: "Retry succeeds", + setup: func(_ *testing.T, _ *aibridged.Server, _ context.CancelFunc) (aibridged.ProviderReloader, *atomic.Int32) { + reloader := &failThenSucceedReloader{failUntil: 2} + return reloader, &reloader.calls + }, + wantCalls: 3, + wantLoaded: true, + }, + { + name: "Daemon stops retry", + setup: func(t *testing.T, daemon *aibridged.Server, _ context.CancelFunc) (aibridged.ProviderReloader, *atomic.Int32) { + reloader := &failingReloader{ + after: func() { + require.NoError(t, daemon.Close()) + }, + err: reloadErr, + } + return reloader, &reloader.calls + }, + wantErr: reloadErr, + wantCalls: 1, + }, + { + name: "Context cancellation stops retry", + setup: func(_ *testing.T, _ *aibridged.Server, cancel context.CancelFunc) (aibridged.ProviderReloader, *atomic.Int32) { + reloader := &failingReloader{after: cancel, err: reloadErr} + return reloader, &reloader.calls + }, + wantErr: context.Canceled, + wantCalls: 1, + }, } - require.NoError(t, gateway.loadProviders(ctx)) - require.True(t, gateway.providersLoaded.Load()) - require.Equal(t, int32(3), reloader.calls.Load()) -} + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() -func TestStandaloneGatewayLoadProviders_DaemonDoneStopsRetry(t *testing.T) { - t.Parallel() + ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitShort)) + defer cancel() + logger := slog.Make() + daemon := newTestStandaloneDaemon(t, logger) + reloader, calls := tc.setup(t, daemon, cancel) + gateway := &standaloneGateway{ + daemon: daemon, + providerLogger: logger, + reloader: reloader, + } - ctx := testutil.Context(t, testutil.WaitShort) - logger := slog.Make() - daemon := newTestStandaloneDaemon(t, logger) - reloadErr := xerrors.New("reload failed") - reloader := &failingReloader{ - after: func() { - require.NoError(t, daemon.Close()) - }, - err: reloadErr, - } - gateway := &standaloneGateway{ - daemon: daemon, - providerLogger: logger, - reloader: reloader, + err := gateway.loadProviders(ctx) + if tc.wantErr == nil { + require.NoError(t, err) + } else { + require.ErrorIs(t, err, tc.wantErr) + } + require.Equal(t, tc.wantCalls, calls.Load()) + require.Equal(t, tc.wantLoaded, gateway.providersLoaded.Load()) + }) } - - require.ErrorIs(t, gateway.loadProviders(ctx), reloadErr) - require.Equal(t, int32(1), reloader.calls.Load()) } func TestStandaloneGatewayHealthAndReadiness(t *testing.T) { @@ -414,11 +421,9 @@ func TestStandaloneGatewayServe_ShutdownOrder(t *testing.T) { require.NoError(t, err) httpAddress := fmt.Sprintf("127.0.0.1:%d", testutil.RandomPort(t)) - reloader := &blockingReloader{ - canceled: make(chan struct{}, 1), - started: make(chan struct{}, 1), - } + reloader := &failThenSucceedReloader{} handlerStarted := make(chan struct{}, 1) + httpShutdownStarted := make(chan struct{}, 1) releaseHandler := make(chan struct{}) gateway := &standaloneGateway{ daemon: daemon, @@ -438,6 +443,9 @@ func TestStandaloneGatewayServe_ShutdownOrder(t *testing.T) { providerLogger: logger, reloader: reloader, } + gateway.httpServer.RegisterOnShutdown(func() { + httpShutdownStarted <- struct{}{} + }) serveCtx, cancelServe := context.WithCancel(testCtx) serveDone := make(chan error, 1) @@ -446,7 +454,7 @@ func TestStandaloneGatewayServe_ShutdownOrder(t *testing.T) { }() dialCtx := testutil.RequireReceive(testCtx, t, dialCtxCh) - testutil.RequireReceive(testCtx, t, reloader.started) + require.Eventually(t, gateway.providersLoaded.Load, testutil.WaitShort, testutil.IntervalFast) requireListenerReady(t, httpAddress) requestDone := make(chan error, 1) @@ -467,9 +475,9 @@ func TestStandaloneGatewayServe_ShutdownOrder(t *testing.T) { }() testutil.RequireReceive(testCtx, t, handlerStarted) - // Trigger shutdown while the HTTP request is still in flight. + // Trigger shutdown after the initial load enters the provider watch loop. cancelServe() - testutil.RequireReceive(testCtx, t, reloader.canceled) + testutil.RequireReceive(testCtx, t, httpShutdownStarted) select { case <-dialCtx.Done(): t.Fatal("daemon context canceled before the in-flight HTTP request drained") @@ -478,7 +486,7 @@ func TestStandaloneGatewayServe_ShutdownOrder(t *testing.T) { default: } - // Expect provider synchronization to stop while HTTP draining keeps the daemon alive. + // Expect provider reload to stop while HTTP draining keeps the daemon alive. close(releaseHandler) require.NoError(t, testutil.RequireReceive(testCtx, t, requestDone)) require.NoError(t, testutil.RequireReceive(testCtx, t, serveDone)) diff --git a/helm/ai-gateway/values.yaml b/helm/ai-gateway/values.yaml index 151a9459d0c..4d8b43d4744 100644 --- a/helm/ai-gateway/values.yaml +++ b/helm/ai-gateway/values.yaml @@ -188,7 +188,7 @@ aigateway: name: "" certKey: tls.crt keyKey: tls.key - # Allow 5 seconds for provider synchronization, 300 seconds for HTTP + # Allow 5 seconds for provider reload shutdown, 300 seconds for HTTP # draining, 5 seconds for daemon shutdown, 5 seconds for tracing shutdown, # and 15 seconds of termination headroom. terminationGracePeriodSeconds: 330 From 494e3ce76a8d0cc2848661ea127c04a0fb7a7df8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Wed, 22 Jul 2026 17:11:31 +0000 Subject: [PATCH 5/6] move resolveAIGatewayKey to end of file for less diff --- enterprise/cli/aigatewaystart.go | 38 ++++++++++++++++---------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/enterprise/cli/aigatewaystart.go b/enterprise/cli/aigatewaystart.go index 31e0812d865..26b4ffcc774 100644 --- a/enterprise/cli/aigatewaystart.go +++ b/enterprise/cli/aigatewaystart.go @@ -235,25 +235,6 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command { return cmd } -// resolveAIGatewayKey resolves key from --key or --key-file flags. -// If both are set, an error is returned. If neither is set, an empty string is returned. -func resolveAIGatewayKey(key string, keyFile string) (string, error) { - if key != "" && keyFile != "" { - return "", xerrors.New(keyFlagsExclusiveErr) - } - if key == "" && keyFile == "" { - return "", xerrors.New(keyFlagsMissingErr) - } - if keyFile == "" { - return key, nil - } - data, err := os.ReadFile(keyFile) - if err != nil { - return "", xerrors.Errorf("read AI Gateway key file %q: %w", keyFile, err) - } - return strings.TrimSpace(string(data)), nil -} - type standaloneGatewayParams struct { // Configuration. bridgeConfig codersdk.AIBridgeConfig @@ -509,3 +490,22 @@ func newGatewayMux(aibridgedHandler http.Handler, aibridgedReady func() bool, mi return mux } + +// resolveAIGatewayKey resolves key from --key or --key-file flags. +// If both are set, an error is returned. If neither is set, an empty string is returned. +func resolveAIGatewayKey(key string, keyFile string) (string, error) { + if key != "" && keyFile != "" { + return "", xerrors.New(keyFlagsExclusiveErr) + } + if key == "" && keyFile == "" { + return "", xerrors.New(keyFlagsMissingErr) + } + if keyFile == "" { + return key, nil + } + data, err := os.ReadFile(keyFile) + if err != nil { + return "", xerrors.Errorf("read AI Gateway key file %q: %w", keyFile, err) + } + return strings.TrimSpace(string(data)), nil +} From 8ac17f10f71ebd045d41b9d1e74318de6fa48f2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Thu, 23 Jul 2026 09:35:06 +0000 Subject: [PATCH 6/6] review 1: test cleanup --- .../cli/aigatewaystart_internal_test.go | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/enterprise/cli/aigatewaystart_internal_test.go b/enterprise/cli/aigatewaystart_internal_test.go index 2d55b7d7ae7..4ebb14d5951 100644 --- a/enterprise/cli/aigatewaystart_internal_test.go +++ b/enterprise/cli/aigatewaystart_internal_test.go @@ -500,16 +500,8 @@ func TestStandaloneGatewayServe_ShutdownOrder(t *testing.T) { // Expect the runtime owner to shut down the daemon after HTTP serving stops. require.NoError(t, shutdownWithTimeout(daemon.Shutdown, daemonShutdownTimeout)) - select { - case <-dialCtx.Done(): - case <-testCtx.Done(): - t.Fatal("daemon dial context was not canceled") - } - select { - case <-daemon.Done(): - case <-testCtx.Done(): - t.Fatal("daemon did not stop") - } + testutil.TryReceive(testCtx, t, dialCtx.Done()) + testutil.TryReceive(testCtx, t, daemon.Done()) requireListenerAvailable(t, httpAddress, "HTTP listener must be closed before serve returns") } @@ -517,14 +509,15 @@ func TestStandaloneGatewayServe_ShutdownOrder(t *testing.T) { func requireListenerReady(t *testing.T, address string) { t.Helper() - require.Eventually(t, func() bool { - conn, err := net.Dial("tcp", address) + ctx := testutil.Context(t, testutil.WaitShort) + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", address) if err != nil { return false } _ = conn.Close() return true - }, testutil.WaitShort, testutil.IntervalFast) + }, testutil.IntervalFast) } func requireListenerAvailable(t *testing.T, address, message string) {