feat: add /healtz and /readyz endpoints to standalone Gateway (#26988)

Adds `/healthz` and `/readyz` endpoints to standalone AI Gateway.
* `/healthz`: returns 200 once the AI Gateways HTTP server is listening.
* `/readyz`: returns 200 when the DRPC connection to `coderd` is
established.

Cleanup:
Removed `initConnection*` fields from `aibridged.Server` as they where
not used anywhere.

coderd/aibridged/aibridged.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"net/http"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"go.opentelemetry.io/otel/trace"
@@ -41,10 +42,8 @@ type Server struct {
 	tracer trace.Tracer
 	wg     sync.WaitGroup
 
-	// initConnectionCh will receive when the daemon connects to coderd for the
-	// first time.
-	initConnectionCh   chan struct{}
-	initConnectionOnce sync.Once
+	// connected tracks whether the DRPC connection to coderd is currently active.
+	connected atomic.Bool
 
 	// lifecycleCtx is canceled when we start closing or when the
 	// connection loop exits permanently.
@@ -62,13 +61,12 @@ func New(ctx context.Context, pool Pooler, rpcDialer Dialer, logger slog.Logger,
 
 	ctx, cancel := context.WithCancelCause(ctx)
 	daemon := &Server{
-		logger:           logger,
-		tracer:           tracer,
-		clientDialer:     rpcDialer,
-		clientCh:         make(chan DRPCClient),
-		lifecycleCtx:     ctx,
-		cancelFn:         cancel,
-		initConnectionCh: make(chan struct{}),
+		logger:       logger,
+		tracer:       tracer,
+		clientDialer: rpcDialer,
+		clientCh:     make(chan DRPCClient),
+		lifecycleCtx: ctx,
+		cancelFn:     cancel,
 
 		requestBridgePool: pool,
 	}
@@ -135,17 +133,17 @@ connectLoop:
 		// failure (paired with the warning logged above).
 		s.logger.Info(s.lifecycleCtx, "successfully connected to coderd")
 		retrier.Reset()
-		s.initConnectionOnce.Do(func() {
-			close(s.initConnectionCh)
-		})
+		s.connected.Store(true)
 
 		// Serve the client until we are closed or it disconnects.
 		for {
 			select {
 			case <-s.lifecycleCtx.Done():
+				s.connected.Store(false)
 				client.DRPCConn().Close()
 				return
 			case <-client.DRPCConn().Closed():
+				s.connected.Store(false)
 				logConnect(s.lifecycleCtx, "connection to coderd closed")
 				continue connectLoop
 			case s.clientCh <- client:
@@ -201,6 +199,11 @@ func (s *Server) GetRequestHandler(ctx context.Context, req Request) (http.Handl
 	return reqBridge, nil
 }
 
+// Ready reports whether the server currently has an active DRPC connection to coderd.
+func (s *Server) Ready() bool {
+	return s.connected.Load()
+}
+
 // isShutdown returns whether the Server is shutdown or not.
 func (s *Server) isShutdown() bool {
 	select {

coderd/aibridged/aibridged_test.go

@@ -70,18 +70,28 @@ func newTestServerWithDialer(t *testing.T, dialer aibridged.Dialer, loggerOption
 	return srv, client, pool
 }
 
-// mockDRPCConn is a mock implementation of drpc.Conn
-type mockDRPCConn struct{}
+// mockDRPCConn is a mock implementation of drpc.Conn.
+// If closedCh is set, Closed() returns it so the caller can trigger a
+// disconnect by closing the channel. Otherwise a fresh never-closed
+// channel is returned on each call.
+type mockDRPCConn struct {
+	closedCh chan struct{}
+}
 
-func (*mockDRPCConn) Close() error              { return nil }
-func (*mockDRPCConn) Closed() <-chan struct{}   { ch := make(chan struct{}); return ch }
+func (*mockDRPCConn) Close() error { return nil }
+func (c *mockDRPCConn) Closed() <-chan struct{} {
+	if c.closedCh != nil {
+		return c.closedCh
+	}
+	return make(chan struct{})
+}
 func (*mockDRPCConn) Transport() drpc.Transport { return nil }
-func (*mockDRPCConn) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, in, out drpc.Message) error {
+func (*mockDRPCConn) Invoke(_ context.Context, _ string, _ drpc.Encoding, _, _ drpc.Message) error {
 	return nil
 }
 
-func (*mockDRPCConn) NewStream(ctx context.Context, rpc string, enc drpc.Encoding) (drpc.Stream, error) {
-	// nolint:nilnil // Chillchill.
+func (*mockDRPCConn) NewStream(_ context.Context, _ string, _ drpc.Encoding) (drpc.Stream, error) {
+	//nolint:nilnil // test stub
 	return nil, nil
 }
 
@@ -985,3 +995,98 @@ func TestServeHTTP_StripInternalHeaders(t *testing.T) {
 		})
 	}
 }
+
+func TestReady(t *testing.T) {
+	t.Parallel()
+
+	t.Run("FalseBeforeConnection", func(t *testing.T) {
+		t.Parallel()
+
+		logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
+		ctrl := gomock.NewController(t)
+		pool := mock.NewMockPooler(ctrl)
+		pool.EXPECT().Shutdown(gomock.Any()).MinTimes(1).Return(nil)
+
+		dialerCalled := make(chan struct{})
+		blockDialer := func(ctx context.Context) (aibridged.DRPCClient, error) {
+			select {
+			case dialerCalled <- struct{}{}:
+			default:
+			}
+			<-ctx.Done()
+			return nil, ctx.Err()
+		}
+
+		srv, err := aibridged.New(t.Context(), pool, blockDialer, logger, testTracer)
+		require.NoError(t, err)
+		t.Cleanup(func() { srv.Close() })
+
+		testutil.RequireReceive(t.Context(), t, dialerCalled)
+		require.False(t, srv.Ready(), "expected not ready before first connection")
+	})
+
+	t.Run("TrueAfterConnection", func(t *testing.T) {
+		t.Parallel()
+
+		srv, _, _ := newTestServer(t)
+
+		// newTestServer uses an immediate dialer, server should become ready quickly.
+		require.Eventually(t, srv.Ready, testutil.WaitShort, testutil.IntervalFast,
+			"expected ready after first connection")
+	})
+
+	t.Run("DisconnectAndReconnect", func(t *testing.T) {
+		t.Parallel()
+
+		logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
+		ctrl := gomock.NewController(t)
+		pool := mock.NewMockPooler(ctrl)
+		pool.EXPECT().Shutdown(gomock.Any()).MinTimes(1).Return(nil)
+
+		// allowDial gates the dialer. When open, dials succeed
+		// immediately. Replace with a fresh channel to block dials.
+		allowDial := make(chan struct{})
+		// firstConn lets the test trigger a disconnect on the
+		// initial connection.
+		firstConn := &mockDRPCConn{closedCh: make(chan struct{})}
+		var dialCount atomic.Int32
+		dialer := func(ctx context.Context) (aibridged.DRPCClient, error) {
+			select {
+			case <-allowDial:
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			}
+			var conn *mockDRPCConn
+			if dialCount.Add(1) == 1 {
+				conn = firstConn
+			} else {
+				conn = &mockDRPCConn{}
+			}
+			c := mock.NewMockDRPCClient(ctrl)
+			c.EXPECT().DRPCConn().AnyTimes().Return(conn)
+			return c, nil
+		}
+
+		// Start with dialer unblocked.
+		close(allowDial)
+
+		srv, err := aibridged.New(t.Context(), pool, dialer, logger, testTracer)
+		require.NoError(t, err)
+		t.Cleanup(func() { srv.Close() })
+		srvNotReady := func() bool { return !srv.Ready() }
+
+		// Wait for the initial connection.
+		require.Eventually(t, srv.Ready, testutil.WaitShort, testutil.IntervalFast, "expected ready after first connection")
+
+		// Block future dials, then trigger disconnect.
+		allowDial = make(chan struct{})
+		close(firstConn.closedCh)
+
+		require.Eventually(t, srvNotReady, testutil.WaitShort, testutil.IntervalFast, "expected not ready after disconnect")
+
+		// Unblock the dialer and wait for reconnect.
+		close(allowDial)
+		require.Eventually(t, srv.Ready, testutil.WaitShort, testutil.IntervalFast,
+			"expected ready after reconnect")
+	})
+}

enterprise/cli/aigatewaystart.go

@@ -136,6 +136,20 @@ func (r *RootCmd) aiGatewayStart() *serpent.Command {
 			mux.Handle("/api/v2/ai-gateway/", mw(http.StripPrefix("/api/v2/ai-gateway", srv)))
 			mux.Handle("/", mw(srv))
 
+			// healthz: returns 200 once the HTTP server is listening.
+			mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) {
+				w.WriteHeader(http.StatusOK)
+			})
+
+			// readyz: returns 200 only when the DRPC connection to coderd is established.
+			mux.HandleFunc("/readyz", func(w http.ResponseWriter, _ *http.Request) {
+				if srv.Ready() {
+					w.WriteHeader(http.StatusOK)
+					return
+				}
+				w.WriteHeader(http.StatusServiceUnavailable)
+			})
+
 			listener, err := net.Listen("tcp", httpAddress)
 			if err != nil {
 				return xerrors.Errorf("listen on %q: %w", httpAddress, err)