diff --git a/docs/features/telemetry.md b/docs/features/telemetry.md index 6e8d1a2c9..e0c437820 100644 --- a/docs/features/telemetry.md +++ b/docs/features/telemetry.md @@ -49,6 +49,16 @@ The `docker_cli_source` field is likewise a **fixed enum** (`path`, `bundled`, ` Docker isolation failures surface in `error_code_counts_24h` via three stable diagnostic codes (schema v5): `MCPX_DOCKER_CLI_NOT_FOUND` (isolation requested but the `docker` binary is unresolved — issue #696), `MCPX_DOCKER_EXEC_NOT_FOUND` (the image lacks the interpreter the server needs, e.g. `uvx` missing in `python:3.11`), and `MCPX_DOCKER_OCI_RUNTIME` (OCI runtime / architecture-mismatch failures). +## When heartbeats are sent + +A heartbeat goes out 5 minutes after start (so short-lived processes stay quiet), then once every 24 hours, and **once more on graceful shutdown**. + +The shutdown heartbeat exists because the windowed counters live in memory and are cleared only after the endpoint accepts a send. Without it, everything an install did after its 5-minute heartbeat was discarded whenever the process exited before the 24-hour tick — which is most desktop sessions. It is skipped when telemetry is disabled or opted out, and when nothing has been recorded since the last accepted heartbeat. It is hard-bounded to a few seconds, so an unreachable endpoint can never delay quitting, and because counters are cleared only on an accepted send, a failed shutdown flush simply leaves the counts for the next run rather than dropping them. + +Delivery is therefore **at-least-once**, not exactly-once. If a send is cut off after the endpoint accepted it but before mcpproxy sees the response, the counts stay pending locally and are re-sent — a lost response is indistinguishable from a lost request. Receivers should treat a heartbeat as idempotent per `(anonymous_id, timestamp)`. Retrying is the deliberate choice: the alternative is silently discarding the window, which is the problem the shutdown heartbeat was added to solve. + +Two housekeeping steps are deliberately *not* performed by the shutdown heartbeat — the 365-day anonymous-ID rotation and the one-shot `installer` launch-source attribution. Both are consumed when the payload is built rather than when it is accepted, so running them for a send that may never land would spend state a normal run still has. The next start performs both. + ## What is NOT collected The following is **never** collected: diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index e48f196da..da1f60614 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -776,6 +776,13 @@ func (r *Runtime) Close() error { // ActivityService.Stop above — terminally stops the service so a Start // goroutine not yet scheduled (lifecycle.go launches it via `go`) becomes // a no-op instead of writing after this point. + // + // Stop also performs the graceful-shutdown heartbeat flush — one bounded + // (4s) final send of the counters recorded since the last accepted + // heartbeat, which would otherwise die with the process. It runs after the + // loop is joined, so it is the only sender at that moment, and it happens + // HERE — while the BBolt handle is still open and before the marker + // resolves below — because its buildHeartbeat writes funnel activity. if r.telemetryService != nil { r.telemetryService.Stop() } diff --git a/internal/telemetry/activation_test.go b/internal/telemetry/activation_test.go index 9db5fd8de..35400dde7 100644 --- a/internal/telemetry/activation_test.go +++ b/internal/telemetry/activation_test.go @@ -337,14 +337,28 @@ func TestInstallerPending_ClearedAfterFirstHeartbeat(t *testing.T) { activationDB: db, } - // First call → "installer" and the pending flag must be cleared. - ls1 := s.resolveLaunchSource() + // The shutdown flush must NOT consume the one-shot: it builds a payload that + // may never be accepted, so a non-consuming resolve has to leave the flag + // armed for the next run and report the runtime-detected source instead. + if ls := s.resolveLaunchSource(false); ls == LaunchSourceInstaller { + t.Fatal("non-consuming resolveLaunchSource claimed installer; the shutdown flush must not spend the one-shot") + } + pendingAfterFlush, err := store.IsInstallerPending(db) + if err != nil { + t.Fatalf("IsInstallerPending after non-consuming resolve: %v", err) + } + if !pendingAfterFlush { + t.Fatal("non-consuming resolveLaunchSource cleared installer_heartbeat_pending; a failed shutdown flush would destroy the attribution") + } + + // First consuming call → "installer" and the pending flag must be cleared. + ls1 := s.resolveLaunchSource(true) if ls1 != LaunchSourceInstaller { t.Fatalf("first resolveLaunchSource = %q, want %q", ls1, LaunchSourceInstaller) } pending, err := store.IsInstallerPending(db) if err != nil { - t.Fatalf("IsInstallerPending after first call: %v", err) + t.Fatalf("IsInstallerPending after first consuming call: %v", err) } if pending { t.Fatalf("installer_heartbeat_pending should be cleared after first heartbeat") @@ -352,7 +366,7 @@ func TestInstallerPending_ClearedAfterFirstHeartbeat(t *testing.T) { // Second call → no longer installer; should return the runtime // detector's value (whatever it is in the test env), not installer. - ls2 := s.resolveLaunchSource() + ls2 := s.resolveLaunchSource(true) if ls2 == LaunchSourceInstaller { t.Fatalf("second resolveLaunchSource still returned installer — one-shot broken") } diff --git a/internal/telemetry/feedback.go b/internal/telemetry/feedback.go index 8f273e66e..8fc7c0116 100644 --- a/internal/telemetry/feedback.go +++ b/internal/telemetry/feedback.go @@ -133,7 +133,7 @@ func (s *Service) SubmitFeedback(ctx context.Context, req *FeedbackRequest) (*Fe return nil, fmt.Errorf("failed to marshal feedback: %w", err) } - url := s.endpoint + "/feedback" + url := s.liveEndpoint() + "/feedback" httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(data)) if err != nil { return nil, fmt.Errorf("failed to create feedback request: %w", err) diff --git a/internal/telemetry/optout.go b/internal/telemetry/optout.go index 679625f8e..aa5cf011d 100644 --- a/internal/telemetry/optout.go +++ b/internal/telemetry/optout.go @@ -79,7 +79,7 @@ func (s *Service) EmitOptOutBeacon(ctx context.Context) bool { if disabled, _ := IsDisabledByEnv(); disabled { return false } - if s.config.GetAnonymousID() == "" { + if s.liveAnonymousID() == "" { return false } if err := s.SendOptOutBeacon(ctx); err != nil { @@ -98,7 +98,7 @@ func (s *Service) EmitOptOutBeacon(ctx context.Context) bool { // caller-supplied URL, so this never sends to an arbitrary, request-derived // host. func (s *Service) SendOptOutBeacon(ctx context.Context) error { - anonID := s.config.GetAnonymousID() + anonID := s.liveAnonymousID() if anonID == "" { // Nothing to dedup on — never send an identity-less beacon. return errors.New("opt-out beacon skipped: no anonymous_id") @@ -122,7 +122,7 @@ func (s *Service) SendOptOutBeacon(ctx context.Context) error { // require a host, and issue the request against the re-serialized URL. This // mirrors validateRegistryURL and guarantees a malformed/non-http endpoint // can never aim the beacon at, e.g., file:// or a schemeless host. - beaconURL, err := validateTelemetryURL(strings.TrimRight(s.endpoint, "/") + "/heartbeat") + beaconURL, err := validateTelemetryURL(strings.TrimRight(s.liveEndpoint(), "/") + "/heartbeat") if err != nil { return err } diff --git a/internal/telemetry/registry.go b/internal/telemetry/registry.go index 2f799ecb6..9ddc32476 100644 --- a/internal/telemetry/registry.go +++ b/internal/telemetry/registry.go @@ -490,6 +490,50 @@ func (r *CounterRegistry) Snapshot() RegistrySnapshot { return snap } +// HasPendingCounters reports whether anything has been recorded since the last +// Reset — i.e. whether a heartbeat would carry usage data that has not yet been +// accepted by the endpoint. +// +// The counter set inspected here is EXACTLY the set Reset zeroes: those are the +// counters whose window is "since the last accepted send", so they are the ones +// that would be lost if the process exited without transmitting. Lifetime +// tallies that Reset deliberately preserves (anonymityViolations) are excluded — +// they are never transmitted and must not, on their own, justify a send. +// +// Used by the graceful-shutdown flush to skip a final heartbeat that would carry +// no new information. +func (r *CounterRegistry) HasPendingCounters() bool { + if r == nil { + return false + } + for i := range r.surfaceCounts { + if r.surfaceCounts[i].Load() > 0 { + return true + } + } + if r.upstreamTotal.Load() > 0 { + return true + } + + r.mu.RLock() + defer r.mu.RUnlock() + + if len(r.builtinCalls) > 0 || len(r.restEndpoints) > 0 || + len(r.errorCategories) > 0 || len(r.doctorChecks) > 0 { + return true + } + if r.tpaScansCompleted > 0 || r.tpaScansFailed > 0 || r.tpaScansWithFindings > 0 || + r.tpaToolChangeGateScans > 0 || r.tpaPromptScans > 0 { + return true + } + for _, v := range r.tpaFindings { + if v > 0 { + return true + } + } + return false +} + // Reset zeros all counters. Called only after a successful heartbeat send. // // Deliberate, registry-wide trade-off: an event recorded between Snapshot() diff --git a/internal/telemetry/shutdown_flush_test.go b/internal/telemetry/shutdown_flush_test.go new file mode 100644 index 000000000..e2a095bea --- /dev/null +++ b/internal/telemetry/shutdown_flush_test.go @@ -0,0 +1,539 @@ +package telemetry + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "go.uber.org/zap" + + "github.com/smart-mcp-proxy/mcpproxy-go/internal/config" +) + +// Graceful-shutdown flush: telemetry counters live in memory and are reset only +// after an ACCEPTED (2xx) send, so before this everything recorded after the +// 5-minute first heartbeat was dropped unless the process survived to the 24h +// tick. Stop now performs ONE bounded final send after joining the loop. +// +// The invariants these tests pin down: +// - flush sends exactly once and resets the counters (2xx), +// - a dead endpoint cannot hang shutdown, and its counters are RETAINED, +// - opted-out / disabled installs send nothing at all, +// - an ACCEPTED periodic tick leaves the flush nothing to send, and an +// unconfirmed one is joined and re-sent (delivery is at-least-once). + +// heartbeatRecorder is a test endpoint that captures every usage heartbeat. +// +// It must distinguish heartbeats from opt-out beacons: the beacon deliberately +// reuses the SAME /heartbeat ingest path (MCP-2482) and is told apart only by +// its `event` field. Counting it as a heartbeat would make the opt-out test +// assert on whichever of two goroutines happened to win. +type heartbeatRecorder struct { + mu sync.Mutex + payloads []HeartbeatPayload + beacons []OptOutBeacon + status int // response status; 0 means 200 +} + +func (h *heartbeatRecorder) ServeHTTP(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + + var beacon OptOutBeacon + isBeacon := json.Unmarshal(body, &beacon) == nil && beacon.Event == OptOutEvent + + var p HeartbeatPayload + if !isBeacon { + _ = json.Unmarshal(body, &p) + } + + h.mu.Lock() + if isBeacon { + h.beacons = append(h.beacons, beacon) + } else { + h.payloads = append(h.payloads, p) + } + status := h.status + h.mu.Unlock() + + if status == 0 { + status = http.StatusOK + } + w.WriteHeader(status) +} + +func (h *heartbeatRecorder) count() int { + h.mu.Lock() + defer h.mu.Unlock() + return len(h.payloads) +} + +func (h *heartbeatRecorder) beaconCount() int { + h.mu.Lock() + defer h.mu.Unlock() + return len(h.beacons) +} + +// mcpTotal sums surface_requests.mcp across every received payload. Because +// counters reset only on an accepted send, this total must equal the number of +// MCP events recorded — anything higher proves a double-send. +func (h *heartbeatRecorder) mcpTotal() int64 { return h.surfaceTotal("mcp") } + +func (h *heartbeatRecorder) surfaceTotal(key string) int64 { + h.mu.Lock() + defer h.mu.Unlock() + var total int64 + for i := range h.payloads { + total += h.payloads[i].SurfaceRequests[key] + } + return total +} + +// newFlushTestService builds an enabled service whose heartbeat loop parks in +// the initial delay (no periodic send) unless the caller shortens it, with a +// snappy shutdown-flush timeout. +func newFlushTestService(t *testing.T, endpoint string) *Service { + t.Helper() + cfg := &config.Config{ + Telemetry: &config.TelemetryConfig{ + AnonymousID: "test-uuid-flush", + Endpoint: endpoint, + }, + RoutingMode: "retrieve_tools", + } + svc := New(cfg, "", "v1.0.0", "personal", zap.NewNop()) + svc.initialDelay = time.Hour // no periodic send unless overridden + svc.heartbeatInterval = time.Hour + svc.flushTimeout = 2 * time.Second + svc.SetRuntimeStats(&mockRuntimeStats{}) + return svc +} + +// startAndArm launches the heartbeat loop and waits until it has passed every +// emission gate (the point at which the shutdown flush is armed). +func startAndArm(t *testing.T, svc *Service, ctx context.Context) { + t.Helper() + go svc.Start(ctx) + deadline := time.Now().Add(5 * time.Second) + for !svc.flushEligible.Load() { + if time.Now().After(deadline) { + t.Fatal("telemetry Start never armed the shutdown flush") + } + time.Sleep(time.Millisecond) + } +} + +// TestFlushOnStopSendsFinalHeartbeatAndResetsCounters: activity recorded before +// any heartbeat went out must be transmitted exactly once on graceful shutdown, +// and the 2xx must reset the counters through the normal send+reset path. +func TestFlushOnStopSendsFinalHeartbeatAndResetsCounters(t *testing.T) { + clearTelemetryEnv(t) + + rec := &heartbeatRecorder{} + server := httptest.NewServer(rec) + defer server.Close() + + svc := newFlushTestService(t, server.URL) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + startAndArm(t, svc, ctx) + + // Usage the process would otherwise take to the grave. + svc.Registry().RecordSurface(SurfaceMCP) + svc.Registry().RecordUpstreamTool() + + cancel() + svc.Stop() + + if got := rec.count(); got != 1 { + t.Fatalf("heartbeats received: got %d, want 1 (the shutdown flush)", got) + } + if got := rec.mcpTotal(); got != 1 { + t.Fatalf("surface_requests.mcp transmitted: got %d, want 1", got) + } + if svc.Registry().HasPendingCounters() { + t.Fatal("counters were not reset after the accepted shutdown flush") + } + + // Stop is idempotent and the flush is single-shot. + svc.Stop() + svc.Stop() + if got := rec.count(); got != 1 { + t.Fatalf("repeated Stop re-sent the flush: got %d heartbeats, want 1", got) + } +} + +// TestFlushOnStopTimeoutDoesNotBlockShutdown: a blackholed endpoint must not +// turn shutdown into a hang, and because no 2xx arrived the counters must be +// RETAINED for the next run rather than silently dropped. +func TestFlushOnStopTimeoutDoesNotBlockShutdown(t *testing.T) { + clearTelemetryEnv(t) + + release := make(chan struct{}) + var releaseOnce sync.Once + unblock := func() { releaseOnce.Do(func() { close(release) }) } + defer unblock() + + var received atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + received.Add(1) + select { + case <-release: + case <-r.Context().Done(): + } + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + svc := newFlushTestService(t, server.URL) + svc.flushTimeout = 100 * time.Millisecond + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + startAndArm(t, svc, ctx) + svc.Registry().RecordSurface(SurfaceMCP) + + cancel() + stopReturned := make(chan struct{}) + start := time.Now() + go func() { + svc.Stop() + close(stopReturned) + }() + select { + case <-stopReturned: + case <-time.After(5 * time.Second): + t.Fatal("Stop hung on an unresponsive telemetry endpoint") + } + if elapsed := time.Since(start); elapsed > 2*time.Second { + t.Fatalf("shutdown flush was not hard-bounded: took %v", elapsed) + } + + if received.Load() == 0 { + t.Fatal("shutdown flush never attempted a send") + } + if !svc.Registry().HasPendingCounters() { + t.Fatal("counters were dropped by a flush that was never accepted") + } + + unblock() +} + +// TestFlushOnStopSkippedWhenOptedOut: once the opt-out latch is set, shutdown +// must not become a loophole that ships one last usage payload. +func TestFlushOnStopSkippedWhenOptedOut(t *testing.T) { + clearTelemetryEnv(t) + + rec := &heartbeatRecorder{} + server := httptest.NewServer(rec) + defer server.Close() + + svc := newFlushTestService(t, server.URL) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + startAndArm(t, svc, ctx) + svc.Registry().RecordSurface(SurfaceMCP) + + // The user turns telemetry off mid-run: NotifyConfigChanged latches + // optedOut and swaps in the disabled config. + disabled := false + svc.NotifyConfigChanged(&config.Config{ + Telemetry: &config.TelemetryConfig{ + AnonymousID: "test-uuid-flush", + Endpoint: server.URL, + Enabled: &disabled, + }, + RoutingMode: "retrieve_tools", + }) + + cancel() + svc.Stop() + + // NotifyConfigChanged fires the opt-out beacon from its own goroutine, and + // that beacon POSTs to the SAME /heartbeat ingest path (MCP-2482) — so + // asserting straight after Stop would race it. Wait for the beacon to land + // first; only then is "no usage heartbeat arrived" a settled fact rather + // than a statement about which goroutine was faster. + deadline := time.Now().Add(5 * time.Second) + for rec.beaconCount() == 0 { + if time.Now().After(deadline) { + t.Fatal("opt-out beacon never reached the endpoint") + } + time.Sleep(time.Millisecond) + } + + // The beacon carries ONLY the anonymous id + event marker. No usage payload + // may ever follow an opt-out, shutdown flush included. + if got := rec.count(); got != 0 { + t.Fatalf("usage heartbeats sent after opt-out: got %d, want 0", got) + } +} + +// TestFlushOnStopSkippedWhenTelemetryDisabled: a service whose loop never +// started (telemetry disabled by config) has no armed flush, so Stop sends +// nothing even though counters were recorded. +func TestFlushOnStopSkippedWhenTelemetryDisabled(t *testing.T) { + clearTelemetryEnv(t) + + rec := &heartbeatRecorder{} + server := httptest.NewServer(rec) + defer server.Close() + + disabled := false + cfg := &config.Config{ + Telemetry: &config.TelemetryConfig{ + AnonymousID: "test-uuid-flush", + Endpoint: server.URL, + Enabled: &disabled, + }, + RoutingMode: "retrieve_tools", + } + svc := New(cfg, "", "v1.0.0", "personal", zap.NewNop()) + svc.flushTimeout = time.Second + svc.SetRuntimeStats(&mockRuntimeStats{}) + svc.Registry().RecordSurface(SurfaceMCP) + + svc.Start(context.Background()) // returns immediately: disabled by config + svc.Stop() + + if got := rec.count(); got != 0 { + t.Fatalf("heartbeats sent by a disabled install: got %d, want 0", got) + } + if svc.flushEligible.Load() { + t.Fatal("shutdown flush was armed although telemetry is disabled") + } +} + +// TestFlushOnStopSkippedAfterAnAcceptedTick: once a periodic tick has been +// ACCEPTED (2xx ⇒ Reset), the shutdown flush must find nothing left to report +// and send nothing. This is the no-double-send guarantee the flush actually +// makes — it is a property of the counter contract, not of send timing. +func TestFlushOnStopSkippedAfterAnAcceptedTick(t *testing.T) { + clearTelemetryEnv(t) + + rec := &heartbeatRecorder{} + server := httptest.NewServer(rec) + defer server.Close() + + svc := newFlushTestService(t, server.URL) + svc.initialDelay = 5 * time.Millisecond // the first heartbeat fires promptly + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + svc.Registry().RecordSurface(SurfaceMCP) + svc.Registry().RecordSurface(SurfaceMCP) + go svc.Start(ctx) + + // Wait for the tick AND its post-2xx reset — that reset is what makes the + // rest of this test deterministic. + deadline := time.Now().Add(5 * time.Second) + for rec.count() == 0 || svc.Registry().HasPendingCounters() { + if time.Now().After(deadline) { + t.Fatal("first heartbeat never arrived (or never reset the registry)") + } + time.Sleep(time.Millisecond) + } + + cancel() + svc.Stop() + + if got := rec.count(); got != 1 { + t.Fatalf("heartbeats received: got %d, want 1 (the flush must skip after an accepted tick)", got) + } + if got := rec.mcpTotal(); got != 2 { + t.Fatalf("surface_requests.mcp transmitted: got %d, want 2", got) + } +} + +// TestFlushOnStopJoinsAndResendsAnAbortedTick pins the two things that DO hold +// when Stop lands on a send that is still in flight: +// +// - Stop joins it (the loop is never left running behind shutdown), and +// - delivery is AT-LEAST-ONCE. The aborted send never observed a 2xx, so the +// counters stay pending and the flush re-sends them. If the endpoint had in +// fact accepted that aborted request, it would see the window twice — that +// is undecidable client-side and is the deliberate trade against the +// alternative, which is losing the window entirely (the bug the flush +// exists to fix). Receivers dedup per (anonymous_id, timestamp). +// +// The blocked first request records NOTHING, standing in for exactly that case: +// a request the client could not confirm. +func TestFlushOnStopJoinsAndResendsAnAbortedTick(t *testing.T) { + clearTelemetryEnv(t) + + rec := &heartbeatRecorder{} + inFlight := make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + unblock := func() { releaseOnce.Do(func() { close(release) }) } + defer unblock() + + var seen atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if seen.Add(1) == 1 { + // The tick's send: hold it open until the test releases it, and + // never record it — its outcome is unobservable to the client. + close(inFlight) + <-release + w.WriteHeader(http.StatusOK) + return + } + rec.ServeHTTP(w, r) + })) + defer server.Close() + + svc := newFlushTestService(t, server.URL) + svc.initialDelay = 5 * time.Millisecond + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + svc.Registry().RecordSurface(SurfaceMCP) + svc.Registry().RecordSurface(SurfaceMCP) + go svc.Start(ctx) + + select { + case <-inFlight: + case <-time.After(5 * time.Second): + t.Fatal("first heartbeat never reached the endpoint") + } + + // Stop while that send is still in flight (Runtime.Close cancels first). + stopReturned := make(chan struct{}) + go func() { + cancel() + svc.Stop() + close(stopReturned) + }() + + select { + case <-stopReturned: + case <-time.After(10 * time.Second): + t.Fatal("Stop did not return while a tick was in flight") + } + unblock() + + if got := rec.count(); got != 1 { + t.Fatalf("heartbeats recorded: got %d, want 1 (the flush re-sending the unconfirmed window)", got) + } + if got := rec.mcpTotal(); got != 2 { + t.Fatalf("surface_requests.mcp re-sent by the flush: got %d, want 2", got) + } + if svc.Registry().HasPendingCounters() { + t.Fatal("counters were not reset after the accepted shutdown flush") + } +} + +// TestOptOutBeaconReadsAnonymousIDUnderLock covers the second half of the +// s.endpoint drive-by. NotifyConfigChanged replaces s.config under s.mu, and +// GetAnonymousID dereferences cfg.Telemetry — a pointer that +// ensureAnonymousIDOnce / advanceUpgradeFunnelOnce also install under s.mu. The +// opt-out beacon read both unlocked, two lines above the now-locked endpoint +// read. Fails under -race without liveAnonymousID(). +func TestOptOutBeaconReadsAnonymousIDUnderLock(t *testing.T) { + clearTelemetryEnv(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + svc := newFlushTestService(t, server.URL) + + // Every config stays ENABLED, so no enabled->disabled transition fires and + // this test drives the reader/writer pair directly rather than through the + // beacon's own goroutine. + stop := make(chan struct{}) + var reloads sync.WaitGroup + reloads.Add(1) + go func() { + defer reloads.Done() + for i := 0; ; i++ { + select { + case <-stop: + return + default: + } + svc.NotifyConfigChanged(&config.Config{ + Telemetry: &config.TelemetryConfig{ + AnonymousID: fmt.Sprintf("test-uuid-%d", i), + Endpoint: server.URL, + }, + RoutingMode: "retrieve_tools", + }) + } + }() + + ctx := context.Background() + for i := 0; i < 100; i++ { + if err := svc.SendOptOutBeacon(ctx); err != nil { + t.Fatalf("SendOptOutBeacon: %v", err) + } + if !svc.EmitOptOutBeacon(ctx) { + t.Fatal("EmitOptOutBeacon reported no send attempt despite a live anonymous id") + } + } + + close(stop) + reloads.Wait() +} + +// TestFlushOnStopCapturesActivityAfterFirstHeartbeat is the regression the +// whole change exists for: work done AFTER the 5-minute first heartbeat used to +// be lost unless the process lived to the 24h tick. +func TestFlushOnStopCapturesActivityAfterFirstHeartbeat(t *testing.T) { + clearTelemetryEnv(t) + + rec := &heartbeatRecorder{} + server := httptest.NewServer(rec) + defer server.Close() + + svc := newFlushTestService(t, server.URL) + svc.initialDelay = 5 * time.Millisecond + + // A marker recorded BEFORE the first heartbeat: once it has been reported + // and the registry reset, the initial send is provably complete, so the + // activity recorded below cannot be swallowed by that send's Reset. + svc.Registry().RecordSurface(SurfaceCLI) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go svc.Start(ctx) + + // Wait for the first (initial-delay) heartbeat AND its post-2xx reset. + deadline := time.Now().Add(5 * time.Second) + for rec.count() == 0 || svc.Registry().HasPendingCounters() { + if time.Now().After(deadline) { + t.Fatal("first heartbeat never arrived (or never reset the registry)") + } + time.Sleep(time.Millisecond) + } + + // Post-heartbeat usage — the data that used to be dropped. + svc.Registry().RecordSurface(SurfaceMCP) + svc.Registry().RecordSurface(SurfaceMCP) + svc.Registry().RecordSurface(SurfaceMCP) + + cancel() + svc.Stop() + + if got := rec.count(); got != 2 { + t.Fatalf("heartbeats received: got %d, want 2 (initial + shutdown flush)", got) + } + if got := rec.mcpTotal(); got != 3 { + t.Fatalf("surface_requests.mcp transmitted: got %d, want 3", got) + } + if got := rec.surfaceTotal("cli"); got != 1 { + t.Fatalf("surface_requests.cli transmitted: got %d, want 1 (no double-count)", got) + } + if svc.Registry().HasPendingCounters() { + t.Fatal("counters were not reset after the accepted shutdown flush") + } +} diff --git a/internal/telemetry/telemetry.go b/internal/telemetry/telemetry.go index a585b1238..dc740d49a 100644 --- a/internal/telemetry/telemetry.go +++ b/internal/telemetry/telemetry.go @@ -462,12 +462,36 @@ type Service struct { startMu sync.Mutex started bool stopped bool + + // Graceful-shutdown flush. Counters live in memory and are only reset after + // an ACCEPTED send, so everything recorded after the last accepted heartbeat + // (in practice: everything after the 5-minute first heartbeat, for any + // process that does not survive to the 24h tick) was previously dropped on + // exit. Stop now performs ONE final send after joining the loop. + // + // flushEligible is set by Start once it has passed every emission gate + // (env / config / semver) and ensured the anonymous id — i.e. exactly when a + // heartbeat from this process would have been legitimate. Stop consults it + // so a disabled or dev build never sends on shutdown. + // + // flushOnce keeps the flush single-shot across the idempotent Stop. + // flushTimeout is overridable in tests; see shutdownFlushTimeout. + flushEligible atomic.Bool + flushOnce sync.Once + flushTimeout time.Duration } // optOutBeaconTimeout bounds the best-effort opt-out beacon send so a slow or // unreachable endpoint never delays the config save that triggered it. const optOutBeaconTimeout = 5 * time.Second +// shutdownFlushTimeout hard-bounds the final heartbeat sent on graceful +// shutdown. It is deliberately much shorter than the HTTP client's own 10s +// timeout: a dead or blackholed endpoint must never turn a user's quit into a +// visible hang. The send is best-effort — on timeout the counters are simply +// retained (no 2xx ⇒ no Reset) and the next run reports them. +const shutdownFlushTimeout = 4 * time.Second + // New creates a new telemetry service. func New(cfg *config.Config, cfgPath, version, edition string, logger *zap.Logger) *Service { _, envReason := IsDisabledByEnv() @@ -485,6 +509,7 @@ func New(cfg *config.Config, cfgPath, version, edition string, logger *zap.Logge envDisabledReason: envReason, initialDelay: 5 * time.Minute, heartbeatInterval: 24 * time.Hour, + flushTimeout: shutdownFlushTimeout, resolvedEnabled: EffectiveTelemetryEnabled(cfg), done: make(chan struct{}), } @@ -737,8 +762,18 @@ func (s *Service) SetOnboardingProvider(fn func() *OnboardingSnapshot) { // falls through to the runtime detector — this preserves liveness of the // heartbeat pipeline at the cost of losing the "installer" classification // for this one cycle. -func (s *Service) resolveLaunchSource() LaunchSource { - if s.activationStore != nil && s.activationDB != nil { +// +// consumeInstallerPending is false ONLY on the graceful-shutdown flush. Clearing +// the flag happens at BUILD time, before the send is known to have been +// accepted, which is a deliberate trade-off for the long-running loop (a crash +// mid-send must not re-emit "installer" forever). On the shutdown flush that +// trade-off inverts: the process is about to exit, so a flush that fails — +// offline machine, 4s timeout — would destroy an attribution that previously +// survived intact into the next run. The flush therefore leaves the flag ALONE +// and reports the runtime-detected source; the next run's first heartbeat still +// claims "installer" exactly once. +func (s *Service) resolveLaunchSource(consumeInstallerPending bool) LaunchSource { + if consumeInstallerPending && s.activationStore != nil && s.activationDB != nil { pending, err := s.activationStore.IsInstallerPending(s.activationDB) if err == nil && pending { // Clear the flag synchronously so a crash before the HTTP POST @@ -813,8 +848,15 @@ func (s *Service) Start(ctx context.Context) { // leaks before they leave the machine. Idempotent. PopulateBlockedValues() + // Every emission gate above has passed and the anonymous id exists, so a + // heartbeat from this process is legitimate from here on. Arm the + // graceful-shutdown flush (see flushFinalHeartbeat): even if the process + // exits before the initial delay elapses, whatever it recorded is worth one + // bounded final send. + s.flushEligible.Store(true) + s.logger.Info("Telemetry service starting", - zap.String("endpoint", s.endpoint), + zap.String("endpoint", s.liveEndpoint()), zap.Duration("initial_delay", s.initialDelay), zap.Duration("interval", s.heartbeatInterval)) @@ -856,6 +898,15 @@ func (s *Service) Start(ctx context.Context) { // Runtime.Close relies on this so no telemetry BBolt write can land after // the clean-shutdown marker resolves or after the DB closes (Spec 080 // FR-010, review round 6). +// +// After the loop has been joined, Stop performs ONE bounded final heartbeat +// (flushFinalHeartbeat) so counters recorded since the last accepted send are +// not lost on exit. Doing it AFTER the join — not from a second goroutine +// racing the loop — is what keeps the "the loop owns sends" invariant intact: +// at that point there is provably no other sender, so the send+reset path +// cannot interleave with a tick. Runtime.Close calls Stop while the BBolt +// handle is still open and before the clean-shutdown marker resolves, so the +// flush's buildHeartbeat writes are still legal there. func (s *Service) Stop() { s.startMu.Lock() s.stopped = true @@ -867,16 +918,87 @@ func (s *Service) Stop() { // Start ran (or is running): its defer guarantees done closes on every // exit path, including the disabled-by-env/config early returns. <-s.done + + s.flushFinalHeartbeat() +} + +// flushFinalHeartbeat sends at most one final heartbeat on the graceful +// shutdown path. Called by Stop AFTER the heartbeat loop has exited, so it is +// never concurrent with a tick. +// +// It reuses sendHeartbeat's send+reset path verbatim, which is what preserves +// the counter contract: counters are zeroed ONLY on a 2xx, so a failed or +// timed-out flush drops nothing (the next run reports the same counts) and an +// accepted flush cannot double-count with a later start. +// +// Delivery is AT-LEAST-ONCE, exactly as it already is between ticks. When a send +// is aborted after the endpoint accepted it but before the client observed the +// 2xx — the app context is cancelled a moment into an in-flight tick, say — the +// counters stay pending and this flush re-sends them, so the endpoint can see +// the same window twice. That is undecidable client-side (a lost response is +// indistinguishable from a lost request) and is the pre-existing retry semantic +// of Reset-on-2xx; the alternative is dropping the window entirely, which is the +// bug this flush exists to fix. Receivers must treat heartbeats as idempotent +// per (anonymous_id, timestamp) rather than assume exactly-once. +// +// Skipped when: Start never armed it (telemetry disabled by env/config, dev +// build), the user opted out, telemetry was turned off mid-run, or nothing has +// been recorded since the last accepted send. +// +// The send runs on a FRESH context, not the loop's: Runtime.Close cancels the +// app context before calling Stop, so a derived context would abort the flush +// instantly. It is hard-bounded by flushTimeout so a dead endpoint cannot hang +// shutdown. +// +// Build-time one-shot side effects are suppressed here (consumeOneShots=false): +// the 365-day anonymous-id rotation, which rewrites the whole config file at the +// worst possible moment (the daemon's other config writers are winding down), +// and the installer_heartbeat_pending flag, which is cleared at build time and +// would be destroyed by a flush that never lands. Both fire on the next run's +// first heartbeat instead. See buildHeartbeatWithOneShots. +func (s *Service) flushFinalHeartbeat() { + s.flushOnce.Do(func() { + if !s.flushEligible.Load() { + return + } + if s.optedOut.Load() { + return + } + if !s.telemetryEnabledLive() { + return + } + if !s.registry.HasPendingCounters() { + s.logger.Debug("Skipping telemetry shutdown flush: no counters recorded since the last accepted heartbeat") + return + } + + timeout := s.flushTimeout + if timeout <= 0 { + timeout = shutdownFlushTimeout + } + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + s.logger.Debug("Flushing final telemetry heartbeat on shutdown", + zap.Duration("timeout", timeout)) + s.sendHeartbeatWithOneShots(ctx, false) + }) } func (s *Service) sendHeartbeat(ctx context.Context) { + s.sendHeartbeatWithOneShots(ctx, true) +} + +// sendHeartbeatWithOneShots is sendHeartbeat's body. consumeOneShots is false +// only on the shutdown-flush path; see flushFinalHeartbeat. +func (s *Service) sendHeartbeatWithOneShots(ctx context.Context, consumeOneShots bool) { // MCP-2482: once the user has opted out, no further telemetry is emitted — // even if the long-running heartbeat loop is still ticking. if s.optedOut.Load() { return } - payload := s.buildHeartbeat() + payload := s.buildHeartbeatWithOneShots(consumeOneShots) data, err := json.Marshal(payload) if err != nil { @@ -905,7 +1027,7 @@ func (s *Service) sendHeartbeat(ctx context.Context) { return } - url := s.endpoint + "/heartbeat" + url := s.liveEndpoint() + "/heartbeat" req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(data)) if err != nil { s.logger.Debug("Failed to create heartbeat request", zap.Error(err)) @@ -1008,6 +1130,43 @@ func (s *Service) liveConfig() *config.Config { return s.config } +// liveEndpoint returns the service's current telemetry endpoint, read under +// s.mu. NotifyConfigChanged rewrites s.endpoint on a live config reload, so +// every read must take the same lock the write does — an unlocked read is a +// data race, not merely a stale value. The shutdown flush made this reachable in +// practice (it sends after the loop has been joined, i.e. from whichever +// goroutine called Stop), and it is equally reachable from the feedback and +// opt-out-beacon paths. +func (s *Service) liveEndpoint() string { + s.mu.Lock() + defer s.mu.Unlock() + return s.endpoint +} + +// liveAnonymousID returns the anonymous install id from the service's current +// config, read under s.mu — the same lock NotifyConfigChanged takes to swap +// s.config, and the same one ensureAnonymousIDOnce / advanceUpgradeFunnelOnce +// take to install cfg.Telemetry on a config that arrived without one. +// +// Snapshotting the pointer and dereferencing it after unlocking would NOT be +// enough, for exactly the reason telemetryEnabledLive documents: +// GetAnonymousID reads cfg.Telemetry, so the whole read has to happen inside +// the critical section. The opt-out beacon read it unlocked in both of its +// eligibility gates — two lines above the endpoint read this PR moved under the +// lock — which the race detector confirms (see +// TestOptOutBeaconReadsAnonymousIDUnderLock). +// +// config.GetAnonymousID only reads config fields, so calling it under the lock +// cannot re-enter the Service. +func (s *Service) liveAnonymousID() string { + s.mu.Lock() + defer s.mu.Unlock() + if s.config == nil { + return "" + } + return s.config.GetAnonymousID() +} + // telemetryCursor reads the four cfg.Telemetry scalars the heartbeat reports, // in one hold of s.mu. Every writer of these fields — maybeRotateAnonymousID // and advanceUpgradeFunnelOnce — mutates them under the same mutex, so the read @@ -1033,6 +1192,25 @@ func (s *Service) telemetryCursor(cfg *config.Config) (anonID, createdAt, previo } func (s *Service) buildHeartbeat() HeartbeatPayload { + return s.buildHeartbeatWithOneShots(true) +} + +// buildHeartbeatWithOneShots is buildHeartbeat's body. +// +// consumeOneShots gates every side effect this build performs on state that +// only fires ONCE and is consumed at build time rather than on a 2xx: +// - the 365-day anonymous-id rotation (rewrites the config FILE), and +// - the installer_heartbeat_pending flag (BBolt one-shot, see +// resolveLaunchSource). +// +// It is false only on the shutdown-flush path. Both effects are irreversible +// the moment they run, so performing them for a send that may never be accepted +// — the flush runs while the daemon tears down, on a 4s budget, possibly +// offline — would destroy state that previously survived intact into the next +// run. The next run's first heartbeat performs both instead. Everything else in +// this build (BBolt reads, funnel activity marking, counter snapshots) is +// idempotent or reset-on-2xx and is therefore unconditional. +func (s *Service) buildHeartbeatWithOneShots(consumeOneShots bool) HeartbeatPayload { // Take ONE snapshot of the live config pointer under s.mu and read only that // below. NotifyConfigChanged swaps s.config wholesale on a reload, so an // unsynchronized read of the field races that swap; snapshotting also means a @@ -1043,7 +1221,9 @@ func (s *Service) buildHeartbeat() HeartbeatPayload { // Spec 042: rotate the anonymous ID if it's older than 365 days. Runs on the // snapshot, so the rotated ID is the one this payload reports. - s.maybeRotateAnonymousID(cfg, time.Now().UTC()) + if consumeOneShots { + s.maybeRotateAnonymousID(cfg, time.Now().UTC()) + } // Read every cfg.Telemetry-derived scalar in ONE locked pass. These four // fields are written under s.mu by maybeRotateAnonymousID (anonymous id + @@ -1139,7 +1319,7 @@ func (s *Service) buildHeartbeat() HeartbeatPayload { // the installer_heartbeat_pending flag set at process startup when // MCPPROXY_LAUNCHED_BY=installer was observed. Otherwise the runtime // detector result (tray/login_item/cli/unknown) is emitted. - payload.LaunchSource = string(s.resolveLaunchSource()) + payload.LaunchSource = string(s.resolveLaunchSource(consumeOneShots)) // Spec 044 (T051): AutostartEnabled. Tri-state; nil when the tray sidecar // is absent/unreachable/malformed (Linux always falls here).