fix(coglet): return 503 from /health-check when not ready

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: bfirsh

crates/coglet/src/transport/http/routes.rs

@@ -99,7 +99,9 @@ fn generate_prediction_id() -> String {
     format!("pred_{:x}", timestamp)
 }
 
-async fn health_check(State(service): State<Arc<PredictionService>>) -> Json<HealthCheckResponse> {
+async fn health_check(
+    State(service): State<Arc<PredictionService>>,
+) -> (StatusCode, Json<HealthCheckResponse>) {
     let snapshot = service.health().await;
 
     // Run user healthcheck if ready and not busy
@@ -116,10 +118,13 @@ async fn health_check(State(service): State<Arc<PredictionService>>) -> Json<Hea
         None
     };
 
-    Json(HealthCheckResponse::from_snapshot(
-        snapshot,
-        user_healthcheck_error,
-    ))
+    let response = HealthCheckResponse::from_snapshot(snapshot, user_healthcheck_error);
+    let status_code = match response.status {
+        HealthResponse::Ready => StatusCode::OK,
+        _ => StatusCode::SERVICE_UNAVAILABLE,
+    };
+
+    (status_code, Json(response))
 }
 
 /// Write /var/run/cog/ready for K8s readiness probe.
@@ -549,7 +554,7 @@ mod tests {
             .await
             .unwrap();
 
-        assert_eq!(response.status(), StatusCode::OK);
+        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
 
         let json = response_json(response).await;
         assert_eq!(json["status"], "STARTING");
@@ -566,6 +571,7 @@ mod tests {
             .await
             .unwrap();
 
+        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
         let json = response_json(response).await;
         assert_eq!(json["status"], "UNKNOWN");
     }
@@ -920,6 +926,7 @@ mod tests {
             .await
             .unwrap();
 
+        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
         let json = response_json(response).await;
         assert_eq!(json["status"], "BUSY");
     }

docs/http.md

@@ -214,8 +214,8 @@ If the upload fails, the server responds with an error.
 ### `GET /health-check`
 
 Returns the current health status of the model container.
-This endpoint always responds with `200 OK` —
-check the `status` field in the response body to determine readiness.
+Responds with `200 OK` when the model is ready,
+or `503 Service Unavailable` otherwise.
 
 The response body is a JSON object with the following fields:
 

docs/llms.txt

@@ -458,7 +458,8 @@ func (h *Harness) cmdCurl(ts *testscript.TestScript, neg bool, args []string) {
 
 		if neg {
 			if !statusOK {
-				// Expected to fail - success!
+				// Expected to fail - write body to stdout and return
+				_, _ = ts.Stdout().Write([]byte(respBody))
 				return
 			}
 		} else {
@@ -551,9 +552,9 @@ type healthCheckResponse struct {
 	Status string `json:"status"`
 }
 
-// waitForServer polls the server's health-check endpoint until it returns READY status.
-// The server may return HTTP 200 while still in STARTING state (during setup),
-// so we must check the actual status field in the response.
+// waitForServer polls the server's health-check endpoint until it returns a
+// post-setup status. It checks the status field in the JSON response body
+// regardless of HTTP status code.
 func waitForServer(serverURL string, timeout time.Duration) bool {
 	client := &http.Client{Timeout: 5 * time.Second}
 	deadline := time.Now().Add(timeout)
@@ -565,34 +566,30 @@ func waitForServer(serverURL string, timeout time.Duration) bool {
 			continue
 		}
 
-		if resp.StatusCode == http.StatusOK {
-			body, err := io.ReadAll(resp.Body)
-			_ = resp.Body.Close()
-			if err != nil {
-				time.Sleep(200 * time.Millisecond)
-				continue
-			}
+		body, err := io.ReadAll(resp.Body)
+		_ = resp.Body.Close()
+		if err != nil {
+			time.Sleep(200 * time.Millisecond)
+			continue
+		}
 
-			var health healthCheckResponse
-			if err := json.Unmarshal(body, &health); err != nil {
-				time.Sleep(200 * time.Millisecond)
-				continue
-			}
+		var health healthCheckResponse
+		if err := json.Unmarshal(body, &health); err != nil {
+			time.Sleep(200 * time.Millisecond)
+			continue
+		}
 
-			// Return success when the server has completed setup
-			// READY = setup completed, healthcheck passed (or no healthcheck)
-			// UNHEALTHY = setup completed, but user healthcheck failed
-			// BUSY = setup completed, prediction in progress
-			if health.Status == "READY" || health.Status == "UNHEALTHY" || health.Status == "BUSY" {
-				return true
-			}
+		// Return success when the server has completed setup
+		// READY = setup completed, healthcheck passed (or no healthcheck)
+		// UNHEALTHY = setup completed, but user healthcheck failed
+		// BUSY = setup completed, prediction in progress
+		if health.Status == "READY" || health.Status == "UNHEALTHY" || health.Status == "BUSY" {
+			return true
+		}
 
-			// If setup failed, no point waiting
-			if health.Status == "SETUP_FAILED" || health.Status == "DEFUNCT" {
-				return false
-			}
-		} else {
-			_ = resp.Body.Close()
+		// If setup failed, no point waiting
+		if health.Status == "SETUP_FAILED" || health.Status == "DEFUNCT" {
+			return false
 		}
 
 		time.Sleep(200 * time.Millisecond)

integration-tests/harness/harness.go

@@ -7,8 +7,8 @@ cog build -t $TEST_IMAGE
 # Start the server
 cog serve
 
-# Test: Async healthcheck raising exception gives UNHEALTHY with error
-curl GET /health-check
+# Test: Async healthcheck raising exception gives UNHEALTHY with error (503)
+! curl GET /health-check
 stdout '"status":"UNHEALTHY"'
 stdout 'user_healthcheck_error'
 stdout 'Async healthcheck error'

integration-tests/tests/healthcheck_async_exception.txtar

@@ -15,8 +15,8 @@ stdout '"status":"READY"'
 curl POST /predictions '{"input":{"text":"trigger_slow"}}'
 stdout '"status":"succeeded"'
 
-# Now healthcheck should timeout and return UNHEALTHY
-curl GET /health-check
+# Now healthcheck should timeout and return UNHEALTHY (503)
+! curl GET /health-check
 stdout '"status":"UNHEALTHY"'
 stdout 'user_healthcheck_error'
 stdout 'timed out after 5.0 seconds'

integration-tests/tests/healthcheck_async_timeout.txtar

@@ -7,8 +7,8 @@ cog build -t $TEST_IMAGE
 # Start the server
 cog serve
 
-# Test: Async healthcheck returning False gives UNHEALTHY
-curl GET /health-check
+# Test: Async healthcheck returning False gives UNHEALTHY (503)
+! curl GET /health-check
 stdout '"status":"UNHEALTHY"'
 stdout 'user_healthcheck_error'
 stdout 'returned False'

integration-tests/tests/healthcheck_async_unhealthy.txtar

@@ -7,8 +7,8 @@ cog build -t $TEST_IMAGE
 # Start the server
 cog serve
 
-# Exception in healthcheck should return UNHEALTHY with error message
-curl GET /health-check
+# Exception in healthcheck should return UNHEALTHY with error message (503)
+! curl GET /health-check
 stdout '"status":"UNHEALTHY"'
 stdout 'user_healthcheck_error'
 stdout 'Critical system error'

integration-tests/tests/healthcheck_exception.txtar

@@ -15,8 +15,8 @@ stdout '"status":"READY"'
 curl POST /predictions '{"input":{"text":"trigger_slow"}}'
 stdout '"status":"succeeded"'
 
-# Now healthcheck should timeout and return UNHEALTHY
-curl GET /health-check
+# Now healthcheck should timeout and return UNHEALTHY (503)
+! curl GET /health-check
 stdout '"status":"UNHEALTHY"'
 stdout 'user_healthcheck_error'
 stdout 'timed out after 5.0 seconds'

integration-tests/tests/healthcheck_timeout.txtar

@@ -7,8 +7,8 @@ cog build -t $TEST_IMAGE
 # Start the server
 cog serve
 
-# Unhealthy healthcheck should return UNHEALTHY status
-curl GET /health-check
+# Unhealthy healthcheck should return UNHEALTHY status (503)
+! curl GET /health-check
 stdout '"status":"UNHEALTHY"'
 stdout 'user_healthcheck_error'
 stdout 'user-defined healthcheck returned False'