fix(telemetry): drain remaining BBolt writers before shutdown-marker resolve (Spec 080, review round 4)

Two writers could still hit BBolt after Close() resolved the clean-shutdown
marker (FR-010: the marker resolve must be the final DB write of a graceful
shutdown) or even after the DB handle closed:

1. Error-code notifier (runtime.go SetTelemetry wiring): review round 1 made
   the notifier synchronous for FR-012 but kept the 24h diagnostics counter
   write in an untracked goroutine — no better tracked than the pre-branch
   `go notifier(code)` in the supervisor, and now able to race the marker
   resolve. The counter write is now synchronous too: it completes inside the
   supervisor's call stack, so Supervisor.Stop() (which joins its goroutines
   before Close touches storage) is a hard barrier. Safe: sub-ms BBolt
   Update, the callback never re-enters the supervisor (documented at
   notifyErrorCode), and RecordLastErrorCode already writes synchronously
   under identical locking. SetErrorCodeNotifier's contract now forbids
   callback goroutines that outlive the call. The untracked-goroutine
   exposure itself pre-dates the branch; the marker invariant it violates is
   new here, and the specific goroutine was introduced in round 1 (1ec8109).

2. ActivityService (activity records, retention pruning, usage-snapshot
   flushes, async sensitive-data detection — all BBolt writes): Close() only
   context-cancelled it and never awaited it, so the flush-on-shutdown
   (persistUsage) and any in-flight worker could land after the marker
   resolve or be lost at DB close. Pre-existing on main (Stop() existed but
   was never called from Close); the Spec 080 marker invariant makes it
   observable, so it is fixed on this branch. ActivityService.Stop() is now
   idempotent, returns immediately if Start never ran, and waits for the
   event loop's final flush plus a WaitGroup over the retention loop, the
   usage-flush loop, and per-event detection goroutines. Close() calls it
   after appCancel (so the final flush is captured) and BEFORE
   StopAsync/ResolveCleanShutdown/storage Close.

Tests: TestRuntimeCloseWaitsForActivityWritersBeforeMarkerResolve proves the
shutdown flush lands before the DB closes AND the marker still resolves to
clean; TestActivityServiceStopSafeWhenNeverStartedAndIdempotent covers the
never-started and repeated-Stop paths (every existing prechurn test now
exercises the never-started branch through Close).

Author: Dumbris

internal/runtime/activity_service.go

@@ -3,6 +3,7 @@ package runtime
 import (
 	"context"
 	"encoding/json"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -39,8 +40,18 @@ type ActivityService struct {
 
 	// Channel for receiving events
 	eventCh chan Event
-	// Done channel for graceful shutdown
-	done chan struct{}
+
+	// Shutdown coordination (Spec 080 FR-010): Runtime.Close must be able to
+	// await every BBolt writer this service owns BEFORE the clean-shutdown
+	// marker resolves and the DB closes. done signals the main event loop's
+	// exit (the final flush-on-shutdown included); workersWG tracks the
+	// background loops (retention, usage flush) and the per-event async
+	// detection goroutines, all of which write to BBolt. startMu/started make
+	// Stop return immediately when Start never ran (done would never close).
+	done      chan struct{}
+	workersWG sync.WaitGroup
+	startMu   sync.Mutex
+	started   bool
 
 	// Retention configuration
 	maxAge        time.Duration
@@ -116,17 +127,37 @@ func (s *ActivityService) SetRetentionConfig(maxAge time.Duration, maxRecords in
 // Start begins listening for activity events and persisting them.
 // It should be called as a goroutine: go svc.Start(ctx, runtime)
 func (s *ActivityService) Start(ctx context.Context, rt *Runtime) {
+	// Mark started so Stop knows the done channel WILL close; refuse a second
+	// Start (the done/WaitGroup bookkeeping is single-shot).
+	s.startMu.Lock()
+	if s.started {
+		s.startMu.Unlock()
+		s.logger.Warn("Activity service Start called twice; ignoring")
+		return
+	}
+	s.started = true
+	s.startMu.Unlock()
+
 	// Subscribe to runtime events
 	eventCh := rt.SubscribeEvents()
 	defer rt.UnsubscribeEvents(eventCh)
 
-	// Start retention loop in a separate goroutine
-	go s.runRetentionLoop(ctx)
+	// Start retention loop in a separate goroutine. Tracked in workersWG: it
+	// prunes activity records (BBolt writes), so Stop must await it.
+	s.workersWG.Add(1)
+	go func() {
+		defer s.workersWG.Done()
+		s.runRetentionLoop(ctx)
+	}()
 
 	// Spec 069 A2: load/rebuild the usage aggregate before processing events,
-	// then start the periodic snapshot flush loop.
+	// then start the periodic snapshot flush loop (tracked: it writes BBolt).
 	s.initUsageFromStorage()
-	go s.runUsageFlushLoop(ctx)
+	s.workersWG.Add(1)
+	go func() {
+		defer s.workersWG.Done()
+		s.runUsageFlushLoop(ctx)
+	}()
 
 	s.logger.Info("Activity service started")
 
@@ -213,9 +244,30 @@ func (s *ActivityService) runRetentionCleanup() {
 	}
 }
 
-// Stop gracefully shuts down the activity service.
+// Stop gracefully shuts down the activity service, waiting for every BBolt
+// writer it owns to finish: the main event loop (including its final
+// flush-on-shutdown of the usage snapshot), the retention and usage-flush
+// loops, and any in-flight async detection goroutines (Spec 080 FR-010: no
+// activity write may land after Runtime.Close resolves the clean-shutdown
+// marker or closes the DB).
+//
+// Callers must cancel the context passed to Start FIRST — the final usage
+// flush runs on ctx.Done inside the event loop, and Stop waits for it, so the
+// flush is captured before the marker resolves. Idempotent, and returns
+// immediately when Start never ran.
 func (s *ActivityService) Stop() {
+	s.startMu.Lock()
+	started := s.started
+	s.startMu.Unlock()
+	if !started {
+		return
+	}
+	// Main loop exit (closes done AFTER the shutdown flush). All workersWG.Add
+	// calls happen before done closes — the loop goroutines are registered at
+	// the top of Start and detection goroutines are only spawned from the event
+	// loop — so Wait below cannot race an Add.
 	<-s.done
+	s.workersWG.Wait()
 }
 
 // handleEvent processes an activity event and persists it to storage.
@@ -355,9 +407,14 @@ func (s *ActivityService) handleToolCallCompleted(evt Event) {
 			s.usage.Apply(record)
 		}
 
-		// Run async sensitive data detection (Spec 026)
+		// Run async sensitive data detection (Spec 026). Tracked in workersWG:
+		// it updates the record's metadata in BBolt, so Stop must await it.
 		if s.detector != nil {
-			go s.runAsyncDetection(record.ID, arguments, response)
+			s.workersWG.Add(1)
+			go func() {
+				defer s.workersWG.Done()
+				s.runAsyncDetection(record.ID, arguments, response)
+			}()
 		}
 	}
 }

internal/runtime/prechurn_lifecycle_test.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"go.etcd.io/bbolt"
 	berrors "go.etcd.io/bbolt/errors"
@@ -139,6 +140,107 @@ func TestRuntimeCloseCleanupBranchStillResolvesAndClosesStorage(t *testing.T) {
 	}
 }
 
+// TestRuntimeCloseWaitsForActivityWritersBeforeMarkerResolve guards the Close
+// ordering added in review round 4 (Spec 080 FR-010): the ActivityService owns
+// BBolt writers (activity records, usage-snapshot flushes, retention pruning,
+// async detection), and before this fix Close only context-cancelled it —
+// its final flush-on-shutdown (persistUsage) raced the shutdown-marker resolve
+// and the DB close, so the flush could be lost or land AFTER the marker
+// claimed "no writes remain". Now Close awaits ActivityService.Stop() before
+// StopAsync/ResolveCleanShutdown/db.Close, so the flush must both survive and
+// precede the clean marker.
+func TestRuntimeCloseWaitsForActivityWritersBeforeMarkerResolve(t *testing.T) {
+	t.Setenv("MCPPROXY_LAUNCHED_BY", "")
+	dataDir := t.TempDir()
+
+	rt, _ := previousShutdownVia(t, dataDir)
+
+	// Start the activity service exactly as StartBackgroundInitialization does.
+	rt.ActivityService().SetEventEmitter(rt)
+	go rt.ActivityService().Start(rt.AppContext(), rt)
+
+	// Emit a completed tool call and wait until the event loop has persisted it
+	// and folded it into the in-memory usage aggregate (Apply runs only after a
+	// successful SaveActivity). Re-emit each attempt: an event published before
+	// Start's SubscribeEvents registers is silently dropped, so a single early
+	// emit could race the service's startup.
+	deadline := time.Now().Add(5 * time.Second)
+	for {
+		rt.EmitActivityToolCallCompleted(
+			"prechurn-srv", "prechurn-tool", "sess-1", "req-1", "mcp",
+			"success", "", 7, nil, "ok", false, "", nil, "", "", 0, 0)
+		if snap := rt.ActivityService().UsageSnapshot(); snap != nil && len(snap.Tools) > 0 {
+			break
+		}
+		if time.Now().After(deadline) {
+			t.Fatal("timed out waiting for the activity event to reach the usage aggregate")
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+
+	// The default flush cadence (30s) cannot fire within this test, so the ONLY
+	// path that persists the aggregate is the flush-on-shutdown — which Close
+	// must wait for before resolving the marker and closing the DB.
+	if err := rt.Close(); err != nil {
+		t.Fatalf("close: %v", err)
+	}
+
+	// Reopen: the flush landed AND the marker resolved to clean afterwards.
+	rt2, prev := previousShutdownVia(t, dataDir)
+	defer func() { _ = rt2.Close() }()
+	if prev != "clean" {
+		t.Fatalf("after graceful close: expected clean, got %q", prev)
+	}
+	data, err := rt2.StorageManager().LoadUsageSnapshot()
+	if err != nil {
+		t.Fatalf("load usage snapshot: %v", err)
+	}
+	if len(data) == 0 {
+		t.Fatal("usage snapshot missing: shutdown flush did not land before the DB closed")
+	}
+	agg, err := decodeUsageAggregate(data)
+	if err != nil {
+		t.Fatalf("decode usage snapshot: %v", err)
+	}
+	if _, ok := agg.Tools[toolKey("prechurn-srv", "prechurn-tool")]; !ok {
+		t.Fatalf("persisted usage snapshot lacks the recorded tool call; tools = %d", len(agg.Tools))
+	}
+}
+
+// TestActivityServiceStopSafeWhenNeverStartedAndIdempotent: Close() now calls
+// ActivityService.Stop() unconditionally, including on runtimes whose
+// background initialization never ran (every other test in this file) and on
+// double Close. Stop must return immediately when Start never ran, and be
+// idempotent after a normal shutdown.
+func TestActivityServiceStopSafeWhenNeverStartedAndIdempotent(t *testing.T) {
+	t.Setenv("MCPPROXY_LAUNCHED_BY", "")
+	dataDir := t.TempDir()
+
+	rt, _ := previousShutdownVia(t, dataDir)
+
+	// Never started: must not block.
+	doneCh := make(chan struct{})
+	go func() {
+		rt.ActivityService().Stop()
+		close(doneCh)
+	}()
+	select {
+	case <-doneCh:
+	case <-time.After(2 * time.Second):
+		t.Fatal("ActivityService.Stop blocked although Start never ran")
+	}
+
+	// Started, then Close (which stops it), then Stop again: idempotent.
+	// (Runtime.Close itself is not re-runnable — cache.Manager.Close panics on
+	// a second call, a pre-existing constraint — so only Stop is re-invoked.)
+	go rt.ActivityService().Start(rt.AppContext(), rt)
+	if err := rt.Close(); err != nil {
+		t.Fatalf("close: %v", err)
+	}
+	rt.ActivityService().Stop() // second stop must not hang or panic
+	rt.ActivityService().Stop() // and a third
+}
+
 // TestRuntimeCloseAfterExternalStopAsyncStillResolvesClean guards the split
 // Close sequence (Spec 080 FR-010, review round 3): Close now runs
 // storage.StopAsync (stop + drain queued async DB ops) BEFORE resolving the

internal/runtime/runtime.go

@@ -653,6 +653,20 @@ func (r *Runtime) Close() error {
 		}
 	}
 
+	// Spec 080 (FR-010, review round 4): the ActivityService owns BBolt
+	// writers — activity records, retention pruning, usage-snapshot flushes,
+	// async sensitive-data detection. The appCancel at the top of Close
+	// triggered its flush-on-shutdown; await that flush AND all its worker
+	// goroutines here, BEFORE the async-op drain and the shutdown-marker
+	// resolve below, so no activity write can land after the marker claims
+	// the shutdown was clean (or after the DB closes). This writer-vs-close
+	// race pre-dates Spec 080 (Stop existed but was never called); the marker
+	// invariant makes it observable, so it is closed here. Stop returns
+	// immediately when Start never ran and is itself idempotent.
+	if r.activityService != nil {
+		r.activityService.Stop()
+	}
+
 	// Spec 080 (US3, FR-010): resolve the shutdown marker to "clean" at the
 	// LAST point the DB is still open — i.e. after the async storage manager
 	// has stopped AND drained its queue (those queued operations perform BBolt
@@ -2524,9 +2538,18 @@ func (r *Runtime) SetTelemetry(version, edition string) {
 					if prechurnStore != nil {
 						_ = prechurnStore.RecordLastErrorCode(db, code)
 					}
-					// The 24h aggregate counter keeps its pre-existing async
-					// posture: loss-tolerant, off the supervisor's path.
-					go func() { _ = diagStore.RecordErrorCode(db, code) }()
+					// The 24h aggregate counter write is synchronous too
+					// (review round 4): an untracked goroutine here could run
+					// after Close resolves the shutdown marker (FR-010) or
+					// after the DB handle closes. Synchronous means the write
+					// completes inside the supervisor's call stack, so
+					// supervisor.Stop() — which joins its goroutines before
+					// Close touches storage — is a hard barrier: after it
+					// returns, no notifier-driven DB write remains. Same
+					// safety argument as above: sub-ms BBolt Update, no
+					// supervisor re-entry, so no lock cycle even when the
+					// caller holds stateMu.
+					_ = diagStore.RecordErrorCode(db, code)
 				})
 			}
 

internal/runtime/supervisor/supervisor.go

@@ -810,9 +810,12 @@ func (s *Supervisor) SetOnServerConnectedCallback(callback func(serverName strin
 // callback receives the stable MCPX_* code string and is invoked
 // SYNCHRONOUSLY at the classification site (Spec 080 FR-012: the pre-churn
 // last_error_code must be durable before a crash can follow the
-// classification). The callback must be fast (sub-ms) and must not call back
-// into the Supervisor; anything slow or loss-tolerant (e.g. aggregate
-// counters) should offload to a goroutine inside the callback itself.
+// classification). The callback must be fast (sub-ms), must not call back
+// into the Supervisor, and must NOT spawn goroutines that outlive the call:
+// Runtime.Close relies on Supervisor.Stop() as a barrier — once Stop returns,
+// no callback-driven DB write may remain in flight, or it could land after
+// the clean-shutdown marker resolves or after the DB closes (Spec 080
+// FR-010).
 func (s *Supervisor) SetErrorCodeNotifier(fn func(code string)) {
 	s.callbackMu.Lock()
 	defer s.callbackMu.Unlock()