diff --git a/workspace-server/internal/handlers/hibernation_test.go b/workspace-server/internal/handlers/hibernation_test.go index da5f8df3..9d0c99c9 100644 --- a/workspace-server/internal/handlers/hibernation_test.go +++ b/workspace-server/internal/handlers/hibernation_test.go @@ -173,6 +173,17 @@ func hibernateRequest(t *testing.T, handler *WorkspaceHandler, wsID string) *htt return w } +// hibernateRequestWithQuery is like hibernateRequest but appends a query string. +func hibernateRequestWithQuery(t *testing.T, handler *WorkspaceHandler, wsID, query string) *httptest.ResponseRecorder { + t.Helper() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: wsID}} + c.Request = httptest.NewRequest(http.MethodPost, "/workspaces/"+wsID+"/hibernate?"+query, nil) + handler.Hibernate(c) + return w +} + // TestHibernateHandler_Online_Returns200 verifies that an online workspace // that is eligible for hibernation returns 200 {"status":"hibernated"}. // With the 3-step fix: handler SELECT → atomic claim UPDATE → name/tier SELECT @@ -186,9 +197,9 @@ func TestHibernateHandler_Online_Returns200(t *testing.T) { wsID := "ws-handler-online" // Hibernate() handler eligibility SELECT — checks status IN ('online','degraded'). - mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id = .* AND status IN`). + mock.ExpectQuery(`SELECT name, tier, active_tasks FROM workspaces WHERE id = .* AND status IN`). WithArgs(wsID). - WillReturnRows(sqlmock.NewRows([]string{"name", "tier"}).AddRow("Online Bot", 1)) + WillReturnRows(sqlmock.NewRows([]string{"name", "tier", "active_tasks"}).AddRow("Online Bot", 1, 0)) // HibernateWorkspace() step 1: atomic claim. mock.ExpectExec(`UPDATE workspaces`). @@ -198,7 +209,7 @@ func TestHibernateHandler_Online_Returns200(t *testing.T) { // Post-claim SELECT for name/tier. mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id`). WithArgs(wsID). - WillReturnRows(sqlmock.NewRows([]string{"name", "tier"}).AddRow("Online Bot", 1)) + WillReturnRows(sqlmock.NewRows([]string{"name", "tier", "active_tasks"}).AddRow("Online Bot", 1, 0)) // Step 3: final UPDATE. mock.ExpectExec(`UPDATE workspaces SET status = 'hibernated'`). @@ -239,7 +250,7 @@ func TestHibernateHandler_NotActive_Returns404(t *testing.T) { wsID := "ws-handler-paused" // Handler's eligibility SELECT returns no rows — workspace is not online/degraded. - mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id = .* AND status IN`). + mock.ExpectQuery(`SELECT name, tier, active_tasks FROM workspaces WHERE id = .* AND status IN`). WithArgs(wsID). WillReturnError(sql.ErrNoRows) @@ -262,6 +273,75 @@ func TestHibernateHandler_NotActive_Returns404(t *testing.T) { } } +// TestHibernateHandler_ActiveTasks_Returns409 verifies that hibernating a +// workspace with active_tasks > 0 returns 409 unless ?force=true is passed. +// (#822) +func TestHibernateHandler_ActiveTasks_Returns409(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + broadcaster := newTestBroadcaster() + handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) + + wsID := "ws-busy" + + mock.ExpectQuery(`SELECT name, tier, active_tasks FROM workspaces WHERE id = .* AND status IN`). + WithArgs(wsID). + WillReturnRows(sqlmock.NewRows([]string{"name", "tier", "active_tasks"}).AddRow("Busy Bot", 1, 3)) + + w := hibernateRequest(t, handler, wsID) + + if w.Code != http.StatusConflict { + t.Fatalf("expected 409, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if active, _ := resp["active_tasks"].(float64); active != 3 { + t.Errorf("expected active_tasks=3 in response, got %v", resp["active_tasks"]) + } +} + +// TestHibernateHandler_ActiveTasks_ForceTrue_Returns200 verifies that +// ?force=true overrides the 409 guard and proceeds with hibernation. (#822) +func TestHibernateHandler_ActiveTasks_ForceTrue_Returns200(t *testing.T) { + mock := setupTestDB(t) + setupTestRedis(t) + broadcaster := newTestBroadcaster() + handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) + + wsID := "ws-force-hibernate" + + mock.ExpectQuery(`SELECT name, tier, active_tasks FROM workspaces WHERE id = .* AND status IN`). + WithArgs(wsID). + WillReturnRows(sqlmock.NewRows([]string{"name", "tier", "active_tasks"}).AddRow("Force Bot", 1, 2)) + + // HibernateWorkspace claim + mock.ExpectExec(`UPDATE workspaces`). + WithArgs(wsID). + WillReturnResult(sqlmock.NewResult(0, 1)) + + // Post-claim SELECT + mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id`). + WithArgs(wsID). + WillReturnRows(sqlmock.NewRows([]string{"name", "tier"}).AddRow("Force Bot", 1)) + + // Final UPDATE to hibernated + mock.ExpectExec(`UPDATE workspaces SET status = 'hibernated'`). + WithArgs(wsID). + WillReturnResult(sqlmock.NewResult(0, 1)) + + // Broadcaster + mock.ExpectExec(`INSERT INTO structure_events`). + WillReturnResult(sqlmock.NewResult(0, 1)) + + w := hibernateRequestWithQuery(t, handler, wsID, "force=true") + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } +} + // TestHibernateHandler_DBError_Returns500 verifies that an unexpected DB error // on the eligibility SELECT returns 500. func TestHibernateHandler_DBError_Returns500(t *testing.T) { @@ -272,7 +352,7 @@ func TestHibernateHandler_DBError_Returns500(t *testing.T) { wsID := "ws-handler-dberror" - mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id = .* AND status IN`). + mock.ExpectQuery(`SELECT name, tier, active_tasks FROM workspaces WHERE id = .* AND status IN`). WithArgs(wsID). WillReturnError(fmt.Errorf("db: connection reset")) diff --git a/workspace-server/internal/handlers/workspace_restart.go b/workspace-server/internal/handlers/workspace_restart.go index 711e2c77..686f0596 100644 --- a/workspace-server/internal/handlers/workspace_restart.go +++ b/workspace-server/internal/handlers/workspace_restart.go @@ -190,10 +190,10 @@ func (h *WorkspaceHandler) Hibernate(c *gin.Context) { ctx := c.Request.Context() var wsName string - var tier int + var tier, activeTasks int err := db.DB.QueryRowContext(ctx, - `SELECT name, tier FROM workspaces WHERE id = $1 AND status IN ('online', 'degraded')`, id, - ).Scan(&wsName, &tier) + `SELECT name, tier, active_tasks FROM workspaces WHERE id = $1 AND status IN ('online', 'degraded')`, id, + ).Scan(&wsName, &tier, &activeTasks) if err == sql.ErrNoRows { c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found or not in a hibernatable state (must be online or degraded)"}) return @@ -203,6 +203,20 @@ func (h *WorkspaceHandler) Hibernate(c *gin.Context) { return } + // #822: reject hibernation when active tasks are in flight unless caller + // passes ?force=true. Prevents operator from accidentally killing a + // mid-task agent. + if activeTasks > 0 && c.Query("force") != "true" { + c.JSON(http.StatusConflict, gin.H{ + "error": "workspace has active tasks; use ?force=true to terminate them", + "active_tasks": activeTasks, + }) + return + } + if activeTasks > 0 { + log.Printf("[WARN] force-hibernating workspace %s (%s) with %d active tasks", id, wsName, activeTasks) + } + h.HibernateWorkspace(ctx, id) c.JSON(http.StatusOK, gin.H{"status": "hibernated"}) }