DO NOT MERGE / DO NOT DEPLOY: Refactor vm-agent lifecycle shutdown ownership

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,15 +97,23 @@ 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.
 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
 }
 

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
 

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
 }

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
 }
 

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)

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
 }

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")
+	}
+}