diff --git a/.gitignore b/.gitignore index 8ebb177..21f8f37 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,4 @@ Thumbs.db # Solr data (managed by Docker volume) solr-data/ /solr-mem-indexer +/solr-mem-bench diff --git a/Makefile b/Makefile index e776d44..11f7bd6 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: build build-indexer install install-indexer test tidy run dev up down logs reset docker-build docker-up config systemd-install systemd-uninstall indexer-install indexer-uninstall launchd-install launchd-install-server launchd-install-indexer launchd-uninstall skill-install skill-uninstall +.PHONY: build build-indexer build-bench install install-indexer test tidy run dev up down logs reset docker-build docker-up config systemd-install systemd-uninstall indexer-install indexer-uninstall launchd-install launchd-install-server launchd-install-indexer launchd-uninstall skill-install skill-uninstall bench # Go targets build: @@ -7,7 +7,19 @@ build: build-indexer: go build -o bin/solr-mem-indexer ./cmd/solr-mem-indexer -build-all: build build-indexer +build-all: build build-indexer build-bench + +build-bench: + go build -o bin/solr-mem-bench ./cmd/solr-mem-bench + +# Retrieval benchmark. Seeds the memories collection with namespaced bench-* +# docs (safe to run against a live collection — only touches bench-* IDs) and +# runs the shipped query set. Override BENCH_URL to point elsewhere. +BENCH_URL ?= http://pax89.local:8983/solr/memories +bench: build-bench + ./bin/solr-mem-bench -solr-url $(BENCH_URL) -seed \ + -corpus cmd/solr-mem-bench/testdata/corpus.jsonl \ + -queries cmd/solr-mem-bench/testdata/queries.jsonl install: go install ./cmd/solr-mem-server diff --git a/cmd/solr-mem-bench/main.go b/cmd/solr-mem-bench/main.go new file mode 100644 index 0000000..d08f52f --- /dev/null +++ b/cmd/solr-mem-bench/main.go @@ -0,0 +1,230 @@ +// solr-mem-bench is a small harness for measuring memory retrieval quality. +// +// Usage: +// +// solr-mem-bench -solr-url http://localhost:8983/solr -collection memories_bench \ +// -corpus testdata/corpus.jsonl -queries testdata/queries.jsonl -seed +// +// The -seed flag clears the target collection and writes the corpus before +// running queries. Without -seed, the harness assumes the collection is +// already populated with the gold IDs referenced by queries. +// +// Gold format (JSONL, one per line): +// +// corpus: {"id":"mem-001","title":"...","content":"...","tags":["..."]} +// queries: {"id":"q-001","text":"database N+1 fix","gold":["mem-007"]} +package main + +import ( + "bufio" + "context" + "encoding/json" + "flag" + "fmt" + "log" + "os" + "time" + + "github.com/arreyder/solr-mem/internal/solr" +) + +type corpusItem struct { + ID string `json:"id"` + Title string `json:"title"` + Content string `json:"content"` + Tags []string `json:"tags,omitempty"` +} + +type queryItem struct { + ID string `json:"id"` + Text string `json:"text"` + Gold []string `json:"gold"` +} + +func main() { + var ( + solrURL = flag.String("solr-url", "http://localhost:8983/solr/memories", "Base URL of the memories collection") + corpusPath = flag.String("corpus", "cmd/solr-mem-bench/testdata/corpus.jsonl", "Path to corpus JSONL") + queryPath = flag.String("queries", "cmd/solr-mem-bench/testdata/queries.jsonl", "Path to queries JSONL") + seed = flag.Bool("seed", false, "Clear the collection and write the corpus before running queries") + topK = flag.Int("topk", 10, "Number of hits to retrieve per query") + variant = flag.String("variant", "bm25", "Retrieval variant name (for labeling output)") + ) + flag.Parse() + + corpus, err := loadCorpus(*corpusPath) + if err != nil { + log.Fatalf("load corpus: %v", err) + } + queries, err := loadQueries(*queryPath) + if err != nil { + log.Fatalf("load queries: %v", err) + } + log.Printf("loaded %d corpus items, %d queries", len(corpus), len(queries)) + + client := solr.NewClient(*solrURL) + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + if *seed { + log.Printf("seeding collection at %s (clearing first)", *solrURL) + if err := seedCorpus(ctx, client, corpus); err != nil { + log.Fatalf("seed: %v", err) + } + // Give Solr a moment to commit. + time.Sleep(1 * time.Second) + } + + agg := NewAggregate(1, 3, 5, 10) + perQuery := make([]queryResult, 0, len(queries)) + + start := time.Now() + for _, q := range queries { + hits, err := runQuery(ctx, client, q.Text, *topK) + if err != nil { + log.Printf("query %s failed: %v", q.ID, err) + continue + } + agg.Add(q.Gold, hits) + perQuery = append(perQuery, queryResult{ID: q.ID, Text: q.Text, Hits: hits, Gold: q.Gold}) + } + elapsed := time.Since(start) + + agg.Finalize() + printReport(*variant, agg, perQuery, elapsed) +} + +type queryResult struct { + ID string + Text string + Hits []string + Gold []string +} + +func loadCorpus(path string) ([]corpusItem, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + var out []corpusItem + s := bufio.NewScanner(f) + s.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for s.Scan() { + line := s.Bytes() + if len(line) == 0 { + continue + } + var item corpusItem + if err := json.Unmarshal(line, &item); err != nil { + return nil, fmt.Errorf("parse corpus line: %w", err) + } + out = append(out, item) + } + return out, s.Err() +} + +func loadQueries(path string) ([]queryItem, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + var out []queryItem + s := bufio.NewScanner(f) + s.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for s.Scan() { + line := s.Bytes() + if len(line) == 0 { + continue + } + var item queryItem + if err := json.Unmarshal(line, &item); err != nil { + return nil, fmt.Errorf("parse query line: %w", err) + } + out = append(out, item) + } + return out, s.Err() +} + +// seedCorpus writes the bench corpus to Solr. It scopes its cleanup to +// IDs starting with "bench-" so it is safe to run against a live memories +// collection — it will only ever touch its own namespaced docs. +func seedCorpus(ctx context.Context, client *solr.Client, items []corpusItem) error { + for _, it := range items { + if !isBenchID(it.ID) { + return fmt.Errorf("corpus id %q must start with 'bench-' for safe seeding", it.ID) + } + } + if err := client.DeleteByQuery(ctx, "id:bench-*"); err != nil { + return fmt.Errorf("clear bench docs: %w", err) + } + docs := make([]solr.Document, 0, len(items)) + now := time.Now().UTC() + for _, it := range items { + docs = append(docs, solr.Document{ + ID: it.ID, + Title: it.Title, + Content: it.Content, + Tags: it.Tags, + CreatedAt: now, + UpdatedAt: now, + Lifetime: "permanent", + Format: "prose", + }) + } + return client.Add(ctx, docs...) +} + +func isBenchID(s string) bool { + return len(s) > 6 && s[:6] == "bench-" +} + +func runQuery(ctx context.Context, client *solr.Client, text string, topK int) ([]string, error) { + params := solr.QueryParams{ + Query: text, + Rows: topK, + Fields: []string{"id"}, + } + resp, err := client.Query(ctx, params) + if err != nil { + return nil, err + } + out := make([]string, 0, len(resp.Docs)) + for _, d := range resp.Docs { + if id, ok := d["id"].(string); ok { + out = append(out, id) + } + } + return out, nil +} + +func printReport(variant string, agg *Aggregate, results []queryResult, elapsed time.Duration) { + fmt.Printf("# solr-mem retrieval benchmark\n\n") + fmt.Printf("Variant: `%s` | Queries: %d | Elapsed: %s\n\n", variant, agg.Queries, elapsed.Round(time.Millisecond)) + fmt.Printf("| Metric | Value |\n|---|---|\n") + for _, k := range []int{1, 3, 5, 10} { + if v, ok := agg.RecallAt[k]; ok { + fmt.Printf("| R@%d | %.3f |\n", k, v) + } + } + fmt.Printf("| MRR | %.3f |\n\n", agg.MRR) + + fmt.Printf("## Per-query breakdown\n\n") + fmt.Printf("| Query | Gold | R@5 | MRR | Top hits |\n|---|---|---|---|---|\n") + for _, r := range results { + top := r.Hits + if len(top) > 5 { + top = top[:5] + } + fmt.Printf("| %s (%s) | %v | %.2f | %.2f | %v |\n", + r.ID, truncate(r.Text, 40), r.Gold, RecallAtK(r.Gold, r.Hits, 5), MRR(r.Gold, r.Hits), top) + } +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n-1] + "…" +} diff --git a/cmd/solr-mem-bench/metrics.go b/cmd/solr-mem-bench/metrics.go new file mode 100644 index 0000000..6853dcc --- /dev/null +++ b/cmd/solr-mem-bench/metrics.go @@ -0,0 +1,78 @@ +package main + +// RecallAtK returns the fraction of gold IDs that appear in the top-K hits. +// gold is the set of relevant IDs for a query; hits is the ranked result list. +func RecallAtK(gold, hits []string, k int) float64 { + if len(gold) == 0 { + return 0 + } + if k <= 0 || k > len(hits) { + k = len(hits) + } + goldSet := make(map[string]struct{}, len(gold)) + for _, g := range gold { + goldSet[g] = struct{}{} + } + found := 0 + for i := 0; i < k; i++ { + if _, ok := goldSet[hits[i]]; ok { + found++ + } + } + return float64(found) / float64(len(gold)) +} + +// MRR returns the reciprocal rank of the first gold hit in the ranked list, +// or 0 if no gold item was retrieved. Ranks are 1-indexed. +func MRR(gold, hits []string) float64 { + if len(gold) == 0 || len(hits) == 0 { + return 0 + } + goldSet := make(map[string]struct{}, len(gold)) + for _, g := range gold { + goldSet[g] = struct{}{} + } + for i, h := range hits { + if _, ok := goldSet[h]; ok { + return 1.0 / float64(i+1) + } + } + return 0 +} + +// Aggregate holds per-run averaged metrics. +type Aggregate struct { + Queries int + RecallAt map[int]float64 // k -> mean recall + MRR float64 +} + +// NewAggregate accumulates per-query metrics over a set of queries. +func NewAggregate(ks ...int) *Aggregate { + m := make(map[int]float64, len(ks)) + for _, k := range ks { + m[k] = 0 + } + return &Aggregate{RecallAt: m} +} + +// Add records one query's hits against its gold. +func (a *Aggregate) Add(gold, hits []string) { + a.Queries++ + for k := range a.RecallAt { + a.RecallAt[k] += RecallAtK(gold, hits, k) + } + a.MRR += MRR(gold, hits) +} + +// Finalize divides sums by query count, producing mean metrics. +func (a *Aggregate) Finalize() { + if a.Queries == 0 { + return + } + n := float64(a.Queries) + for k, v := range a.RecallAt { + a.RecallAt[k] = v / n + } + a.MRR /= n +} diff --git a/cmd/solr-mem-bench/metrics_test.go b/cmd/solr-mem-bench/metrics_test.go new file mode 100644 index 0000000..a447ee5 --- /dev/null +++ b/cmd/solr-mem-bench/metrics_test.go @@ -0,0 +1,111 @@ +package main + +import ( + "math" + "testing" +) + +func approx(a, b float64) bool { return math.Abs(a-b) < 1e-9 } + +func TestRecallAtKPerfect(t *testing.T) { + gold := []string{"a", "b"} + hits := []string{"a", "b", "c", "d", "e"} + if got := RecallAtK(gold, hits, 5); !approx(got, 1.0) { + t.Errorf("R@5 should be 1.0, got %f", got) + } +} + +func TestRecallAtKPartial(t *testing.T) { + gold := []string{"a", "b", "c"} + hits := []string{"a", "x", "y", "z", "b"} + // 2 of 3 gold found in top 5. + if got := RecallAtK(gold, hits, 5); !approx(got, 2.0/3.0) { + t.Errorf("expected 2/3, got %f", got) + } + // Only 1 in top 2. + if got := RecallAtK(gold, hits, 2); !approx(got, 1.0/3.0) { + t.Errorf("expected 1/3, got %f", got) + } +} + +func TestRecallAtKNoHits(t *testing.T) { + gold := []string{"a"} + hits := []string{"x", "y", "z"} + if got := RecallAtK(gold, hits, 3); !approx(got, 0) { + t.Errorf("expected 0, got %f", got) + } +} + +func TestRecallAtKKLargerThanHits(t *testing.T) { + gold := []string{"a"} + hits := []string{"a", "b"} + if got := RecallAtK(gold, hits, 10); !approx(got, 1.0) { + t.Errorf("k > len(hits) should not error, got %f", got) + } +} + +func TestRecallAtKEmptyGold(t *testing.T) { + if got := RecallAtK(nil, []string{"a"}, 5); got != 0 { + t.Errorf("no gold → 0, got %f", got) + } +} + +func TestMRRFirstPosition(t *testing.T) { + gold := []string{"a"} + hits := []string{"a", "b", "c"} + if got := MRR(gold, hits); !approx(got, 1.0) { + t.Errorf("first-rank gold → 1.0, got %f", got) + } +} + +func TestMRRThirdPosition(t *testing.T) { + gold := []string{"c"} + hits := []string{"a", "b", "c"} + if got := MRR(gold, hits); !approx(got, 1.0/3.0) { + t.Errorf("rank 3 → 1/3, got %f", got) + } +} + +func TestMRRNoMatch(t *testing.T) { + gold := []string{"z"} + hits := []string{"a", "b", "c"} + if got := MRR(gold, hits); got != 0 { + t.Errorf("no match → 0, got %f", got) + } +} + +func TestMRRMultipleGoldTakesBest(t *testing.T) { + // Two gold items; MRR uses the best (earliest) rank. + gold := []string{"c", "a"} + hits := []string{"a", "b", "c"} + if got := MRR(gold, hits); !approx(got, 1.0) { + t.Errorf("best rank should win: expected 1.0, got %f", got) + } +} + +func TestAggregateEndToEnd(t *testing.T) { + agg := NewAggregate(5, 10) + agg.Add([]string{"a"}, []string{"a", "b", "c"}) // R@5=1, R@10=1, MRR=1 + agg.Add([]string{"x"}, []string{"a", "b", "c"}) // all zero + agg.Add([]string{"c"}, []string{"a", "b", "c"}) // R@5=1, R@10=1, MRR=1/3 + agg.Finalize() + + // Means across 3 queries: R@5 = 2/3, R@10 = 2/3, MRR = (1 + 0 + 1/3)/3 = 4/9 + if !approx(agg.RecallAt[5], 2.0/3.0) { + t.Errorf("R@5 mean expected 2/3, got %f", agg.RecallAt[5]) + } + if !approx(agg.RecallAt[10], 2.0/3.0) { + t.Errorf("R@10 mean expected 2/3, got %f", agg.RecallAt[10]) + } + if !approx(agg.MRR, 4.0/9.0) { + t.Errorf("MRR mean expected 4/9, got %f", agg.MRR) + } +} + +func TestAggregateEmpty(t *testing.T) { + agg := NewAggregate(5) + agg.Finalize() // should not panic on zero queries + if agg.MRR != 0 || agg.RecallAt[5] != 0 { + t.Errorf("empty aggregate should stay zero") + } +} diff --git a/cmd/solr-mem-bench/testdata/corpus.jsonl b/cmd/solr-mem-bench/testdata/corpus.jsonl new file mode 100644 index 0000000..b71784f --- /dev/null +++ b/cmd/solr-mem-bench/testdata/corpus.jsonl @@ -0,0 +1,30 @@ +{"id":"bench-001","title":"JWT auth via jose middleware","content":"Project uses jose (not jsonwebtoken) for Edge runtime compatibility. Middleware in src/middleware/auth.ts validates Authorization: Bearer tokens, decodes JWT claims, attaches user to request context. Token signing uses RS256 with keys rotated weekly.","tags":["auth","jwt","middleware","edge"]} +{"id":"bench-002","title":"N+1 query fix on orders list","content":"OrdersController.list was fetching each order's customer in a separate SELECT, causing N+1. Switched to a single JOIN with customers table. p95 latency dropped from 1200ms to 140ms on the orders index page. Related to ticket ORD-413.","tags":["database","performance","n+1","orders"]} +{"id":"bench-003","title":"Rate limiting with Redis token bucket","content":"API gateway uses a sliding-window token-bucket algorithm backed by Redis. 100 requests per minute per API key, 1000 per org. Implementation in pkg/ratelimit/bucket.go. Tokens refill in 600ms increments to avoid thundering herd at minute boundaries.","tags":["rate-limit","redis","api","gateway"]} +{"id":"bench-004","title":"Postgres connection pool sizing","content":"Production pool maxed at 100 connections per service instance, 5 instances → 500 total. PgBouncer in transaction mode in front. Lowering pool size to 50 per instance reduced idle connection overhead by 30% with no throughput loss.","tags":["postgres","connection-pool","pgbouncer","performance"]} +{"id":"bench-005","title":"React context provider pattern for feature flags","content":"FeatureFlagProvider wraps the app root, hydrates flags once from /api/flags at mount, exposes useFeatureFlag(name) hook. Flag changes from admin UI push via WebSocket and trigger provider re-render. Stale-while-revalidate fallback to localStorage for offline.","tags":["react","feature-flags","websocket","frontend"]} +{"id":"bench-006","title":"Memory leak in WebSocket reconnect handler","content":"The reconnect exponential-backoff timer kept a closure reference to the old socket, preventing GC. Fix: explicitly null socket.onmessage and socket.onclose in cleanup, store the current timer ID outside the closure so successive retries replace it.","tags":["websocket","memory-leak","javascript","frontend"]} +{"id":"bench-007","title":"Slow Solr range facet on timestamp field","content":"facet.range on created_at with gap=+1DAY over 90 days was taking 3s. Switched to docValues=true on the pdate field and enabled facet.method=dv. Latency dropped to 80ms. Applies to any range faceting on high-cardinality date fields.","tags":["solr","facet","performance","date"]} +{"id":"bench-008","title":"Docker Compose healthcheck for Solr","content":"solr service in docker-compose.yml needs healthcheck: curl -sf http://localhost:8983/solr/admin/ping. Without it, dependent services start before the configset is loaded, causing schema-not-found errors on first run. 10s interval, 5 retries.","tags":["docker","solr","healthcheck","devops"]} +{"id":"bench-009","title":"TypeScript strict mode migration","content":"Migrated tsconfig with strict: true in stages: strictNullChecks first (revealed 180 null-deref bugs), then noImplicitAny (caught 40 places passing any to third-party APIs). Suppress with // @ts-expect-error initially, burn down incrementally.","tags":["typescript","migration","strict","types"]} +{"id":"bench-010","title":"Kafka consumer lag alerting","content":"Alert fires when consumer group lag exceeds 10k for 5 minutes on any partition. Check by: kafka-consumer-groups --describe --group X. Most common root cause: slow downstream DB writes. Fix is usually raising batch size, not scaling consumers.","tags":["kafka","alerting","lag","monitoring"]} +{"id":"bench-011","title":"GraphQL N+1 via DataLoader","content":"Nested resolvers for User.posts and Post.comments were firing one query per row. Wrapped each resolver in a DataLoader that batches IDs per request. Cut DB round-trips on the feed query from 300 to 3.","tags":["graphql","n+1","dataloader","database"]} +{"id":"bench-012","title":"AWS S3 presigned URL expiry","content":"Upload URLs signed with 15-min expiry. Users on slow networks occasionally fail right at the boundary. Bumped to 60 min for uploads, kept 5 min for downloads. Monitoring the RequestExpired 4xx count on the bucket metrics dashboard.","tags":["aws","s3","presigned-url","upload"]} +{"id":"bench-013","title":"Go generics for repository pattern","content":"Replaced a dozen Find/Create/Delete methods across 6 repos with one generic[T any] Repository struct. Used constraints.Ordered for sortable fields. Reduced LOC by ~40% and caught two places where type mismatches were silently compiling.","tags":["go","generics","repository","refactor"]} +{"id":"bench-014","title":"Flaky integration test on CI","content":"auth_integration_test.go failed 1-in-20 on GitHub Actions, always passed locally. Root cause: test reused a Redis key across parallel goroutines; local runs happened to serialize. Fix: use t.Name() as key prefix to isolate.","tags":["flaky","test","redis","ci"]} +{"id":"bench-015","title":"CORS preflight cache tuning","content":"Access-Control-Max-Age defaulted to 5s; browsers were firing OPTIONS every request. Raised to 86400 (1 day). Preflight traffic dropped 98%. Works well because our CORS config is stable.","tags":["cors","http","performance","browser"]} +{"id":"bench-016","title":"Pulumi stack for ephemeral preview envs","content":"Each PR gets a preview environment: Pulumi stack named pr-, S3 bucket, RDS serverless, domain cnamed to preview.example.com. Cleanup lambda removes stacks after PR close or 7 days inactivity. Cost per stack: ~$2/day.","tags":["pulumi","infra","preview","aws"]} +{"id":"bench-017","title":"Prometheus histogram bucket tuning","content":"Default buckets are tuned for web latency (le=0.005..10). For background job duration we used le=1,5,10,30,60,120,300. Too-coarse buckets hide p99 changes; too-fine explodes time-series cardinality. Revisit quarterly.","tags":["prometheus","metrics","histogram","monitoring"]} +{"id":"bench-018","title":"ElasticSearch to Solr migration","content":"Migrated search from ES 7.x to Solr 9. Key differences: Solr uses managed-schema.xml vs ES mappings, edismax is the edismax parser (no dis_max clause), field boosting syntax is title^2 not title:(query)^2. Faceting syntax maps cleanly 1:1.","tags":["elasticsearch","solr","migration","search"]} +{"id":"bench-019","title":"Stripe webhook idempotency","content":"Store the webhook event ID in a processed_events table with a unique constraint. Reject duplicates with 200 (not 4xx, Stripe will retry). 90-day retention then expire. Saw duplicate deliveries ~0.3% of the time.","tags":["stripe","webhook","idempotency","payments"]} +{"id":"bench-020","title":"pgx vs database/sql connection handling","content":"Switched from lib/pq to pgx/v5. pgx.Conn is stateful per-connection (prepared statements cached locally) so it doesn't play with PgBouncer in transaction mode. Use pgx.Pool with simple_protocol=true or stay with database/sql.","tags":["postgres","pgx","go","connection"]} +{"id":"bench-021","title":"nginx buffer size for large POST bodies","content":"Uploads >1MB were hitting client_body_buffer_size limits and spilling to /tmp on nginx hosts. Raised to 10m and moved temp dir to tmpfs. 413 errors went to zero. Also bumped client_max_body_size to 100m for the upload endpoint.","tags":["nginx","upload","buffer","http"]} +{"id":"bench-022","title":"Redis key prefix conventions","content":"We use :: (e.g. auth:session:abc123). Makes KEYS/SCAN across a namespace easy and lets us purge per-service with a single pattern. Avoid ::: double-colons; Redis pattern matching is happier with single separators.","tags":["redis","conventions","cache","keys"]} +{"id":"bench-023","title":"Go pprof on production pods","content":"Enabled net/http/pprof on /debug/pprof, bound to localhost only, exposed via kubectl port-forward on demand. Used pprof -http to analyze heap dumps. Found a string-append loop in JSON marshaling that was 40% of allocations.","tags":["go","pprof","performance","profiling"]} +{"id":"bench-024","title":"Datadog log pattern cardinality","content":"Log messages containing user-supplied UUIDs blew up pattern aggregation. Now we scrub UUIDs to in log formatter before shipping. Cardinality on the /patterns query dropped 100x, making anomaly detection useful again.","tags":["datadog","logs","cardinality","observability"]} +{"id":"bench-025","title":"Linear API rate limit","content":"Linear's GraphQL API caps at 1500 requests per hour per API key. Issue sync job batches 50 issues per request. Retry-After header is respected. Hit the cap once during a full re-import — now we paginate with 1s jitter between pages.","tags":["linear","api","rate-limit","integration"]} +{"id":"bench-026","title":"macOS launchd user agent for background service","content":"Created ~/Library/LaunchAgents/com.example.service.plist with KeepAlive and RunAtLoad. Logs to /tmp, environment vars set under EnvironmentVariables dict. launchctl load starts immediately and on login. Beats writing a custom systemd wrapper.","tags":["macos","launchd","service","daemon"]} +{"id":"bench-027","title":"Frontend bundle size budget enforcement","content":"bundle-analyzer in CI fails if main.js > 400KB gzipped. Broke twice this quarter — once from an accidental import of all of lodash, once from a moment.js locale bundle. Now we lint for barrel imports and use dayjs exclusively.","tags":["frontend","bundle","webpack","ci"]} +{"id":"bench-028","title":"PostgreSQL partial index for sparse column","content":"created_at IS NOT NULL on a 50M row table; only ~200k rows have it set. Regular btree took 800MB; partial index (WHERE created_at IS NOT NULL) is 12MB. Query planner picks it correctly after ANALYZE.","tags":["postgres","index","partial","performance"]} +{"id":"bench-029","title":"Claude Code hooks for auto-capture","content":"settings.json PostToolUse hook runs a shell script that parses the tool event JSON, filters for file-touching tools, and POSTs to a local observer endpoint. UserPromptSubmit mirrors prompts to a journaling log. Zero-effort session recording.","tags":["claude-code","hooks","observability","agent"]} +{"id":"bench-030","title":"Solr more-like-this for related memories","content":"MLT on title,content with mintf=2 mindf=2 gives decent recall for 'memories similar to X'. Filter out the source doc id. For memory system: raise mintf for short memories so we don't match on single common tokens.","tags":["solr","mlt","similarity","search"]} diff --git a/cmd/solr-mem-bench/testdata/queries.jsonl b/cmd/solr-mem-bench/testdata/queries.jsonl new file mode 100644 index 0000000..df1a2b5 --- /dev/null +++ b/cmd/solr-mem-bench/testdata/queries.jsonl @@ -0,0 +1,25 @@ +{"id":"q-01","text":"JWT authentication middleware","gold":["bench-001"]} +{"id":"q-02","text":"rate limiting API gateway","gold":["bench-003"]} +{"id":"q-03","text":"Postgres connection pool sizing","gold":["bench-004"]} +{"id":"q-04","text":"feature flag React provider","gold":["bench-005"]} +{"id":"q-05","text":"Stripe webhook idempotency","gold":["bench-019"]} +{"id":"q-06","text":"Kafka consumer lag alerts","gold":["bench-010"]} +{"id":"q-07","text":"database performance optimization","gold":["bench-002","bench-011","bench-028","bench-007"]} +{"id":"q-08","text":"secure token validation edge runtime","gold":["bench-001"]} +{"id":"q-09","text":"fixing slow search facets on dates","gold":["bench-007"]} +{"id":"q-10","text":"preventing duplicate payment processing","gold":["bench-019"]} +{"id":"q-11","text":"browser preflight caching","gold":["bench-015"]} +{"id":"q-12","text":"monitoring job duration histograms","gold":["bench-017"]} +{"id":"q-13","text":"find related records efficiently","gold":["bench-011","bench-030"]} +{"id":"q-14","text":"fixing a memory leak in reconnection logic","gold":["bench-006"]} +{"id":"q-15","text":"ephemeral environments per pull request","gold":["bench-016"]} +{"id":"q-16","text":"upload large files HTTP body limit","gold":["bench-021","bench-012"]} +{"id":"q-17","text":"type safety migration staged rollout","gold":["bench-009"]} +{"id":"q-18","text":"flaky tests on continuous integration","gold":["bench-014"]} +{"id":"q-19","text":"moving from Elasticsearch","gold":["bench-018"]} +{"id":"q-20","text":"launchd plist for a background service","gold":["bench-026"]} +{"id":"q-21","text":"Go profiling heap allocations","gold":["bench-023"]} +{"id":"q-22","text":"N+1 query with GraphQL resolvers","gold":["bench-011","bench-002"]} +{"id":"q-23","text":"automatic observation capture from agent hooks","gold":["bench-029"]} +{"id":"q-24","text":"generic repository pattern in Go","gold":["bench-013"]} +{"id":"q-25","text":"scrubbing high-cardinality identifiers from logs","gold":["bench-024"]} diff --git a/cmd/solr-mem-server/broker.go b/cmd/solr-mem-server/broker.go index 13bfca4..75a938d 100644 --- a/cmd/solr-mem-server/broker.go +++ b/cmd/solr-mem-server/broker.go @@ -36,6 +36,11 @@ const ( // defaultRunSweeperInterval is how often the stale-run sweeper runs. defaultRunSweeperInterval = 5 * time.Minute + + // brokerSessionCap is the max number of memory items per session_id in a + // single packet. Code items (no session) are unaffected. With maxItems=5 + // this leaves room for at least 3 distinct sessions (or a mix with code). + brokerSessionCap = 2 ) // WorkObservation is a single structured report from a worker agent. @@ -65,6 +70,7 @@ type MemoryPacketItem struct { MemoryType string `json:"memory_type,omitempty"` FilePath string `json:"file_path,omitempty"` SymbolName string `json:"symbol_name,omitempty"` + SessionID string `json:"session_id,omitempty"` // populated for memory items; used for session-diversified ranking } // MemoryPacket is a precomputed bundle of relevant context for a worker agent. @@ -320,13 +326,17 @@ func (b *Broker) doBuild(obs WorkObservation, buildSeq int) *MemoryPacket { // 3. Score and dedupe. scored := scoreCandidates(candidates, obs) - // 4. Pick top items (max 5). + // 4. Cap per session so one chatty session can't dominate the packet. + // Code items have empty SessionID and are unaffected. + scored = diversifyBySession(scored, func(i MemoryPacketItem) string { return i.SessionID }, brokerSessionCap) + + // 5. Pick top items (max 5). const maxItems = 5 if len(scored) > maxItems { scored = scored[:maxItems] } - // 5. Determine delivery class. + // 6. Determine delivery class. delivery := DeliveryCheckpoint for _, item := range scored { // Promote to interrupt if we found a high-relevance hazard or prior solution. @@ -336,7 +346,7 @@ func (b *Broker) doBuild(obs WorkObservation, buildSeq int) *MemoryPacket { } } - // 6. Build summary. + // 7. Build summary. summary := buildPacketSummary(scored, obs) return &MemoryPacket{ @@ -371,6 +381,7 @@ func (b *Broker) searchMemories(ctx context.Context, queryTerms string, obs Work title, _ := doc["title"].(string) content, _ := doc["content"].(string) memType, _ := doc["memory_type"].(string) + sessionID, _ := doc["session_id"].(string) tags := getStringSliceFromDoc(doc, "tags") // Truncate content for summary. @@ -388,6 +399,7 @@ func (b *Broker) searchMemories(ctx context.Context, queryTerms string, obs Work Reason: "memory search hit", Tags: tags, MemoryType: memType, + SessionID: sessionID, }) } return items diff --git a/cmd/solr-mem-server/broker_tool.go b/cmd/solr-mem-server/broker_tool.go index 4b5de62..0f1ee69 100644 --- a/cmd/solr-mem-server/broker_tool.go +++ b/cmd/solr-mem-server/broker_tool.go @@ -14,17 +14,24 @@ func observeWorkTool(broker *Broker) ToolHandler { return nil, fmt.Errorf("run_id is required") } + // Scrub free-text fields; entities/code_refs/repo/phase are symbolic + // and unlikely to carry secrets. + task, _ := scrubString(getString(args, "task")) + subgoal, _ := scrubString(getString(args, "subgoal")) + uncertainty, _ := scrubString(getString(args, "uncertainty")) + nextAction, _ := scrubString(getString(args, "next_action")) + obs := WorkObservation{ RunID: runID, AgentID: getString(args, "agent_id"), Repo: getString(args, "repo"), Phase: getString(args, "phase"), - Task: getString(args, "task"), - Subgoal: getString(args, "subgoal"), + Task: task, + Subgoal: subgoal, Entities: getStringSlice(args, "entities"), CodeRefs: getStringSlice(args, "code_refs"), - Uncertainty: getString(args, "uncertainty"), - NextAction: getString(args, "next_action"), + Uncertainty: uncertainty, + NextAction: nextAction, } result := broker.Observe(obs) diff --git a/cmd/solr-mem-server/bulk_store_tool.go b/cmd/solr-mem-server/bulk_store_tool.go index df8ebfb..6f64a75 100644 --- a/cmd/solr-mem-server/bulk_store_tool.go +++ b/cmd/solr-mem-server/bulk_store_tool.go @@ -5,6 +5,7 @@ import ( "fmt" "time" + "github.com/arreyder/solr-mem/internal/privacy" "github.com/arreyder/solr-mem/internal/solr" "github.com/google/uuid" ) @@ -26,6 +27,7 @@ func bulkStoreMemoriesTool(ctx context.Context, args map[string]any) (any, error now := time.Now().UTC() var docs []solr.Document var ids []string + totalScrubbed := 0 for i, raw := range memories { m, ok := raw.(map[string]any) @@ -49,16 +51,26 @@ func bulkStoreMemoriesTool(ctx context.Context, args map[string]any) (any, error format = "prose" } + title := getString(m, "title") + metadata := getString(m, "metadata") + scrubbedContent, contentHits := scrubString(content) + scrubbedTitle, titleHits := scrubString(title) + allHits := privacy.MergeHits(contentHits, titleHits) + metadata = privacy.MergeMetadata(metadata, allHits) + for _, v := range allHits { + totalScrubbed += v + } + docs = append(docs, solr.Document{ ID: id, AgentID: getString(m, "agent_id"), MemoryType: getString(m, "memory_type"), - Content: content, - Title: getString(m, "title"), + Content: scrubbedContent, + Title: scrubbedTitle, Tags: getStringSlice(m, "tags"), Source: getString(m, "source"), Importance: getFloat(m, "importance", 0.5), - Metadata: getString(m, "metadata"), + Metadata: metadata, CreatedAt: now, UpdatedAt: now, ExpiresAt: expiresAt, @@ -73,11 +85,18 @@ func bulkStoreMemoriesTool(ctx context.Context, args map[string]any) (any, error return nil, fmt.Errorf("failed to bulk store memories: %w", err) } + text := fmt.Sprintf("Successfully stored %d memories.\nIDs: %v", len(docs), ids) + structured := map[string]any{ + "count": len(docs), + "ids": ids, + } + if totalScrubbed > 0 { + text += fmt.Sprintf("\nPrivacy: redacted %d secret(s) across all memories", totalScrubbed) + structured["scrub_count"] = totalScrubbed + } + return ToolOutput{ - Text: fmt.Sprintf("Successfully stored %d memories.\nIDs: %v", len(docs), ids), - Structured: map[string]any{ - "count": len(docs), - "ids": ids, - }, + Text: text, + Structured: structured, }, nil } diff --git a/cmd/solr-mem-server/diversify.go b/cmd/solr-mem-server/diversify.go new file mode 100644 index 0000000..2d4a71c --- /dev/null +++ b/cmd/solr-mem-server/diversify.go @@ -0,0 +1,28 @@ +package main + +// diversifyBySession caps the number of items per session in a ranked slice. +// Items with an empty group key are passed through without counting — callers +// use this for hits that have no natural grouping (e.g. code docs without a +// session_id). +// +// cap <= 0 disables diversification and returns the input unchanged. +func diversifyBySession[T any](items []T, key func(T) string, cap int) []T { + if cap <= 0 || len(items) == 0 { + return items + } + counts := make(map[string]int, len(items)) + out := make([]T, 0, len(items)) + for _, item := range items { + k := key(item) + if k == "" { + out = append(out, item) + continue + } + if counts[k] >= cap { + continue + } + counts[k]++ + out = append(out, item) + } + return out +} diff --git a/cmd/solr-mem-server/diversify_test.go b/cmd/solr-mem-server/diversify_test.go new file mode 100644 index 0000000..668efbc --- /dev/null +++ b/cmd/solr-mem-server/diversify_test.go @@ -0,0 +1,96 @@ +package main + +import "testing" + +type divItem struct { + id string + session string +} + +func ids(items []divItem) []string { + out := make([]string, len(items)) + for i, it := range items { + out[i] = it.id + } + return out +} + +func sessionKey(it divItem) string { return it.session } + +func TestDiversifyBySessionCapsPerSession(t *testing.T) { + in := []divItem{ + {"a", "s1"}, + {"b", "s1"}, + {"c", "s1"}, // should be dropped at cap=2 + {"d", "s2"}, + {"e", "s1"}, // dropped + {"f", "s2"}, + {"g", "s3"}, + } + got := diversifyBySession(in, sessionKey, 2) + want := []string{"a", "b", "d", "f", "g"} + if !equalStrings(ids(got), want) { + t.Errorf("cap=2: got %v, want %v", ids(got), want) + } +} + +func TestDiversifyBySessionPreservesOrder(t *testing.T) { + // Top-ranked items should be kept over later same-session items. + in := []divItem{ + {"top", "s1"}, + {"mid", "s2"}, + {"also1", "s1"}, + {"low1", "s1"}, // dropped at cap=2 + {"low2", "s1"}, // dropped + } + got := diversifyBySession(in, sessionKey, 2) + want := []string{"top", "mid", "also1"} + if !equalStrings(ids(got), want) { + t.Errorf("order: got %v, want %v", ids(got), want) + } +} + +func TestDiversifyBySessionEmptyKeyNotCapped(t *testing.T) { + // Items with empty key (e.g. code docs) pass through without counting. + in := []divItem{ + {"a", ""}, + {"b", ""}, + {"c", ""}, + {"d", "s1"}, + {"e", "s1"}, + {"f", "s1"}, // dropped + } + got := diversifyBySession(in, sessionKey, 2) + want := []string{"a", "b", "c", "d", "e"} + if !equalStrings(ids(got), want) { + t.Errorf("empty key: got %v, want %v", ids(got), want) + } +} + +func TestDiversifyBySessionDisabled(t *testing.T) { + in := []divItem{{"a", "s1"}, {"b", "s1"}, {"c", "s1"}} + if got := diversifyBySession(in, sessionKey, 0); len(got) != 3 { + t.Errorf("cap=0 should passthrough, got %d items", len(got)) + } + if got := diversifyBySession(in, sessionKey, -1); len(got) != 3 { + t.Errorf("cap<0 should passthrough, got %d items", len(got)) + } +} + +func TestDiversifyBySessionEmpty(t *testing.T) { + if got := diversifyBySession[divItem](nil, sessionKey, 2); got != nil { + t.Errorf("nil input: got %v, want nil", got) + } +} + +func equalStrings(a, b []string) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} diff --git a/cmd/solr-mem-server/scrub.go b/cmd/solr-mem-server/scrub.go new file mode 100644 index 0000000..6d0695e --- /dev/null +++ b/cmd/solr-mem-server/scrub.go @@ -0,0 +1,42 @@ +package main + +import ( + "log" + "os" + "strings" + "sync" + + "github.com/arreyder/solr-mem/internal/privacy" +) + +// privacyScrubDisabled is resolved once from the environment. +var ( + privacyScrubOnce sync.Once + privacyScrubOff bool +) + +// privacyScrubEnabled reports whether secret scrubbing is on. Controlled by +// SOLR_MEM_PRIVACY_SCRUB=off (default: on). +func privacyScrubEnabled() bool { + privacyScrubOnce.Do(func() { + v := strings.ToLower(strings.TrimSpace(os.Getenv("SOLR_MEM_PRIVACY_SCRUB"))) + if v == "off" || v == "false" || v == "0" || v == "disabled" { + privacyScrubOff = true + log.Println("privacy: scrub disabled by SOLR_MEM_PRIVACY_SCRUB") + } + }) + return !privacyScrubOff +} + +// scrubString returns a scrubbed copy of s and the hit map; if scrubbing is +// disabled or s is empty, the hit map is nil. +func scrubString(s string) (string, map[string]int) { + if !privacyScrubEnabled() || s == "" { + return s, nil + } + r := privacy.Scrub(s) + if r.Count() == 0 { + return r.Content, nil + } + return r.Content, r.Hits +} diff --git a/cmd/solr-mem-server/search_tool.go b/cmd/solr-mem-server/search_tool.go index a18ae80..2efee40 100644 --- a/cmd/solr-mem-server/search_tool.go +++ b/cmd/solr-mem-server/search_tool.go @@ -74,6 +74,16 @@ func searchMemoriesTool(ctx context.Context, args map[string]any) (any, error) { return nil, fmt.Errorf("search failed: %w", err) } + // Cap per session so one chatty session can't dominate results. + // Default 3; pass 0 to disable. + sessionCap := getInt(args, "session_cap", 3) + if sessionCap > 0 { + resp.Docs = diversifyBySession(resp.Docs, func(d map[string]any) string { + s, _ := d["session_id"].(string) + return s + }, sessionCap) + } + return ToolOutput{ Text: formatSearchResults(resp), Structured: resp, diff --git a/cmd/solr-mem-server/store_tool.go b/cmd/solr-mem-server/store_tool.go index 88ddf47..c17459c 100644 --- a/cmd/solr-mem-server/store_tool.go +++ b/cmd/solr-mem-server/store_tool.go @@ -5,6 +5,7 @@ import ( "fmt" "time" + "github.com/arreyder/solr-mem/internal/privacy" "github.com/arreyder/solr-mem/internal/solr" "github.com/google/uuid" ) @@ -24,16 +25,23 @@ func storeMemoryTool(ctx context.Context, args map[string]any) (any, error) { format = "prose" } + title := getString(args, "title") + metadata := getString(args, "metadata") + scrubbedContent, contentHits := scrubString(content) + scrubbedTitle, titleHits := scrubString(title) + allHits := privacy.MergeHits(contentHits, titleHits) + metadata = privacy.MergeMetadata(metadata, allHits) + doc := solr.Document{ ID: uuid.New().String(), AgentID: getString(args, "agent_id"), MemoryType: getString(args, "memory_type"), - Content: content, - Title: getString(args, "title"), + Content: scrubbedContent, + Title: scrubbedTitle, Tags: getStringSlice(args, "tags"), Source: getString(args, "source"), Importance: getFloat(args, "importance", 0.5), - Metadata: getString(args, "metadata"), + Metadata: metadata, CreatedAt: now, UpdatedAt: now, ExpiresAt: expiresAt, @@ -53,13 +61,23 @@ func storeMemoryTool(ctx context.Context, args map[string]any) (any, error) { result += fmt.Sprintf("\nExpires: %s", expiresAt) } + structured := map[string]any{ + "id": doc.ID, + "lifetime": doc.Lifetime, + "expires_at": expiresAt, + "created_at": doc.CreatedAt.Format(time.RFC3339), + } + if n := len(allHits); n > 0 { + total := 0 + for _, v := range allHits { + total += v + } + structured["scrub_count"] = total + result += fmt.Sprintf("\nPrivacy: redacted %d secret(s) across %d kind(s)", total, n) + } + return ToolOutput{ - Text: result, - Structured: map[string]any{ - "id": doc.ID, - "lifetime": doc.Lifetime, - "expires_at": expiresAt, - "created_at": doc.CreatedAt.Format(time.RFC3339), - }, + Text: result, + Structured: structured, }, nil } diff --git a/cmd/solr-mem-server/tools.go b/cmd/solr-mem-server/tools.go index 44d5d0b..f74b853 100644 --- a/cmd/solr-mem-server/tools.go +++ b/cmd/solr-mem-server/tools.go @@ -57,7 +57,7 @@ func ToolSchemas() []ToolDefinition { **When to use**: Find relevant memories by content, tags, type, agent, or time range. Uses edismax with field boosting (content^3, title^2, tags^1.5) and recency boost. **Required**: query (search text) -**Optional**: agent_id, memory_type, tags, source, importance_min, from, to, limit, highlight, facet, session_id, lifetime`, +**Optional**: agent_id, memory_type, tags, source, importance_min, from, to, limit, highlight, facet, session_id, lifetime, session_cap`, InputSchema: NewObjectSchema(map[string]any{ "query": prop("string", "Full-text search query (required)"), "agent_id": prop("string", "Filter by agent ID"), @@ -72,6 +72,7 @@ func ToolSchemas() []ToolDefinition { "facet": prop("boolean", "Include facet counts (default: false)"), "session_id": prop("string", "Filter by session ID"), "lifetime": prop("string", "Filter by lifetime (permanent, session, ephemeral, temporary)"), + "session_cap": integerProp("Max hits per session_id after ranking (default 3, 0 = disable)", intPtr(0), intPtr(100)), }, "query"), }, Handler: searchMemoriesTool, diff --git a/cmd/solr-mem-server/update_tool.go b/cmd/solr-mem-server/update_tool.go index a8e5b24..26d25af 100644 --- a/cmd/solr-mem-server/update_tool.go +++ b/cmd/solr-mem-server/update_tool.go @@ -4,6 +4,8 @@ import ( "context" "fmt" "time" + + "github.com/arreyder/solr-mem/internal/privacy" ) func updateMemoryTool(ctx context.Context, args map[string]any) (any, error) { @@ -13,12 +15,17 @@ func updateMemoryTool(ctx context.Context, args map[string]any) (any, error) { } fields := make(map[string]any) + scrubHits := map[string]int{} if v := getString(args, "content"); v != "" { - fields["content"] = v + scrubbed, hits := scrubString(v) + fields["content"] = scrubbed + scrubHits = privacy.MergeHits(scrubHits, hits) } if v := getString(args, "title"); v != "" { - fields["title"] = v + scrubbed, hits := scrubString(v) + fields["title"] = scrubbed + scrubHits = privacy.MergeHits(scrubHits, hits) } if v := getString(args, "memory_type"); v != "" { fields["memory_type"] = v @@ -35,6 +42,13 @@ func updateMemoryTool(ctx context.Context, args map[string]any) (any, error) { if v := getString(args, "metadata"); v != "" { fields["metadata"] = v } + // If secrets were scrubbed, merge the tally into metadata (existing + // metadata arg takes precedence over the stored doc's metadata, which + // matches atomic-update semantics). + if len(scrubHits) > 0 { + existing, _ := fields["metadata"].(string) + fields["metadata"] = privacy.MergeMetadata(existing, scrubHits) + } if v := getString(args, "session_id"); v != "" { fields["session_id"] = v } diff --git a/internal/privacy/scrub.go b/internal/privacy/scrub.go new file mode 100644 index 0000000..3fb38f9 --- /dev/null +++ b/internal/privacy/scrub.go @@ -0,0 +1,130 @@ +// Package privacy scrubs secrets from strings before they are stored as memory +// content. It is conservative: a missed redaction is preferred over mangling +// legitimate content, so patterns are anchored and stand-alone. +package privacy + +import ( + "encoding/json" + "regexp" + "sort" +) + +// Result is the outcome of a Scrub call. +type Result struct { + // Content is the input string with any matched secrets replaced. + Content string + // Hits maps pattern name -> number of matches redacted. + Hits map[string]int +} + +// Count returns the total number of redactions applied. +func (r Result) Count() int { + n := 0 + for _, v := range r.Hits { + n += v + } + return n +} + +// Kinds returns a sorted list of the pattern names that matched. +func (r Result) Kinds() []string { + out := make([]string, 0, len(r.Hits)) + for k := range r.Hits { + out = append(out, k) + } + sort.Strings(out) + return out +} + +type pattern struct { + name string + re *regexp.Regexp + // replace overrides the default "[REDACTED:]" template. + // Use Go regexp replacement syntax ($1, $2, ...) if needed. + replace string +} + +// Patterns are applied in order. Multi-line / block patterns run first so +// they subsume any single-line keys that would otherwise match inside the +// block. The sk-ant- pattern runs before the generic sk- OpenAI pattern so +// Anthropic keys aren't mislabeled. +var patterns = []pattern{ + {name: "private_key_block", re: regexp.MustCompile(`(?s)-----BEGIN (?:RSA |EC |DSA |OPENSSH |PGP )?PRIVATE KEY(?: BLOCK)?-----.*?-----END[^-]*-----`)}, + {name: "private_tag", re: regexp.MustCompile(`(?is).*?`)}, + {name: "secret_tag", re: regexp.MustCompile(`(?is).*?`)}, + {name: "anthropic_key", re: regexp.MustCompile(`sk-ant-[A-Za-z0-9_\-]{20,}`)}, + {name: "openai_key", re: regexp.MustCompile(`\bsk-[A-Za-z0-9]{48}\b`)}, + {name: "github_pat", re: regexp.MustCompile(`\bgithub_pat_[A-Za-z0-9_]{82}\b`)}, + {name: "github_token", re: regexp.MustCompile(`\bghp_[A-Za-z0-9]{36}\b`)}, + {name: "aws_access_key", re: regexp.MustCompile(`\bAKIA[0-9A-Z]{16}\b`)}, + {name: "aws_secret_key", re: regexp.MustCompile(`(?i)aws_secret_access_key\s*[=:]\s*[A-Za-z0-9/+=]{40}`)}, + {name: "slack_token", re: regexp.MustCompile(`\bxox[baprs]-[A-Za-z0-9\-]+`)}, + {name: "bearer_token", re: regexp.MustCompile(`(?i)bearer\s+[A-Za-z0-9_\-\.=]{20,}`)}, + {name: "url_creds", re: regexp.MustCompile(`(https?://)[^:/@\s]+:[^@\s]+@`), replace: "${1}[REDACTED:url_creds]@"}, +} + +// Scrub replaces recognized secrets in s with [REDACTED:] and returns +// the scrubbed content plus a tally of what was hit. +func Scrub(s string) Result { + res := Result{Content: s, Hits: map[string]int{}} + if s == "" { + return res + } + for _, p := range patterns { + matches := p.re.FindAllStringIndex(res.Content, -1) + if len(matches) == 0 { + continue + } + res.Hits[p.name] = len(matches) + replace := p.replace + if replace == "" { + replace = "[REDACTED:" + p.name + "]" + } + res.Content = p.re.ReplaceAllString(res.Content, replace) + } + return res +} + +// MergeHits combines two hit maps into a new one. +func MergeHits(a, b map[string]int) map[string]int { + out := make(map[string]int, len(a)+len(b)) + for k, v := range a { + out[k] += v + } + for k, v := range b { + out[k] += v + } + return out +} + +// MergeMetadata augments an existing metadata JSON string with scrub_count and +// scrub_kinds fields. If existing is empty or invalid JSON, a fresh object is +// produced. If the result has no hits, existing is returned unchanged. +func MergeMetadata(existing string, hits map[string]int) string { + total := 0 + kinds := make([]string, 0, len(hits)) + for k, v := range hits { + total += v + kinds = append(kinds, k) + } + if total == 0 { + return existing + } + sort.Strings(kinds) + + var obj map[string]any + if existing != "" { + _ = json.Unmarshal([]byte(existing), &obj) + } + if obj == nil { + obj = map[string]any{} + } + obj["scrub_count"] = total + obj["scrub_kinds"] = kinds + + b, err := json.Marshal(obj) + if err != nil { + return existing + } + return string(b) +} diff --git a/internal/privacy/scrub_test.go b/internal/privacy/scrub_test.go new file mode 100644 index 0000000..8248cf5 --- /dev/null +++ b/internal/privacy/scrub_test.go @@ -0,0 +1,249 @@ +package privacy + +import ( + "encoding/json" + "strings" + "testing" +) + +func TestScrubAnthropicKey(t *testing.T) { + // Anthropic keys start with sk-ant-. + in := "My key is sk-ant-api03-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA done" + r := Scrub(in) + if strings.Contains(r.Content, "sk-ant-") { + t.Errorf("expected sk-ant- to be redacted, got: %s", r.Content) + } + if !strings.Contains(r.Content, "[REDACTED:anthropic_key]") { + t.Errorf("expected REDACTED marker, got: %s", r.Content) + } + if r.Hits["anthropic_key"] != 1 { + t.Errorf("expected 1 anthropic_key hit, got %d", r.Hits["anthropic_key"]) + } +} + +func TestScrubAnthropicBeforeOpenAI(t *testing.T) { + // sk-ant- keys must not be labeled as OpenAI keys. + in := "sk-ant-api03-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" + r := Scrub(in) + if r.Hits["openai_key"] != 0 { + t.Errorf("anthropic key should not match openai pattern, got %+v", r.Hits) + } + if r.Hits["anthropic_key"] != 1 { + t.Errorf("expected anthropic match, got %+v", r.Hits) + } +} + +func TestScrubOpenAIKey(t *testing.T) { + // OpenAI keys are sk- followed by exactly 48 alphanumeric chars. + in := "use sk-" + strings.Repeat("a", 48) + " done" + r := Scrub(in) + if !strings.Contains(r.Content, "[REDACTED:openai_key]") { + t.Errorf("expected openai redaction, got: %s", r.Content) + } +} + +func TestScrubGithubToken(t *testing.T) { + in := "token=ghp_1234567890abcdefghijKLMNOPQRSTUVwxyz extra" + r := Scrub(in) + if r.Hits["github_token"] != 1 { + t.Errorf("expected github_token hit, got %+v", r.Hits) + } + if !strings.Contains(r.Content, "[REDACTED:github_token]") { + t.Errorf("expected redaction, got: %s", r.Content) + } +} + +func TestScrubGithubPAT(t *testing.T) { + // github_pat_ has exactly 82 chars of [A-Za-z0-9_] after the prefix. + suffix := strings.Repeat("a", 82) + in := "token: github_pat_" + suffix + r := Scrub(in) + if r.Hits["github_pat"] != 1 { + t.Errorf("expected github_pat hit, got %+v", r.Hits) + } +} + +func TestScrubAWSAccessKey(t *testing.T) { + in := "AWS_KEY=AKIAIOSFODNN7EXAMPLE here" + r := Scrub(in) + if r.Hits["aws_access_key"] != 1 { + t.Errorf("expected aws_access_key hit, got %+v", r.Hits) + } +} + +func TestScrubAWSSecretKey(t *testing.T) { + in := "aws_secret_access_key=wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" + r := Scrub(in) + if r.Hits["aws_secret_key"] != 1 { + t.Errorf("expected aws_secret_key hit, got %+v", r.Hits) + } +} + +func TestScrubSlackToken(t *testing.T) { + in := "slack=xoxb-12345-67890-abcdef leaks" + r := Scrub(in) + if r.Hits["slack_token"] != 1 { + t.Errorf("expected slack_token hit, got %+v", r.Hits) + } +} + +func TestScrubBearerToken(t *testing.T) { + in := "Authorization: Bearer eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.abc" + r := Scrub(in) + if r.Hits["bearer_token"] != 1 { + t.Errorf("expected bearer_token hit, got %+v", r.Hits) + } + if strings.Contains(r.Content, "eyJ") { + t.Errorf("JWT should be redacted, got: %s", r.Content) + } +} + +func TestScrubPrivateKeyBlock(t *testing.T) { + in := "note:\n-----BEGIN RSA PRIVATE KEY-----\nMIIEowIBAAKCAQEAvZ\nnextline\n-----END RSA PRIVATE KEY-----\nend" + r := Scrub(in) + if r.Hits["private_key_block"] != 1 { + t.Errorf("expected private_key_block hit, got %+v", r.Hits) + } + if strings.Contains(r.Content, "MIIEowIBAAK") { + t.Errorf("key body should be gone, got: %s", r.Content) + } +} + +func TestScrubURLCreds(t *testing.T) { + in := "connect https://admin:s3cr3t@db.example.com:5432/foo" + r := Scrub(in) + if r.Hits["url_creds"] != 1 { + t.Errorf("expected url_creds hit, got %+v", r.Hits) + } + if !strings.Contains(r.Content, "https://[REDACTED:url_creds]@db.example.com") { + t.Errorf("expected host preserved with redacted creds, got: %s", r.Content) + } +} + +func TestScrubPrivateTag(t *testing.T) { + in := "before hidden stuff after" + r := Scrub(in) + if r.Hits["private_tag"] != 1 { + t.Errorf("expected private_tag hit, got %+v", r.Hits) + } + if strings.Contains(r.Content, "hidden stuff") { + t.Errorf("private content should be gone, got: %s", r.Content) + } +} + +func TestScrubSecretTag(t *testing.T) { + in := "foo bar" + r := Scrub(in) + if r.Hits["secret_tag"] != 1 { + t.Errorf("expected secret_tag hit, got %+v", r.Hits) + } +} + +func TestScrubEmpty(t *testing.T) { + r := Scrub("") + if r.Count() != 0 { + t.Errorf("empty input should have no hits") + } +} + +func TestScrubNoMatches(t *testing.T) { + in := "this is plain text with no secrets" + r := Scrub(in) + if r.Content != in { + t.Errorf("no matches should leave content unchanged") + } + if r.Count() != 0 { + t.Errorf("no matches should have zero count") + } +} + +func TestScrubMultipleKinds(t *testing.T) { + in := "AKIAIOSFODNN7EXAMPLE and ghp_1234567890abcdefghijKLMNOPQRSTUVwxyz" + r := Scrub(in) + if r.Hits["aws_access_key"] != 1 || r.Hits["github_token"] != 1 { + t.Errorf("expected both hits, got %+v", r.Hits) + } + if r.Count() != 2 { + t.Errorf("expected total count 2, got %d", r.Count()) + } + if got := r.Kinds(); len(got) != 2 || got[0] != "aws_access_key" || got[1] != "github_token" { + t.Errorf("expected sorted kinds, got %v", got) + } +} + +func TestScrubIdempotent(t *testing.T) { + // Scrubbing already-scrubbed content should be a no-op. + in := "sk-ant-api03-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" + first := Scrub(in) + second := Scrub(first.Content) + if second.Count() != 0 { + t.Errorf("re-scrubbing should produce no hits, got %+v", second.Hits) + } + if second.Content != first.Content { + t.Errorf("re-scrubbing changed content: %q -> %q", first.Content, second.Content) + } +} + +func TestMergeMetadataFresh(t *testing.T) { + hits := map[string]int{"anthropic_key": 2, "github_token": 1} + got := MergeMetadata("", hits) + var obj map[string]any + if err := json.Unmarshal([]byte(got), &obj); err != nil { + t.Fatalf("invalid JSON: %v", err) + } + if c, _ := obj["scrub_count"].(float64); int(c) != 3 { + t.Errorf("expected scrub_count=3, got %v", obj["scrub_count"]) + } + kinds, _ := obj["scrub_kinds"].([]any) + if len(kinds) != 2 { + t.Errorf("expected 2 kinds, got %v", kinds) + } +} + +func TestMergeMetadataExisting(t *testing.T) { + existing := `{"source":"chat","importance":0.7}` + hits := map[string]int{"slack_token": 1} + got := MergeMetadata(existing, hits) + var obj map[string]any + if err := json.Unmarshal([]byte(got), &obj); err != nil { + t.Fatalf("invalid JSON: %v", err) + } + if obj["source"] != "chat" { + t.Errorf("existing keys should be preserved, got %v", obj) + } + if c, _ := obj["scrub_count"].(float64); int(c) != 1 { + t.Errorf("expected scrub_count=1, got %v", obj["scrub_count"]) + } +} + +func TestMergeMetadataNoHits(t *testing.T) { + existing := `{"foo":"bar"}` + got := MergeMetadata(existing, map[string]int{}) + if got != existing { + t.Errorf("no hits should leave metadata untouched, got %q", got) + } +} + +func TestMergeMetadataInvalidExisting(t *testing.T) { + // Garbage existing metadata should be replaced by a fresh object. + got := MergeMetadata("not-json", map[string]int{"github_token": 1}) + var obj map[string]any + if err := json.Unmarshal([]byte(got), &obj); err != nil { + t.Fatalf("expected valid JSON: %v", err) + } + if _, ok := obj["scrub_count"]; !ok { + t.Errorf("expected scrub_count key, got %v", obj) + } +} + +func TestMergeHits(t *testing.T) { + a := map[string]int{"anthropic_key": 1, "github_token": 1} + b := map[string]int{"anthropic_key": 2, "slack_token": 1} + got := MergeHits(a, b) + if got["anthropic_key"] != 3 { + t.Errorf("expected sum for overlapping key, got %d", got["anthropic_key"]) + } + if got["github_token"] != 1 || got["slack_token"] != 1 { + t.Errorf("expected non-overlapping keys kept, got %+v", got) + } +}