diff --git a/packages/vm-agent/internal/errorreport/reporter.go b/packages/vm-agent/internal/errorreport/reporter.go index 2134ad5f5..bd1cdc324 100644 --- a/packages/vm-agent/internal/errorreport/reporter.go +++ b/packages/vm-agent/internal/errorreport/reporter.go @@ -47,6 +47,9 @@ type Reporter struct { queue []ErrorEntry stopC chan struct{} doneC chan struct{} + + startOnce sync.Once + shutdownOnce sync.Once } // New creates a Reporter with the given configuration. @@ -94,7 +97,9 @@ func (r *Reporter) Start() { if r == nil { return } - go r.flushLoop() + r.startOnce.Do(func() { + go r.flushLoop() + }) } // Shutdown flushes any remaining entries and stops the background goroutine. @@ -102,7 +107,13 @@ func (r *Reporter) Shutdown() { if r == nil { return } - close(r.stopC) + r.shutdownOnce.Do(func() { + close(r.stopC) + r.startOnce.Do(func() { + r.flush() + close(r.doneC) + }) + }) <-r.doneC } diff --git a/packages/vm-agent/internal/errorreport/reporter_test.go b/packages/vm-agent/internal/errorreport/reporter_test.go index 02d334d64..6b13b094c 100644 --- a/packages/vm-agent/internal/errorreport/reporter_test.go +++ b/packages/vm-agent/internal/errorreport/reporter_test.go @@ -171,6 +171,73 @@ func TestShutdownFlushesRemaining(t *testing.T) { } } +func TestShutdownIsIdempotentBeforeStart(t *testing.T) { + var mu sync.Mutex + var received []ErrorEntry + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var payload struct { + Errors []ErrorEntry `json:"errors"` + } + json.Unmarshal(body, &payload) + + mu.Lock() + received = append(received, payload.Errors...) + mu.Unlock() + + w.WriteHeader(204) + })) + defer srv.Close() + + r := New(srv.URL, "node-1", "token", Config{ + FlushInterval: time.Hour, + MaxBatchSize: 100, + MaxQueueSize: 50, + HTTPTimeout: 5 * time.Second, + }) + r.Report(ErrorEntry{Message: "queued-before-start", Source: "test"}) + + r.Shutdown() + r.Shutdown() + + mu.Lock() + defer mu.Unlock() + if len(received) != 1 { + t.Fatalf("expected queued entry to flush on Shutdown before Start, got %d", len(received)) + } +} + +func TestShutdownIsConcurrentSafe(t *testing.T) { + r := New("http://localhost", "node-1", "token", Config{ + FlushInterval: time.Hour, + MaxBatchSize: 100, + MaxQueueSize: 50, + }) + r.Start() + + var wg sync.WaitGroup + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + r.Shutdown() + }() + } + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("concurrent Shutdown calls did not complete") + } +} + func TestSendIncludesAuthHeader(t *testing.T) { var authHeader string diff --git a/packages/vm-agent/internal/eventstore/store.go b/packages/vm-agent/internal/eventstore/store.go index 95d03b179..93bbfb481 100644 --- a/packages/vm-agent/internal/eventstore/store.go +++ b/packages/vm-agent/internal/eventstore/store.go @@ -28,9 +28,11 @@ type EventRecord struct { // Store is a SQLite-backed event store. type Store struct { - db *sql.DB - dbPath string - mu sync.Mutex // serializes writes + db *sql.DB + dbPath string + mu sync.Mutex // serializes writes + closeOnce sync.Once + closeErr error } // New opens (or creates) a SQLite event store at the given path. @@ -198,5 +200,8 @@ func (s *Store) DBPath() string { // Close closes the underlying database connection. func (s *Store) Close() error { - return s.db.Close() + s.closeOnce.Do(func() { + s.closeErr = s.db.Close() + }) + return s.closeErr } diff --git a/packages/vm-agent/internal/messagereport/reporter.go b/packages/vm-agent/internal/messagereport/reporter.go index 9c5e166d8..ee1ad1748 100644 --- a/packages/vm-agent/internal/messagereport/reporter.go +++ b/packages/vm-agent/internal/messagereport/reporter.go @@ -50,6 +50,8 @@ type Reporter struct { stopC chan struct{} doneC chan struct{} + + shutdownOnce sync.Once } // New creates a Reporter backed by the given SQLite database. @@ -291,7 +293,9 @@ func (r *Reporter) Shutdown() { if r == nil { return } - close(r.stopC) + r.shutdownOnce.Do(func() { + close(r.stopC) + }) <-r.doneC } diff --git a/packages/vm-agent/internal/messagereport/reporter_test.go b/packages/vm-agent/internal/messagereport/reporter_test.go index f96bf2959..de37c02e6 100644 --- a/packages/vm-agent/internal/messagereport/reporter_test.go +++ b/packages/vm-agent/internal/messagereport/reporter_test.go @@ -74,6 +74,48 @@ func TestNew_ValidConfig(t *testing.T) { } } +func TestShutdownIsIdempotent(t *testing.T) { + db := openTestDB(t) + cfg := testConfig("http://localhost", "ws-1") + r, err := New(db, cfg) + if err != nil { + t.Fatalf("new reporter: %v", err) + } + + r.Shutdown() + r.Shutdown() +} + +func TestShutdownIsConcurrentSafe(t *testing.T) { + db := openTestDB(t) + cfg := testConfig("http://localhost", "ws-1") + r, err := New(db, cfg) + if err != nil { + t.Fatalf("new reporter: %v", err) + } + + var wg sync.WaitGroup + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + r.Shutdown() + }() + } + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("concurrent Shutdown calls did not complete") + } +} + func TestNew_NilDB(t *testing.T) { cfg := testConfig("http://localhost", "ws-1") _, err := New(nil, cfg) diff --git a/packages/vm-agent/internal/resourcemon/monitor.go b/packages/vm-agent/internal/resourcemon/monitor.go index 802863448..f442c48cc 100644 --- a/packages/vm-agent/internal/resourcemon/monitor.go +++ b/packages/vm-agent/internal/resourcemon/monitor.go @@ -13,6 +13,7 @@ import ( "runtime" "strconv" "strings" + "sync" "syscall" "time" @@ -36,10 +37,12 @@ type Snapshot struct { // Monitor collects and stores resource metrics. type Monitor struct { - db *sql.DB - dbPath string - cancel context.CancelFunc - done chan struct{} + db *sql.DB + dbPath string + cancel context.CancelFunc + done chan struct{} + closeOnce sync.Once + closeErr error } // New creates a resource monitor that writes metrics to the given SQLite path. @@ -222,7 +225,10 @@ func (m *Monitor) DBPath() string { // Close stops the collection loop and closes the database. func (m *Monitor) Close() error { - m.cancel() - <-m.done - return m.db.Close() + m.closeOnce.Do(func() { + m.cancel() + <-m.done + m.closeErr = m.db.Close() + }) + return m.closeErr } diff --git a/packages/vm-agent/internal/server/lifecycle_shutdown_test.go b/packages/vm-agent/internal/server/lifecycle_shutdown_test.go new file mode 100644 index 000000000..84a3529f9 --- /dev/null +++ b/packages/vm-agent/internal/server/lifecycle_shutdown_test.go @@ -0,0 +1,129 @@ +package server + +import ( + "context" + "net/http" + "path/filepath" + "testing" + "time" + + "github.com/workspace/vm-agent/internal/acp" + "github.com/workspace/vm-agent/internal/config" + "github.com/workspace/vm-agent/internal/container" + "github.com/workspace/vm-agent/internal/eventstore" + "github.com/workspace/vm-agent/internal/messagereport" + "github.com/workspace/vm-agent/internal/persistence" + "github.com/workspace/vm-agent/internal/ports" + "github.com/workspace/vm-agent/internal/resourcemon" +) + +func TestServerStopIsIdempotentAndClosesOwnedResources(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + store, err := persistence.Open(filepath.Join(dir, "persistence.db")) + if err != nil { + t.Fatalf("open persistence store: %v", err) + } + evStore, err := eventstore.New(filepath.Join(dir, "events.db")) + if err != nil { + t.Fatalf("open event store: %v", err) + } + resMon, err := resourcemon.New(filepath.Join(dir, "metrics.db"), time.Hour) + if err != nil { + t.Fatalf("open resource monitor: %v", err) + } + + s := &Server{ + config: &config.Config{}, + httpServer: &http.Server{}, + done: make(chan struct{}), + sessionHosts: make(map[string]*acp.SessionHost), + sessionMcpServers: make(map[string][]acp.McpServerEntry), + sessionProfileOvr: make(map[string]profileOverrides), + sessionTaskCtx: make(map[string]taskCallbackContext), + messageReporters: make(map[string]*messagereport.Reporter), + store: store, + eventStore: evStore, + resourceMonitor: resMon, + workspaces: make(map[string]*WorkspaceRuntime), + portScanners: make(map[string]*ports.Scanner), + portDiscoveries: make(map[string]*container.Discovery), + } + + if err := s.Stop(context.Background()); err != nil { + t.Fatalf("first Stop returned error: %v", err) + } + + select { + case <-s.done: + default: + t.Fatal("expected Stop to close done channel") + } + + if err := s.Stop(context.Background()); err != nil { + t.Fatalf("second Stop returned error: %v", err) + } + if err := evStore.Close(); err != nil { + t.Fatalf("event store Close should be idempotent after Server.Stop: %v", err) + } + if err := resMon.Close(); err != nil { + t.Fatalf("resource monitor Close should be idempotent after Server.Stop: %v", err) + } +} + +func TestRemoveSessionHostsForWorkspaceReleasesLockAndCleansMetadata(t *testing.T) { + t.Parallel() + + s := &Server{ + sessionHosts: map[string]*acp.SessionHost{ + "ws-1:sess-1": {}, + "ws-1:sess-2": {}, + "ws-2:sess-1": {}, + }, + sessionMcpServers: map[string][]acp.McpServerEntry{ + "ws-1:sess-1": {{URL: "https://example.com/a"}}, + "ws-2:sess-1": {{URL: "https://example.com/b"}}, + }, + sessionProfileOvr: map[string]profileOverrides{ + "ws-1:sess-1": {Model: "model-a"}, + "ws-2:sess-1": {Model: "model-b"}, + }, + sessionTaskCtx: map[string]taskCallbackContext{ + "ws-1:sess-2": {TaskID: "task-a"}, + "ws-2:sess-1": {TaskID: "task-b"}, + }, + } + + hosts := s.removeSessionHostsForWorkspace("ws-1") + if len(hosts) != 2 { + t.Fatalf("expected 2 removed hosts, got %d", len(hosts)) + } + + if !s.sessionHostMu.TryLock() { + t.Fatal("removeSessionHostsForWorkspace returned while sessionHostMu was still locked") + } + defer s.sessionHostMu.Unlock() + + if _, ok := s.sessionHosts["ws-1:sess-1"]; ok { + t.Fatal("expected ws-1:sess-1 host to be removed") + } + if _, ok := s.sessionHosts["ws-1:sess-2"]; ok { + t.Fatal("expected ws-1:sess-2 host to be removed") + } + if _, ok := s.sessionHosts["ws-2:sess-1"]; !ok { + t.Fatal("expected other workspace host to remain") + } + if _, ok := s.sessionMcpServers["ws-1:sess-1"]; ok { + t.Fatal("expected ws-1 MCP metadata to be removed") + } + if _, ok := s.sessionMcpServers["ws-2:sess-1"]; !ok { + t.Fatal("expected other workspace MCP metadata to remain") + } + if _, ok := s.sessionProfileOvr["ws-1:sess-1"]; ok { + t.Fatal("expected ws-1 profile override to be removed") + } + if _, ok := s.sessionTaskCtx["ws-1:sess-2"]; ok { + t.Fatal("expected ws-1 task context to be removed") + } +} diff --git a/packages/vm-agent/internal/server/server.go b/packages/vm-agent/internal/server/server.go index db22796ef..dbe22815d 100644 --- a/packages/vm-agent/internal/server/server.go +++ b/packages/vm-agent/internal/server/server.go @@ -6,6 +6,7 @@ import ( "context" "database/sql" "encoding/json" + "errors" "fmt" "io" "log/slog" @@ -95,6 +96,8 @@ type Server struct { callbackToken string httpClient *http.Client // shared HTTP client with timeout for control-plane callbacks done chan struct{} + stopOnce sync.Once + stopErr error // Deployment mode — set via SetDeployEngine() when Role=deployment deployEngine *deploy.Engine @@ -762,16 +765,28 @@ func (s *Server) Start() error { // StopAllWorkspacesAndSessions transitions all local workloads to stopped state. // This is invoked during node shutdown to ensure no child workloads are left active. func (s *Server) StopAllWorkspacesAndSessions() { + var ptyManagers []*pty.Manager + s.workspaceMu.Lock() workspaceIDs := make([]string, 0, len(s.workspaces)) + seenPTYManagers := make(map[*pty.Manager]struct{}, len(s.workspaces)) for id, runtime := range s.workspaces { - runtime.PTY.CloseAllSessions() + if runtime.PTY != nil { + if _, ok := seenPTYManagers[runtime.PTY]; !ok { + seenPTYManagers[runtime.PTY] = struct{}{} + ptyManagers = append(ptyManagers, runtime.PTY) + } + } runtime.Status = "stopped" runtime.UpdatedAt = nowUTC() workspaceIDs = append(workspaceIDs, id) } s.workspaceMu.Unlock() + for _, manager := range ptyManagers { + manager.CloseAllSessions() + } + for _, workspaceID := range workspaceIDs { if s.agentSessions != nil { sessions := s.agentSessions.List(workspaceID) @@ -790,46 +805,109 @@ func (s *Server) StopAllWorkspacesAndSessions() { // Stop gracefully stops the server. func (s *Server) Stop(ctx context.Context) error { + s.stopOnce.Do(func() { + s.stopErr = s.stop(ctx) + }) + return s.stopErr +} + +func (s *Server) stop(ctx context.Context) error { + var errs []error + // Signal background goroutines to stop. - close(s.done) + if s.done != nil { + close(s.done) + } // Stop all port scanners s.stopAllPortScanners() // Close JWT validator - s.jwtValidator.Close() + if s.jwtValidator != nil { + s.jwtValidator.Close() + } - s.sessionHostMu.Lock() - for key, host := range s.sessionHosts { + for _, host := range s.removeAllSessionHosts() { if host != nil { host.Stop() } - delete(s.sessionHosts, key) } - s.sessionHostMu.Unlock() // Close all workspace PTY sessions. - s.workspaceMu.Lock() - for _, runtime := range s.workspaces { - runtime.PTY.CloseAllSessions() + for _, manager := range s.collectWorkspacePTYManagers() { + manager.CloseAllSessions() } - s.workspaceMu.Unlock() - // Flush and stop error reporter + // Flush and stop error reporter. s.errorReporter.Shutdown() - // Flush and stop all per-workspace message reporters + // Flush and stop all per-workspace message reporters. s.shutdownAllReporters() - // Close persistence store + // Close server-owned persistence and telemetry stores. if s.store != nil { if err := s.store.Close(); err != nil { + errs = append(errs, fmt.Errorf("close persistence store: %w", err)) slog.Warn("Failed to close persistence store", "error", err) } } + if s.eventStore != nil { + if err := s.eventStore.Close(); err != nil { + errs = append(errs, fmt.Errorf("close event store: %w", err)) + slog.Warn("Failed to close event store", "error", err) + } + } + if s.resourceMonitor != nil { + if err := s.resourceMonitor.Close(); err != nil { + errs = append(errs, fmt.Errorf("close resource monitor: %w", err)) + slog.Warn("Failed to close resource monitor", "error", err) + } + } - // Shutdown HTTP server - return s.httpServer.Shutdown(ctx) + // Shutdown HTTP server. + if s.httpServer != nil { + if err := s.httpServer.Shutdown(ctx); err != nil { + errs = append(errs, fmt.Errorf("shutdown http server: %w", err)) + } + } + + return errors.Join(errs...) +} + +func (s *Server) collectWorkspacePTYManagers() []*pty.Manager { + s.workspaceMu.RLock() + defer s.workspaceMu.RUnlock() + + seen := make(map[*pty.Manager]struct{}, len(s.workspaces)) + managers := make([]*pty.Manager, 0, len(s.workspaces)) + for _, runtime := range s.workspaces { + if runtime.PTY == nil { + continue + } + if _, ok := seen[runtime.PTY]; ok { + continue + } + seen[runtime.PTY] = struct{}{} + managers = append(managers, runtime.PTY) + } + return managers +} + +func (s *Server) removeAllSessionHosts() map[string]*acp.SessionHost { + s.sessionHostMu.Lock() + defer s.sessionHostMu.Unlock() + + hosts := make(map[string]*acp.SessionHost, len(s.sessionHosts)) + for key, host := range s.sessionHosts { + if host != nil { + hosts[key] = host + } + delete(s.sessionHosts, key) + delete(s.sessionMcpServers, key) + delete(s.sessionProfileOvr, key) + delete(s.sessionTaskCtx, key) + } + return hosts } // setupRoutes configures the HTTP routes. diff --git a/packages/vm-agent/internal/server/workspaces.go b/packages/vm-agent/internal/server/workspaces.go index 95e4b3eb2..bfcd90d67 100644 --- a/packages/vm-agent/internal/server/workspaces.go +++ b/packages/vm-agent/internal/server/workspaces.go @@ -20,17 +20,10 @@ import ( ) func (s *Server) stopSessionHost(workspaceID, sessionID string) { - hostKey := workspaceID + ":" + sessionID - s.sessionHostMu.Lock() - existing := s.sessionHosts[hostKey] + existing := s.removeSessionHost(workspaceID, sessionID) if existing != nil { existing.Stop() - delete(s.sessionHosts, hostKey) } - delete(s.sessionMcpServers, hostKey) - delete(s.sessionProfileOvr, hostKey) - delete(s.sessionTaskCtx, hostKey) - s.sessionHostMu.Unlock() // Clean up persisted MCP servers (best-effort). if s.store != nil { @@ -41,6 +34,19 @@ func (s *Server) stopSessionHost(workspaceID, sessionID string) { } } +func (s *Server) removeSessionHost(workspaceID, sessionID string) *acp.SessionHost { + hostKey := workspaceID + ":" + sessionID + s.sessionHostMu.Lock() + defer s.sessionHostMu.Unlock() + + existing := s.sessionHosts[hostKey] + delete(s.sessionHosts, hostKey) + delete(s.sessionMcpServers, hostKey) + delete(s.sessionProfileOvr, hostKey) + delete(s.sessionTaskCtx, hostKey) + return existing +} + // removeWorkspaceContainer stops and removes the Docker container associated with // a workspace. This must be called before removing the workspace's Docker volume // (Docker won't remove a volume that's in use by a container). @@ -77,28 +83,43 @@ func (s *Server) removeWorkspaceContainer(workspaceID string) { } func (s *Server) stopSessionHostsForWorkspace(workspaceID string) { + hosts := s.removeSessionHostsForWorkspace(workspaceID) + for _, host := range hosts { + if host != nil { + host.Stop() + } + } + + // Clean up all persisted MCP servers for this workspace (best-effort). + if s.store != nil { + if err := s.store.DeleteWorkspaceMcpServers(workspaceID); err != nil { + slog.Warn("Failed to delete workspace MCP servers from SQLite", + "workspace", workspaceID, "error", err) + } + } +} + +func (s *Server) removeSessionHostsForWorkspace(workspaceID string) map[string]*acp.SessionHost { prefix := workspaceID + ":" + hosts := make(map[string]*acp.SessionHost) s.sessionHostMu.Lock() + defer s.sessionHostMu.Unlock() + for key, host := range s.sessionHosts { if !strings.HasPrefix(key, prefix) { continue } - host.Stop() + if host != nil { + hosts[key] = host + } delete(s.sessionHosts, key) delete(s.sessionMcpServers, key) delete(s.sessionProfileOvr, key) delete(s.sessionTaskCtx, key) } - s.sessionHostMu.Unlock() - // Clean up all persisted MCP servers for this workspace (best-effort). - if s.store != nil { - if err := s.store.DeleteWorkspaceMcpServers(workspaceID); err != nil { - slog.Warn("Failed to delete workspace MCP servers from SQLite", - "workspace", workspaceID, "error", err) - } - } + return hosts } func (s *Server) requireNodeManagementAuth(w http.ResponseWriter, r *http.Request, workspaceID string) bool { diff --git a/tasks/backlog/2026-06-18-vm-agent-lifecycle-idempotent-shutdown.md b/tasks/archive/2026-06-18-vm-agent-lifecycle-idempotent-shutdown.md similarity index 85% rename from tasks/backlog/2026-06-18-vm-agent-lifecycle-idempotent-shutdown.md rename to tasks/archive/2026-06-18-vm-agent-lifecycle-idempotent-shutdown.md index aaa5d7ca2..465c3c1b2 100644 --- a/tasks/backlog/2026-06-18-vm-agent-lifecycle-idempotent-shutdown.md +++ b/tasks/archive/2026-06-18-vm-agent-lifecycle-idempotent-shutdown.md @@ -50,18 +50,18 @@ The idea and audit context identify these concrete risks: ## Implementation Checklist -- [ ] Make `Server.Stop` idempotent and return the first shutdown result on repeated calls. -- [ ] Ensure `Server.Stop` closes/stops server-owned event store and resource monitor. -- [ ] Stop ACP session hosts with collect-then-stop semantics outside `sessionHostMu`. -- [ ] Remove PTY session cleanup from broad `workspaceMu` critical sections where feasible. -- [ ] Make `errorreport.Reporter.Shutdown` idempotent and safe under concurrent callers. -- [ ] Make `messagereport.Reporter.Shutdown` idempotent and safe under concurrent callers. -- [ ] Add focused tests for `Server.Stop` idempotency and owned resource shutdown. -- [ ] Add focused tests for session-host shutdown lock discipline. -- [ ] Add reporter idempotency/concurrency tests. -- [ ] Run focused Go tests from `packages/vm-agent`. -- [ ] Run `go test ./...` from `packages/vm-agent` if environment supports it. -- [ ] Run race tests for touched packages if feasible. +- [x] Make `Server.Stop` idempotent and return the first shutdown result on repeated calls. +- [x] Ensure `Server.Stop` closes/stops server-owned event store and resource monitor. +- [x] Stop ACP session hosts with collect-then-stop semantics outside `sessionHostMu`. +- [x] Remove PTY session cleanup from broad `workspaceMu` critical sections where feasible. +- [x] Make `errorreport.Reporter.Shutdown` idempotent and safe under concurrent callers. +- [x] Make `messagereport.Reporter.Shutdown` idempotent and safe under concurrent callers. +- [x] Add focused tests for `Server.Stop` idempotency and owned resource shutdown. +- [x] Add focused tests for session-host shutdown lock discipline. +- [x] Add reporter idempotency/concurrency tests. +- [x] Run focused Go tests from `packages/vm-agent`. +- [x] Run `go test ./...` from `packages/vm-agent` if environment supports it. +- [x] Run race tests for touched packages if feasible. ## Acceptance Criteria