From ba52527e1eb3be2c028a5fa7be58bd8287f3b6c2 Mon Sep 17 00:00:00 2001 From: marlian <45340800+marlian@users.noreply.github.com> Date: Tue, 14 Apr 2026 22:36:08 +0200 Subject: [PATCH 01/11] feat: add internal/telemetry package MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port of the Node telemetry design to Go with three refinements: - Nil-tolerant *Client: InitIfEnabled returns nil when disabled, every method no-ops on nil receiver. Replaces per-callsite guards. - No globals: client is plumbed via constructor, never module-level. - Privacy-strict mode: new MEMORY_TELEMETRY_PRIVACY=strict option hashes entity/query/label values with sha256 before storage. Observation content is always reduced to regardless of mode. Package contents: - schema.go — DDL for tool_calls + search_metrics + indexes - client.go — Client, InitIfEnabled, FromEnv, LogToolCall, LogSearchMetrics - sanitize.go — args/result summarizers, HTML escape disabled for readable sentinels - hash.go — sha256 helpers for strict mode - detect.go — client identity (MCP handshake primary, env fallback) All ported env signals from Node preserved: KILO, CLAUDE_CODE_SSE_PORT, CURSOR_TRACE_ID, WINDSURF_EXTENSION_ID, VSCODE_MCP_HTTP_PREFER, TERM_PROGRAM. Added non-empty guard on VSCODE_MCP_HTTP_PREFER vs the Node presence-only check — simpler semantics, same real-world detection. Uses modernc.org/sqlite for the telemetry DB to preserve the pure-Go single-binary invariant (no CGO added for analytics). --- internal/telemetry/client.go | 210 ++++++++++++++++++++++++++++ internal/telemetry/client_test.go | 158 +++++++++++++++++++++ internal/telemetry/detect.go | 39 ++++++ internal/telemetry/detect_test.go | 70 ++++++++++ internal/telemetry/hash.go | 25 ++++ internal/telemetry/hash_test.go | 39 ++++++ internal/telemetry/sanitize.go | 115 +++++++++++++++ internal/telemetry/sanitize_test.go | 110 +++++++++++++++ internal/telemetry/schema.go | 44 ++++++ 9 files changed, 810 insertions(+) create mode 100644 internal/telemetry/client.go create mode 100644 internal/telemetry/client_test.go create mode 100644 internal/telemetry/detect.go create mode 100644 internal/telemetry/detect_test.go create mode 100644 internal/telemetry/hash.go create mode 100644 internal/telemetry/hash_test.go create mode 100644 internal/telemetry/sanitize.go create mode 100644 internal/telemetry/sanitize_test.go create mode 100644 internal/telemetry/schema.go diff --git a/internal/telemetry/client.go b/internal/telemetry/client.go new file mode 100644 index 0000000..2b676ae --- /dev/null +++ b/internal/telemetry/client.go @@ -0,0 +1,210 @@ +// Package telemetry provides opt-in usage analytics for the workmem MCP server. +// +// The package is a no-op when MEMORY_TELEMETRY_PATH is unset — InitIfEnabled +// returns nil, and every method of *Client returns immediately when the +// receiver is nil. There is no global state, no side channel: the server +// passes a *Client down to wherever logging happens. +// +// When MEMORY_TELEMETRY_PRIVACY=strict is set, entity/query/label values in +// args_summary and search_metrics.query are sha256-hashed. observation +// content and structured array values are always reduced to counts, in any +// mode. +package telemetry + +import ( + "database/sql" + "encoding/json" + "fmt" + "os" + + _ "modernc.org/sqlite" +) + +// Client is a telemetry sink. A nil *Client is valid and represents the +// disabled state — every method below checks for nil and returns harmlessly. +type Client struct { + db *sql.DB + insertCall *sql.Stmt + insertSearch *sql.Stmt + strict bool +} + +// FromEnv reads MEMORY_TELEMETRY_PATH and MEMORY_TELEMETRY_PRIVACY from the +// environment and delegates to InitIfEnabled. Returns nil when the path is +// unset (telemetry disabled). Strict mode is enabled when +// MEMORY_TELEMETRY_PRIVACY is exactly "strict"; any other value (including +// empty) means permissive. +func FromEnv() *Client { + return InitIfEnabled( + os.Getenv("MEMORY_TELEMETRY_PATH"), + os.Getenv("MEMORY_TELEMETRY_PRIVACY") == "strict", + ) +} + +// InitIfEnabled opens (or creates) the telemetry SQLite database at the given +// path. If path is empty, returns nil (telemetry disabled). If init fails at +// any step, logs a single warning to stderr and returns nil — telemetry is +// never allowed to break the tool call path. +func InitIfEnabled(path string, strict bool) *Client { + if path == "" { + return nil + } + db, err := sql.Open("sqlite", path) + if err != nil { + initWarn(err) + return nil + } + if _, err := db.Exec("PRAGMA journal_mode=WAL"); err != nil { + _ = db.Close() + initWarn(err) + return nil + } + if _, err := db.Exec(schemaSQL); err != nil { + _ = db.Close() + initWarn(err) + return nil + } + insertCall, err := db.Prepare(insertCallSQL) + if err != nil { + _ = db.Close() + initWarn(err) + return nil + } + insertSearch, err := db.Prepare(insertSearchSQL) + if err != nil { + _ = insertCall.Close() + _ = db.Close() + initWarn(err) + return nil + } + return &Client{db: db, insertCall: insertCall, insertSearch: insertSearch, strict: strict} +} + +func initWarn(err error) { + fmt.Fprintf(os.Stderr, "[memory] telemetry init failed (disabled for this session): %v\n", err) +} + +// Close releases the prepared statements and the underlying database +// connection. Safe to call on a nil receiver. +func (c *Client) Close() error { + if c == nil { + return nil + } + if c.insertCall != nil { + _ = c.insertCall.Close() + } + if c.insertSearch != nil { + _ = c.insertSearch.Close() + } + return c.db.Close() +} + +// Strict reports whether privacy-strict mode is active. +func (c *Client) Strict() bool { + if c == nil { + return false + } + return c.strict +} + +// ToolCallInput captures a single tool invocation for telemetry. +type ToolCallInput struct { + Tool string + Client ClientInfo + DBScope string // "global" or "project" + ProjectPath string + DurationMs float64 + ArgsSummary string + ResultSummary string + IsError bool +} + +// LogToolCall inserts a tool_calls row. Returns the new row id on success or +// 0 on failure / disabled client. The returned id is used by LogSearchMetrics +// to link the search_metrics row back to its parent tool call. +func (c *Client) LogToolCall(in ToolCallInput) int64 { + if c == nil { + return 0 + } + dbScope := in.DBScope + if dbScope == "" { + dbScope = "global" + } + res, err := c.insertCall.Exec( + in.Tool, + nullIfEmpty(in.Client.Name), + nullIfEmpty(in.Client.Version), + nullIfEmpty(in.Client.Source), + dbScope, + nullIfEmpty(in.ProjectPath), + in.DurationMs, + nullIfEmpty(in.ArgsSummary), + nullIfEmpty(in.ResultSummary), + boolToInt(in.IsError), + ) + if err != nil { + fmt.Fprintf(os.Stderr, "[memory] telemetry log failed: %v\n", err) + return 0 + } + id, err := res.LastInsertId() + if err != nil { + return 0 + } + return id +} + +// SearchMetricsInput captures ranking-pipeline metrics for a single recall call. +type SearchMetricsInput struct { + ToolCallID int64 + Query string + Channels map[string]int + CandidatesTotal int + ResultsReturned int + LimitRequested int + ScoreMin float64 + ScoreMax float64 + ScoreMedian float64 + Compact bool +} + +// LogSearchMetrics inserts a search_metrics row linked to the tool_call id. +// No-op when client is nil or ToolCallID is 0 (the linking parent failed). +// In strict mode, Query is hashed before insertion. +func (c *Client) LogSearchMetrics(in SearchMetricsInput) { + if c == nil || in.ToolCallID == 0 { + return + } + channelsJSON, err := json.Marshal(in.Channels) + if err != nil { + channelsJSON = []byte("{}") + } + query := hashIfStrict(in.Query, c.strict) + if _, err := c.insertSearch.Exec( + in.ToolCallID, + nullIfEmpty(query), + string(channelsJSON), + in.CandidatesTotal, + in.ResultsReturned, + in.LimitRequested, + in.ScoreMin, + in.ScoreMax, + in.ScoreMedian, + boolToInt(in.Compact), + ); err != nil { + fmt.Fprintf(os.Stderr, "[memory] telemetry search log failed: %v\n", err) + } +} + +func nullIfEmpty(s string) any { + if s == "" { + return nil + } + return s +} + +func boolToInt(b bool) int { + if b { + return 1 + } + return 0 +} diff --git a/internal/telemetry/client_test.go b/internal/telemetry/client_test.go new file mode 100644 index 0000000..62ccf85 --- /dev/null +++ b/internal/telemetry/client_test.go @@ -0,0 +1,158 @@ +package telemetry + +import ( + "database/sql" + "path/filepath" + "strings" + "testing" +) + +func TestNilClientMethodsAreSafe(t *testing.T) { + var c *Client + if c.Strict() { + t.Fatalf("nil client should not be strict") + } + if got := c.LogToolCall(ToolCallInput{Tool: "noop"}); got != 0 { + t.Fatalf("nil client LogToolCall should return 0, got %d", got) + } + c.LogSearchMetrics(SearchMetricsInput{ToolCallID: 123, Query: "anything"}) + if err := c.Close(); err != nil { + t.Fatalf("nil client Close should be nil, got %v", err) + } +} + +func TestInitIfEnabledEmptyPathReturnsNil(t *testing.T) { + if got := InitIfEnabled("", false); got != nil { + t.Fatalf("InitIfEnabled(\"\", false) should return nil, got %+v", got) + } +} + +func TestInitIfEnabledInvalidPathReturnsNil(t *testing.T) { + // A path inside a nonexistent directory should fail open + bad := "/nonexistent-directory-for-telemetry-test/telemetry.db" + if got := InitIfEnabled(bad, false); got != nil { + _ = got.Close() + t.Fatalf("InitIfEnabled on bad path should return nil, got client") + } +} + +func TestInitIfEnabledCreatesSchemaAndInserts(t *testing.T) { + path := filepath.Join(t.TempDir(), "telemetry.db") + c := InitIfEnabled(path, false) + if c == nil { + t.Fatalf("InitIfEnabled returned nil on valid path") + } + t.Cleanup(func() { _ = c.Close() }) + + id := c.LogToolCall(ToolCallInput{ + Tool: "remember", + Client: ClientInfo{Name: "claude-code", Source: "env"}, + DBScope: "global", + DurationMs: 1.23, + ArgsSummary: `{"entity":"Alice"}`, + ResultSummary: `{"stored":true}`, + }) + if id == 0 { + t.Fatalf("LogToolCall returned 0 on valid client") + } + + c.LogSearchMetrics(SearchMetricsInput{ + ToolCallID: id, + Query: "alice", + Channels: map[string]int{"fts": 3}, + CandidatesTotal: 3, + ResultsReturned: 1, + LimitRequested: 5, + ScoreMin: 0.1, + ScoreMax: 0.9, + ScoreMedian: 0.5, + }) + + // Read back from disk to prove persistence + rdb, err := sql.Open("sqlite", path) + if err != nil { + t.Fatalf("open telemetry db for readback: %v", err) + } + defer rdb.Close() + + var toolCount, searchCount int + if err := rdb.QueryRow(`SELECT COUNT(*) FROM tool_calls`).Scan(&toolCount); err != nil { + t.Fatalf("count tool_calls: %v", err) + } + if err := rdb.QueryRow(`SELECT COUNT(*) FROM search_metrics`).Scan(&searchCount); err != nil { + t.Fatalf("count search_metrics: %v", err) + } + if toolCount != 1 { + t.Fatalf("tool_calls count = %d, want 1", toolCount) + } + if searchCount != 1 { + t.Fatalf("search_metrics count = %d, want 1", searchCount) + } + + var argsSummary, resultSummary, clientName sql.NullString + if err := rdb.QueryRow(`SELECT client_name, args_summary, result_summary FROM tool_calls WHERE id = ?`, id).Scan(&clientName, &argsSummary, &resultSummary); err != nil { + t.Fatalf("readback tool_calls: %v", err) + } + if clientName.String != "claude-code" { + t.Fatalf("client_name = %q, want claude-code", clientName.String) + } + if !strings.Contains(argsSummary.String, "Alice") { + t.Fatalf("args_summary missing entity: %q", argsSummary.String) + } +} + +func TestStrictModeHashesSearchQuery(t *testing.T) { + path := filepath.Join(t.TempDir(), "strict.db") + c := InitIfEnabled(path, true) + if c == nil { + t.Fatalf("InitIfEnabled strict returned nil") + } + t.Cleanup(func() { _ = c.Close() }) + + if !c.Strict() { + t.Fatalf("expected strict mode active") + } + + id := c.LogToolCall(ToolCallInput{Tool: "recall", DBScope: "project"}) + c.LogSearchMetrics(SearchMetricsInput{ + ToolCallID: id, + Query: "sensitive therapy question", + Channels: map[string]int{"fts": 1}, + CandidatesTotal: 1, + ResultsReturned: 1, + }) + + rdb, err := sql.Open("sqlite", path) + if err != nil { + t.Fatalf("open db: %v", err) + } + defer rdb.Close() + var storedQuery sql.NullString + if err := rdb.QueryRow(`SELECT query FROM search_metrics WHERE tool_call_id = ?`, id).Scan(&storedQuery); err != nil { + t.Fatalf("readback: %v", err) + } + if strings.Contains(storedQuery.String, "sensitive") || strings.Contains(storedQuery.String, "therapy") { + t.Fatalf("strict mode leaked plaintext query: %q", storedQuery.String) + } + if !strings.HasPrefix(storedQuery.String, "sha256:") { + t.Fatalf("strict mode did not hash query, got %q", storedQuery.String) + } +} + +func TestLogSearchMetricsZeroToolCallIDIsNoop(t *testing.T) { + path := filepath.Join(t.TempDir(), "noop.db") + c := InitIfEnabled(path, false) + if c == nil { + t.Fatalf("init failed") + } + t.Cleanup(func() { _ = c.Close() }) + + c.LogSearchMetrics(SearchMetricsInput{ToolCallID: 0, Query: "anything"}) + var count int + if err := c.db.QueryRow(`SELECT COUNT(*) FROM search_metrics`).Scan(&count); err != nil { + t.Fatalf("count: %v", err) + } + if count != 0 { + t.Fatalf("LogSearchMetrics with ToolCallID=0 inserted row, count=%d", count) + } +} diff --git a/internal/telemetry/detect.go b/internal/telemetry/detect.go new file mode 100644 index 0000000..1887053 --- /dev/null +++ b/internal/telemetry/detect.go @@ -0,0 +1,39 @@ +package telemetry + +import "os" + +// ClientInfo identifies the MCP client currently calling the server. +type ClientInfo struct { + Name string + Version string + Source string // "protocol" | "env" | "none" +} + +// DetectClient resolves the active client identity. Priority: +// 1. protocolName / protocolVersion from the MCP initialize handshake +// 2. environment variables set by known MCP clients +// 3. ClientInfo{Name: "unknown", Source: "none"} +func DetectClient(protocolName, protocolVersion string) ClientInfo { + if protocolName != "" { + return ClientInfo{Name: protocolName, Version: protocolVersion, Source: "protocol"} + } + if os.Getenv("KILO") != "" { + return ClientInfo{Name: "kilo", Version: os.Getenv("KILOCODE_VERSION"), Source: "env"} + } + if os.Getenv("CLAUDE_CODE_SSE_PORT") != "" { + return ClientInfo{Name: "claude-code", Source: "env"} + } + if os.Getenv("CURSOR_TRACE_ID") != "" { + return ClientInfo{Name: "cursor", Source: "env"} + } + if os.Getenv("WINDSURF_EXTENSION_ID") != "" { + return ClientInfo{Name: "windsurf", Source: "env"} + } + if os.Getenv("VSCODE_MCP_HTTP_PREFER") != "" { + return ClientInfo{Name: "vscode-copilot", Source: "env"} + } + if os.Getenv("TERM_PROGRAM") == "vscode" { + return ClientInfo{Name: "vscode-unknown", Source: "env"} + } + return ClientInfo{Name: "unknown", Source: "none"} +} diff --git a/internal/telemetry/detect_test.go b/internal/telemetry/detect_test.go new file mode 100644 index 0000000..5042fd1 --- /dev/null +++ b/internal/telemetry/detect_test.go @@ -0,0 +1,70 @@ +package telemetry + +import ( + "testing" +) + +func TestDetectClientProtocolWins(t *testing.T) { + t.Setenv("KILO", "1") + info := DetectClient("claude-desktop", "0.9.1") + if info.Name != "claude-desktop" || info.Version != "0.9.1" || info.Source != "protocol" { + t.Fatalf("protocol should win, got %+v", info) + } +} + +func TestDetectClientEnvFallbackKilo(t *testing.T) { + t.Setenv("KILO", "1") + t.Setenv("KILOCODE_VERSION", "0.43.6") + unsetMCPClientEnv(t) + t.Setenv("KILO", "1") + t.Setenv("KILOCODE_VERSION", "0.43.6") + info := DetectClient("", "") + if info.Name != "kilo" || info.Version != "0.43.6" || info.Source != "env" { + t.Fatalf("kilo env not detected: %+v", info) + } +} + +func TestDetectClientEnvFallbackClaudeCode(t *testing.T) { + unsetMCPClientEnv(t) + t.Setenv("CLAUDE_CODE_SSE_PORT", "7123") + info := DetectClient("", "") + if info.Name != "claude-code" || info.Source != "env" { + t.Fatalf("claude-code not detected: %+v", info) + } +} + +func TestDetectClientVSCodeCopilotRequiresNonEmptyVar(t *testing.T) { + unsetMCPClientEnv(t) + t.Setenv("VSCODE_MCP_HTTP_PREFER", "auto") + info := DetectClient("", "") + if info.Name != "vscode-copilot" || info.Source != "env" { + t.Fatalf("vscode-copilot non-empty var not detected: %+v", info) + } +} + +func TestDetectClientNoSignals(t *testing.T) { + unsetMCPClientEnv(t) + info := DetectClient("", "") + if info.Name != "unknown" || info.Source != "none" { + t.Fatalf("no-signal detection wrong: %+v", info) + } +} + +func unsetMCPClientEnv(t *testing.T) { + t.Helper() + for _, v := range []string{ + "KILO", "KILOCODE_VERSION", + "CLAUDE_CODE_SSE_PORT", + "CURSOR_TRACE_ID", + "WINDSURF_EXTENSION_ID", + "VSCODE_MCP_HTTP_PREFER", + "TERM_PROGRAM", + } { + t.Setenv(v, "") + } + // t.Setenv with "" sets to empty, but LookupEnv still reports present. + // For the ones we need genuinely absent, unset via os.Unsetenv is not + // t.Setenv-compatible; we accept that "" is how the tests model absence + // for signals whose presence-only matters (VSCODE_MCP_HTTP_PREFER handled + // in its own test). +} diff --git a/internal/telemetry/hash.go b/internal/telemetry/hash.go new file mode 100644 index 0000000..37634e0 --- /dev/null +++ b/internal/telemetry/hash.go @@ -0,0 +1,25 @@ +package telemetry + +import ( + "crypto/sha256" + "encoding/hex" +) + +// hashString returns sha256: for a non-empty input, or "" for empty input. +func hashString(s string) string { + if s == "" { + return "" + } + h := sha256.Sum256([]byte(s)) + return "sha256:" + hex.EncodeToString(h[:]) +} + +// hashIfStrict returns s unchanged unless strict is true, in which case s is +// replaced with its sha256 hash (prefixed with "sha256:"). Empty input always +// returns empty. +func hashIfStrict(s string, strict bool) string { + if !strict { + return s + } + return hashString(s) +} diff --git a/internal/telemetry/hash_test.go b/internal/telemetry/hash_test.go new file mode 100644 index 0000000..c65f532 --- /dev/null +++ b/internal/telemetry/hash_test.go @@ -0,0 +1,39 @@ +package telemetry + +import ( + "strings" + "testing" +) + +func TestHashStringDeterministicAndPrefixed(t *testing.T) { + a := hashString("hello") + b := hashString("hello") + if a != b { + t.Fatalf("hashString not deterministic: %q vs %q", a, b) + } + if !strings.HasPrefix(a, "sha256:") { + t.Fatalf("hashString missing sha256: prefix: %q", a) + } + if len(a) != len("sha256:")+64 { + t.Fatalf("hashString unexpected length %d: %q", len(a), a) + } +} + +func TestHashStringEmptyInputReturnsEmpty(t *testing.T) { + if got := hashString(""); got != "" { + t.Fatalf("hashString(\"\") = %q, want empty", got) + } +} + +func TestHashIfStrict(t *testing.T) { + if got := hashIfStrict("alpha", false); got != "alpha" { + t.Fatalf("hashIfStrict(strict=false) should pass through, got %q", got) + } + got := hashIfStrict("alpha", true) + if !strings.HasPrefix(got, "sha256:") { + t.Fatalf("hashIfStrict(strict=true) should hash, got %q", got) + } + if got == "alpha" { + t.Fatalf("hashIfStrict(strict=true) returned plaintext") + } +} diff --git a/internal/telemetry/sanitize.go b/internal/telemetry/sanitize.go new file mode 100644 index 0000000..97d4374 --- /dev/null +++ b/internal/telemetry/sanitize.go @@ -0,0 +1,115 @@ +package telemetry + +import ( + "bytes" + "encoding/json" + "fmt" + "strings" +) + +// encodeNoEscape marshals v as JSON without escaping <, >, & — telemetry +// rows are never rendered as HTML, and the escapes obscure sentinel markers +// like "<51 chars>" that must remain readable in the DB. +func encodeNoEscape(v any) (string, error) { + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + enc.SetEscapeHTML(false) + if err := enc.Encode(v); err != nil { + return "", err + } + return strings.TrimRight(buf.String(), "\n"), nil +} + +// SanitizeArgs produces the args_summary string for a tool call. +// +// Rules: +// - "observation", "content", "context" string values are replaced with +// "" (length only, never content) +// - "facts" and "observations" arrays are replaced with "" / +// "" (count only, never contents) +// - "entity", "from", "to", "label", "query" string values are hashed when +// strict is true; left unchanged otherwise +// - all other fields pass through unchanged +func SanitizeArgs(args map[string]any, strict bool) string { + if args == nil { + return "" + } + safe := make(map[string]any, len(args)) + for k, v := range args { + switch k { + case "observation", "content", "context": + if s, ok := v.(string); ok { + safe[k] = fmt.Sprintf("<%d chars>", len(s)) + continue + } + safe[k] = v + case "facts": + if arr, ok := v.([]any); ok { + safe[k] = fmt.Sprintf("<%d facts>", len(arr)) + continue + } + safe[k] = v + case "observations": + if arr, ok := v.([]any); ok { + safe[k] = fmt.Sprintf("<%d observations>", len(arr)) + continue + } + safe[k] = v + case "entity", "from", "to", "label", "query": + if s, ok := v.(string); ok { + safe[k] = hashIfStrict(s, strict) + continue + } + safe[k] = v + default: + safe[k] = v + } + } + out, err := encodeNoEscape(safe) + if err != nil { + return "" + } + return out +} + +// SummarizeResult extracts count-only telemetry fields from a tool result. +// It never leaks content: only counts, booleans, and structural indicators. +// +// Uses JSON round-trip to avoid direct coupling to the store package result +// types — the telemetry package stays domain-agnostic. +func SummarizeResult(result any) string { + if result == nil { + return "" + } + raw, err := json.Marshal(result) + if err != nil { + return "" + } + var m map[string]any + if err := json.Unmarshal(raw, &m); err != nil { + return "" + } + out := make(map[string]any, 8) + for _, k := range []string{"total_facts", "total", "stored", "deleted", "created", "found", "observations_attached", "compact"} { + if v, ok := m[k]; ok { + out[k] = v + } + } + if arr, ok := m["results"].([]any); ok { + out["entity_groups"] = len(arr) + } + if arr, ok := m["entities"].([]any); ok { + out["entities"] = len(arr) + } + if arr, ok := m["events"].([]any); ok { + out["events"] = len(arr) + } + if arr, ok := m["facts"].([]any); ok { + out["facts"] = len(arr) + } + encoded, err := encodeNoEscape(out) + if err != nil { + return "" + } + return encoded +} diff --git a/internal/telemetry/sanitize_test.go b/internal/telemetry/sanitize_test.go new file mode 100644 index 0000000..a76348d --- /dev/null +++ b/internal/telemetry/sanitize_test.go @@ -0,0 +1,110 @@ +package telemetry + +import ( + "strings" + "testing" +) + +func TestSanitizeArgsStripsObservationContent(t *testing.T) { + got := SanitizeArgs(map[string]any{ + "entity": "Alice", + "observation": "this is sensitive content that must never be logged", + }, false) + if strings.Contains(got, "sensitive content") { + t.Fatalf("observation content leaked in args_summary: %s", got) + } + if !strings.Contains(got, "<51 chars>") { + t.Fatalf("observation length marker missing: %s", got) + } + if !strings.Contains(got, "Alice") { + t.Fatalf("entity should be plaintext in permissive mode: %s", got) + } +} + +func TestSanitizeArgsStripsFactsAndObservationsArrays(t *testing.T) { + got := SanitizeArgs(map[string]any{ + "facts": []any{map[string]any{}, map[string]any{}, map[string]any{}}, + "observations": []any{map[string]any{}, map[string]any{}}, + }, false) + if !strings.Contains(got, "<3 facts>") { + t.Fatalf("facts count marker missing: %s", got) + } + if !strings.Contains(got, "<2 observations>") { + t.Fatalf("observations count marker missing: %s", got) + } +} + +func TestSanitizeArgsStrictModeHashesIdentifiers(t *testing.T) { + got := SanitizeArgs(map[string]any{ + "entity": "Alice", + "from": "Bob", + "to": "Carol", + "label": "therapy session", + "query": "anxiety", + "limit": 10, + }, true) + for _, leak := range []string{"Alice", "Bob", "Carol", "therapy session", "anxiety"} { + if strings.Contains(got, leak) { + t.Fatalf("strict mode leaked %q in args_summary: %s", leak, got) + } + } + if !strings.Contains(got, "sha256:") { + t.Fatalf("strict mode did not hash identifiers: %s", got) + } + if !strings.Contains(got, "\"limit\":10") { + t.Fatalf("non-identifier fields should pass through in strict mode: %s", got) + } +} + +func TestSanitizeArgsPermissiveModeKeepsIdentifiers(t *testing.T) { + got := SanitizeArgs(map[string]any{ + "entity": "Alice", + "query": "preferences", + }, false) + if !strings.Contains(got, "Alice") { + t.Fatalf("permissive mode should keep entity plaintext: %s", got) + } + if !strings.Contains(got, "preferences") { + t.Fatalf("permissive mode should keep query plaintext: %s", got) + } +} + +func TestSanitizeArgsNilReturnsEmpty(t *testing.T) { + if got := SanitizeArgs(nil, false); got != "" { + t.Fatalf("SanitizeArgs(nil) = %q, want empty", got) + } +} + +func TestSummarizeResultReturnsCountsOnly(t *testing.T) { + result := map[string]any{ + "stored": true, + "results": []any{ + map[string]any{"entity": "A", "observations": []any{"secret content 1"}}, + map[string]any{"entity": "B", "observations": []any{"secret content 2"}}, + }, + "total": 5, + "compact": false, + } + got := SummarizeResult(result) + if strings.Contains(got, "secret content") { + t.Fatalf("SummarizeResult leaked content: %s", got) + } + if strings.Contains(got, "\"A\"") || strings.Contains(got, "\"B\"") { + t.Fatalf("SummarizeResult leaked entity names: %s", got) + } + if !strings.Contains(got, "\"entity_groups\":2") { + t.Fatalf("SummarizeResult missing entity_groups count: %s", got) + } + if !strings.Contains(got, "\"stored\":true") { + t.Fatalf("SummarizeResult missing stored: %s", got) + } + if !strings.Contains(got, "\"total\":5") { + t.Fatalf("SummarizeResult missing total: %s", got) + } +} + +func TestSummarizeResultNilReturnsEmpty(t *testing.T) { + if got := SummarizeResult(nil); got != "" { + t.Fatalf("SummarizeResult(nil) = %q, want empty", got) + } +} diff --git a/internal/telemetry/schema.go b/internal/telemetry/schema.go new file mode 100644 index 0000000..8d862d6 --- /dev/null +++ b/internal/telemetry/schema.go @@ -0,0 +1,44 @@ +package telemetry + +const schemaSQL = ` +CREATE TABLE IF NOT EXISTS tool_calls ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + ts TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%f', 'now')), + tool TEXT NOT NULL, + client_name TEXT, + client_version TEXT, + client_source TEXT, + db_scope TEXT NOT NULL DEFAULT 'global', + project_path TEXT, + duration_ms REAL, + args_summary TEXT, + result_summary TEXT, + is_error INTEGER NOT NULL DEFAULT 0 +); +CREATE TABLE IF NOT EXISTS search_metrics ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + tool_call_id INTEGER REFERENCES tool_calls(id), + query TEXT, + channels TEXT, + candidates_total INTEGER, + results_returned INTEGER, + limit_requested INTEGER, + score_min REAL, + score_max REAL, + score_median REAL, + compact INTEGER DEFAULT 0 +); +CREATE INDEX IF NOT EXISTS idx_tool_calls_ts ON tool_calls(ts); +CREATE INDEX IF NOT EXISTS idx_tool_calls_tool ON tool_calls(tool); +CREATE INDEX IF NOT EXISTS idx_search_metrics_tool_call ON search_metrics(tool_call_id); +` + +const ( + insertCallSQL = `INSERT INTO tool_calls + (tool, client_name, client_version, client_source, db_scope, project_path, duration_ms, args_summary, result_summary, is_error) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)` + + insertSearchSQL = `INSERT INTO search_metrics + (tool_call_id, query, channels, candidates_total, results_returned, limit_requested, score_min, score_max, score_median, compact) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)` +) From 48f2296ec5f2f845054bd1eb2b9ccb81c72eebff Mon Sep 17 00:00:00 2001 From: marlian <45340800+marlian@users.noreply.github.com> Date: Tue, 14 Apr 2026 22:36:25 +0200 Subject: [PATCH 02/11] refactor: SearchMemory returns SearchMetrics for telemetry hook MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SearchMemory now returns (results, metrics, err). The metrics struct captures channel counts, candidates_total, results_returned, score distribution, and limit_requested — exactly what the search_metrics telemetry table needs, with no coupling to the telemetry package. Replaces the Node pattern of a module-level _lastSearchMetrics side channel. Go callers get an explicit value and tests can assert on it without any global state. HandleTool keeps its original (any, error) signature for backward compatibility. A new HandleToolWithMetrics sibling returns (any, *SearchMetrics, error) and is called by the mcpserver dispatch when it wants the search_metrics payload. The shared body lives in an inner dispatchTool helper that writes metrics via an out-pointer — avoids rewriting every case's return statement. Callers updated: internal/store/tools.go recall handler, internal/store/parity_test.go (5 call sites). No behavior change for any existing test. --- internal/store/parity_test.go | 10 ++--- internal/store/search.go | 80 ++++++++++++++++++++++++++++++++--- internal/store/tools.go | 23 +++++++++- 3 files changed, 100 insertions(+), 13 deletions(-) diff --git a/internal/store/parity_test.go b/internal/store/parity_test.go index 7e3fc1f..8ee1912 100644 --- a/internal/store/parity_test.go +++ b/internal/store/parity_test.go @@ -32,7 +32,7 @@ func TestCoreMemoryParity(t *testing.T) { if _, err := AddObservation(db, entityID, "test observation", "user", 1.0); err != nil { t.Fatalf("AddObservation() error = %v", err) } - results, err := SearchMemory(db, "test", 5, 12) + results, _, err := SearchMemory(db, "test", 5, 12) if err != nil { t.Fatalf("SearchMemory() error = %v", err) } @@ -65,7 +65,7 @@ func TestCoreMemoryParity(t *testing.T) { if !tombstoned { t.Fatalf("observation tombstone not set") } - recalled, err := SearchMemory(db, "observation to forget", 5, 12) + recalled, _, err := SearchMemory(db, "observation to forget", 5, 12) if err != nil { t.Fatalf("SearchMemory() error = %v", err) } @@ -353,7 +353,7 @@ func TestProjectIsolationParity(t *testing.T) { if _, err := AddObservation(projectDB, entityID, "project fact", "user", 1.0); err != nil { t.Fatalf("AddObservation(project) error = %v", err) } - globalResults, err := SearchMemory(db, "ProjectOnly", 5, 12) + globalResults, _, err := SearchMemory(db, "ProjectOnly", 5, 12) if err != nil { t.Fatalf("SearchMemory(global) error = %v", err) } @@ -402,7 +402,7 @@ func TestRankingAndRecallParity(t *testing.T) { } t.Run("composite score present on results", func(t *testing.T) { - results, err := SearchMemory(db, "Zara", 10, 12) + results, _, err := SearchMemory(db, "Zara", 10, 12) if err != nil { t.Fatalf("SearchMemory() error = %v", err) } @@ -417,7 +417,7 @@ func TestRankingAndRecallParity(t *testing.T) { }) t.Run("entity exact outranks entity like", func(t *testing.T) { - results, err := SearchMemory(db, "Zara", 10, 12) + results, _, err := SearchMemory(db, "Zara", 10, 12) if err != nil { t.Fatalf("SearchMemory() error = %v", err) } diff --git a/internal/store/search.go b/internal/store/search.go index acc908a..9717896 100644 --- a/internal/store/search.go +++ b/internal/store/search.go @@ -29,6 +29,21 @@ type candidate struct { FTSPosition *int } +// SearchMetrics captures per-call ranking-pipeline observations, suitable for +// telemetry. SearchMemory populates every field except Compact (which depends +// on tool args and is set by the caller). +type SearchMetrics struct { + Query string + Channels map[string]int + CandidatesTotal int + ResultsReturned int + LimitRequested int + ScoreMin float64 + ScoreMax float64 + ScoreMedian float64 + Compact bool +} + type SearchObservation struct { ID int64 EntityID int64 @@ -315,14 +330,16 @@ func ScoreCandidates(observations []SearchObservation, candidateMap map[int64]*c return ranked } -func SearchMemory(db *sql.DB, query string, limit int, halfLifeWeeks float64) ([]SearchObservation, error) { +func SearchMemory(db *sql.DB, query string, limit int, halfLifeWeeks float64) ([]SearchObservation, SearchMetrics, error) { + requestedLimit := limit limit = SanitizeSearchLimit(limit) + emptyMetrics := SearchMetrics{Query: query, LimitRequested: requestedLimit, Channels: map[string]int{}} if limit <= 0 { - return []SearchObservation{}, nil + return []SearchObservation{}, emptyMetrics, nil } trimmed := strings.TrimSpace(query) if trimmed == "" { - return []SearchObservation{}, nil + return []SearchObservation{}, emptyMetrics, nil } collectLimit := limit * collectionMultiplier @@ -332,11 +349,12 @@ func SearchMemory(db *sql.DB, query string, limit int, halfLifeWeeks float64) ([ } candidates, err := CollectCandidates(db, trimmed, collectLimit, maxCandidates) if err != nil { - return nil, err + return nil, SearchMetrics{}, err } + channelCounts := countCandidateChannels(candidates) hydrated, err := HydrateCandidates(db, candidates) if err != nil { - return nil, err + return nil, SearchMetrics{}, err } ranked := ScoreCandidates(hydrated, candidates, halfLifeWeeks, limit) returnedIDs := make([]int64, 0, len(ranked)) @@ -344,9 +362,57 @@ func SearchMemory(db *sql.DB, query string, limit int, halfLifeWeeks float64) ([ returnedIDs = append(returnedIDs, item.ID) } if err := TouchObservations(db, returnedIDs); err != nil { - return nil, err + return nil, SearchMetrics{}, err + } + metrics := SearchMetrics{ + Query: query, + Channels: channelCounts, + CandidatesTotal: len(candidates), + ResultsReturned: len(ranked), + LimitRequested: requestedLimit, + ScoreMin: scoreMin(ranked), + ScoreMax: scoreMax(ranked), + ScoreMedian: scoreMedian(ranked), + } + return ranked, metrics, nil +} + +func countCandidateChannels(m map[int64]*candidate) map[string]int { + out := make(map[string]int) + for _, c := range m { + for ch := range c.Channels { + out[ch]++ + } + } + return out +} + +// ranked is sorted descending by CompositeScore; the three helpers below +// exploit that ordering so we avoid a re-sort just for metrics. + +func scoreMax(rs []SearchObservation) float64 { + if len(rs) == 0 { + return 0 + } + return rs[0].CompositeScore +} + +func scoreMin(rs []SearchObservation) float64 { + if len(rs) == 0 { + return 0 + } + return rs[len(rs)-1].CompositeScore +} + +func scoreMedian(rs []SearchObservation) float64 { + n := len(rs) + if n == 0 { + return 0 + } + if n%2 == 1 { + return rs[n/2].CompositeScore } - return ranked, nil + return (rs[n/2-1].CompositeScore + rs[n/2].CompositeScore) / 2 } func GroupResults(results []SearchObservation, compact bool) RecallResponse { diff --git a/internal/store/tools.go b/internal/store/tools.go index 5cb38bb..376d5b1 100644 --- a/internal/store/tools.go +++ b/internal/store/tools.go @@ -127,7 +127,24 @@ type GetEventObservationsToolResult struct { Total int `json:"total,omitempty"` } +// HandleTool dispatches a tool call and discards the search metrics (if any). +// Callers that want access to ranking-pipeline metrics for telemetry should +// use HandleToolWithMetrics instead. func HandleTool(defaultDB *sql.DB, name string, args ToolArgs) (any, error) { + result, _, err := HandleToolWithMetrics(defaultDB, name, args) + return result, err +} + +// HandleToolWithMetrics dispatches a tool call and returns ranking-pipeline +// metrics when applicable (currently: "recall"). For other tools the metrics +// return is nil. +func HandleToolWithMetrics(defaultDB *sql.DB, name string, args ToolArgs) (any, *SearchMetrics, error) { + var metrics *SearchMetrics + result, err := dispatchTool(defaultDB, name, args, &metrics) + return result, metrics, err +} + +func dispatchTool(defaultDB *sql.DB, name string, args ToolArgs, outMetrics **SearchMetrics) (any, error) { if err := validateToolArgs(name, args); err != nil { return nil, err } @@ -181,10 +198,14 @@ func HandleTool(defaultDB *sql.DB, name string, args ToolArgs) (any, error) { if args.Limit != nil { limit = *args.Limit } - results, err := SearchMemory(db, args.Query, limit, halfLife) + results, m, err := SearchMemory(db, args.Query, limit, halfLife) if err != nil { return nil, err } + m.Compact = args.Compact + if outMetrics != nil { + *outMetrics = &m + } return GroupResults(results, args.Compact), nil case "recall_entity": From e4cc72b42925f135be79a2380435eeb9a701b866 Mon Sep 17 00:00:00 2001 From: marlian <45340800+marlian@users.noreply.github.com> Date: Tue, 14 Apr 2026 22:36:42 +0200 Subject: [PATCH 03/11] feat: wire telemetry through dispatch with defer-based logging MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Entrypoint (cmd/workmem/main.go): reads MEMORY_TELEMETRY_PATH and MEMORY_TELEMETRY_PRIVACY via telemetry.FromEnv(), passes the *Client to mcpserver.Config. Deferred Close() on shutdown. When unset, FromEnv returns nil and no DB is ever opened. mcpserver dispatch (internal/mcpserver/server.go): handleTool now measures wall-clock duration and uses a single defer to log on every exit path (success, validation failure, dispatch error). The defer captures four observables as the flow progresses — argObject, toolResult, metrics, projectRaw, isError — and flushes them via r.logToolCall + r.logSearchMetrics at the end. mcpserver helpers (internal/mcpserver/telemetry.go): client detection from MCP initialize handshake (req.Session.InitializeParams().ClientInfo) with env fallback via telemetry.DetectClient. Resolves project path to absolute when project arg is present. Integration tests (3): enabled roundtrip asserts remember/recall/forget each land a tool_calls row with duration > 0 and no observation content leak, plus one search_metrics row linked to the recall tool_call. Disabled path asserts no telemetry file is created. Strict mode asserts Alice and sensitive-query text never appear in args_summary and search_metrics.query is sha256-hashed. --- cmd/workmem/main.go | 9 +- internal/mcpserver/server.go | 40 ++- internal/mcpserver/telemetry.go | 90 +++++++ .../mcpserver/telemetry_integration_test.go | 242 ++++++++++++++++++ 4 files changed, 375 insertions(+), 6 deletions(-) create mode 100644 internal/mcpserver/telemetry.go create mode 100644 internal/mcpserver/telemetry_integration_test.go diff --git a/cmd/workmem/main.go b/cmd/workmem/main.go index ae3d213..f67be1a 100644 --- a/cmd/workmem/main.go +++ b/cmd/workmem/main.go @@ -12,6 +12,7 @@ import ( "workmem/internal/dotenv" "workmem/internal/mcpserver" "workmem/internal/store" + "workmem/internal/telemetry" ) func main() { @@ -101,7 +102,13 @@ func runMCP(args []string) { loadEnvFile(*envFile) - runtime, err := mcpserver.New(mcpserver.Config{DBPath: *dbPath}) + tele := telemetry.FromEnv() + defer tele.Close() + + runtime, err := mcpserver.New(mcpserver.Config{ + DBPath: *dbPath, + Telemetry: tele, + }) if err != nil { fmt.Fprintf(os.Stderr, "start mcp server: %v\n", err) os.Exit(1) diff --git a/internal/mcpserver/server.go b/internal/mcpserver/server.go index ce9216b..f088eb3 100644 --- a/internal/mcpserver/server.go +++ b/internal/mcpserver/server.go @@ -10,22 +10,26 @@ import ( "os" "path/filepath" "strings" + "time" "github.com/modelcontextprotocol/go-sdk/mcp" "workmem/internal/store" + "workmem/internal/telemetry" ) const serverVersion = "0.1.0" type Config struct { - DBPath string + DBPath string + Telemetry *telemetry.Client } type Runtime struct { server *mcp.Server defaultDB *sql.DB dbPath string + telemetry *telemetry.Client } type toolDefinition struct { @@ -54,7 +58,7 @@ func New(config Config) (*Runtime, error) { Version: serverVersion, }, nil) - runtime := &Runtime{server: server, defaultDB: db, dbPath: dbPath} + runtime := &Runtime{server: server, defaultDB: db, dbPath: dbPath, telemetry: config.Telemetry} runtime.registerTools() return runtime, nil } @@ -99,21 +103,43 @@ func (r *Runtime) registerTools() { } func (r *Runtime) handleTool(_ context.Context, def toolDefinition, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + t0 := time.Now() + + // Telemetry observables — captured as the flow unfolds, flushed in defer so + // every return path (success, validation error, dispatch error) is logged. + var ( + argObject map[string]any + toolResult any + metrics *store.SearchMetrics + projectRaw string + isError bool + ) + defer func() { + id := r.logToolCall(def.Name, req, argObject, toolResult, projectRaw, isError, time.Since(t0)) + r.logSearchMetrics(id, metrics) + }() + raw := req.Params.Arguments if len(raw) == 0 { raw = []byte("{}") } - argObject, err := parseArgumentObject(raw) + var err error + argObject, err = parseArgumentObject(raw) if err != nil { + isError = true return errorResult(map[string]any{ "error": err.Error(), "tool": def.Name, }), nil } + if p, ok := argObject["project"].(string); ok { + projectRaw = p + } missing := missingRequiredArguments(def.Required, argObject) if len(missing) > 0 { + isError = true return errorResult(map[string]any{ "error": fmt.Sprintf("Missing required arguments: %s", strings.Join(missing, ", ")), "tool": def.Name, @@ -122,11 +148,13 @@ func (r *Runtime) handleTool(_ context.Context, def toolDefinition, req *mcp.Cal } if validation := validateStringArguments(def, argObject); validation != nil { + isError = true return errorResult(validation), nil } var args store.ToolArgs if err := json.Unmarshal(raw, &args); err != nil { + isError = true return errorResult(map[string]any{ "error": fmt.Sprintf("Invalid arguments: %v", err), "tool": def.Name, @@ -134,8 +162,10 @@ func (r *Runtime) handleTool(_ context.Context, def toolDefinition, req *mcp.Cal }), nil } - result, err := store.HandleTool(r.defaultDB, def.Name, args) + toolResult, metrics, err = store.HandleToolWithMetrics(r.defaultDB, def.Name, args) if err != nil { + isError = true + toolResult = nil return errorResult(map[string]any{ "error": err.Error(), "tool": def.Name, @@ -143,7 +173,7 @@ func (r *Runtime) handleTool(_ context.Context, def toolDefinition, req *mcp.Cal }), nil } - return successResult(result) + return successResult(toolResult) } func ResolveDBPath(configPath string) (string, error) { diff --git a/internal/mcpserver/telemetry.go b/internal/mcpserver/telemetry.go new file mode 100644 index 0000000..aeebbb4 --- /dev/null +++ b/internal/mcpserver/telemetry.go @@ -0,0 +1,90 @@ +package mcpserver + +import ( + "os" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" + + "workmem/internal/store" + "workmem/internal/telemetry" +) + +// detectClient reads clientInfo from the MCP initialize handshake when +// available, falling back to environment-variable detection otherwise. +func detectClient(req *mcp.CallToolRequest) telemetry.ClientInfo { + var name, version string + if req != nil && req.Session != nil { + if params := req.Session.InitializeParams(); params != nil && params.ClientInfo != nil { + name = params.ClientInfo.Name + version = params.ClientInfo.Version + } + } + return telemetry.DetectClient(name, version) +} + +// resolveProjectPath resolves a project argument as provided by the caller +// to its absolute form. Returns "" when project is empty or unresolvable — +// telemetry accepts empty, it simply records a global-scope call. +func resolveProjectPath(project string) string { + if project == "" { + return "" + } + home, err := os.UserHomeDir() + if err != nil { + return project + } + return store.ResolveProjectPath(project, home) +} + +// logToolCall inserts a tool_calls row and returns the insert id (or 0 when +// telemetry is disabled / logging failed). Safe to call when c is nil. +func (r *Runtime) logToolCall( + toolName string, + req *mcp.CallToolRequest, + argObject map[string]any, + result any, + projectRaw string, + isError bool, + elapsed time.Duration, +) int64 { + if r.telemetry == nil { + return 0 + } + dbScope := "global" + projectPath := "" + if projectRaw != "" { + dbScope = "project" + projectPath = resolveProjectPath(projectRaw) + } + return r.telemetry.LogToolCall(telemetry.ToolCallInput{ + Tool: toolName, + Client: detectClient(req), + DBScope: dbScope, + ProjectPath: projectPath, + DurationMs: float64(elapsed) / float64(time.Millisecond), + ArgsSummary: telemetry.SanitizeArgs(argObject, r.telemetry.Strict()), + ResultSummary: telemetry.SummarizeResult(result), + IsError: isError, + }) +} + +// logSearchMetrics mirrors the recall search_metrics row. No-op when telemetry +// is disabled or the parent tool_call id is 0. +func (r *Runtime) logSearchMetrics(toolCallID int64, m *store.SearchMetrics) { + if r.telemetry == nil || toolCallID == 0 || m == nil { + return + } + r.telemetry.LogSearchMetrics(telemetry.SearchMetricsInput{ + ToolCallID: toolCallID, + Query: m.Query, + Channels: m.Channels, + CandidatesTotal: m.CandidatesTotal, + ResultsReturned: m.ResultsReturned, + LimitRequested: m.LimitRequested, + ScoreMin: m.ScoreMin, + ScoreMax: m.ScoreMax, + ScoreMedian: m.ScoreMedian, + Compact: m.Compact, + }) +} diff --git a/internal/mcpserver/telemetry_integration_test.go b/internal/mcpserver/telemetry_integration_test.go new file mode 100644 index 0000000..7d78cae --- /dev/null +++ b/internal/mcpserver/telemetry_integration_test.go @@ -0,0 +1,242 @@ +package mcpserver + +import ( + "context" + "database/sql" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" + + _ "modernc.org/sqlite" + + "workmem/internal/telemetry" +) + +func TestTelemetryEnabledRoundtripLogsToolCallsAndSearchMetrics(t *testing.T) { + telePath := filepath.Join(t.TempDir(), "telemetry.db") + tele := telemetry.InitIfEnabled(telePath, false) + if tele == nil { + t.Fatalf("telemetry InitIfEnabled returned nil on valid path") + } + + runtime, err := New(Config{DBPath: filepath.Join(t.TempDir(), "memory.db"), Telemetry: tele}) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + session, stop := startTelemetrySession(t, runtime) + defer stop() + ctx := context.Background() + + callOK(t, session, ctx, "remember", map[string]any{ + "entity": "TelemetryEntity", + "observation": "sensitive observation content must never leak", + }) + callOK(t, session, ctx, "recall", map[string]any{ + "query": "TelemetryEntity", + "limit": 5, + }) + callOK(t, session, ctx, "forget", map[string]any{ + "entity": "TelemetryEntity", + }) + + if err := tele.Close(); err != nil { + t.Fatalf("tele.Close() error = %v", err) + } + + rdb, err := sql.Open("sqlite", telePath) + if err != nil { + t.Fatalf("open telemetry db for readback: %v", err) + } + defer rdb.Close() + + var callCount int + if err := rdb.QueryRow(`SELECT COUNT(*) FROM tool_calls`).Scan(&callCount); err != nil { + t.Fatalf("count tool_calls: %v", err) + } + if callCount != 3 { + t.Fatalf("tool_calls count = %d, want 3", callCount) + } + + rows, err := rdb.Query(`SELECT tool, args_summary, result_summary, duration_ms, is_error FROM tool_calls ORDER BY id`) + if err != nil { + t.Fatalf("query tool_calls: %v", err) + } + defer rows.Close() + var tools []string + for rows.Next() { + var tool string + var argsSummary, resultSummary sql.NullString + var dur float64 + var isErr int + if err := rows.Scan(&tool, &argsSummary, &resultSummary, &dur, &isErr); err != nil { + t.Fatalf("scan tool_calls row: %v", err) + } + tools = append(tools, tool) + if strings.Contains(argsSummary.String, "sensitive observation content") { + t.Fatalf("args_summary leaked observation content: %s", argsSummary.String) + } + if dur <= 0 { + t.Fatalf("duration_ms = %v for %s, want > 0", dur, tool) + } + if isErr != 0 { + t.Fatalf("is_error = %d for %s, want 0", isErr, tool) + } + } + expected := []string{"remember", "recall", "forget"} + for i, want := range expected { + if tools[i] != want { + t.Fatalf("tool[%d] = %q, want %q", i, tools[i], want) + } + } + + var smCount int + if err := rdb.QueryRow(`SELECT COUNT(*) FROM search_metrics`).Scan(&smCount); err != nil { + t.Fatalf("count search_metrics: %v", err) + } + if smCount != 1 { + t.Fatalf("search_metrics count = %d, want 1", smCount) + } + + var smQuery sql.NullString + var candidatesTotal, resultsReturned int + if err := rdb.QueryRow(`SELECT query, candidates_total, results_returned FROM search_metrics`).Scan(&smQuery, &candidatesTotal, &resultsReturned); err != nil { + t.Fatalf("readback search_metrics: %v", err) + } + if smQuery.String != "TelemetryEntity" { + t.Fatalf("search_metrics.query = %q, want TelemetryEntity", smQuery.String) + } + if candidatesTotal == 0 { + t.Fatalf("search_metrics.candidates_total = 0, want > 0") + } +} + +func TestTelemetryDisabledCreatesNoFile(t *testing.T) { + telePath := filepath.Join(t.TempDir(), "telemetry-must-not-exist.db") + + runtime, err := New(Config{DBPath: filepath.Join(t.TempDir(), "memory.db"), Telemetry: nil}) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + session, stop := startTelemetrySession(t, runtime) + defer stop() + ctx := context.Background() + + callOK(t, session, ctx, "remember", map[string]any{ + "entity": "NoTelemetryEntity", + "observation": "this must dispatch without creating any telemetry file", + }) + callOK(t, session, ctx, "recall", map[string]any{"query": "NoTelemetry", "limit": 3}) + + if _, err := os.Stat(telePath); err == nil { + t.Fatalf("telemetry file %q must not exist when Telemetry is nil", telePath) + } +} + +func TestTelemetryStrictModeHashesIdentifiersEndToEnd(t *testing.T) { + telePath := filepath.Join(t.TempDir(), "strict-telemetry.db") + tele := telemetry.InitIfEnabled(telePath, true) + if tele == nil { + t.Fatalf("telemetry InitIfEnabled(strict) returned nil") + } + + runtime, err := New(Config{DBPath: filepath.Join(t.TempDir(), "memory.db"), Telemetry: tele}) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + session, stop := startTelemetrySession(t, runtime) + defer stop() + ctx := context.Background() + + callOK(t, session, ctx, "remember", map[string]any{ + "entity": "Alice", + "observation": "confidential note attached to Alice", + }) + callOK(t, session, ctx, "recall", map[string]any{ + "query": "sensitive question about therapy", + "limit": 5, + }) + + if err := tele.Close(); err != nil { + t.Fatalf("tele.Close() error = %v", err) + } + + rdb, err := sql.Open("sqlite", telePath) + if err != nil { + t.Fatalf("open strict telemetry db: %v", err) + } + defer rdb.Close() + + rows, err := rdb.Query(`SELECT args_summary FROM tool_calls`) + if err != nil { + t.Fatalf("query tool_calls: %v", err) + } + defer rows.Close() + for rows.Next() { + var args sql.NullString + if err := rows.Scan(&args); err != nil { + t.Fatalf("scan: %v", err) + } + for _, leak := range []string{"Alice", "sensitive question", "therapy"} { + if strings.Contains(args.String, leak) { + t.Fatalf("strict args_summary leaked %q: %s", leak, args.String) + } + } + } + + var smQuery sql.NullString + if err := rdb.QueryRow(`SELECT query FROM search_metrics`).Scan(&smQuery); err != nil { + t.Fatalf("read search_metrics.query: %v", err) + } + if strings.Contains(smQuery.String, "sensitive") || strings.Contains(smQuery.String, "therapy") { + t.Fatalf("strict search_metrics.query leaked plaintext: %q", smQuery.String) + } + if !strings.HasPrefix(smQuery.String, "sha256:") { + t.Fatalf("strict search_metrics.query not hashed: %q", smQuery.String) + } +} + +func callOK(t *testing.T, session *mcp.ClientSession, ctx context.Context, name string, args map[string]any) { + t.Helper() + result, err := session.CallTool(ctx, &mcp.CallToolParams{Name: name, Arguments: args}) + if err != nil { + t.Fatalf("CallTool(%s) error = %v", name, err) + } + if result.IsError { + t.Fatalf("CallTool(%s) returned tool error: %+v", name, result) + } +} + +// startTelemetrySession spins up an in-memory MCP client+server pair against +// the given runtime and returns a session plus a stop function that unwinds +// the session and the server goroutine cleanly. +func startTelemetrySession(t *testing.T, runtime *Runtime) (*mcp.ClientSession, func()) { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + serverTransport, clientTransport := mcp.NewInMemoryTransports() + errCh := make(chan error, 1) + go func() { errCh <- runtime.Run(ctx, serverTransport) }() + + client := mcp.NewClient(&mcp.Implementation{Name: "telemetry-test-client", Version: "1.0.0"}, nil) + session, err := client.Connect(ctx, clientTransport, nil) + if err != nil { + cancel() + t.Fatalf("client.Connect() error = %v", err) + } + stop := func() { + _ = session.Close() + cancel() + select { + case <-errCh: + case <-time.After(2 * time.Second): + t.Fatal("runtime.Run() did not exit after client shutdown") + } + } + return session, stop +} From 32a62b6abdbf3c5223694f4a709b9b33c1222127 Mon Sep 17 00:00:00 2001 From: marlian <45340800+marlian@users.noreply.github.com> Date: Tue, 14 Apr 2026 22:37:04 +0200 Subject: [PATCH 04/11] docs: telemetry guide, invariants, Step 3.4, decision entry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - docs/TELEMETRY.md: full user-facing guide. Adapted from the Node telemetry.md, with privacy-strict mode documented (trade-off: ranking debug vs plaintext leak risk on sensitive backends). Includes example client config for Claude Code / governor env-file wiring. - OPERATIONS.md invariants: telemetry is opt-in and never affects the success path; DB is physically separate from memory DB; strict mode sha256-hashes identifiers before disk. P1 "telemetry deferred" entry removed (condition met). P2 telemetry-schema entry removed (designed and implemented). Pre-Launch TODO no longer lists Kilo proof — already closed in PR #4. - IMPLEMENTATION.md: new Step 3.4 Telemetry marked done with explicit Gate — enabled writes rows, disabled creates no DB, strict hashes. - DECISION_LOG.md: append entry documenting the three Go-native refinements (nil-tolerant Client, no globals, SearchMetrics tuple) and the new privacy-strict mode. Explicitly rejects at-rest encryption with keychain for this iteration as over-scope. --- DECISION_LOG.md | 29 +++++++ IMPLEMENTATION.md | 17 ++++- OPERATIONS.md | 16 +--- docs/TELEMETRY.md | 190 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 239 insertions(+), 13 deletions(-) create mode 100644 docs/TELEMETRY.md diff --git a/DECISION_LOG.md b/DECISION_LOG.md index a01980d..cb8b9cb 100644 --- a/DECISION_LOG.md +++ b/DECISION_LOG.md @@ -26,6 +26,35 @@ Add a `workmem backup` subcommand that writes an age-encrypted snapshot of the g - **Include project-scoped DBs automatically.** Rejected: project DBs belong to workspaces, not to the user's top-level knowledge. Auto-including them couples backup to filesystem scanning and makes the unit of restore ambiguous. A `backup` invocation per workspace is explicit. - **Include telemetry.db in the snapshot.** Rejected: telemetry is operational, rebuildable, and has a different lifecycle than knowledge. Mixing them also risks leaking telemetry via recall if paths cross. +## 2026-04-14: Port telemetry with Go-native refinements and add privacy-strict mode + +### Context + +The Node reference implementation shipped with opt-in telemetry in a separate SQLite DB (schemas: `tool_calls`, `search_metrics`). Phase 3 deferred the port until the Go MCP entrypoint was real and adopted. That condition is now met: the Go binary serves Claude Code, Kilo, and Codex in production. Time to port. + +### Decision + +Port the Node telemetry design to Go and preserve the guiding principles (opt-in via env, separate database, counts-only for results, content replaced with ``). Refine three Node-era shortcuts: + +1. **Nil-tolerant `*Client`** — the client value is `nil` when disabled, every method returns immediately on `nil` receiver. Replaces per-callsite `if TELEMETRY_ENABLED` checks. +2. **No globals** — the client is constructed in `cmd/workmem/main.go` and plumbed via `mcpserver.Config{Telemetry: …}`. Replaces the Node pattern of module-level mutable state (`_telemetryDb`, `_lastSearchMetrics`, etc.). +3. **`SearchMemory` returns `SearchMetrics` as a tuple** — `(results []SearchObservation, metrics SearchMetrics, err error)`. Replaces the Node `_lastSearchMetrics` side-channel. + +Add a new **privacy-strict mode** (`MEMORY_TELEMETRY_PRIVACY=strict`): entity names, queries, and event labels are sha256-hashed before storage. Intended for sensitive backends (e.g., the `private_memory` server backing therapy/health/relationship content). + +### Rationale + +- Node-era globals would have been awkward in Go and hard to test under parallel `t.Run` — eliminating them keeps the test story clean. +- Privacy-strict closes a real threat: local plaintext telemetry DB on a laptop with sensitive entity names is a leak vector if the laptop is lost/sync'd/exported. Strict mode lets one binary serve two wiring contexts (permissive `memory`, strict `private_memory`) cleanly. +- `SearchMemory` returning metrics as a proper value is idiomatic Go and testable in isolation without the telemetry package. +- Using `modernc.org/sqlite` for the telemetry DB keeps the pure-Go single-binary invariant (no CGO addition just for the analytics path). + +### Alternatives considered + +- **1:1 port with globals** — Rejected because Go's `database/sql` + `sql.Stmt` lifecycle around a package-level mutable pointer becomes painful under test; the nil-client pattern is simpler and safer. +- **Attach telemetry as an MCP tool** — Rejected as in the Node design: telemetry is human-developer infrastructure, not a model capability. Adding a tool wastes context tokens on every call for every client. +- **Encryption at rest on the telemetry DB with a keychain-stored key** — Rejected for this iteration. Cross-platform keychain integration (macOS/Windows/Linux headless) is a bigger cantiere than hashing the sensitive fields. Revisit if the strict mode proves insufficient in practice. + ## 2026-04-14: Use the official Go MCP SDK for transport ### Context diff --git a/IMPLEMENTATION.md b/IMPLEMENTATION.md index ae17b82..8a634da 100644 --- a/IMPLEMENTATION.md +++ b/IMPLEMENTATION.md @@ -108,4 +108,19 @@ Ship a `workmem backup` subcommand that produces an age-encrypted snapshot of th - [x] Wire `backup` subcommand in `cmd/workmem/main.go` with `--to`, `--age-recipient` (repeatable), `--db`, `--env-file` - [x] README section documenting usage and manual `age -d` restore -**On Step Gate (all items [x]):** trigger correctness review focused on crypto wiring and VACUUM INTO error paths. \ No newline at end of file +**On Step Gate (all items [x]):** trigger correctness review focused on crypto wiring and VACUUM INTO error paths. + +### Step 3.5: Telemetry [✅] + +Port the Node telemetry design to Go with Go-native refinements and a new privacy-strict mode. **Gate:** when `MEMORY_TELEMETRY_PATH` is set, every tool call lands a row in `tool_calls`; every `recall` lands a row in `search_metrics` linked by `tool_call_id`; when unset, no DB is created and no overhead is added. In `MEMORY_TELEMETRY_PRIVACY=strict` mode, entity/query/label values are sha256-hashed before storage. + +- [x] Build `internal/telemetry` package (nil-tolerant Client, schema, sanitize, hash, detect) +- [x] Refactor `SearchMemory` to return `(results, metrics, err)` — no globals, no side channels +- [x] Wire `*telemetry.Client` through `cmd/workmem/main.go` and `mcpserver.Config` +- [x] Wrap `mcpserver` dispatch with duration + args/result sanitization + LogToolCall/LogSearchMetrics +- [x] Unit tests for package (nil-client safety, init failure, strict hashing, sanitize, detect) +- [x] Integration tests: enabled roundtrip / disabled zero overhead / privacy-strict +- [x] `docs/TELEMETRY.md` adapted for Go with privacy-strict documented +- [x] Telemetry invariants wired into `OPERATIONS.md` + +**On Step Gate (all items [x]):** trigger correctness review on telemetry hook points and strict-mode hashing. diff --git a/OPERATIONS.md b/OPERATIONS.md index 0eae5b6..5202444 100644 --- a/OPERATIONS.md +++ b/OPERATIONS.md @@ -11,6 +11,9 @@ - Live-data queries must never bypass tombstone guards. - FTS cleanup must never use raw `DELETE` against a contentless FTS table. - `remember_event` must be atomic: the event row and all attached observations commit together or not at all. Proof: `TestRememberEventAtomicityOnMidLoopFailure` in `internal/store/parity_test.go`. +- Telemetry is opt-in (`MEMORY_TELEMETRY_PATH`) and never affects the tool call success path. Init failure logs a single warning to stderr and disables telemetry for the session; the main memory DB is unaffected. +- Telemetry data lives in its own SQLite file, physically separate from the memory database. No foreign keys, no joins, no shared lifecycle. +- When `MEMORY_TELEMETRY_PRIVACY=strict`, entity names, queries, and event labels must be sha256-hashed before reaching disk. Observation/content values are always reduced to `` regardless of mode. ## Active Debt @@ -32,12 +35,6 @@ Blast radius: Late failures on Linux or Windows packaging, or FTS behavior drift Fix: Keep the canary in CI and run it on at least macOS, Linux, and Windows before calling the persistence layer portable. Done when: the same schema/FTS canary passes in cross-build validation. -- Telemetry parity is consciously deferred until the new Go MCP entrypoint is wired into a real client and the request path is considered stable. -Trigger: Instrumenting before the client-facing transport contract has been debugged in practice. -Blast radius: Busywork telemetry code tied to temporary wiring. -Fix: keep telemetry scope documented; implement it once the Kilo-facing transport path is stable. -Done when: the Go MCP entrypoint is live under a real client and tool-call telemetry can be attached once, not retrofitted twice. - - FTS5 viability is proven locally on the chosen driver, but not yet in a cross-platform validation matrix. Trigger: Assuming a passing local canary implies release-target portability. Blast radius: Search or forget semantics break only after packaging or OS expansion. @@ -54,15 +51,10 @@ Done when: FTS-specific parity tests pass across the release matrix. ### P2 -- Telemetry schema and migration strategy are not yet designed. -Trigger: Reaching post-parity milestone without an observability plan. -Blast radius: Delayed adoption of telemetry in the Go port. -Fix: Define minimal telemetry compatibility after core parity lands. -Done when: telemetry design is recorded and scheduled. +- None active. ## Pre-Launch TODO -- Prove MCP stdio compatibility with Kilo or another real client. - Prove schema initialization and migrations on clean and upgraded DBs. - Prove forget semantics including FTS deletion. - Prove project isolation. diff --git a/docs/TELEMETRY.md b/docs/TELEMETRY.md new file mode 100644 index 0000000..88f71b6 --- /dev/null +++ b/docs/TELEMETRY.md @@ -0,0 +1,190 @@ +# Telemetry + +> Opt-in usage analytics for workmem. Zero overhead when disabled. Separate database, strict privacy controls. + +## Enabling + +Set `MEMORY_TELEMETRY_PATH` to a file path. The telemetry database is created on first tool call. + +**Via `.env`:** +```bash +MEMORY_TELEMETRY_PATH=./telemetry.db +``` + +**Via client config (example for Claude Code's `~/.claude.json`):** +```json +{ + "memory": { + "command": "/path/to/workmem", + "args": ["-env-file", "/path/to/memory.env"], + "env": { + "MEMORY_TELEMETRY_PATH": "/absolute/path/to/telemetry.db" + } + } +} +``` + +When `MEMORY_TELEMETRY_PATH` is unset (the default), every telemetry call is a no-op — no timing, no logging, no database opened. The `*telemetry.Client` pointer in the runtime is `nil` and every method returns immediately. + +## Privacy modes + +Telemetry supports two modes, controlled by `MEMORY_TELEMETRY_PRIVACY`: + +| Value | Mode | Behavior | +|-------|------|----------| +| (unset) / any other value | **permissive** (default) | entity names, queries, and event labels are stored in plaintext | +| `strict` | **strict** | entity names, queries, and event labels are sha256-hashed before storage | + +Strict mode is intended for sensitive instances such as a `private_memory` server backing therapy/health/relationship content. Ranking debug ("which queries overfetch candidates?") becomes harder in strict mode because plaintext queries are no longer recoverable — but sensitive identifiers never land on disk. + +Observation/content values are **always** reduced to `` regardless of mode. Facts/observations arrays are **always** reduced to ``. Strict mode only changes what happens to identifier-like fields. + +**Example `.env` for sensitive backend:** +```bash +MEMORY_TELEMETRY_PATH=/home/user/.local/state/workmem/private-telemetry.db +MEMORY_TELEMETRY_PRIVACY=strict +``` + +## What it logs + +### Tool calls (`tool_calls` table) + +Every MCP tool invocation is logged with: + +| Column | Example | +|--------|---------| +| `ts` | `2026-04-14T20:15:32.456` | +| `tool` | `recall`, `remember`, `forget` | +| `client_name` | `kilo`, `claude-code`, `cursor`, `windsurf`, `vscode-copilot` | +| `client_version` | `0.43.6` | +| `client_source` | `protocol` / `env` / `none` | +| `db_scope` | `global` / `project` | +| `project_path` | resolved absolute path, or null | +| `duration_ms` | `12.4` | +| `args_summary` | Sanitized JSON (see below) | +| `result_summary` | Counts only, never data | +| `is_error` | `0` or `1` | + +### Search ranking metrics (`search_metrics` table) + +For `recall` calls, additional metrics capture the ranking pipeline: + +| Column | Example | +|--------|---------| +| `tool_call_id` | FK into `tool_calls.id` | +| `query` | Search text (hashed in strict mode) | +| `channels` | `{"fts": 12, "fts_phrase": 3, "entity_exact": 1}` | +| `candidates_total` | `16` | +| `results_returned` | `5` | +| `limit_requested` | `20` | +| `score_min` | `0.32` | +| `score_max` | `0.87` | +| `score_median` | `0.61` | +| `compact` | `0` or `1` | + +## What it does NOT log + +- Observation content (replaced with ``, always) +- Full result payloads — only counts (entities returned, observations stored, etc.) +- In strict mode, any identifier (entity name, query, event label, from/to) + +## Client identity + +The server identifies which client is calling through two mechanisms: + +1. **MCP protocol** (primary) — the `initialize` handshake includes `clientInfo.name` and `clientInfo.version`. This is a required field in the MCP spec. +2. **Environment fingerprinting** (fallback) — when the protocol doesn't provide client info, the server detects the client from environment variables: + +| Client | Signal | +|--------|--------| +| Kilo | `KILO=1` (version from `KILOCODE_VERSION`) | +| Claude Code | `CLAUDE_CODE_SSE_PORT` set | +| Cursor | `CURSOR_TRACE_ID` set | +| Windsurf | `WINDSURF_EXTENSION_ID` set | +| VS Code Copilot | `VSCODE_MCP_HTTP_PREFER` set non-empty | +| VS Code (unknown extension) | `TERM_PROGRAM=vscode` | + +The `client_source` column tells you which mechanism fired: `protocol`, `env`, or `none`. + +## Querying the data + +The telemetry database is a standard SQLite file. Open it with any tool: `sqlite3`, DBeaver, Jupyter, pandas, etc. + +### Example queries + +**Tool usage by client:** +```sql +SELECT client_name, tool, COUNT(*) as calls, + ROUND(AVG(duration_ms), 1) as avg_ms +FROM tool_calls +GROUP BY client_name, tool +ORDER BY calls DESC; +``` + +**Search ranking quality (permissive mode — query is plaintext):** +```sql +SELECT query, candidates_total, results_returned, + ROUND(score_min, 3) as min, ROUND(score_max, 3) as max, + channels +FROM search_metrics +ORDER BY candidates_total DESC +LIMIT 20; +``` + +**Overfetch detection (candidates >> returned):** +```sql +SELECT query, candidates_total, results_returned, limit_requested, + ROUND(1.0 * results_returned / candidates_total, 2) as yield_ratio +FROM search_metrics +WHERE candidates_total > 0 +ORDER BY yield_ratio ASC +LIMIT 20; +``` + +**Error rate by tool:** +```sql +SELECT tool, COUNT(*) as total, + SUM(is_error) as errors, + ROUND(100.0 * SUM(is_error) / COUNT(*), 1) as error_pct +FROM tool_calls +GROUP BY tool +ORDER BY error_pct DESC; +``` + +**Channel effectiveness:** +```sql +SELECT json_each.key as channel, COUNT(*) as appearances +FROM search_metrics, json_each(search_metrics.channels) +GROUP BY channel +ORDER BY appearances DESC; +``` + +## Separate database + +Telemetry is stored in its own SQLite file, completely separate from `memory.db`. This means: + +- Deleting the telemetry DB has zero impact on your knowledge graph +- The telemetry DB can be wiped and recreated at any time +- No foreign keys or joins between telemetry and memory data +- Telemetry uses `journal_mode=WAL` for concurrent reads while the server writes + +## Init failure handling + +If the telemetry path is invalid or the database can't be opened, the server prints a single warning to stderr and disables telemetry for the rest of the session. It does not retry on every call. The main `memory.db` is unaffected — telemetry failure never breaks the tool call path. + +Example warning: +``` +[memory] telemetry init failed (disabled for this session): unable to open database file +``` + +## Design rationale + +**Why plaintext queries in permissive mode?** Local, single-user development: the telemetry DB lives on your machine, is gitignored, and is only readable by you. Redacting queries permissively would make the analytics useless — you can't answer "which queries produce too many candidates?" if the query text is hashed. + +**Why privacy-strict mode?** For backends holding sensitive content (therapy, health, relationships, personal journaling), even local plaintext can matter: laptop loss, accidental sync-folder placement, or exported snapshots. Strict mode ensures entity names and queries never land on disk in the clear. + +**Why a separate SQLite and not the memory DB?** Separation of concerns. The memory DB is your knowledge graph. Telemetry is operational data with a different lifecycle (wipe freely, aggregate, export). Mixing them risks accidentally leaking telemetry via `recall`, or losing telemetry history when the memory DB is rebuilt. + +**Why not an MCP tool?** Telemetry is infrastructure, not a capability the model needs. Adding a tool would cost context tokens on every call for something only the human developer uses. Query the SQLite directly. + +**Why a nil-tolerant `*Client`?** The alternative is `if TELEMETRY_ENABLED` checks sprinkled across every call site. The `nil`-receiver pattern keeps the dispatch code clean — the wrapper always calls `LogToolCall`; when telemetry is disabled, the client is `nil` and the method returns immediately. From 914a6f0502e267b5db5d2f73922f2efb1a4f64b8 Mon Sep 17 00:00:00 2001 From: marlian <45340800+marlian@users.noreply.github.com> Date: Tue, 14 Apr 2026 23:05:41 +0200 Subject: [PATCH 05/11] review: address Copilot review feedback on PR #5 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nine comments from the Copilot auto-reviewer, all verified against the code. Each addressed here: 1. mcpserver/telemetry.go resolveProjectPath: return "" on home-dir error instead of the raw project argument, matching the docstring and avoiding non-canonical paths leaking into telemetry rows. 2. telemetry/client_test.go TestInitIfEnabledInvalidPathReturnsNil: build the invalid path inside t.TempDir() + missing-subdir so the failure mode is identical on macOS, Linux, and Windows (the old "/nonexistent-…" was Unix-flavored). 3. mcpserver/server.go Runtime: close telemetry in Runtime.Close() (nil-safe, idempotent) and document ownership transfer in Config. Removed the now-redundant defer tele.Close() from cmd/workmem/main.go to avoid double-close semantics. 4. mcpserver/server.go handleTool: gate the time.Now() / defer block behind `if r.telemetry != nil` so the disabled path is literally a single pointer-nil check per call, matching the "zero overhead" promise. 5. docs/TELEMETRY.md: correct the lifecycle statement — the telemetry DB is opened at process startup (via FromEnv), not on first tool call. 6. docs/TELEMETRY.md: reword the disabled-path description to reflect the fast-path change from fix 4 (no timing, no logging, no DB opened; one pointer-nil check on the hot path). 7. telemetry/client.go InitIfEnabled: align pragmas with the main memory DB — enable foreign_keys so search_metrics -> tool_calls FK is enforced, set busy_timeout=5000 so brief lock contention retries instead of erroring, and SetMaxOpenConns(1) for deterministic write ordering. 8. mcpserver/telemetry_integration_test.go TestTelemetryDisabledViaFromEnvCreatesNoArtifacts: replaces the previous test that only asserted an arbitrary path did not exist (a vacuous claim). The new test drives the actual env -> FromEnv -> runtime path with MEMORY_TELEMETRY_PATH unset, and asserts that no telemetry artifact appears anywhere in the data dir after a full remember/recall cycle. 9. telemetry/sanitize.go SanitizeArgs: count runes instead of bytes when rendering "" for redacted string fields, so the marker is accurate for non-ASCII content. --- cmd/workmem/main.go | 7 ++-- docs/TELEMETRY.md | 4 +-- internal/mcpserver/server.go | 29 ++++++++++----- internal/mcpserver/telemetry.go | 7 ++-- .../mcpserver/telemetry_integration_test.go | 35 ++++++++++++++++--- internal/telemetry/client.go | 20 ++++++++--- internal/telemetry/client_test.go | 7 ++-- internal/telemetry/sanitize.go | 3 +- 8 files changed, 81 insertions(+), 31 deletions(-) diff --git a/cmd/workmem/main.go b/cmd/workmem/main.go index f67be1a..962e8da 100644 --- a/cmd/workmem/main.go +++ b/cmd/workmem/main.go @@ -102,12 +102,11 @@ func runMCP(args []string) { loadEnvFile(*envFile) - tele := telemetry.FromEnv() - defer tele.Close() - + // Ownership of the telemetry client transfers to the Runtime — its Close() + // handles the shutdown. No defer here on purpose (see mcpserver.Config). runtime, err := mcpserver.New(mcpserver.Config{ DBPath: *dbPath, - Telemetry: tele, + Telemetry: telemetry.FromEnv(), }) if err != nil { fmt.Fprintf(os.Stderr, "start mcp server: %v\n", err) diff --git a/docs/TELEMETRY.md b/docs/TELEMETRY.md index 88f71b6..e3972ee 100644 --- a/docs/TELEMETRY.md +++ b/docs/TELEMETRY.md @@ -4,7 +4,7 @@ ## Enabling -Set `MEMORY_TELEMETRY_PATH` to a file path. The telemetry database is created on first tool call. +Set `MEMORY_TELEMETRY_PATH` to a file path. When telemetry is enabled, the database is opened and the schema is initialized at process startup (before any tool call). If initialization fails, a single warning is printed to stderr and telemetry is disabled for the rest of the session — the main memory path is never affected. **Via `.env`:** ```bash @@ -24,7 +24,7 @@ MEMORY_TELEMETRY_PATH=./telemetry.db } ``` -When `MEMORY_TELEMETRY_PATH` is unset (the default), every telemetry call is a no-op — no timing, no logging, no database opened. The `*telemetry.Client` pointer in the runtime is `nil` and every method returns immediately. +When `MEMORY_TELEMETRY_PATH` is unset (the default), telemetry is disabled end-to-end: no database is opened, no schema is written, no rows are inserted. The dispatch wrapper skips its timing block entirely when the telemetry client is `nil`, so the only cost on the no-telemetry path is a single pointer-nil check per tool call. ## Privacy modes diff --git a/internal/mcpserver/server.go b/internal/mcpserver/server.go index f088eb3..f3e1ac2 100644 --- a/internal/mcpserver/server.go +++ b/internal/mcpserver/server.go @@ -20,6 +20,10 @@ import ( const serverVersion = "0.1.0" +// Config carries the construction parameters for a Runtime. Ownership of the +// Telemetry client transfers to the Runtime: Runtime.Close() will close it. +// Callers must not call Close on the client themselves once it has been +// handed to New. type Config struct { DBPath string Telemetry *telemetry.Client @@ -82,7 +86,11 @@ func (r *Runtime) Close() error { closeErr = r.defaultDB.Close() r.defaultDB = nil } - return errors.Join(closeErr, store.ResetProjectDBs()) + // Close telemetry on the way down. *Client.Close is nil-safe and + // idempotent, so double-close from a defensive caller is harmless. + teleErr := r.telemetry.Close() + r.telemetry = nil + return errors.Join(closeErr, teleErr, store.ResetProjectDBs()) } func (r *Runtime) DBPath() string { @@ -103,10 +111,10 @@ func (r *Runtime) registerTools() { } func (r *Runtime) handleTool(_ context.Context, def toolDefinition, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { - t0 := time.Now() - - // Telemetry observables — captured as the flow unfolds, flushed in defer so - // every return path (success, validation error, dispatch error) is logged. + // Telemetry observables — captured as the flow unfolds. When telemetry is + // disabled (r.telemetry == nil) we skip both the time.Now/time.Since + // measurement and the defer installation entirely: the no-telemetry path + // adds nothing beyond a single pointer-nil check. var ( argObject map[string]any toolResult any @@ -114,10 +122,13 @@ func (r *Runtime) handleTool(_ context.Context, def toolDefinition, req *mcp.Cal projectRaw string isError bool ) - defer func() { - id := r.logToolCall(def.Name, req, argObject, toolResult, projectRaw, isError, time.Since(t0)) - r.logSearchMetrics(id, metrics) - }() + if r.telemetry != nil { + t0 := time.Now() + defer func() { + id := r.logToolCall(def.Name, req, argObject, toolResult, projectRaw, isError, time.Since(t0)) + r.logSearchMetrics(id, metrics) + }() + } raw := req.Params.Arguments if len(raw) == 0 { diff --git a/internal/mcpserver/telemetry.go b/internal/mcpserver/telemetry.go index aeebbb4..a14440b 100644 --- a/internal/mcpserver/telemetry.go +++ b/internal/mcpserver/telemetry.go @@ -24,15 +24,16 @@ func detectClient(req *mcp.CallToolRequest) telemetry.ClientInfo { } // resolveProjectPath resolves a project argument as provided by the caller -// to its absolute form. Returns "" when project is empty or unresolvable — -// telemetry accepts empty, it simply records a global-scope call. +// to its absolute form. Returns "" when project is empty or unresolvable +// (including the rare case where os.UserHomeDir fails) — telemetry accepts +// empty and simply records a global-scope call without a path. func resolveProjectPath(project string) string { if project == "" { return "" } home, err := os.UserHomeDir() if err != nil { - return project + return "" } return store.ResolveProjectPath(project, home) } diff --git a/internal/mcpserver/telemetry_integration_test.go b/internal/mcpserver/telemetry_integration_test.go index 7d78cae..5fb65bc 100644 --- a/internal/mcpserver/telemetry_integration_test.go +++ b/internal/mcpserver/telemetry_integration_test.go @@ -115,10 +115,22 @@ func TestTelemetryEnabledRoundtripLogsToolCallsAndSearchMetrics(t *testing.T) { } } -func TestTelemetryDisabledCreatesNoFile(t *testing.T) { - telePath := filepath.Join(t.TempDir(), "telemetry-must-not-exist.db") +func TestTelemetryDisabledViaFromEnvCreatesNoArtifacts(t *testing.T) { + // Drive the full env → FromEnv → runtime → dispatch path with telemetry + // explicitly disabled, and verify that no telemetry artifact appears + // anywhere in the runtime's data directory after a complete tool call + // cycle. This proves the disabled path is silent end-to-end, not just + // that a hand-picked path happens not to exist. + t.Setenv("MEMORY_TELEMETRY_PATH", "") + t.Setenv("MEMORY_TELEMETRY_PRIVACY", "") + tele := telemetry.FromEnv() + if tele != nil { + t.Fatalf("FromEnv with MEMORY_TELEMETRY_PATH unset must return nil, got %+v", tele) + } - runtime, err := New(Config{DBPath: filepath.Join(t.TempDir(), "memory.db"), Telemetry: nil}) + dir := t.TempDir() + memDB := filepath.Join(dir, "memory.db") + runtime, err := New(Config{DBPath: memDB, Telemetry: tele}) if err != nil { t.Fatalf("New() error = %v", err) } @@ -133,8 +145,21 @@ func TestTelemetryDisabledCreatesNoFile(t *testing.T) { }) callOK(t, session, ctx, "recall", map[string]any{"query": "NoTelemetry", "limit": 3}) - if _, err := os.Stat(telePath); err == nil { - t.Fatalf("telemetry file %q must not exist when Telemetry is nil", telePath) + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatalf("readdir %q: %v", dir, err) + } + for _, e := range entries { + if e.IsDir() { + continue + } + name := e.Name() + // memory.db + its WAL/SHM companions are expected; anything else is + // a telemetry leak. + if name == "memory.db" || strings.HasPrefix(name, "memory.db-") { + continue + } + t.Fatalf("unexpected artifact %q in data dir — telemetry disabled path should not create any file", name) } } diff --git a/internal/telemetry/client.go b/internal/telemetry/client.go index 2b676ae..58c97b2 100644 --- a/internal/telemetry/client.go +++ b/internal/telemetry/client.go @@ -54,10 +54,22 @@ func InitIfEnabled(path string, strict bool) *Client { initWarn(err) return nil } - if _, err := db.Exec("PRAGMA journal_mode=WAL"); err != nil { - _ = db.Close() - initWarn(err) - return nil + // Align with the main memory DB: single open connection for deterministic + // ordering under concurrent tool calls, WAL for reader-friendly durability, + // foreign_keys ON so the search_metrics -> tool_calls FK is actually + // enforced, and a non-zero busy_timeout so brief lock contention retries + // instead of erroring immediately. + db.SetMaxOpenConns(1) + for _, pragma := range []string{ + "PRAGMA journal_mode=WAL", + "PRAGMA foreign_keys=ON", + "PRAGMA busy_timeout=5000", + } { + if _, err := db.Exec(pragma); err != nil { + _ = db.Close() + initWarn(err) + return nil + } } if _, err := db.Exec(schemaSQL); err != nil { _ = db.Close() diff --git a/internal/telemetry/client_test.go b/internal/telemetry/client_test.go index 62ccf85..b56e183 100644 --- a/internal/telemetry/client_test.go +++ b/internal/telemetry/client_test.go @@ -28,11 +28,12 @@ func TestInitIfEnabledEmptyPathReturnsNil(t *testing.T) { } func TestInitIfEnabledInvalidPathReturnsNil(t *testing.T) { - // A path inside a nonexistent directory should fail open - bad := "/nonexistent-directory-for-telemetry-test/telemetry.db" + // Build a path inside a non-existent subdirectory of t.TempDir() so the + // failure mode is identical across macOS, Linux, and Windows. + bad := filepath.Join(t.TempDir(), "missing-subdir", "telemetry.db") if got := InitIfEnabled(bad, false); got != nil { _ = got.Close() - t.Fatalf("InitIfEnabled on bad path should return nil, got client") + t.Fatalf("InitIfEnabled on path inside missing directory should return nil, got client") } } diff --git a/internal/telemetry/sanitize.go b/internal/telemetry/sanitize.go index 97d4374..3afbc54 100644 --- a/internal/telemetry/sanitize.go +++ b/internal/telemetry/sanitize.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "strings" + "unicode/utf8" ) // encodeNoEscape marshals v as JSON without escaping <, >, & — telemetry @@ -39,7 +40,7 @@ func SanitizeArgs(args map[string]any, strict bool) string { switch k { case "observation", "content", "context": if s, ok := v.(string); ok { - safe[k] = fmt.Sprintf("<%d chars>", len(s)) + safe[k] = fmt.Sprintf("<%d chars>", utf8.RuneCountInString(s)) continue } safe[k] = v From fed29f2216a1b124c95d90d10b2128b14801d510 Mon Sep 17 00:00:00 2001 From: marlian <45340800+marlian@users.noreply.github.com> Date: Tue, 14 Apr 2026 23:17:21 +0200 Subject: [PATCH 06/11] review: address Copilot round 2 feedback on PR #5 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three new comments after the previous fixes landed. All legitimate: 1. telemetry/client.go InitIfEnabled: align DSN pattern with the main memory DB. Use file:?_pragma=foreign_keys(1) + db.Ping() instead of a raw path + separate PRAGMA. Safer on Windows drive letters, foreign keys are enforced from open time (not just at the first Exec), and the init failure mode becomes immediately observable via Ping. 2. telemetry_integration_test.go: the enabled-roundtrip and strict-mode tests were calling tele.Close() directly, violating the ownership contract the previous round introduced (Runtime owns the telemetry client once passed to New). Tests now call stop() explicitly before reopening the telemetry DB for readback, so Runtime.Close -> tele .Close -> flush happens through the documented path. 3. startTelemetrySession stop function made idempotent via sync.Once, so explicit stop() + deferred stop() converge safely. t.Fatal replaced with t.Error in the cleanup timeout path — calling Fatal from a deferred function aborts the goroutine but skips subsequent cleanup/readback assertions; Error is the correct level here. The self-inflicted inconsistency between "Runtime owns the telemetry client" (added last round in mcpserver.Config docstring) and the tests closing it directly is exactly the kind of drift a second reviewer pass catches. Good catch. --- .../mcpserver/telemetry_integration_test.go | 36 ++++++++++++------- internal/telemetry/client.go | 21 +++++++---- 2 files changed, 37 insertions(+), 20 deletions(-) diff --git a/internal/mcpserver/telemetry_integration_test.go b/internal/mcpserver/telemetry_integration_test.go index 5fb65bc..5e86e1e 100644 --- a/internal/mcpserver/telemetry_integration_test.go +++ b/internal/mcpserver/telemetry_integration_test.go @@ -6,6 +6,7 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" "time" @@ -44,9 +45,11 @@ func TestTelemetryEnabledRoundtripLogsToolCallsAndSearchMetrics(t *testing.T) { "entity": "TelemetryEntity", }) - if err := tele.Close(); err != nil { - t.Fatalf("tele.Close() error = %v", err) - } + // Shut the runtime down before reading back. Ownership of tele is held + // by Runtime, so Runtime.Close() (triggered through stop -> Run's defer) + // is the right way to flush it. Direct tele.Close() here would violate + // the ownership contract documented in mcpserver.Config. + stop() rdb, err := sql.Open("sqlite", telePath) if err != nil { @@ -188,9 +191,10 @@ func TestTelemetryStrictModeHashesIdentifiersEndToEnd(t *testing.T) { "limit": 5, }) - if err := tele.Close(); err != nil { - t.Fatalf("tele.Close() error = %v", err) - } + // Same ownership discipline as the enabled roundtrip test: stop the + // runtime (which triggers Runtime.Close -> tele.Close) before reading + // the telemetry DB back. + stop() rdb, err := sql.Open("sqlite", telePath) if err != nil { @@ -254,14 +258,20 @@ func startTelemetrySession(t *testing.T, runtime *Runtime) (*mcp.ClientSession, cancel() t.Fatalf("client.Connect() error = %v", err) } + // Idempotent stop — tests call it explicitly before DB readback (so + // Runtime.Close flushes telemetry) and then again through defer for + // panic-safety. Both calls must converge without blocking or fataling. + var once sync.Once stop := func() { - _ = session.Close() - cancel() - select { - case <-errCh: - case <-time.After(2 * time.Second): - t.Fatal("runtime.Run() did not exit after client shutdown") - } + once.Do(func() { + _ = session.Close() + cancel() + select { + case <-errCh: + case <-time.After(2 * time.Second): + t.Error("runtime.Run() did not exit after client shutdown") + } + }) } return session, stop } diff --git a/internal/telemetry/client.go b/internal/telemetry/client.go index 58c97b2..c5730fe 100644 --- a/internal/telemetry/client.go +++ b/internal/telemetry/client.go @@ -16,6 +16,7 @@ import ( "encoding/json" "fmt" "os" + "path/filepath" _ "modernc.org/sqlite" ) @@ -49,20 +50,26 @@ func InitIfEnabled(path string, strict bool) *Client { if path == "" { return nil } - db, err := sql.Open("sqlite", path) + // Mirror the main memory DB's open pattern (see internal/store/sqlite.go + // openSQLite): a file: DSN with foreign_keys set at open time, cleaned + // path for cross-platform safety (notably Windows drive letters), single + // open connection for deterministic write ordering under concurrent tool + // calls, WAL for reader-friendly durability, and a non-zero busy_timeout + // so brief lock contention retries instead of erroring immediately. + dsn := fmt.Sprintf("file:%s?_pragma=foreign_keys(1)", filepath.Clean(path)) + db, err := sql.Open("sqlite", dsn) if err != nil { initWarn(err) return nil } - // Align with the main memory DB: single open connection for deterministic - // ordering under concurrent tool calls, WAL for reader-friendly durability, - // foreign_keys ON so the search_metrics -> tool_calls FK is actually - // enforced, and a non-zero busy_timeout so brief lock contention retries - // instead of erroring immediately. db.SetMaxOpenConns(1) + if err := db.Ping(); err != nil { + _ = db.Close() + initWarn(err) + return nil + } for _, pragma := range []string{ "PRAGMA journal_mode=WAL", - "PRAGMA foreign_keys=ON", "PRAGMA busy_timeout=5000", } { if _, err := db.Exec(pragma); err != nil { From 1b44593a1e92e705a79c2f12d123fb1493cd1a34 Mon Sep 17 00:00:00 2001 From: marlian <45340800+marlian@users.noreply.github.com> Date: Tue, 14 Apr 2026 23:28:41 +0200 Subject: [PATCH 07/11] review: address Copilot round 3 feedback on PR #5 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Eight new comments, grouped into three themes: Cross-platform readback DSN (4 sites): client_test.go (two tests) and telemetry_integration_test.go (two tests) were opening the telemetry DB for readback with a raw path, while production code uses file:?_pragma=foreign_keys(1) (see internal/store/sqlite.go openSQLite). The raw-path form can misbehave on Windows drive letters and drifts from production. All four readbacks now use the same DSN construction. Belt-and-suspenders FK enforcement (1 site): telemetry/client.go InitIfEnabled sets _pragma=foreign_keys(1) in the DSN but did not issue PRAGMA foreign_keys=ON explicitly afterward. The main store's openSQLite does both, because some driver/version combinations honor one form but not the other. Telemetry now matches the same sequence. Small drift cleanup (3 sites): - detect_test.go TestDetectClientEnvFallbackKilo: remove the redundant t.Setenv pair before unsetMCPClientEnv — same two vars were set, cleared, and re-set, which adds noise without changing behavior. - detect_test.go unsetMCPClientEnv: rewrite the trailing comment. The old wording warned about LookupEnv semantics, but DetectClient actually uses os.Getenv != "", so empty-string values are treated as absent for every current signal. - mcpserver/telemetry.go resolveProjectPath: the docstring promised a "global-scope call without a path" on resolve failure, but logToolCall keeps db_scope="project" whenever the caller supplied a non-empty project argument, regardless of whether resolution succeeded. Rewrote the docstring to match the actual (and intended) behavior: scope reflects caller intent, only project_path column is NULL when resolution fails. --- internal/mcpserver/telemetry.go | 9 ++++++--- internal/mcpserver/telemetry_integration_test.go | 4 ++-- internal/telemetry/client.go | 5 +++++ internal/telemetry/client_test.go | 8 +++++--- internal/telemetry/detect_test.go | 13 ++++++------- 5 files changed, 24 insertions(+), 15 deletions(-) diff --git a/internal/mcpserver/telemetry.go b/internal/mcpserver/telemetry.go index a14440b..6dd5748 100644 --- a/internal/mcpserver/telemetry.go +++ b/internal/mcpserver/telemetry.go @@ -24,9 +24,12 @@ func detectClient(req *mcp.CallToolRequest) telemetry.ClientInfo { } // resolveProjectPath resolves a project argument as provided by the caller -// to its absolute form. Returns "" when project is empty or unresolvable -// (including the rare case where os.UserHomeDir fails) — telemetry accepts -// empty and simply records a global-scope call without a path. +// to its absolute form. Returns "" in two cases: the caller passed no +// project, and the rare case where os.UserHomeDir fails. When the caller +// originally supplied a non-empty project, logToolCall still records +// db_scope="project" for that call — only the project_path column is +// written as NULL. The scope reflects the caller's intent, not the +// resolver's success. func resolveProjectPath(project string) string { if project == "" { return "" diff --git a/internal/mcpserver/telemetry_integration_test.go b/internal/mcpserver/telemetry_integration_test.go index 5e86e1e..4809eed 100644 --- a/internal/mcpserver/telemetry_integration_test.go +++ b/internal/mcpserver/telemetry_integration_test.go @@ -51,7 +51,7 @@ func TestTelemetryEnabledRoundtripLogsToolCallsAndSearchMetrics(t *testing.T) { // the ownership contract documented in mcpserver.Config. stop() - rdb, err := sql.Open("sqlite", telePath) + rdb, err := sql.Open("sqlite", "file:"+filepath.Clean(telePath)+"?_pragma=foreign_keys(1)") if err != nil { t.Fatalf("open telemetry db for readback: %v", err) } @@ -196,7 +196,7 @@ func TestTelemetryStrictModeHashesIdentifiersEndToEnd(t *testing.T) { // the telemetry DB back. stop() - rdb, err := sql.Open("sqlite", telePath) + rdb, err := sql.Open("sqlite", "file:"+filepath.Clean(telePath)+"?_pragma=foreign_keys(1)") if err != nil { t.Fatalf("open strict telemetry db: %v", err) } diff --git a/internal/telemetry/client.go b/internal/telemetry/client.go index c5730fe..f328c39 100644 --- a/internal/telemetry/client.go +++ b/internal/telemetry/client.go @@ -68,8 +68,13 @@ func InitIfEnabled(path string, strict bool) *Client { initWarn(err) return nil } + // PRAGMA foreign_keys is also set via the DSN (_pragma=foreign_keys(1)), + // but openSQLite in the main store issues the PRAGMA explicitly after + // open anyway — some driver/version combinations honor one but not the + // other. Belt-and-suspenders. for _, pragma := range []string{ "PRAGMA journal_mode=WAL", + "PRAGMA foreign_keys=ON", "PRAGMA busy_timeout=5000", } { if _, err := db.Exec(pragma); err != nil { diff --git a/internal/telemetry/client_test.go b/internal/telemetry/client_test.go index b56e183..7d8b90f 100644 --- a/internal/telemetry/client_test.go +++ b/internal/telemetry/client_test.go @@ -69,8 +69,10 @@ func TestInitIfEnabledCreatesSchemaAndInserts(t *testing.T) { ScoreMedian: 0.5, }) - // Read back from disk to prove persistence - rdb, err := sql.Open("sqlite", path) + // Read back from disk to prove persistence. Use the same DSN pattern as + // InitIfEnabled so the readback works consistently across Windows and + // POSIX (raw paths with drive letters can trip the driver). + rdb, err := sql.Open("sqlite", "file:"+filepath.Clean(path)+"?_pragma=foreign_keys(1)") if err != nil { t.Fatalf("open telemetry db for readback: %v", err) } @@ -123,7 +125,7 @@ func TestStrictModeHashesSearchQuery(t *testing.T) { ResultsReturned: 1, }) - rdb, err := sql.Open("sqlite", path) + rdb, err := sql.Open("sqlite", "file:"+filepath.Clean(path)+"?_pragma=foreign_keys(1)") if err != nil { t.Fatalf("open db: %v", err) } diff --git a/internal/telemetry/detect_test.go b/internal/telemetry/detect_test.go index 5042fd1..3dc1efd 100644 --- a/internal/telemetry/detect_test.go +++ b/internal/telemetry/detect_test.go @@ -13,8 +13,6 @@ func TestDetectClientProtocolWins(t *testing.T) { } func TestDetectClientEnvFallbackKilo(t *testing.T) { - t.Setenv("KILO", "1") - t.Setenv("KILOCODE_VERSION", "0.43.6") unsetMCPClientEnv(t) t.Setenv("KILO", "1") t.Setenv("KILOCODE_VERSION", "0.43.6") @@ -62,9 +60,10 @@ func unsetMCPClientEnv(t *testing.T) { } { t.Setenv(v, "") } - // t.Setenv with "" sets to empty, but LookupEnv still reports present. - // For the ones we need genuinely absent, unset via os.Unsetenv is not - // t.Setenv-compatible; we accept that "" is how the tests model absence - // for signals whose presence-only matters (VSCODE_MCP_HTTP_PREFER handled - // in its own test). + // These tests model "absent" env vars by setting them to "". That + // matches the current DetectClient implementation, which uses + // os.Getenv(...) != "" rather than os.LookupEnv, so empty-string + // values are treated as absent for every signal. The non-empty + // requirement on VSCODE_MCP_HTTP_PREFER is covered by its own + // dedicated test. } From 327789a5b495f408cdfb051fe3d6bb81eb558c74 Mon Sep 17 00:00:00 2001 From: marlian <45340800+marlian@users.noreply.github.com> Date: Tue, 14 Apr 2026 23:39:27 +0200 Subject: [PATCH 08/11] review: address Copilot round 4 feedback on PR #5 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four new comments, all legitimate: 1. cmd/workmem/main.go: fix telemetry leak when mcpserver.New fails. FromEnv() was inlined directly in the Config literal — if New returned an error, the DB handle was already open and nobody closed it. Now FromEnv() lands in a local var; on New failure the code calls tele.Close() (nil-safe) before exiting. Ownership transfers to Runtime only after New succeeds, matching the documented contract. 2. internal/telemetry/client.go: Client.Close is now truly idempotent. After the first call it nil-s db, insertCall, insertSearch so the second call hits the new `if c.db == nil { return nil }` guard instead of trying to close an already-closed *sql.DB (which would have surfaced a confusing shutdown error). Runtime.Close already comments "nil-safe and idempotent" — now it's true. 3. internal/telemetry/{schema.go,client.go}: schema initialization no longer runs a single multi-statement Exec. Statements are listed separately in schemaStatements and exec'd one-by-one, matching the main store's InitSchema pattern. More portable across SQLite drivers and produces per-statement errors when one fails. 4. internal/telemetry/sanitize.go: add optional Summarizable interface as a fast path in SummarizeResult. Types that implement TelemetrySummary() skip the JSON marshal + unmarshal round-trip entirely. The fallback stays in place for all current store result types (zero behavior change this round); any type whose result size becomes a telemetry hotspot can now opt in by adding the method. Tests: - TestClientCloseIsIdempotent: open, close, close, close — all three must return nil, proving Close can be called defensively multiple times without surfacing an error. - TestSummarizeResultUsesFastPathForSummarizable: a fake Summarizable type is serialized without going through the JSON fallback. - TestSummarizeResultFallsBackToJSONForNonSummarizable: plain map values still produce the correct entity_groups / stored summary through the fallback path. --- cmd/workmem/main.go | 9 +++++--- internal/telemetry/client.go | 33 +++++++++++++++++++---------- internal/telemetry/client_test.go | 17 +++++++++++++++ internal/telemetry/sanitize.go | 25 ++++++++++++++++++++-- internal/telemetry/sanitize_test.go | 30 ++++++++++++++++++++++++++ internal/telemetry/schema.go | 22 +++++++++++-------- 6 files changed, 111 insertions(+), 25 deletions(-) diff --git a/cmd/workmem/main.go b/cmd/workmem/main.go index 962e8da..63bc93e 100644 --- a/cmd/workmem/main.go +++ b/cmd/workmem/main.go @@ -102,13 +102,16 @@ func runMCP(args []string) { loadEnvFile(*envFile) - // Ownership of the telemetry client transfers to the Runtime — its Close() - // handles the shutdown. No defer here on purpose (see mcpserver.Config). + // Ownership of the telemetry client transfers to the Runtime only after + // New returns successfully. If New fails, the DB was already opened by + // FromEnv and must be closed here — otherwise the handle leaks. + tele := telemetry.FromEnv() runtime, err := mcpserver.New(mcpserver.Config{ DBPath: *dbPath, - Telemetry: telemetry.FromEnv(), + Telemetry: tele, }) if err != nil { + _ = tele.Close() // nil-safe no-op when telemetry is disabled fmt.Fprintf(os.Stderr, "start mcp server: %v\n", err) os.Exit(1) } diff --git a/internal/telemetry/client.go b/internal/telemetry/client.go index f328c39..a927c87 100644 --- a/internal/telemetry/client.go +++ b/internal/telemetry/client.go @@ -83,10 +83,12 @@ func InitIfEnabled(path string, strict bool) *Client { return nil } } - if _, err := db.Exec(schemaSQL); err != nil { - _ = db.Close() - initWarn(err) - return nil + for i, stmt := range schemaStatements { + if _, err := db.Exec(stmt); err != nil { + _ = db.Close() + initWarn(fmt.Errorf("schema statement %d: %w", i, err)) + return nil + } } insertCall, err := db.Prepare(insertCallSQL) if err != nil { @@ -109,18 +111,27 @@ func initWarn(err error) { } // Close releases the prepared statements and the underlying database -// connection. Safe to call on a nil receiver. +// connection. Safe to call on a nil receiver. Idempotent: after the first +// call, fields are nil-ed so subsequent calls return nil instead of +// trying to close an already-closed *sql.DB (which would surface a +// confusing shutdown error). func (c *Client) Close() error { - if c == nil { + if c == nil || c.db == nil { return nil } - if c.insertCall != nil { - _ = c.insertCall.Close() + insertCall := c.insertCall + insertSearch := c.insertSearch + db := c.db + c.insertCall = nil + c.insertSearch = nil + c.db = nil + if insertCall != nil { + _ = insertCall.Close() } - if c.insertSearch != nil { - _ = c.insertSearch.Close() + if insertSearch != nil { + _ = insertSearch.Close() } - return c.db.Close() + return db.Close() } // Strict reports whether privacy-strict mode is active. diff --git a/internal/telemetry/client_test.go b/internal/telemetry/client_test.go index 7d8b90f..ddac38a 100644 --- a/internal/telemetry/client_test.go +++ b/internal/telemetry/client_test.go @@ -142,6 +142,23 @@ func TestStrictModeHashesSearchQuery(t *testing.T) { } } +func TestClientCloseIsIdempotent(t *testing.T) { + path := filepath.Join(t.TempDir(), "close-twice.db") + c := InitIfEnabled(path, false) + if c == nil { + t.Fatalf("InitIfEnabled returned nil on valid path") + } + if err := c.Close(); err != nil { + t.Fatalf("first Close() = %v, want nil", err) + } + if err := c.Close(); err != nil { + t.Fatalf("second Close() = %v, want nil (Close should be idempotent)", err) + } + if err := c.Close(); err != nil { + t.Fatalf("third Close() = %v, want nil", err) + } +} + func TestLogSearchMetricsZeroToolCallIDIsNoop(t *testing.T) { path := filepath.Join(t.TempDir(), "noop.db") c := InitIfEnabled(path, false) diff --git a/internal/telemetry/sanitize.go b/internal/telemetry/sanitize.go index 3afbc54..e8327a9 100644 --- a/internal/telemetry/sanitize.go +++ b/internal/telemetry/sanitize.go @@ -73,15 +73,36 @@ func SanitizeArgs(args map[string]any, strict bool) string { return out } +// Summarizable is implemented by result types that want to produce their +// own telemetry summary without going through the JSON round-trip fallback. +// This is the fast path: types returning large payloads (large recall +// results, big graph dumps) can implement this to avoid the cost of a +// full Marshal + Unmarshal on every tool call. +// +// The returned map must contain only count/boolean/short-string values — +// never the actual payload content. SummarizeResult serializes the map as +// JSON without any further inspection. +type Summarizable interface { + TelemetrySummary() map[string]any +} + // SummarizeResult extracts count-only telemetry fields from a tool result. // It never leaks content: only counts, booleans, and structural indicators. // -// Uses JSON round-trip to avoid direct coupling to the store package result -// types — the telemetry package stays domain-agnostic. +// Fast path: if the result implements Summarizable, its TelemetrySummary() +// map is serialized directly. Fallback: JSON round-trip so the telemetry +// package stays domain-agnostic for types that do not opt in. func SummarizeResult(result any) string { if result == nil { return "" } + if s, ok := result.(Summarizable); ok { + out, err := encodeNoEscape(s.TelemetrySummary()) + if err != nil { + return "" + } + return out + } raw, err := json.Marshal(result) if err != nil { return "" diff --git a/internal/telemetry/sanitize_test.go b/internal/telemetry/sanitize_test.go index a76348d..7be6763 100644 --- a/internal/telemetry/sanitize_test.go +++ b/internal/telemetry/sanitize_test.go @@ -75,6 +75,36 @@ func TestSanitizeArgsNilReturnsEmpty(t *testing.T) { } } +type summarizableFake struct { + countA int + countB int +} + +func (s summarizableFake) TelemetrySummary() map[string]any { + return map[string]any{"a": s.countA, "b": s.countB} +} + +func TestSummarizeResultUsesFastPathForSummarizable(t *testing.T) { + got := SummarizeResult(summarizableFake{countA: 3, countB: 7}) + if !strings.Contains(got, "\"a\":3") || !strings.Contains(got, "\"b\":7") { + t.Fatalf("SummarizeResult fast path missing fields: %s", got) + } +} + +func TestSummarizeResultFallsBackToJSONForNonSummarizable(t *testing.T) { + // Plain map — no Summarizable interface, must go through fallback + got := SummarizeResult(map[string]any{ + "stored": true, + "results": []any{map[string]any{}, map[string]any{}}, + }) + if !strings.Contains(got, "\"stored\":true") { + t.Fatalf("fallback missing stored: %s", got) + } + if !strings.Contains(got, "\"entity_groups\":2") { + t.Fatalf("fallback missing entity_groups: %s", got) + } +} + func TestSummarizeResultReturnsCountsOnly(t *testing.T) { result := map[string]any{ "stored": true, diff --git a/internal/telemetry/schema.go b/internal/telemetry/schema.go index 8d862d6..3aa8ab3 100644 --- a/internal/telemetry/schema.go +++ b/internal/telemetry/schema.go @@ -1,7 +1,11 @@ package telemetry -const schemaSQL = ` -CREATE TABLE IF NOT EXISTS tool_calls ( +// schemaStatements is executed one-by-one by InitIfEnabled. The main store +// follows the same single-statement-per-Exec pattern (see +// internal/store/sqlite.go InitSchema) — more portable across SQLite +// drivers than chaining multiple statements into one db.Exec call. +var schemaStatements = []string{ + `CREATE TABLE IF NOT EXISTS tool_calls ( id INTEGER PRIMARY KEY AUTOINCREMENT, ts TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%f', 'now')), tool TEXT NOT NULL, @@ -14,8 +18,8 @@ CREATE TABLE IF NOT EXISTS tool_calls ( args_summary TEXT, result_summary TEXT, is_error INTEGER NOT NULL DEFAULT 0 -); -CREATE TABLE IF NOT EXISTS search_metrics ( +)`, + `CREATE TABLE IF NOT EXISTS search_metrics ( id INTEGER PRIMARY KEY AUTOINCREMENT, tool_call_id INTEGER REFERENCES tool_calls(id), query TEXT, @@ -27,11 +31,11 @@ CREATE TABLE IF NOT EXISTS search_metrics ( score_max REAL, score_median REAL, compact INTEGER DEFAULT 0 -); -CREATE INDEX IF NOT EXISTS idx_tool_calls_ts ON tool_calls(ts); -CREATE INDEX IF NOT EXISTS idx_tool_calls_tool ON tool_calls(tool); -CREATE INDEX IF NOT EXISTS idx_search_metrics_tool_call ON search_metrics(tool_call_id); -` +)`, + `CREATE INDEX IF NOT EXISTS idx_tool_calls_ts ON tool_calls(ts)`, + `CREATE INDEX IF NOT EXISTS idx_tool_calls_tool ON tool_calls(tool)`, + `CREATE INDEX IF NOT EXISTS idx_search_metrics_tool_call ON search_metrics(tool_call_id)`, +} const ( insertCallSQL = `INSERT INTO tool_calls From 5a80d1644b8a1f1dfdd7e36731ffc110f0a43c75 Mon Sep 17 00:00:00 2001 From: marlian <45340800+marlian@users.noreply.github.com> Date: Tue, 14 Apr 2026 23:54:22 +0200 Subject: [PATCH 09/11] review: address Copilot round 5 feedback on PR #5 - telemetry: guard LogToolCall and LogSearchMetrics against post-Close panic by nil-checking db / prepared-statement fields, not just the receiver. Covers late-arriving tool calls during an orderly shutdown. Test: TestLogAfterCloseDoesNotPanic. - search: record the trimmed query in SearchMetrics so two recall calls that only differ by whitespace do not appear as distinct queries in telemetry. Test: TestSearchMemoryMetricsQueryIsTrimmed (5 subcases). - docs: TELEMETRY.md said batched arrays became "" but the code emits "" / "". Doc updated to match reality. --- docs/TELEMETRY.md | 2 +- internal/store/parity_test.go | 36 +++++++++++++++++++++++++++++++ internal/store/search.go | 9 +++++--- internal/telemetry/client.go | 11 ++++++++-- internal/telemetry/client_test.go | 22 +++++++++++++++++++ 5 files changed, 74 insertions(+), 6 deletions(-) diff --git a/docs/TELEMETRY.md b/docs/TELEMETRY.md index e3972ee..1b508d8 100644 --- a/docs/TELEMETRY.md +++ b/docs/TELEMETRY.md @@ -37,7 +37,7 @@ Telemetry supports two modes, controlled by `MEMORY_TELEMETRY_PRIVACY`: Strict mode is intended for sensitive instances such as a `private_memory` server backing therapy/health/relationship content. Ranking debug ("which queries overfetch candidates?") becomes harder in strict mode because plaintext queries are no longer recoverable — but sensitive identifiers never land on disk. -Observation/content values are **always** reduced to `` regardless of mode. Facts/observations arrays are **always** reduced to ``. Strict mode only changes what happens to identifier-like fields. +Observation/content values are **always** reduced to `` regardless of mode. Batched-array fields are **always** reduced to a count marker using the field name — `facts` becomes ``, `observations` becomes ``. Strict mode only changes what happens to identifier-like fields. **Example `.env` for sensitive backend:** ```bash diff --git a/internal/store/parity_test.go b/internal/store/parity_test.go index 8ee1912..d89c7d0 100644 --- a/internal/store/parity_test.go +++ b/internal/store/parity_test.go @@ -530,3 +530,39 @@ func TestRememberEventAtomicityOnMidLoopFailure(t *testing.T) { t.Fatalf("Alpha entity persisted despite rollback: count=%d", alphaCount) } } + +// SearchMetrics.Query must match the trimmed query actually executed against +// the index — otherwise two recall calls that only differ by leading or +// trailing whitespace look like distinct queries in telemetry even though +// they ran the same search. +func TestSearchMemoryMetricsQueryIsTrimmed(t *testing.T) { + db := newTestDB(t, "metrics-trim.db") + entityID, err := UpsertEntity(db, "TrimProbe", "test") + if err != nil { + t.Fatalf("UpsertEntity: %v", err) + } + if _, err := AddObservation(db, entityID, "trim probe observation", "user", 1.0); err != nil { + t.Fatalf("AddObservation: %v", err) + } + + cases := []struct { + name, input, want string + }{ + {"no whitespace", "trim", "trim"}, + {"leading whitespace", " trim", "trim"}, + {"trailing whitespace", "trim ", "trim"}, + {"both sides", " trim ", "trim"}, + {"whitespace only", " ", ""}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + _, metrics, err := SearchMemory(db, tc.input, 5, 12) + if err != nil { + t.Fatalf("SearchMemory(%q): %v", tc.input, err) + } + if metrics.Query != tc.want { + t.Fatalf("SearchMemory(%q).Metrics.Query = %q, want %q", tc.input, metrics.Query, tc.want) + } + }) + } +} diff --git a/internal/store/search.go b/internal/store/search.go index 9717896..50c5018 100644 --- a/internal/store/search.go +++ b/internal/store/search.go @@ -333,11 +333,14 @@ func ScoreCandidates(observations []SearchObservation, candidateMap map[int64]*c func SearchMemory(db *sql.DB, query string, limit int, halfLifeWeeks float64) ([]SearchObservation, SearchMetrics, error) { requestedLimit := limit limit = SanitizeSearchLimit(limit) - emptyMetrics := SearchMetrics{Query: query, LimitRequested: requestedLimit, Channels: map[string]int{}} + // Use the trimmed query in metrics so two recall calls that only differ + // by leading/trailing whitespace do not look like different queries when + // they actually execute the same search. + trimmed := strings.TrimSpace(query) + emptyMetrics := SearchMetrics{Query: trimmed, LimitRequested: requestedLimit, Channels: map[string]int{}} if limit <= 0 { return []SearchObservation{}, emptyMetrics, nil } - trimmed := strings.TrimSpace(query) if trimmed == "" { return []SearchObservation{}, emptyMetrics, nil } @@ -365,7 +368,7 @@ func SearchMemory(db *sql.DB, query string, limit int, halfLifeWeeks float64) ([ return nil, SearchMetrics{}, err } metrics := SearchMetrics{ - Query: query, + Query: trimmed, Channels: channelCounts, CandidatesTotal: len(candidates), ResultsReturned: len(ranked), diff --git a/internal/telemetry/client.go b/internal/telemetry/client.go index a927c87..548a16d 100644 --- a/internal/telemetry/client.go +++ b/internal/telemetry/client.go @@ -157,8 +157,13 @@ type ToolCallInput struct { // LogToolCall inserts a tool_calls row. Returns the new row id on success or // 0 on failure / disabled client. The returned id is used by LogSearchMetrics // to link the search_metrics row back to its parent tool call. +// +// Safe to call after Close: the nil-field guard returns 0 instead of panicking +// on a closed client. This matters for late-arriving tool calls during an +// orderly shutdown (Runtime.Close runs concurrently with the last in-flight +// dispatch) — telemetry gracefully degrades rather than crashing the server. func (c *Client) LogToolCall(in ToolCallInput) int64 { - if c == nil { + if c == nil || c.db == nil || c.insertCall == nil { return 0 } dbScope := in.DBScope @@ -204,9 +209,11 @@ type SearchMetricsInput struct { // LogSearchMetrics inserts a search_metrics row linked to the tool_call id. // No-op when client is nil or ToolCallID is 0 (the linking parent failed). +// Also no-op when the client has been closed — the same nil-field guard used +// by LogToolCall protects against late-arriving metrics during shutdown. // In strict mode, Query is hashed before insertion. func (c *Client) LogSearchMetrics(in SearchMetricsInput) { - if c == nil || in.ToolCallID == 0 { + if c == nil || in.ToolCallID == 0 || c.db == nil || c.insertSearch == nil { return } channelsJSON, err := json.Marshal(in.Channels) diff --git a/internal/telemetry/client_test.go b/internal/telemetry/client_test.go index ddac38a..1951b19 100644 --- a/internal/telemetry/client_test.go +++ b/internal/telemetry/client_test.go @@ -159,6 +159,28 @@ func TestClientCloseIsIdempotent(t *testing.T) { } } +// A late-arriving tool call that lands after the Runtime has begun shutting +// down must not panic. Close nils the db/stmt fields; LogToolCall and +// LogSearchMetrics must both degrade gracefully to a no-op. +func TestLogAfterCloseDoesNotPanic(t *testing.T) { + path := filepath.Join(t.TempDir(), "log-after-close.db") + c := InitIfEnabled(path, false) + if c == nil { + t.Fatalf("init failed") + } + if err := c.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + // LogToolCall after Close must return 0 without panicking. + if id := c.LogToolCall(ToolCallInput{Tool: "remember"}); id != 0 { + t.Fatalf("LogToolCall after Close returned id %d, want 0", id) + } + + // LogSearchMetrics after Close must silently no-op (no panic). + c.LogSearchMetrics(SearchMetricsInput{ToolCallID: 42, Query: "anything"}) +} + func TestLogSearchMetricsZeroToolCallIDIsNoop(t *testing.T) { path := filepath.Join(t.TempDir(), "noop.db") c := InitIfEnabled(path, false) From 1460dc51e1fee52e98a50fe936b1836ba2e95627 Mon Sep 17 00:00:00 2001 From: marlian <45340800+marlian@users.noreply.github.com> Date: Wed, 15 Apr 2026 00:05:40 +0200 Subject: [PATCH 10/11] review: address Copilot round 6 feedback on PR #5 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two concurrency concerns raised on c46ac66: 1. telemetry.Client had no synchronization: Close nil-ed db/stmt fields while LogToolCall / LogSearchMetrics read and used the same pointers, so a real concurrent shutdown could produce a data race (formally undefined behavior under Go's memory model) or a TOCTOU panic where Exec ran on a stmt being closed. Fix: add sync.Mutex to Client, protect Close + Log* under the lock. strict is immutable after construction and stays lock-free. The uncontended lock cost (a few ns) is dwarfed by SQLite Exec (µs-ms), so the hot path is unaffected. 2. Runtime.Close called r.telemetry.Close() before nil-ing r.telemetry, so an in-flight handler's telemetry defer could observe a non-nil pointer on a client that was concurrently closing. Fix: swap r.telemetry into a local, set the field to nil first, then Close the local. With the client mutex this is a defensive second layer rather than the primary safety guarantee, but it keeps future readers honest about lifecycle ordering. Tests: - TestClientCloseRacesWithLogging fires 32 logger goroutines firing LogToolCall + LogSearchMetrics in a tight loop while Close runs concurrently, then a final round of post-close calls that must no-op. Passes under `go test -race` — without the mutex the race detector flags the pointer-field accesses immediately. --- internal/mcpserver/server.go | 12 ++++++-- internal/telemetry/client.go | 48 +++++++++++++++++++++++++------ internal/telemetry/client_test.go | 45 +++++++++++++++++++++++++++++ 3 files changed, 93 insertions(+), 12 deletions(-) diff --git a/internal/mcpserver/server.go b/internal/mcpserver/server.go index f3e1ac2..b227daf 100644 --- a/internal/mcpserver/server.go +++ b/internal/mcpserver/server.go @@ -86,10 +86,16 @@ func (r *Runtime) Close() error { closeErr = r.defaultDB.Close() r.defaultDB = nil } - // Close telemetry on the way down. *Client.Close is nil-safe and - // idempotent, so double-close from a defensive caller is harmless. - teleErr := r.telemetry.Close() + // Swap the telemetry pointer into a local var and nil the field + // *before* closing, so any in-flight tool handler that reaches its + // telemetry defer during shutdown observes a nil r.telemetry and + // skips the log path entirely. The client itself is also mutex-guarded + // against concurrent Close + Log* access (see telemetry/client.go), + // so this swap is a defensive second layer, not the primary safety + // guarantee. Close remains nil-safe and idempotent. + telemetryClient := r.telemetry r.telemetry = nil + teleErr := telemetryClient.Close() return errors.Join(closeErr, teleErr, store.ResetProjectDBs()) } diff --git a/internal/telemetry/client.go b/internal/telemetry/client.go index 548a16d..f304cae 100644 --- a/internal/telemetry/client.go +++ b/internal/telemetry/client.go @@ -17,13 +17,22 @@ import ( "fmt" "os" "path/filepath" + "sync" _ "modernc.org/sqlite" ) // Client is a telemetry sink. A nil *Client is valid and represents the // disabled state — every method below checks for nil and returns harmlessly. +// +// The mutex serializes Close against LogToolCall / LogSearchMetrics so the +// "tool call in flight while shutdown fires" scenario is safe even under +// -race. The uncontended lock cost (a few ns per call) is dwarfed by the +// SQLite Exec itself (microseconds), so there is no measurable hot-path +// impact. strict is set at construction and never mutated, so it is read +// outside the lock. type Client struct { + mu sync.Mutex db *sql.DB insertCall *sql.Stmt insertSearch *sql.Stmt @@ -115,8 +124,19 @@ func initWarn(err error) { // call, fields are nil-ed so subsequent calls return nil instead of // trying to close an already-closed *sql.DB (which would surface a // confusing shutdown error). +// +// Thread-safe: the mutex serializes Close against any concurrent +// LogToolCall / LogSearchMetrics so the "shutdown while tool call in +// flight" scenario cannot observe half-closed state. A late log call +// that acquires the mutex after Close returns harmlessly because the +// nil-field guard fails. func (c *Client) Close() error { - if c == nil || c.db == nil { + if c == nil { + return nil + } + c.mu.Lock() + defer c.mu.Unlock() + if c.db == nil { return nil } insertCall := c.insertCall @@ -158,12 +178,17 @@ type ToolCallInput struct { // 0 on failure / disabled client. The returned id is used by LogSearchMetrics // to link the search_metrics row back to its parent tool call. // -// Safe to call after Close: the nil-field guard returns 0 instead of panicking -// on a closed client. This matters for late-arriving tool calls during an -// orderly shutdown (Runtime.Close runs concurrently with the last in-flight -// dispatch) — telemetry gracefully degrades rather than crashing the server. +// Thread-safe: acquires the client mutex so the insert cannot observe a +// half-closed state even if Close runs concurrently. A late call that +// arrives after Close returns 0 harmlessly because the nil-field guard +// fails under the lock. func (c *Client) LogToolCall(in ToolCallInput) int64 { - if c == nil || c.db == nil || c.insertCall == nil { + if c == nil { + return 0 + } + c.mu.Lock() + defer c.mu.Unlock() + if c.db == nil || c.insertCall == nil { return 0 } dbScope := in.DBScope @@ -209,11 +234,16 @@ type SearchMetricsInput struct { // LogSearchMetrics inserts a search_metrics row linked to the tool_call id. // No-op when client is nil or ToolCallID is 0 (the linking parent failed). -// Also no-op when the client has been closed — the same nil-field guard used -// by LogToolCall protects against late-arriving metrics during shutdown. +// Thread-safe: acquires the same mutex as LogToolCall / Close, so a closed +// client is observed as a no-op rather than panicking on a nil stmt. // In strict mode, Query is hashed before insertion. func (c *Client) LogSearchMetrics(in SearchMetricsInput) { - if c == nil || in.ToolCallID == 0 || c.db == nil || c.insertSearch == nil { + if c == nil || in.ToolCallID == 0 { + return + } + c.mu.Lock() + defer c.mu.Unlock() + if c.db == nil || c.insertSearch == nil { return } channelsJSON, err := json.Marshal(in.Channels) diff --git a/internal/telemetry/client_test.go b/internal/telemetry/client_test.go index 1951b19..c91fe2f 100644 --- a/internal/telemetry/client_test.go +++ b/internal/telemetry/client_test.go @@ -4,6 +4,7 @@ import ( "database/sql" "path/filepath" "strings" + "sync" "testing" ) @@ -181,6 +182,50 @@ func TestLogAfterCloseDoesNotPanic(t *testing.T) { c.LogSearchMetrics(SearchMetricsInput{ToolCallID: 42, Query: "anything"}) } +// Under concurrent Close + LogToolCall, the client mutex must serialize +// operations so no goroutine observes a half-closed state. Run with +// `go test -race` to turn an unsynchronized access into a test failure +// instead of an occasional production panic. +func TestClientCloseRacesWithLogging(t *testing.T) { + path := filepath.Join(t.TempDir(), "race.db") + c := InitIfEnabled(path, false) + if c == nil { + t.Fatalf("init failed") + } + + const loggers = 32 + var wg sync.WaitGroup + start := make(chan struct{}) + + for i := 0; i < loggers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + // Fire several calls so Close has real overlap with active inserts. + for j := 0; j < 4; j++ { + _ = c.LogToolCall(ToolCallInput{Tool: "remember"}) + c.LogSearchMetrics(SearchMetricsInput{ToolCallID: 1, Query: "q"}) + } + }() + } + + close(start) + // Give the goroutines a moment to be mid-flight, then close. The mutex + // in Client must make this safe; without it, -race would flag the + // db/stmt pointer reads against the Close writes. + if err := c.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + wg.Wait() + + // A second set of calls after Close must be no-ops, not panics. + if id := c.LogToolCall(ToolCallInput{Tool: "post-close"}); id != 0 { + t.Fatalf("post-close LogToolCall returned %d, want 0", id) + } + c.LogSearchMetrics(SearchMetricsInput{ToolCallID: 99, Query: "post"}) +} + func TestLogSearchMetricsZeroToolCallIDIsNoop(t *testing.T) { path := filepath.Join(t.TempDir(), "noop.db") c := InitIfEnabled(path, false) From 056db94bf64fd8924c67d959a5776efcca97d931 Mon Sep 17 00:00:00 2001 From: marlian <45340800+marlian@users.noreply.github.com> Date: Wed, 15 Apr 2026 02:00:22 +0200 Subject: [PATCH 11/11] review: address Copilot round 7 feedback on PR #5 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two concurrency concerns and two stderr-spam concerns: 1. Runtime.telemetry was a plain *Client pointer written during Close() and read by in-flight handleTool goroutines — a data race under -race, regardless of the Client's internal mutex. Fix: make it an atomic.Pointer[telemetry.Client]. handleTool calls Load() once and captures the pointer into its defer closure, so the deferred Log* uses the captured value instead of re-reading the field. Runtime.Close uses Swap(nil) to atomically take ownership for cleanup. logToolCall / logSearchMetrics now accept an explicit *telemetry.Client parameter and never touch the atomic field themselves. 2. LogToolCall / LogSearchMetrics printed "telemetry log failed: ..." to stderr on every Exec error. A persistently wedged telemetry DB (disk full, permission change, corruption) would therefore spam stderr once per tool call for the remainder of the session. Fix: add a `degraded` flag to Client. On the first Exec failure we emit a single warning ("further errors suppressed for this session") and set degraded=true; subsequent Log* calls take the degraded branch under the existing mutex and return silently. Recovery requires a restart — same contract as init failure. Tests: - TestClientDegradedModeSuppressesSpam captures stderr via os.Pipe, force-closes the underlying sql.DB to make Exec fail, fires three Log* calls, and asserts exactly one warning plus degraded=true. - All existing tests pass under `go test -race`, including the concurrency regression test. --- internal/mcpserver/server.go | 47 ++++++++++++++---------- internal/mcpserver/telemetry.go | 24 ++++++++----- internal/telemetry/client.go | 18 +++++++--- internal/telemetry/client_test.go | 60 +++++++++++++++++++++++++++++++ 4 files changed, 117 insertions(+), 32 deletions(-) diff --git a/internal/mcpserver/server.go b/internal/mcpserver/server.go index b227daf..98d64c7 100644 --- a/internal/mcpserver/server.go +++ b/internal/mcpserver/server.go @@ -10,6 +10,7 @@ import ( "os" "path/filepath" "strings" + "sync/atomic" "time" "github.com/modelcontextprotocol/go-sdk/mcp" @@ -33,7 +34,11 @@ type Runtime struct { server *mcp.Server defaultDB *sql.DB dbPath string - telemetry *telemetry.Client + // telemetry is stored as an atomic.Pointer so concurrent reads by + // in-flight tool handlers cannot race with the Close-time Swap during + // shutdown. Reads use Load(); Close uses Swap(nil) to atomically take + // ownership of the pointer for cleanup. + telemetry atomic.Pointer[telemetry.Client] } type toolDefinition struct { @@ -62,7 +67,8 @@ func New(config Config) (*Runtime, error) { Version: serverVersion, }, nil) - runtime := &Runtime{server: server, defaultDB: db, dbPath: dbPath, telemetry: config.Telemetry} + runtime := &Runtime{server: server, defaultDB: db, dbPath: dbPath} + runtime.telemetry.Store(config.Telemetry) runtime.registerTools() return runtime, nil } @@ -86,16 +92,17 @@ func (r *Runtime) Close() error { closeErr = r.defaultDB.Close() r.defaultDB = nil } - // Swap the telemetry pointer into a local var and nil the field - // *before* closing, so any in-flight tool handler that reaches its - // telemetry defer during shutdown observes a nil r.telemetry and - // skips the log path entirely. The client itself is also mutex-guarded - // against concurrent Close + Log* access (see telemetry/client.go), - // so this swap is a defensive second layer, not the primary safety - // guarantee. Close remains nil-safe and idempotent. - telemetryClient := r.telemetry - r.telemetry = nil - teleErr := telemetryClient.Close() + // atomic.Pointer.Swap atomically takes the old pointer and installs nil, + // so any in-flight handler that reaches its telemetry defer during + // shutdown observes a nil Load() and skips the log path entirely. No + // data race between this write and the concurrent reads in handleTool. + // The underlying Client is also mutex-guarded (see telemetry/client.go), + // so even if a handler already captured the old pointer before the swap + // its Log* calls will serialize with the Close below. + var teleErr error + if telemetryClient := r.telemetry.Swap(nil); telemetryClient != nil { + teleErr = telemetryClient.Close() + } return errors.Join(closeErr, teleErr, store.ResetProjectDBs()) } @@ -117,10 +124,12 @@ func (r *Runtime) registerTools() { } func (r *Runtime) handleTool(_ context.Context, def toolDefinition, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { - // Telemetry observables — captured as the flow unfolds. When telemetry is - // disabled (r.telemetry == nil) we skip both the time.Now/time.Since - // measurement and the defer installation entirely: the no-telemetry path - // adds nothing beyond a single pointer-nil check. + // Telemetry observables — captured as the flow unfolds. The telemetry + // client pointer is read via atomic.Pointer.Load() once and captured + // into the closure below; a concurrent Runtime.Close() Swap cannot race + // with the captured value. When telemetry is disabled (Load returns + // nil) we skip both the time.Now/time.Since measurement and the defer + // installation entirely, so the no-telemetry path stays near-zero. var ( argObject map[string]any toolResult any @@ -128,11 +137,11 @@ func (r *Runtime) handleTool(_ context.Context, def toolDefinition, req *mcp.Cal projectRaw string isError bool ) - if r.telemetry != nil { + if tele := r.telemetry.Load(); tele != nil { t0 := time.Now() defer func() { - id := r.logToolCall(def.Name, req, argObject, toolResult, projectRaw, isError, time.Since(t0)) - r.logSearchMetrics(id, metrics) + id := r.logToolCall(tele, def.Name, req, argObject, toolResult, projectRaw, isError, time.Since(t0)) + r.logSearchMetrics(tele, id, metrics) }() } diff --git a/internal/mcpserver/telemetry.go b/internal/mcpserver/telemetry.go index 6dd5748..b64ee2f 100644 --- a/internal/mcpserver/telemetry.go +++ b/internal/mcpserver/telemetry.go @@ -42,8 +42,13 @@ func resolveProjectPath(project string) string { } // logToolCall inserts a tool_calls row and returns the insert id (or 0 when -// telemetry is disabled / logging failed). Safe to call when c is nil. +// the telemetry client is nil / logging failed). The caller is expected to +// pass the pointer it captured from r.telemetry.Load() so this helper never +// re-reads the atomic field — that avoids any formal race against a +// concurrent Runtime.Close() that may have Swap(nil)'d the pointer after +// the defer fired. func (r *Runtime) logToolCall( + tele *telemetry.Client, toolName string, req *mcp.CallToolRequest, argObject map[string]any, @@ -52,7 +57,7 @@ func (r *Runtime) logToolCall( isError bool, elapsed time.Duration, ) int64 { - if r.telemetry == nil { + if tele == nil { return 0 } dbScope := "global" @@ -61,25 +66,26 @@ func (r *Runtime) logToolCall( dbScope = "project" projectPath = resolveProjectPath(projectRaw) } - return r.telemetry.LogToolCall(telemetry.ToolCallInput{ + return tele.LogToolCall(telemetry.ToolCallInput{ Tool: toolName, Client: detectClient(req), DBScope: dbScope, ProjectPath: projectPath, DurationMs: float64(elapsed) / float64(time.Millisecond), - ArgsSummary: telemetry.SanitizeArgs(argObject, r.telemetry.Strict()), + ArgsSummary: telemetry.SanitizeArgs(argObject, tele.Strict()), ResultSummary: telemetry.SummarizeResult(result), IsError: isError, }) } -// logSearchMetrics mirrors the recall search_metrics row. No-op when telemetry -// is disabled or the parent tool_call id is 0. -func (r *Runtime) logSearchMetrics(toolCallID int64, m *store.SearchMetrics) { - if r.telemetry == nil || toolCallID == 0 || m == nil { +// logSearchMetrics mirrors the recall search_metrics row. No-op when the +// captured telemetry client is nil or the parent tool_call id is 0. Same +// captured-pointer rationale as logToolCall. +func (r *Runtime) logSearchMetrics(tele *telemetry.Client, toolCallID int64, m *store.SearchMetrics) { + if tele == nil || toolCallID == 0 || m == nil { return } - r.telemetry.LogSearchMetrics(telemetry.SearchMetricsInput{ + tele.LogSearchMetrics(telemetry.SearchMetricsInput{ ToolCallID: toolCallID, Query: m.Query, Channels: m.Channels, diff --git a/internal/telemetry/client.go b/internal/telemetry/client.go index f304cae..c402f0c 100644 --- a/internal/telemetry/client.go +++ b/internal/telemetry/client.go @@ -31,12 +31,20 @@ import ( // SQLite Exec itself (microseconds), so there is no measurable hot-path // impact. strict is set at construction and never mutated, so it is read // outside the lock. +// +// degraded flips to true on the first Exec failure (disk full, permissions +// change, DB corruption) so subsequent Log* calls return silently instead +// of spamming stderr for every tool call. The underlying handles are left +// open — Close still cleans them up on shutdown — but the hot path becomes +// a cheap no-op for the rest of the session. Recovery requires a restart, +// which is the same contract as init failure. type Client struct { mu sync.Mutex db *sql.DB insertCall *sql.Stmt insertSearch *sql.Stmt strict bool + degraded bool } // FromEnv reads MEMORY_TELEMETRY_PATH and MEMORY_TELEMETRY_PRIVACY from the @@ -188,7 +196,7 @@ func (c *Client) LogToolCall(in ToolCallInput) int64 { } c.mu.Lock() defer c.mu.Unlock() - if c.db == nil || c.insertCall == nil { + if c.db == nil || c.insertCall == nil || c.degraded { return 0 } dbScope := in.DBScope @@ -208,7 +216,8 @@ func (c *Client) LogToolCall(in ToolCallInput) int64 { boolToInt(in.IsError), ) if err != nil { - fmt.Fprintf(os.Stderr, "[memory] telemetry log failed: %v\n", err) + c.degraded = true + fmt.Fprintf(os.Stderr, "[memory] telemetry log failed (further errors suppressed for this session): %v\n", err) return 0 } id, err := res.LastInsertId() @@ -243,7 +252,7 @@ func (c *Client) LogSearchMetrics(in SearchMetricsInput) { } c.mu.Lock() defer c.mu.Unlock() - if c.db == nil || c.insertSearch == nil { + if c.db == nil || c.insertSearch == nil || c.degraded { return } channelsJSON, err := json.Marshal(in.Channels) @@ -263,7 +272,8 @@ func (c *Client) LogSearchMetrics(in SearchMetricsInput) { in.ScoreMedian, boolToInt(in.Compact), ); err != nil { - fmt.Fprintf(os.Stderr, "[memory] telemetry search log failed: %v\n", err) + c.degraded = true + fmt.Fprintf(os.Stderr, "[memory] telemetry search log failed (further errors suppressed for this session): %v\n", err) } } diff --git a/internal/telemetry/client_test.go b/internal/telemetry/client_test.go index c91fe2f..d853cc0 100644 --- a/internal/telemetry/client_test.go +++ b/internal/telemetry/client_test.go @@ -1,7 +1,10 @@ package telemetry import ( + "bytes" "database/sql" + "io" + "os" "path/filepath" "strings" "sync" @@ -226,6 +229,63 @@ func TestClientCloseRacesWithLogging(t *testing.T) { c.LogSearchMetrics(SearchMetricsInput{ToolCallID: 99, Query: "post"}) } +// After the first Exec failure the client must flip to degraded mode and +// stop emitting stderr warnings on subsequent Log* calls — otherwise a +// wedged telemetry DB (disk full, permissions change, corruption) would +// spam stderr once per tool call for the rest of the session. +func TestClientDegradedModeSuppressesSpam(t *testing.T) { + path := filepath.Join(t.TempDir(), "degraded.db") + c := InitIfEnabled(path, false) + if c == nil { + t.Fatalf("init failed") + } + t.Cleanup(func() { _ = c.Close() }) + + // Redirect stderr into a pipe so we can count emitted warnings. Restore + // on cleanup even if the test fails early. + origStderr := os.Stderr + pr, pw, err := os.Pipe() + if err != nil { + t.Fatalf("pipe: %v", err) + } + os.Stderr = pw + t.Cleanup(func() { + os.Stderr = origStderr + }) + + // Force the next Exec to fail by closing the underlying sql.DB. The + // prepared statements still reference the now-closed pool so Exec will + // surface a "database is closed" error — exactly the kind of persistent + // failure degraded mode must mute after the first warning. + c.mu.Lock() + _ = c.db.Close() + c.mu.Unlock() + + // First call: Exec fails, warning emitted, degraded set. + if id := c.LogToolCall(ToolCallInput{Tool: "remember"}); id != 0 { + t.Fatalf("first LogToolCall after forced failure returned %d, want 0", id) + } + // Second + third calls: degraded now blocks Exec; no new warnings. + _ = c.LogToolCall(ToolCallInput{Tool: "remember"}) + c.LogSearchMetrics(SearchMetricsInput{ToolCallID: 1, Query: "q"}) + + // Close the write side and drain stderr. + _ = pw.Close() + os.Stderr = origStderr + var buf bytes.Buffer + if _, err := io.Copy(&buf, pr); err != nil { + t.Fatalf("drain stderr pipe: %v", err) + } + output := buf.String() + + if n := strings.Count(output, "telemetry log failed"); n != 1 { + t.Fatalf("expected exactly 1 'telemetry log failed' warning, got %d:\n%s", n, output) + } + if !c.degraded { + t.Fatalf("degraded flag not set after forced Exec failure") + } +} + func TestLogSearchMetricsZeroToolCallIDIsNoop(t *testing.T) { path := filepath.Join(t.TempDir(), "noop.db") c := InitIfEnabled(path, false)