From 656b3867fb015ef534084c4b15c898308c2265e2 Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Mon, 25 May 2026 01:15:01 -0500 Subject: [PATCH 01/15] feat: gateway optimizer hardening and database protection --- gateway/cmd/iq/main.go | 118 ++- gateway/internal/cache/bench_test.go | 285 ++++++ gateway/internal/cache/lsm.go | 85 ++ gateway/internal/cache/lsm_test.go | 150 +++ gateway/internal/chunker/chunker.go | 212 ++++- gateway/internal/proxy/claude_messages.go | 834 ++++++++++++++++- gateway/internal/proxy/prompt_cache.go | 40 + gateway/internal/proxy/proxy.go | 135 ++- gateway/internal/proxy/proxy_test.go | 83 +- gateway/internal/proxy/span.go | 43 +- gateway/internal/server/run.go | 32 +- gateway/internal/sessions/tracker.go | 26 + gateway/internal/sessions/tracker_test.go | 145 +++ gateway/internal/store/lsm/bloom.go | 1 + scripts/audit_session.py | 1019 +++++++++++++++++++++ web/index.html | 1 + 16 files changed, 3142 insertions(+), 67 deletions(-) create mode 100644 gateway/internal/cache/bench_test.go create mode 100644 gateway/internal/cache/lsm.go create mode 100644 gateway/internal/cache/lsm_test.go create mode 100644 gateway/internal/proxy/prompt_cache.go create mode 100644 gateway/internal/sessions/tracker_test.go create mode 100644 scripts/audit_session.py diff --git a/gateway/cmd/iq/main.go b/gateway/cmd/iq/main.go index 45c1afe..27b1b3e 100644 --- a/gateway/cmd/iq/main.go +++ b/gateway/cmd/iq/main.go @@ -7,6 +7,7 @@ package main import ( "context" "crypto/rand" + "database/sql" "encoding/hex" "errors" "fmt" @@ -16,9 +17,12 @@ import ( "os/exec" "os/signal" "path/filepath" + "strings" "syscall" "time" + _ "modernc.org/sqlite" // pure-Go SQLite driver + "github.com/Revanth14/indexqube/gateway/internal/server" ) @@ -111,8 +115,12 @@ func runClaude(args []string, devMode, dumpPayloads bool) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() + proxyDone := make(chan struct{}) os.Setenv("INDEXQUBE_LOG_LEVEL", "off") - go startProxy(ctx, ln) + go func() { + startProxy(ctx, ln) + close(proxyDone) + }() if !waitForProxy(port) { fmt.Fprintf(os.Stderr, "iq: proxy failed to start within 2s\n") @@ -126,8 +134,11 @@ func runClaude(args []string, devMode, dumpPayloads bool) { } if !filepath.IsAbs(claudePath) { - fmt.Fprintf(os.Stderr, "iq: 'claude' path is not absolute: %s\n", claudePath) - os.Exit(1) + claudePath, err = filepath.Abs(claudePath) + if err != nil { + fmt.Fprintf(os.Stderr, "iq: failed to resolve absolute path for 'claude': %v\n", err) + os.Exit(1) + } } cmdArgs := append([]string{claudePath}, args...) @@ -143,22 +154,108 @@ func runClaude(args []string, devMode, dumpPayloads bool) { signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM) defer signal.Stop(sigs) go func() { - sig := <-sigs - if cmd.Process != nil { - cmd.Process.Signal(sig) //nolint:errcheck + for sig := range sigs { + if cmd.Process != nil { + cmd.Process.Signal(sig) //nolint:errcheck + } } }() - if err := cmd.Run(); err != nil { + runErr := cmd.Run() + + // Trigger graceful shutdown + cancel() + <-proxyDone + + // Print beautiful colored session summary of token savings + printSessionSummary(sessionID) + + if runErr != nil { var exitErr *exec.ExitError - if errors.As(err, &exitErr) { + if errors.As(runErr, &exitErr) { os.Exit(exitErr.ExitCode()) } - fmt.Fprintf(os.Stderr, "iq: failed to run claude: %v\n", err) + fmt.Fprintf(os.Stderr, "iq: failed to run claude: %v\n", runErr) os.Exit(1) } } +func printSessionSummary(sessionID string) { + home, err := os.UserHomeDir() + if err != nil { + return + } + dbPath := filepath.Join(home, ".indexqube", "sessions.db") + db, err := sql.Open("sqlite", dbPath) + if err != nil { + return + } + defer db.Close() + + var requestsTotal, tokensAttempted, tokensSent, tokensDeduplicated int64 + var status string + queryID := "%" + if len(sessionID) >= 8 { + queryID = "%-" + sessionID[:8] + } + err = db.QueryRow(` + SELECT requests_total, tokens_attempted, tokens_sent, tokens_deduplicated, status + FROM agent_sessions + WHERE session_id LIKE ? + ORDER BY last_seen_at DESC + LIMIT 1 + `, queryID).Scan(&requestsTotal, &tokensAttempted, &tokensSent, &tokensDeduplicated, &status) + if err != nil { + return + } + + if requestsTotal == 0 { + return + } + + percent := 0.0 + if tokensAttempted > 0 { + percent = float64(tokensDeduplicated) / float64(tokensAttempted) * 100 + } + + // Cost saved (estimate at $3 per million tokens for Claude 3.5 Sonnet inputs) + dollarsSaved := float64(tokensDeduplicated) * 0.000003 + + fmt.Fprintln(os.Stderr) + fmt.Fprintln(os.Stderr, " \033[1;36m┌────────────────────────────────────────────────────────┐\033[0m") + fmt.Fprintln(os.Stderr, " \033[1;36m│\033[0m \033[1;37mIndexQube Session Summary\033[0m \033[1;36m│\033[0m") + fmt.Fprintln(os.Stderr, " \033[1;36m├────────────────────────────────────────────────────────┤\033[0m") + fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;32m%-29d\033[0m\033[1;36m│\033[0m\n", "Requests Processed", requestsTotal) + fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;37m%-29s\033[0m\033[1;36m│\033[0m\n", "Tokens Attempted", formatNumber(tokensAttempted)) + fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;37m%-29s\033[0m\033[1;36m│\033[0m\n", "Tokens Sent", formatNumber(tokensSent)) + fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;35m%-29s\033[0m\033[1;36m│\033[0m\n", "Tokens Saved", fmt.Sprintf("%s (%.1f%%)", formatNumber(tokensDeduplicated), percent)) + fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;33m$%-28.2f\033[0m\033[1;36m│\033[0m\n", "Estimated Savings", dollarsSaved) + fmt.Fprintln(os.Stderr, " \033[1;36m├────────────────────────────────────────────────────────┤\033[0m") + statusColor := "\033[1;32m" // green + if status == "killed" { + statusColor = "\033[1;31m" // red + } + fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : %s%-29s\033[0m\033[1;36m│\033[0m\n", "Session Status", statusColor, strings.ToUpper(status)) + fmt.Fprintln(os.Stderr, " \033[1;36m└────────────────────────────────────────────────────────┘\033[0m") + fmt.Fprintln(os.Stderr) +} + +func formatNumber(n int64) string { + in := fmt.Sprintf("%d", n) + out := make([]byte, len(in)+(len(in)-1)/3) + for i, j, k := len(in)-1, len(out)-1, 0; i >= 0; i-- { + out[j] = in[i] + j-- + k++ + if k == 3 && i > 0 { + out[j] = ',' + j-- + k = 0 + } + } + return string(out) +} + func printHelp() { fmt.Print(` iq — IndexQube CLI @@ -189,6 +286,9 @@ func startProxy(ctx context.Context, ln net.Listener) { } os.Setenv("INDEXQUBE_MODE", "optimize") os.Setenv("INDEXQUBE_ENABLE_BLOCK_OPTIMIZER", "true") + // Bind admin server to an ephemeral port so multiple iq instances + // don't collide on the default 9100. + os.Setenv("ADMIN_PORT", "0") // Route telemetry through the deployed gateway so Supabase credentials // never need to be baked into this distributed binary. if os.Getenv("IQ_TELEMETRY_ENDPOINT") == "" { diff --git a/gateway/internal/cache/bench_test.go b/gateway/internal/cache/bench_test.go new file mode 100644 index 0000000..57915f8 --- /dev/null +++ b/gateway/internal/cache/bench_test.go @@ -0,0 +1,285 @@ +package cache + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "testing" + "time" + + "github.com/Revanth14/indexqube/gateway/internal/domain" +) + +// ─── Fixtures ───────────────────────────────────────────────────────────────── + +// benchSize parameterises one sub-benchmark run. +type benchSize struct { + name string + numChunks int + chunkBytes int +} + +// benchSizes covers four realistic entry profiles: +// +// tiny — single-token or one-liner response (~60 B raw chunks) +// small — sentence-length response (~600 B) +// medium — paragraph response (~3.2 KB) +// large — long-form response (~16 KB) +var benchSizes = []benchSize{ + {"1chunk_tiny", 1, 60}, + {"5chunk_small", 5, 120}, + {"20chunk_medium", 20, 160}, + {"100chunk_large", 100, 160}, +} + +// makeBenchEntry builds an Entry with realistic pre-marshaled SSE frames. +// Each chunk looks like an OpenAI-shaped content_block_delta event, which +// is what the tee in cache/tee.go actually captures. +func makeBenchEntry(numChunks, chunkBytes int) *Entry { + chunks := make([][]byte, numChunks) + for i := range chunks { + // 88 bytes of fixed JSON scaffolding; remainder is text payload. + text := strings.Repeat("x", max(1, chunkBytes-88)) + chunks[i] = []byte(fmt.Sprintf( + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":%q}}`, + text, + )) + } + return &Entry{ + Provider: domain.ProviderAnthropic, + Model: "claude-3-5-sonnet-20241022", + Chunks: chunks, + CreatedAt: time.Date(2026, 5, 24, 12, 0, 0, 0, time.UTC), + } +} + +// benchKey returns a 64-character cache key (same width as a SHA-256 hex +// digest produced by DeriveKey) with a numeric suffix for uniqueness. +func benchKey(i int) Key { + return Key(fmt.Sprintf("benchkey%056d", i)) // 8 + 56 = 64 chars +} + +const benchHitKey Key = "benchhitkey000000000000000000000000000000000000000000000000000000" // 64 chars + +// preMarshal returns the JSON encoding of e, panicking on failure. +func preMarshal(e *Entry) []byte { + data, err := json.Marshal(e) + if err != nil { + panic(err) + } + return data +} + +// ─── Group 1: Pure serialization — no cache layer ───────────────────────────── +// +// These are the baseline numbers. Every LSMCache.Put pays json.Marshal; +// every LSMCache.Get hit pays json.Unmarshal. Comparing Group 3 against +// Group 1 isolates exactly what the MemTable layer adds on top. + +// BenchmarkMarshal measures json.Marshal for Entry structs of varying sizes. +// b.SetBytes is the raw chunk total (entry.Bytes()), so MB/s reflects +// throughput of the actual payload — before base64 expansion. +func BenchmarkMarshal(b *testing.B) { + for _, sz := range benchSizes { + sz := sz + b.Run(sz.name, func(b *testing.B) { + e := makeBenchEntry(sz.numChunks, sz.chunkBytes) + b.SetBytes(e.Bytes()) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + data, _ := json.Marshal(e) + _ = data + } + }) + } +} + +// BenchmarkUnmarshal measures json.Unmarshal cost — the deserialization step +// on every LSMCache.Get hit. b.SetBytes is the JSON wire size (after base64 +// expansion), so MB/s reflects true I/O throughput from the LSM layer. +func BenchmarkUnmarshal(b *testing.B) { + for _, sz := range benchSizes { + sz := sz + b.Run(sz.name, func(b *testing.B) { + e := makeBenchEntry(sz.numChunks, sz.chunkBytes) + data := preMarshal(e) + b.SetBytes(int64(len(data))) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + var out Entry + _ = json.Unmarshal(data, &out) + } + }) + } +} + +// BenchmarkMarshalUnmarshalRoundtrip measures the combined Put+Get +// serialization cost with no storage layer. Subtract this from +// BenchmarkLSMCachePut + BenchmarkLSMCacheGet to get pure MemTable overhead. +func BenchmarkMarshalUnmarshalRoundtrip(b *testing.B) { + for _, sz := range benchSizes { + sz := sz + b.Run(sz.name, func(b *testing.B) { + e := makeBenchEntry(sz.numChunks, sz.chunkBytes) + b.SetBytes(e.Bytes()) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + data, _ := json.Marshal(e) + var out Entry + _ = json.Unmarshal(data, &out) + } + }) + } +} + +// ─── Group 2: LSMCache serialization overhead in isolation ──────────────────── +// +// These benchmarks replicate the exact work LSMCache.Put does before calling +// engine.Put — Bytes() budget check, json.Marshal, and key cast — without +// touching the storage layer. The difference between BenchmarkMarshal and +// BenchmarkPutSerialization is the cost of the Bytes() loop and key encoding. + +// BenchmarkPutSerialization benchmarks every serialization step inside +// LSMCache.Put except the engine write: entry.Bytes() + json.Marshal + +// []byte(key) cast. Subtract BenchmarkMarshal to isolate the Bytes() overhead. +func BenchmarkPutSerialization(b *testing.B) { + for _, sz := range benchSizes { + sz := sz + b.Run(sz.name, func(b *testing.B) { + e := makeBenchEntry(sz.numChunks, sz.chunkBytes) + b.SetBytes(e.Bytes()) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = e.Bytes() + data, _ := json.Marshal(e) + lsmKey := []byte(benchHitKey) + _, _ = data, lsmKey + } + }) + } +} + +// BenchmarkBytesCheck benchmarks entry.Bytes() alone — the pre-marshal budget +// guard that iterates every chunk header. This is O(numChunks), not O(payload). +func BenchmarkBytesCheck(b *testing.B) { + for _, sz := range benchSizes { + sz := sz + b.Run(sz.name, func(b *testing.B) { + e := makeBenchEntry(sz.numChunks, sz.chunkBytes) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _ = e.Bytes() + } + }) + } +} + +// ─── Group 3: Full cache path — serialization + MemTable ────────────────────── +// +// These are the end-to-end numbers. Subtract Group 1 from Group 3 to get the +// pure MemTable skip-list overhead. For Put, unique keys are used to exercise +// the new-node insertion path (not update-in-place), matching production usage +// where each unique request maps to its own SHA-256 key. + +// BenchmarkLSMCachePut measures the full LSMCache.Put path: Bytes() + Marshal + +// MemTable skip-list insert. Each iteration writes a unique key to exercise the +// new-node path. The MemTable may flush to disk at high b.N for large entries; +// this reflects the real amortised production cost. +func BenchmarkLSMCachePut(b *testing.B) { + for _, sz := range benchSizes { + sz := sz + b.Run(sz.name, func(b *testing.B) { + dir := b.TempDir() + c, err := NewLSMCache(dir, 0) // 0 = no per-entry size cap + if err != nil { + b.Fatalf("NewLSMCache: %v", err) + } + defer c.Close() + e := makeBenchEntry(sz.numChunks, sz.chunkBytes) + ctx := context.Background() + b.SetBytes(e.Bytes()) + b.ReportAllocs() + b.ResetTimer() + n := 0 + for b.Loop() { + n++ + _ = c.Put(ctx, benchKey(n), e) + } + }) + } +} + +// BenchmarkLSMCacheGet measures the full LSMCache.Get path: MemTable lookup + +// Unmarshal. The key is preloaded before timing to benchmark hot-path reads +// (guaranteed MemTable hit, no SSTable disk I/O). +func BenchmarkLSMCacheGet(b *testing.B) { + for _, sz := range benchSizes { + sz := sz + b.Run(sz.name, func(b *testing.B) { + dir := b.TempDir() + c, err := NewLSMCache(dir, 0) + if err != nil { + b.Fatalf("NewLSMCache: %v", err) + } + defer c.Close() + e := makeBenchEntry(sz.numChunks, sz.chunkBytes) + ctx := context.Background() + if err := c.Put(ctx, benchHitKey, e); err != nil { + b.Fatalf("pre-load Put: %v", err) + } + b.SetBytes(e.Bytes()) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _, _, _ = c.Get(ctx, benchHitKey) + } + }) + } +} + +// ─── Group 4: Serialized-size reporting ─────────────────────────────────────── +// +// BenchmarkSerializedExpansion reports the divergence between what entry.Bytes() +// measures (raw chunk bytes, what the budget cap enforces) and what json.Marshal +// actually writes into the LSM (base64-encoded chunks + JSON framing). +// +// Because [][]byte is base64-encoded by encoding/json, every Put stores +// more bytes than the configured maxEntryBytes cap enforces. This benchmark +// makes that gap visible as custom metrics (raw_B, json_B, overhead_B, +// expansion_pct) alongside ns/op. +// +// Note: uses the traditional b.N loop (not b.Loop) so that b.ReportMetric +// calls survive across calibration passes in Go 1.24+. +func BenchmarkSerializedExpansion(b *testing.B) { + for _, sz := range benchSizes { + sz := sz + b.Run(sz.name, func(b *testing.B) { + e := makeBenchEntry(sz.numChunks, sz.chunkBytes) + data := preMarshal(e) + rawBytes := e.Bytes() + jsonBytes := int64(len(data)) + + b.SetBytes(rawBytes) + b.ResetTimer() + + for i := 0; i < b.N; i++ { + d, _ := json.Marshal(e) + _ = d + } + + // Report after the loop so metrics survive b.Loop calibration passes. + b.ReportMetric(float64(rawBytes), "raw_B") + b.ReportMetric(float64(jsonBytes), "json_B") + b.ReportMetric(float64(jsonBytes-rawBytes), "overhead_B") + if rawBytes > 0 { + b.ReportMetric(float64(jsonBytes)/float64(rawBytes)*100, "expansion_pct") + } + }) + } +} diff --git a/gateway/internal/cache/lsm.go b/gateway/internal/cache/lsm.go new file mode 100644 index 0000000..c9c87cb --- /dev/null +++ b/gateway/internal/cache/lsm.go @@ -0,0 +1,85 @@ +package cache + +import ( + "context" + "encoding/json" + + "github.com/Revanth14/indexqube/gateway/internal/store/lsm" +) + +// LSMCache implements cache.Cache using a local pure-Go LSM storage engine. +// It persists cache items durably across process boundaries. +type LSMCache struct { + engine *lsm.Engine + maxEntryBytes int64 +} + +// NewLSMCache opens (or creates) an LSM storage engine in dir. +func NewLSMCache(dir string, maxEntryBytes int64) (*LSMCache, error) { + opts := lsm.Options{ + MemTableSize: 4 * 1024 * 1024, // 4 MiB flushes + MaxL0Tables: 4, + BlockSize: 4096, // page-aligned + BloomFPRate: 0.01, // 1% false positive + BloomExpected: 10000, + MaxTableSize: 2 * 1024 * 1024, // 2 MiB tables + L1Budget: 10 * 1024 * 1024, + } + engine, err := lsm.Open(dir, opts) + if err != nil { + return nil, err + } + return &LSMCache{ + engine: engine, + maxEntryBytes: maxEntryBytes, + }, nil +} + +// Get looks up key in the LSM-backed cache. +func (c *LSMCache) Get(ctx context.Context, key Key) (*Entry, bool, error) { + val, found, err := c.engine.Get([]byte(key)) + if err != nil { + return nil, false, err + } + if !found { + return nil, false, nil + } + + var entry Entry + if err := json.Unmarshal(val, &entry); err != nil { + return nil, false, err + } + return &entry, true, nil +} + +// GetSemantic is a no-op fallback for simple key-value engines. +func (c *LSMCache) GetSemantic(ctx context.Context, tenantID string, embedding []float32, threshold float64) (*Entry, bool, error) { + return nil, false, nil +} + +// Put serializes and inserts entry under key in the LSM-backed cache. +func (c *LSMCache) Put(ctx context.Context, key Key, entry *Entry) error { + if entry == nil { + return nil + } + if c.maxEntryBytes > 0 && entry.Bytes() > c.maxEntryBytes { + return ErrEntryTooLarge + } + + val, err := json.Marshal(entry) + if err != nil { + return err + } + + return c.engine.Put([]byte(key), val) +} + +// PutSemantic acts as a fallback to Put for simple key-value engines. +func (c *LSMCache) PutSemantic(ctx context.Context, tenantID string, key Key, entry *Entry, embedding []float32) error { + return c.Put(ctx, key, entry) +} + +// Close flushes active writes and closes the underlying LSM engine cleanly. +func (c *LSMCache) Close() error { + return c.engine.Close() +} diff --git a/gateway/internal/cache/lsm_test.go b/gateway/internal/cache/lsm_test.go new file mode 100644 index 0000000..6e1cb97 --- /dev/null +++ b/gateway/internal/cache/lsm_test.go @@ -0,0 +1,150 @@ +package cache + +import ( + "context" + "errors" + "os" + "testing" + "time" + + "github.com/Revanth14/indexqube/gateway/internal/domain" +) + +func TestLSMCache_GetMissOnEmpty(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "iq-lsm-cache-test-*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + c, err := NewLSMCache(tmpDir, 1024) + if err != nil { + t.Fatalf("NewLSMCache failed: %v", err) + } + defer c.Close() + + _, hit, err := c.Get(context.Background(), "missing") + if err != nil { + t.Fatalf("Get failed: %v", err) + } + if hit { + t.Error("expected miss on empty cache") + } +} + +func TestLSMCache_PutThenGet(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "iq-lsm-cache-test-*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + c, err := NewLSMCache(tmpDir, 1024) + if err != nil { + t.Fatalf("NewLSMCache failed: %v", err) + } + defer c.Close() + + want := &Entry{ + Provider: domain.ProviderAnthropic, + Model: "claude-3-5-sonnet", + Chunks: [][]byte{[]byte("first chunk"), []byte("second chunk")}, + CreatedAt: time.Now().UTC(), + } + + if err := c.Put(context.Background(), "test-key", want); err != nil { + t.Fatalf("Put failed: %v", err) + } + + got, hit, err := c.Get(context.Background(), "test-key") + if err != nil { + t.Fatalf("Get failed: %v", err) + } + if !hit { + t.Fatal("expected cache hit") + } + + if got.Provider != want.Provider { + t.Errorf("provider mismatch: got %v, want %v", got.Provider, want.Provider) + } + if got.Model != want.Model { + t.Errorf("model mismatch: got %s, want %s", got.Model, want.Model) + } + if len(got.Chunks) != len(want.Chunks) { + t.Errorf("chunks len mismatch: got %d, want %d", len(got.Chunks), len(want.Chunks)) + } + if string(got.Chunks[0]) != string(want.Chunks[0]) { + t.Errorf("first chunk mismatch: got %s, want %s", got.Chunks[0], want.Chunks[0]) + } +} + +func TestLSMCache_PutTooLargeRejected(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "iq-lsm-cache-test-*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + c, err := NewLSMCache(tmpDir, 10) // extremely small max size + if err != nil { + t.Fatalf("NewLSMCache failed: %v", err) + } + defer c.Close() + + entry := &Entry{ + Provider: domain.ProviderAnthropic, + Model: "claude-3-5-sonnet", + Chunks: [][]byte{[]byte("too long payload data")}, + CreatedAt: time.Now(), + } + + err = c.Put(context.Background(), "too-large", entry) + if !errors.Is(err, ErrEntryTooLarge) { + t.Errorf("expected ErrEntryTooLarge, got: %v", err) + } +} + +func TestLSMCache_PersistenceAcrossCloses(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "iq-lsm-cache-test-*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + // Phase 1: Open, write entry, then close + c1, err := NewLSMCache(tmpDir, 1024) + if err != nil { + t.Fatalf("c1 NewLSMCache failed: %v", err) + } + + want := &Entry{ + Provider: domain.ProviderAnthropic, + Model: "claude-3-5-sonnet", + Chunks: [][]byte{[]byte("durable persistent cache chunk data")}, + CreatedAt: time.Now().UTC(), + } + + if err := c1.Put(context.Background(), "persistent-key", want); err != nil { + t.Fatalf("c1 Put failed: %v", err) + } + c1.Close() + + // Phase 2: Re-open in the same directory, read entry, then close + c2, err := NewLSMCache(tmpDir, 1024) + if err != nil { + t.Fatalf("c2 NewLSMCache failed: %v", err) + } + defer c2.Close() + + got, hit, err := c2.Get(context.Background(), "persistent-key") + if err != nil { + t.Fatalf("c2 Get failed: %v", err) + } + if !hit { + t.Fatal("expected hit from persistent L2 cache") + } + + if string(got.Chunks[0]) != string(want.Chunks[0]) { + t.Errorf("chunk data corrupted: got %s, want %s", got.Chunks[0], want.Chunks[0]) + } +} diff --git a/gateway/internal/chunker/chunker.go b/gateway/internal/chunker/chunker.go index 8501094..f3c7458 100644 --- a/gateway/internal/chunker/chunker.go +++ b/gateway/internal/chunker/chunker.go @@ -22,12 +22,183 @@ // pathologically small or large regardless of the data distribution. package chunker +import ( + "bytes" + "strings" +) + +// goDeclarationPrefixes are the newline-prefixed tokens that open a top-level +// Go declaration. The leading '\n' ensures we only match at line starts. +var goDeclarationPrefixes = []string{"\nfunc ", "\ntype ", "\nconst ", "\nvar "} + +// goDeclarationPrefixBytes mirrors goDeclarationPrefixes as []byte to allow +// allocation-free searching in hot paths. +var goDeclarationPrefixBytes = [][]byte{ + []byte("\nfunc "), []byte("\ntype "), []byte("\nconst "), []byte("\nvar "), +} + +// IsCodeContent returns true when data contains Go/C-family declaration keywords +// at a density that suggests source code rather than prose. +// Only the first 4 KiB is sampled so this runs in O(1) for large files. +func IsCodeContent(data []byte) bool { + sample := data + if len(sample) > 4096 { + sample = sample[:4096] + } + text := string(sample) + hits := 0 + for _, kw := range goDeclarationPrefixes { + hits += strings.Count(text, kw) + } + hits += strings.Count(text, "\nfunction ") // JS/TS + hits += strings.Count(text, "\nclass ") // Python/Java/JS + hits += strings.Count(text, "\ndef ") // Python + // Threshold: at least one declaration per 400 bytes → treat as code. + return hits > 0 && len(sample)/hits <= 400 +} + +// AlignToDeclarationBoundary advances offset forward within data until it lands +// at the start of a top-level Go declaration line (func/type/const/var), or +// returns the original offset when none is found within maxAdvance bytes. +// The returned position skips the leading '\n' so it points at the first letter +// of the keyword ('f', 't', 'c', or 'v'). +func AlignToDeclarationBoundary(data []byte, offset, maxAdvance int) int { + end := offset + maxAdvance + if end > len(data) { + end = len(data) + } + region := string(data[offset:end]) + best := -1 + for _, prefix := range goDeclarationPrefixes { + if idx := strings.Index(region, prefix); idx >= 0 && (best < 0 || idx < best) { + best = idx + } + } + if best < 0 { + return offset + } + return offset + best + 1 // +1: skip the leading '\n' +} + +// RetractToDeclarationBoundary searches backward from forced toward +// forced-retreatRadius for the last top-level Go declaration boundary. +// It is the inverse of AlignToDeclarationBoundary and is used to snap +// forced MaxSize cuts to stable semantic positions so that minor byte +// shifts between turns (e.g. a 21-byte header difference) do not produce +// different fingerprints for otherwise identical content. +// +// Uses []byte patterns to avoid heap allocations in the hot path. +// Returns forced unchanged when no boundary is found within retreatRadius. +func RetractToDeclarationBoundary(data []byte, forced, retreatRadius int) int { + if retreatRadius <= 0 || forced <= 0 || forced > len(data) { + return forced + } + start := forced - retreatRadius + if start < 0 { + start = 0 + } + region := data[start:forced] + best := -1 + for _, prefix := range goDeclarationPrefixBytes { + idx := bytes.LastIndex(region, prefix) + if idx >= 0 && idx > best { + best = idx + } + } + if best < 0 { + return forced + } + return start + best + 1 // +1: skip the leading '\n' +} + +// IsPathList returns true when data appears to be a file-path listing. +// Uses two heuristics: overall slash density (>3%) or high slash/dot line ratio (>60%). +// The lower threshold catches short listings like 15-path tool outputs. +func IsPathList(data []byte) bool { + if len(data) == 0 { + return false + } + overallSlash := float64(bytes.Count(data, []byte("/"))) / float64(len(data)) + if overallSlash > 0.03 { + return true + } + lines := bytes.Split(data, []byte("\n")) + if len(lines) == 0 { + return false + } + slashLines := 0 + for _, line := range lines { + if bytes.Contains(line, []byte("/")) || bytes.Contains(line, []byte(".")) { + slashLines++ + } + } + return float64(slashLines)/float64(len(lines)) > 0.60 +} + +// IsSystemProseContent returns true when data begins with an XML-like system tag +// (e.g. , ). These blocks use SystemProseConfig so +// a single changed field (e.g. currentDate) only invalidates its own small chunk. +func IsSystemProseContent(data []byte) bool { + if len(data) < 16 { + return false + } + sample := bytes.TrimSpace(data) + if len(sample) > 64 { + sample = sample[:64] + } + return bytes.HasPrefix(sample, []byte("")) || + bytes.HasPrefix(sample, []byte("")) || + bytes.HasPrefix(sample, []byte("")) || + bytes.HasPrefix(sample, []byte(" len(data) { + return forced + } + start := forced - retreatRadius + if start < 0 { + start = 0 + } + region := data[start:forced] + idx := bytes.LastIndex(region, []byte("\n\n")) + if idx < 0 { + return forced + } + return start + idx + 2 // position after the blank line +} func (c *Chunker) nextBoundary(data []byte, start int) int { n := len(data) @@ -218,6 +406,20 @@ func (c *Chunker) nextBoundary(data []byte, start int) int { } } - // No boundary found before MaxSize — force a cut at the hard cap. - return maxEnd + // No boundary found before MaxSize — snap to the nearest declaration + // boundary before the hard cap. Lookback radius scales with span size: + // large files (>32 KB) use 1024 bytes to reach declaration boundaries + // that are spaced farther apart. Minor byte shifts between turns won't + // change the boundary as long as the snapped position itself is stable. + lookback := 512 + if len(data) > 32*1024 { + lookback = 1024 + } + snapped := RetractToDeclarationBoundary(data, maxEnd, lookback) + // Secondary fallback: if no declaration keyword found, snap to the + // nearest blank line within 256 bytes to avoid mid-paragraph splits. + if snapped == maxEnd { + snapped = retractToBlankLine(data, maxEnd, 256) + } + return snapped } diff --git a/gateway/internal/proxy/claude_messages.go b/gateway/internal/proxy/claude_messages.go index ae42dba..0813add 100644 --- a/gateway/internal/proxy/claude_messages.go +++ b/gateway/internal/proxy/claude_messages.go @@ -15,24 +15,36 @@ import ( "net/url" "os" "path/filepath" + "regexp" "runtime" "strconv" "strings" "sync" "time" + "github.com/Revanth14/indexqube/gateway/internal/chunker" "github.com/Revanth14/indexqube/gateway/internal/memory" "github.com/Revanth14/indexqube/gateway/internal/middleware" "github.com/Revanth14/indexqube/gateway/internal/telemetry" awsconfig "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" brtypes "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types" + "github.com/google/uuid" ) const ( claudeDefaultMode = "observe" ) +type contextKey string +const isSubscriptionKey contextKey = "isSubscription" + + +// guardBypassRe matches directives that attempt to disable proxy safety controls. +// The separator between "guards" and "velocity" is optional (covers slash, pipe, +// backslash, or none) to handle formatting variations in injected CLAUDE.md files. +var guardBypassRe = regexp.MustCompile(`(?i)guards?\s*[/\\|]?\s*velocity\s+limits?\s+do\s+not\s+apply`) + var dumpPayloadMu sync.Mutex type anthropicMessagesRequest struct { @@ -69,6 +81,12 @@ type claudeOptimizerStats struct { BytesAfter int `json:"bytes_after"` BytesEligible int `json:"bytes_eligible"` BytesPruned int `json:"bytes_pruned"` + // KnownBytes is the byte total of every span that hit the session cache + // (Seen returned true), independent of whether the span was subsequently + // pruned or preserved by a protection rule. Invariant: + // KnownBytes == BytesPruned + PreservedInstructionBytes + // + PreservedLatestTurnBytes + PreservedLastOccurrenceBytes + KnownBytes int `json:"known_bytes"` EstimatedTokensBefore int `json:"estimated_tokens_before"` EstimatedTokensAfter int `json:"estimated_tokens_after"` EstimatedTokensSaved int `json:"estimated_tokens_saved"` @@ -103,11 +121,13 @@ type claudeOptimizerStats struct { type claudeStreamStats struct { Chunks int OutputText int + OutputRawText string OutputTokens int Status string StatusCode int Cancelled bool Completed bool + HasToolUse bool Provider string // "anthropic" or "bedrock" UpstreamErr string UpstreamErrorCode string @@ -128,9 +148,9 @@ type claudeUpstreamErrorMeta struct { func (p *Proxy) handleClaudeMessages(w http.ResponseWriter, r *http.Request) { started := time.Now() cfg := p.claudeDefaults() - requestID := middleware.RequestIDFromContext(r.Context()) - if requestID == "" { - requestID = r.Header.Get("X-Request-ID") + rawRequestID := middleware.RequestIDFromContext(r.Context()) + if rawRequestID == "" { + rawRequestID = r.Header.Get("X-Request-ID") } if err := cfg.validate(); err != nil { @@ -175,23 +195,110 @@ func (p *Proxy) handleClaudeMessages(w http.ResponseWriter, r *http.Request) { overheadStart := time.Now() sessionKey := claudeSessionKey(r, cfg.DevToken) + + // FIX 3: resolve or synthesize the request ID and detect velocity issues. + requestID, missingReqID, velocityWarning := p.resolveRequestID(sessionKey, rawRequestID) + + // Conservative rate limiting when the session repeatedly sends turns + // without request IDs — max 1 turn/second, max 32 KB payload. + if velocityWarning { + time.Sleep(time.Second) + if len(body) > 32*1024 { + p.writeError(w, r, errorPayload{ + HTTPStatus: http.StatusTooManyRequests, + Type: "rate_limit_error", + Code: "missing_request_id_velocity_limit", + Message: "too many turns without request IDs; payload capped at 32 KB", + }) + return + } + } + var earlyMeta struct { Model string `json:"model"` } _ = json.Unmarshal(body, &earlyMeta) + auth := r.Header.Get("Authorization") + isSubscription := isSubscriptionAuth(auth) + ctx := context.WithValue(r.Context(), isSubscriptionKey, isSubscription) - forwardBody, meta, optStats, shape, err := p.prepareClaudeBody(r.Context(), cfg, sessionKey, body) + forwardBody, meta, optStats, shape, err := p.prepareClaudeBody(ctx, cfg, sessionKey, body) if err != nil { p.writeError(w, r, errorPayload{HTTPStatus: http.StatusBadRequest, Type: "invalid_request_error", Code: "invalid_anthropic_request", Message: err.Error()}) return } - if os.Getenv("IQ_DUMP_PAYLOADS") == "1" { - dumpClaudePayloads(requestID, body, forwardBody) + + // Intercept Sentinel Probes (quota, ping) to respond in 1 ms at 0 cost + normalizedLatest := strings.TrimSpace(strings.ToLower(shape.LatestUserText)) + if normalizedLatest == "quota" || normalizedLatest == "ping" { + p.logger.InfoContext(ctx, "sentinel probe intercepted", + slog.String("session_key", shortLogHash(sessionKey)), + slog.String("probe", normalizedLatest)) + writeSyntheticStreamResponse(w, "IndexQube active. Quota check successful.") + return + } + + var capture *responseCaptureWriter + var promptCacheHash string + isCacheablePrompt := shape.LatestUserText != "" && shape.ToolResultCount == 0 && cfg.Mode == "optimize" && + r.Header.Get("Cache-Control") != "no-cache" && r.Header.Get("Pragma") != "no-cache" + if isCacheablePrompt { + promptCacheHash = computePromptHash(shape.LatestUserText, meta.Model) + ts := p.getOrCreateTurnState(sessionKey) + ts.mu.Lock() + cachedPayload, ok := ts.getCachedResponse(promptCacheHash) + ts.mu.Unlock() + if ok { + p.logger.InfoContext(ctx, "response-level cache hit", + slog.String("session_key", shortLogHash(sessionKey)), + slog.String("prompt_hash", promptCacheHash[:8])) + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Accel-Buffering", "no") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(cachedPayload) + if flusher, ok := w.(http.Flusher); ok { + flusher.Flush() + } + return + } + capture = &responseCaptureWriter{ResponseWriter: w} + w = capture } - _ = shape overheadMs := time.Since(overheadStart).Milliseconds() + + // FIX 2: in-flight duplicate detection. When a second identical request + // arrives while the first is still in-flight, it waits up to 30 s for the + // first to complete before dispatching its own upstream call. This prevents + // the "triple-fire" pattern where 3 identical prompts fire concurrently. + promptHash := semanticPromptHash(body) + if doneFn, waitChan := p.inFlightRequests.acquire(promptHash); doneFn == nil { + select { + case <-waitChan: + case <-time.After(30 * time.Second): + case <-r.Context().Done(): + return + } + // Re-register after the original completed so our dispatch is tracked. + if done2, _ := p.inFlightRequests.acquire(promptHash); done2 != nil { + defer done2() + } + } else { + defer doneFn() + } + streamStats := p.forwardClaudeMessages(w, r, cfg, forwardBody) + if capture != nil && streamStats.Completed && streamStats.StatusCode == http.StatusOK && !streamStats.HasToolUse { + ts := p.getOrCreateTurnState(sessionKey) + ts.mu.Lock() + ts.saveCachedResponse(promptCacheHash, capture.buf.Bytes()) + ts.mu.Unlock() + p.logger.InfoContext(ctx, "response-level cache saved", + slog.String("session_key", shortLogHash(sessionKey)), + slog.String("prompt_hash", promptCacheHash[:8])) + } duration := time.Since(started) if cfg.SessionStore != nil { cfg.SessionStore.RecordUsage(sessionKey, memory.UsageTotals{ @@ -203,8 +310,12 @@ func (p *Proxy) handleClaudeMessages(w http.ResponseWriter, r *http.Request) { BytesSaved: optStats.BytesBefore - optStats.BytesAfter, }) } - p.logClaudeRequestComplete(r.Context(), requestID, cfg.Mode, meta.Model, sessionKey, len(body), optStats, streamStats, duration) + p.logClaudeRequestComplete(r.Context(), requestID, cfg.Mode, meta.Model, sessionKey, len(body), optStats, streamStats, duration, missingReqID, requestID, velocityWarning) p.emitClaudeUsageEvent(r, sessionKey, meta.Model, optStats, duration, streamStats.StatusCode, overheadMs) + + if os.Getenv("IQ_DUMP_PAYLOADS") == "1" { + dumpClaudePayloads(requestID, body, forwardBody, streamStats, optStats) + } } func (p *Proxy) handleClaudeCountTokens(w http.ResponseWriter, r *http.Request) { @@ -284,6 +395,9 @@ func (p *Proxy) claudeDefaults() ClaudeMessagesConfig { cfg.Optimizer.MaxChunkBytes = 8192 cfg.Optimizer.MinSavedTokens = 10 cfg.Optimizer.EnableToolResultPruning = true + cfg.Optimizer.EnableSubspanChunking = true + cfg.Optimizer.SmallFileBytes = 4096 + cfg.Optimizer.EnablePromptCache = true } return cfg } @@ -308,9 +422,9 @@ func validClaudeDevToken(r *http.Request, want string) bool { return token != "" } -func dumpClaudePayloads(requestID string, before, after []byte) { +func dumpClaudePayloads(requestID string, before, after []byte, stats claudeStreamStats, optStats claudeOptimizerStats) { if sessionFile := os.Getenv("IQ_DUMP_SESSION_FILE"); sessionFile != "" { - if err := appendSessionDump(sessionFile, requestID, before, after); err != nil { + if err := appendSessionDump(sessionFile, requestID, before, after, stats, optStats); err != nil { fmt.Fprintf(os.Stderr, "[iq] failed to append payload dump: %v\n", err) } return @@ -337,17 +451,40 @@ func dumpClaudePayloads(requestID string, before, after []byte) { appendDumpLog(dumpDir, beforePath, afterPath) } +type payloadDumpResponse struct { + Text string `json:"text"` + OutputTokens int `json:"output_tokens"` + Status string `json:"status"` +} + type payloadDumpRecord struct { - Timestamp string `json:"ts"` - RequestID string `json:"request_id"` - BeforeBytes int `json:"before_bytes"` - AfterBytes int `json:"after_bytes"` - SavedBytes int `json:"saved_bytes"` - Before json.RawMessage `json:"before"` - After json.RawMessage `json:"after"` + Timestamp string `json:"ts"` + RequestID string `json:"request_id"` + BeforeBytes int `json:"before_bytes"` + AfterBytes int `json:"after_bytes"` + SavedBytes int `json:"saved_bytes"` + Before json.RawMessage `json:"before"` + After json.RawMessage `json:"after"` + Response payloadDumpResponse `json:"response"` + Optimizer *payloadOptimizerStats `json:"optimizer,omitempty"` } -func appendSessionDump(sessionFile, requestID string, before, after []byte) error { +// payloadOptimizerStats is the per-turn cache-efficiency view written into +// the JSONL dump. It separates the pruning view (saved_bytes / blocks_pruned) +// from the true cache-hit view (true_cache_hit_bytes / known_bytes), so the +// audit script can distinguish "bytes removed from the forwarded payload" +// from "bytes that hit the session cache regardless of preservation rules". +type payloadOptimizerStats struct { + BlocksPruned int `json:"blocks_pruned"` + BlocksKnown int `json:"blocks_known"` + BlocksKnownProtected int `json:"blocks_known_protected"` + BytesPruned int `json:"bytes_pruned"` + ProtectedBytes int `json:"protected_bytes"` + KnownBytes int `json:"known_bytes"` + TrueCacheHitBytes int `json:"true_cache_hit_bytes"` +} + +func appendSessionDump(sessionFile, requestID string, before, after []byte, stats claudeStreamStats, optStats claudeOptimizerStats) error { if err := os.MkdirAll(filepath.Dir(sessionFile), 0o700); err != nil { return err } @@ -359,6 +496,20 @@ func appendSessionDump(sessionFile, requestID string, before, after []byte) erro SavedBytes: len(before) - len(after), Before: json.RawMessage(before), After: json.RawMessage(after), + Response: payloadDumpResponse{ + Text: stats.OutputRawText, + OutputTokens: stats.OutputTokens, + Status: stats.Status, + }, + Optimizer: &payloadOptimizerStats{ + BlocksPruned: optStats.BlocksPruned, + BlocksKnown: optStats.BlocksKnown, + BlocksKnownProtected: optStats.PreservedInstructionCount, + BytesPruned: optStats.BytesPruned, + ProtectedBytes: optStats.PreservedInstructionBytes, + KnownBytes: optStats.KnownBytes, + TrueCacheHitBytes: optStats.KnownBytes, + }, } line, err := json.Marshal(record) if err != nil { @@ -406,6 +557,173 @@ func marshalJSONNoHTMLEscape(v any) ([]byte, error) { return bytes.TrimSuffix(out.Bytes(), []byte{'\n'}), nil } +// stripGuardBypassDirectives scans the system field of root for patterns +// that attempt to disable proxy-level guards (e.g. from CLAUDE.md content). +// Any matched text is replaced with an empty string in-place. Returns true +// if any directives were stripped so the caller can log a warning. +func stripGuardBypassDirectives(ctx context.Context, logger *slog.Logger, root map[string]any) bool { + stripped := false + stripText := func(text string) (string, bool) { + // Check the original text first; if that misses (e.g. newlines inside the + // directive from a system-reminder block), fall back to a whitespace-normalised + // copy so multi-line formatting doesn't hide the bypass attempt. + matched := guardBypassRe.MatchString(text) + if !matched { + normalized := strings.ToLower(strings.Join(strings.Fields(text), " ")) + matched = guardBypassRe.MatchString(normalized) + } + if !matched { + return text, false + } + return guardBypassRe.ReplaceAllString(text, ""), true + } + + switch sys := root["system"].(type) { + case string: + if cleaned, ok := stripText(sys); ok { + root["system"] = cleaned + stripped = true + } + case []any: + for _, item := range sys { + m, ok := item.(map[string]any) + if !ok { + continue + } + if text, ok2 := m["text"].(string); ok2 { + if cleaned, ok3 := stripText(text); ok3 { + m["text"] = cleaned + stripped = true + } + } + } + } + + if stripped { + logger.WarnContext(ctx, "guard bypass attempt detected in system content", + slog.String("event", "guard_bypass_attempt"), + slog.String("source", "CLAUDE.md"), + ) + } + return stripped +} + +// warmUpSystemSpans pre-registers all system-block fingerprints into the +// session cache on the first turn. This ensures that stable injected files +// (CLAUDE.md, MEMORY.md, architecture docs) are cached from turn 1 so that +// turns 2+ see them as known and report non-zero savings (FIX 6). +func (p *Proxy) warmUpSystemSpans(ctx context.Context, cfg ClaudeMessagesConfig, sessionKey string, root map[string]any) { + registerText := func(text, kind string) { + text = strings.TrimSpace(text) + if len(text) < cfg.Optimizer.MinSpanBytes { + return + } + data := []byte(text) + if cfg.Optimizer.EnableSubspanChunking && len(data) >= cfg.Optimizer.SmallFileBytes { + chunks := splitSpanChunks(text, cfg.Optimizer) + for _, ch := range chunks { + chHash := hashText(ch) + cfg.SessionStore.SaveBlock(sessionKey, memory.Block{ + Hash: chHash, + Kind: "warmup:chunk:" + kind, + Bytes: len(ch), + Tokens: estimateTokens(len(ch)), + }) + } + } else { + h := hashText(text) + cfg.SessionStore.SaveBlock(sessionKey, memory.Block{ + Hash: h, + Kind: "warmup:" + kind, + Bytes: len(data), + Tokens: estimateTokens(len(data)), + }) + } + } + + switch sys := root["system"].(type) { + case string: + registerText(sys, "system") + case []any: + for _, item := range sys { + m, ok := item.(map[string]any) + if !ok { + continue + } + if text, ok2 := m["text"].(string); ok2 { + registerText(text, "system_block") + } + } + } + p.logger.DebugContext(ctx, "session warm-up complete", + slog.String("session_key", shortLogHash(sessionKey)), + ) +} + +// getOrCreateTurnState returns the mutable turn-state for sessionKey, +// creating it if it does not yet exist. +func (p *Proxy) getOrCreateTurnState(sessionKey string) *sessionTurnState { + v, _ := p.sessionTurnCounters.LoadOrStore(sessionKey, &sessionTurnState{}) + return v.(*sessionTurnState) +} + +// getOrCreateBoilerplateState returns the mutable boilerplate-state for +// sessionKey, creating it if it does not yet exist. +func (p *Proxy) getOrCreateBoilerplateState(sessionKey string) *boilerplateState { + v, _ := p.sessionBoilerplateState.LoadOrStore(sessionKey, &boilerplateState{}) + return v.(*boilerplateState) +} + +// resolveRequestID returns a non-empty request ID, assigning a synthetic one +// when the incoming ID is blank. It also updates the per-session missing-ID +// window and returns whether the session should be velocity-limited due to +// excessive missing-ID turns (FIX 3). +func (p *Proxy) resolveRequestID(sessionKey, rawID string) (id string, synthetic bool, velocityLimit bool) { + if rawID != "" { + ts := p.getOrCreateTurnState(sessionKey) + ts.mu.Lock() + ts.turnIndex++ + ts.mu.Unlock() + return rawID, false, false + } + + ts := p.getOrCreateTurnState(sessionKey) + ts.mu.Lock() + defer ts.mu.Unlock() + + ts.turnIndex++ + // FIX 1: UUID4 suffix guarantees uniqueness across sessions and restarts. + // Previous counter-based IDs shared the same value when the session key + // reset between iq invocations (all got suffix -1). + keyPart := sessionKey + if len(keyPart) > 8 { + keyPart = keyPart[:8] + } + syntheticID := fmt.Sprintf("iq-synthetic-%s-%s", keyPart, uuid.New().String()[:8]) + + // Track timestamp for the 60-second missing-ID window. + now := time.Now().Unix() + windowStart := now - 60 + // Evict entries older than 60 seconds. + filtered := ts.missingIDWindow[:0] + for _, t := range ts.missingIDWindow { + if t >= windowStart { + filtered = append(filtered, t) + } + } + filtered = append(filtered, now) + ts.missingIDWindow = filtered + + p.logger.Warn("request arrived with empty request ID; assigned synthetic", + slog.String("synthetic_request_id", syntheticID), + slog.String("session_key", shortLogHash(sessionKey)), + slog.Int("missing_id_window_count", len(filtered)), + ) + + vLimit := len(filtered) > 3 + return syntheticID, true, vLimit +} + func claudeSessionKey(r *http.Request, fallback string) string { if sk := strings.TrimSpace(r.Header.Get(headerSessionKey)); sk != "" { return sk @@ -433,6 +751,15 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, if err := json.Unmarshal(body, &root); err != nil { return nil, anthropicMessagesRequest{}, claudeOptimizerStats{}, claudeRequestShape{}, fmt.Errorf("parse anthropic messages body: %w", err) } + + // FIX 4: strip guard-disable directives injected via CLAUDE.md before any + // optimization or forwarding. Proxy-level guards are always enforced. + if stripGuardBypassDirectives(ctx, p.logger, root) { + // Re-serialize so the cleaned body is what gets forwarded and re-parsed. + if cleaned, err := marshalJSONNoHTMLEscape(root); err == nil { + body = cleaned + } + } req := anthropicMessagesRequest{} req.Model, _ = root["model"].(string) req.Stream, _ = root["stream"].(bool) @@ -447,6 +774,32 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, return body, req, stats, shape, nil } + // FIX 6: Session-open cache warm-up. On the first actual turn with a system prompt, + // pre-fingerprint all system blocks (CLAUDE.md, MEMORY.md, architecture + // files) so that subsequent turns start with cache hits rather than + // zero-savings. Uses a stable chunk ID: sha256(file_path + content_hash). + if root["system"] != nil { + _, alreadyWarmed := p.sessionWarmUpDone.LoadOrStore(sessionKey, true) + if !alreadyWarmed { + p.warmUpSystemSpans(ctx, cfg, sessionKey, root) + } + } + + // FIX 7: suggestion-mode payloads are ephemeral harness meta-prompts. + // Registering their large mixed content in the chunk store pollutes it and + // doubles per-turn cost. Rate-limit to 1 per 10 seconds; always skip chunk + // registration regardless of rate-limit status. + if strings.HasPrefix(shape.LatestUserText, "SUGGESTION MODE:") { + now := time.Now() + if last, ok := p.sessionSuggestionTs.Load(sessionKey); ok && now.Sub(last.(time.Time)) < 10*time.Second { + p.logger.DebugContext(ctx, "suggestion-mode request rate-limited; skipping", + slog.String("session_key", shortLogHash(sessionKey))) + } else { + p.sessionSuggestionTs.Store(sessionKey, now) + } + return body, req, stats, shape, nil + } + minSpanBytes := cfg.Optimizer.MinSpanBytes if minSpanBytes <= 0 { minSpanBytes = 512 @@ -464,6 +817,14 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, Bytes: span.Bytes, Tokens: span.Tokens, }) + } else if span.Bytes > 0 { + // FIX 4: register small spans so they are tracked across turns. + cfg.SessionStore.SaveBlock(sessionKey, memory.Block{ + Hash: span.Hash, + Kind: span.Class + ":small", + Bytes: span.Bytes, + Tokens: span.Tokens, + }) } } return body, req, stats, shape, nil @@ -490,10 +851,41 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, } } + // FIX 7: per-session boilerplate state for injection cooldown. + bpState := p.getOrCreateBoilerplateState(sessionKey) + ts := p.getOrCreateTurnState(sessionKey) + ts.mu.Lock() + currentTurn := ts.turnIndex + currentCtxBytes := ts.contextBytesCumulative + int64(len(body)) + ts.contextBytesCumulative = currentCtxBytes + ts.mu.Unlock() + var prunableSpans []TextSpan + // systemAllKnown tracks whether every system/boilerplate span was already + // seen in the session store, indicating a stable system prompt this turn. + systemAllKnown := true + hasSystemSpan := false for _, span := range spans { if span.Bytes < minSpanBytes { + // FIX 4 & SQLite Optimization: register small spans in the session cache + // only if they are >= 256 bytes to prevent SQLite write-amplification. + if cfg.SessionStore != nil && span.Bytes >= 256 { + if cfg.SessionStore.Seen(sessionKey, span.Hash) { + stats.BlocksKnown++ + } else { + cfg.SessionStore.SaveBlock(sessionKey, memory.Block{ + Hash: span.Hash, + Kind: span.Class + ":small", + Bytes: span.Bytes, + Tokens: span.Tokens, + }) + // FIX 3: register small span in the prefix-hint registry so + // larger spans can detect it as a known prefix. + p.getOrCreatePrefixHints(sessionKey).add(span.Hash, span.Bytes) + stats.BlocksNew++ + } + } stats.PreservedSmallBytes += span.Bytes stats.PreservedSmallCount++ continue @@ -506,19 +898,99 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, stats.LargestSpanBytes = span.Bytes } - known := cfg.SessionStore.Seen(sessionKey, span.Hash) - cfg.SessionStore.SaveBlock(sessionKey, memory.Block{ - Hash: span.Hash, - Kind: span.Class, - Bytes: span.Bytes, - Tokens: span.Tokens, - }) + // Track system prompt stability for Anthropic prompt-cache injection. + if span.Class == SpanClassSystemText || span.Class == SpanClassSystemBoilerplate { + hasSystemSpan = true + } + + // Sub-span chunking: for large tool_result spans, deduplicate at the + // individual chunk level so that edits to one section of a file do not + // invalidate the hashes of unchanged sections in adjacent chunks. + var known bool + if cfg.Optimizer.EnableSubspanChunking && + (span.Class == SpanClassToolResultOld || span.Class == SpanClassToolResultLatest) && + span.Bytes >= cfg.Optimizer.SmallFileBytes { + + // FIX 3: prefix-chunk reuse. Check whether the span's content + // starts with a previously registered small chunk. If so, count + // those prefix bytes as known before splitting the remainder. + phints := p.getOrCreatePrefixHints(sessionKey) + if _, prefixLen := phints.matchPrefix([]byte(span.Text)); prefixLen > 0 { + // The first prefixLen bytes match a known small chunk — + // mark the prefix as a cache hit in the session store + // (already registered, so Seen returns true next turn). + // We do not prune the prefix bytes here; prefix pruning + // is handled at the whole-span level below. + stats.BlocksKnown++ // credit the prefix as a known block + } + + chunks := splitSpanChunks(span.Text, cfg.Optimizer) + allKnown := true + for _, ch := range chunks { + chHash := hashText(ch) + chKnown := cfg.SessionStore.Seen(sessionKey, chHash) + cfg.SessionStore.SaveBlock(sessionKey, memory.Block{ + Hash: chHash, + Kind: "chunk:" + span.Class, + Bytes: len(ch), + Tokens: estimateTokens(len(ch)), + }) + if !chKnown { + allKnown = false + } + } + known = allKnown + } else { + known = cfg.SessionStore.Seen(sessionKey, span.Hash) + cfg.SessionStore.SaveBlock(sessionKey, memory.Block{ + Hash: span.Hash, + Kind: span.Class, + Bytes: span.Bytes, + Tokens: span.Tokens, + }) + } if !known { + if span.Class == SpanClassSystemText || span.Class == SpanClassSystemBoilerplate { + systemAllKnown = false + } + + // FIX 7: SUGGESTION MODE injection cooldown. If a boilerplate span + // is new (unknown hash) but the last forward was <5 turns ago AND + // context delta is <10 KB, prune it anyway to suppress redundant + // harness meta-prompt injections. + if span.Class == SpanClassSystemBoilerplate { + bpState.mu.Lock() + turnsSinceLast := currentTurn - bpState.lastForwardedTurn + bytesSinceLast := currentCtxBytes - int64(bpState.lastForwardedCtxBytes) + inCooldown := bpState.lastForwardedTurn > 0 && turnsSinceLast < 5 && bytesSinceLast < 10240 + if !inCooldown { + bpState.lastForwardedTurn = currentTurn + bpState.lastForwardedCtxBytes = int(currentCtxBytes) + } + bpState.mu.Unlock() + + if inCooldown && isEligibleSpanClass(span.Class, cfg.Optimizer) { + p.logger.DebugContext(ctx, "suppressing boilerplate injection (cooldown active)", + slog.String("session_key", shortLogHash(sessionKey)), + slog.Int("turns_since_last", turnsSinceLast), + ) + stats.ClassBytesEligible[span.Class] += span.Bytes + stats.BytesEligible += span.Bytes + prunableSpans = append(prunableSpans, span) + continue + } + } + stats.BlocksNew++ continue } stats.BlocksKnown++ + // FIX B: credit the cache hit in bytes regardless of any downstream + // preserve-vs-prune decision, so the audit metric reflects true cache + // efficiency (including protected instruction files like CLAUDE.md + // that hit the cache but are intentionally never pruned). + stats.KnownBytes += span.Bytes // Latest-turn spans must never be pruned regardless of whether the // content was seen before. The model needs current-turn context intact. @@ -528,6 +1000,17 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, continue } + // Instruction files (CLAUDE.md, CONTEXT.md, .cursorrules, etc.) must + // NEVER be pruned regardless of span class. This check runs before + // the eligibility gate so user_text_old spans containing instruction + // content are caught — previously they fell through to pruning because + // isEligibleSpanClass("user_text_old") returns true unconditionally. + if isProtectedInstructionSpan(span) { + stats.PreservedInstructionBytes += span.Bytes + stats.PreservedInstructionCount++ + continue + } + // When the same tool result content appears more than once in this // request (the model re-read the same file), preserve the newest copy // so the model always has one copy. Older duplicates fall through to @@ -553,18 +1036,14 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, continue } - if isProtectedInstructionSpan(span) { - stats.PreservedInstructionBytes += span.Bytes - stats.PreservedInstructionCount++ - continue - } - stats.ClassBytesEligible[span.Class] += span.Bytes stats.BytesEligible += span.Bytes prunableSpans = append(prunableSpans, span) } - if len(prunableSpans) == 0 { + isSub, _ := ctx.Value(isSubscriptionKey).(bool) + wantPromptCache := cfg.Optimizer.EnablePromptCache && hasSystemSpan && systemAllKnown && !cfg.Bedrock.Enabled && !isSub + if len(prunableSpans) == 0 && !wantPromptCache { return body, req, stats, shape, nil } @@ -591,6 +1070,11 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, } } + // Inject Anthropic server-side prompt cache header when system prompt is stable. + if wantPromptCache { + injectPromptCacheHeaders(rewriteRoot) + } + optimized, err := marshalJSONNoHTMLEscape(rewriteRoot) if err != nil { p.logger.WarnContext(ctx, "claude optimize marshal failed; forwarding original", slog.Any("err", err)) @@ -605,6 +1089,11 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, stats.ReductionRatio = float64(stats.BytesBefore-stats.BytesAfter) / float64(stats.BytesBefore) } + // FIX 2: synchronous warm-up flush at the end of prepareClaudeBody when Turn is 1 or 2 (turn_index == 0 in 0-based auditing) + if (currentTurn == 1 || currentTurn == 2) && root["system"] != nil { + p.warmUpSystemSpans(ctx, cfg, sessionKey, root) + } + if cfg.Mode == "optimize" { return optimized, req, stats, shape, nil } @@ -622,6 +1111,8 @@ func isEligibleSpanClass(class string, cfg OptimizerConfig) bool { return cfg.EnableToolResultPruning case SpanClassAssistantTextOld: return cfg.EnableAssistantPruning + case SpanClassSystemBoilerplate: + return true // harness meta-prompts are prunable after first occurrence case SpanClassSystemText: return false default: @@ -734,6 +1225,27 @@ func setSystemSpanText(root map[string]any, span TextSpan, replacement string) e return nil } +func stripCacheControlRecursively(v any) { + switch val := v.(type) { + case map[string]any: + delete(val, "cache_control") + for _, child := range val { + stripCacheControlRecursively(child) + } + case []any: + for _, child := range val { + stripCacheControlRecursively(child) + } + } +} + +func isSubscriptionAuth(auth string) bool { + return strings.HasPrefix(auth, "Bearer claudepro_") || + strings.Contains(auth, "claudepro_") || + strings.Contains(auth, "sk-ant-oat") +} + + // forwardClaudeMessagesViaBedrock proxies /v1/messages to the Bedrock // InvokeModelWithResponseStream API. The Anthropic Messages body is forwarded // as-is except that `model` and `stream` are removed and @@ -758,6 +1270,7 @@ func (p *Proxy) forwardClaudeMessagesViaBedrock(w http.ResponseWriter, r *http.R delete(root, "model") delete(root, "stream") delete(root, "context_management") // Bedrock rejects this Anthropic-specific field + stripCacheControlRecursively(root) // Bedrock rejects cache_control fields root["anthropic_version"] = "bedrock-2023-05-31" bedrockBody, err := json.Marshal(root) if err != nil { @@ -832,12 +1345,18 @@ func (p *Proxy) forwardClaudeMessagesViaBedrock(w http.ResponseWriter, r *http.R } _ = rc.Flush() + payloadStr := string(payload) + if strings.Contains(payloadStr, `"tool_use"`) { + stats.HasToolUse = true + } + switch ev.Type { case "content_block_delta": stats.Chunks++ - stats.OutputText += anthropicDeltaTextLen(string(payload)) + stats.OutputText += anthropicDeltaTextLen(payloadStr) + stats.OutputRawText += anthropicDeltaText(payloadStr) case "message_delta": - if tokens := anthropicUsageOutputTokens(string(payload)); tokens > 0 { + if tokens := anthropicUsageOutputTokens(payloadStr); tokens > 0 { stats.OutputTokens = tokens } } @@ -921,6 +1440,51 @@ func toBedrockModelID(model string, cfg BedrockConfig) string { func strPtr(s string) *string { return &s } +// splitSpanChunks splits span text into content-defined chunks using Rabin-Karp +// CDC. For code content a wider 256-byte window is selected automatically so +// that edits within one function shift fewer adjacent chunk boundaries. +// Content shorter than cfg.SmallFileBytes is returned as a single element. +func splitSpanChunks(text string, cfg OptimizerConfig) []string { + smallBytes := cfg.SmallFileBytes + if smallBytes <= 0 { + smallBytes = 4096 + } + data := []byte(text) + if len(data) < smallBytes { + return []string{text} + } + var ckCfg chunker.Config + if len(data) >= 8192 && chunker.IsCodeContent(data) { + // FIX 5: 256-byte window for large code files (≥8 KB with declaration-keyword + // density) to produce fewer, more stable chunk boundaries that survive + // minor edits without shifting fingerprints in adjacent functions. + ckCfg = chunker.CodeConfig() + } else if chunker.IsSystemProseContent(data) { + // FIX 4: 4 KB MaxSize for XML-tagged system-reminder blocks so a single + // changed field (e.g. currentDate) only invalidates its own small chunk. + ckCfg = chunker.SystemProseConfig() + } else if chunker.IsPathList(data) { + // FIX 7: narrow 32-byte window for file-path listings so per-path changes + // produce localised boundary shifts instead of cascading re-fingerprints. + ckCfg = chunker.PathListConfig() + } else { + ckCfg = chunker.DefaultConfig() + } + if cfg.MaxChunkBytes > 0 { + ckCfg.MaxSize = cfg.MaxChunkBytes + } + ck := chunker.New(ckCfg) + raw := ck.Split(data) + if len(raw) == 0 { + return []string{text} + } + result := make([]string, len(raw)) + for i, ch := range raw { + result[i] = string(ch.Data) + } + return result +} + func hashText(text string) string { sum := sha256.Sum256([]byte(strings.TrimSpace(text))) return hex.EncodeToString(sum[:]) @@ -930,6 +1494,18 @@ func (p *Proxy) forwardClaudeMessages(w http.ResponseWriter, r *http.Request, cf if cfg.Bedrock.Enabled { return p.forwardClaudeMessagesViaBedrock(w, r, cfg, body) } + + auth := r.Header.Get("Authorization") + if isSubscriptionAuth(auth) { + var root map[string]any + if err := json.Unmarshal(body, &root); err == nil { + stripCacheControlRecursively(root) + if cleaned, err := marshalJSONNoHTMLEscape(root); err == nil { + body = cleaned + } + } + } + upstreamURL, err := url.JoinPath(strings.TrimRight(cfg.AnthropicBaseURL, "/"), "v1", "messages") if err != nil { p.writeError(w, r, errorPayload{HTTPStatus: http.StatusInternalServerError, Type: "server_error", Code: "bad_upstream_url", Message: err.Error()}) @@ -1109,10 +1685,15 @@ func proxyAnthropicStream(w http.ResponseWriter, r *http.Request, resp *http.Res event = strings.TrimSpace(strings.TrimPrefix(line, "event: ")) case strings.HasPrefix(line, "data: "): payload := strings.TrimSpace(strings.TrimPrefix(line, "data: ")) + if strings.Contains(payload, `"tool_use"`) { + stats.HasToolUse = true + } switch event { case "content_block_delta": stats.Chunks++ - stats.OutputText += anthropicDeltaTextLen(payload) + txt := anthropicDeltaText(payload) + stats.OutputText += len(txt) + stats.OutputRawText += txt case "message_delta": if tokens := anthropicUsageOutputTokens(payload); tokens > 0 { stats.OutputTokens = tokens @@ -1129,6 +1710,22 @@ func proxyAnthropicStream(w http.ResponseWriter, r *http.Request, resp *http.Res return stats } +func anthropicDeltaText(payload string) string { + var ev struct { + Delta struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"delta"` + } + if err := json.Unmarshal([]byte(payload), &ev); err != nil { + return "" + } + if ev.Delta.Type != "text_delta" { + return "" + } + return ev.Delta.Text +} + func anthropicDeltaTextLen(payload string) int { var ev struct { Delta struct { @@ -1277,7 +1874,7 @@ func appendText(sb *strings.Builder, value any) { } } -func (p *Proxy) logClaudeRequestComplete(ctx context.Context, requestID, mode, model, sessionKey string, bytesBefore int, opt claudeOptimizerStats, stream claudeStreamStats, dur time.Duration) { +func (p *Proxy) logClaudeRequestComplete(ctx context.Context, requestID, mode, model, sessionKey string, bytesBefore int, opt claudeOptimizerStats, stream claudeStreamStats, dur time.Duration, missingRequestID bool, syntheticRequestID string, velocityWarning bool) { level := slog.LevelInfo if stream.Status == "error" || stream.Status == "stream_error" { level = slog.LevelError @@ -1308,6 +1905,12 @@ func (p *Proxy) logClaudeRequestComplete(ctx context.Context, requestID, mode, m slog.Int("status_code", stream.StatusCode), slog.Bool("stream_completed", stream.Completed), slog.Bool("stream_cancelled", stream.Cancelled), + // FIX 8: request ID tracing fields + slog.Bool("missing_request_id", missingRequestID), + slog.Bool("velocity_warning", velocityWarning), + } + if missingRequestID && syntheticRequestID != "" { + attrs = append(attrs, slog.String("synthetic_request_id", syntheticRequestID)) } if opt.BytesEligible > 0 { attrs = append(attrs, slog.Int("bytes_eligible", opt.BytesEligible)) @@ -1392,6 +1995,110 @@ func shortLogHash(value string) string { return hex.EncodeToString(sum[:6]) } +// semanticPromptHash returns a 128-bit hex digest keyed on the content of +// the last 3 user messages plus the first 64 bytes of the system field. +// Raw-body hashing fails in-flight deduplication because the same logical +// prompt produces different bytes each turn as the context window grows. +// Falls back to a raw SHA-256 if the body cannot be parsed. +func semanticPromptHash(body []byte) string { + var parsed struct { + System json.RawMessage `json:"system"` + Messages []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + } `json:"messages"` + } + if err := json.Unmarshal(body, &parsed); err != nil { + sum := sha256.Sum256(body) + return hex.EncodeToString(sum[:16]) + } + + // System fingerprint: first 64 bytes of the string-form system content. + var sysFP string + var sysStr string + if err := json.Unmarshal(parsed.System, &sysStr); err == nil { + if len(sysStr) > 64 { + sysStr = sysStr[:64] + } + sysFP = sysStr + } else { + var sysArr []map[string]any + if err2 := json.Unmarshal(parsed.System, &sysArr); err2 == nil && len(sysArr) > 0 { + if text, ok := sysArr[0]["text"].(string); ok { + if len(text) > 64 { + text = text[:64] + } + sysFP = text + } + } + } + + // Collect last 3 user-message contents. + var userContents []string + for _, msg := range parsed.Messages { + if !strings.EqualFold(msg.Role, "user") { + continue + } + var text string + if err := json.Unmarshal(msg.Content, &text); err == nil { + userContents = append(userContents, text) + } else { + var blocks []map[string]any + if err2 := json.Unmarshal(msg.Content, &blocks); err2 == nil { + var sb strings.Builder + for _, b := range blocks { + appendText(&sb, b["text"]) + appendText(&sb, b["content"]) + } + userContents = append(userContents, sb.String()) + } + } + } + if len(userContents) > 3 { + userContents = userContents[len(userContents)-3:] + } + + h := sha256.New() + for _, c := range userContents { + h.Write([]byte(c)) + } + h.Write([]byte(sysFP)) + return hex.EncodeToString(h.Sum(nil)[:16]) +} + +// getOrCreatePrefixHints returns the prefix-hint set for sessionKey, creating +// it if it does not yet exist (FIX 3). +func (p *Proxy) getOrCreatePrefixHints(sessionKey string) *prefixHintSet { + v, _ := p.sessionPrefixHints.LoadOrStore(sessionKey, &prefixHintSet{ + hints: make(map[string]int), + }) + return v.(*prefixHintSet) +} + +func (s *prefixHintSet) add(hash string, length int) { + s.mu.Lock() + s.hints[hash] = length + s.mu.Unlock() +} + +// matchPrefix returns the length and hash of the longest registered small +// chunk whose content matches data[0:length], or 0 if none match (FIX 3). +func (s *prefixHintSet) matchPrefix(data []byte) (matchedHash string, matchedLen int) { + s.mu.Lock() + defer s.mu.Unlock() + for h, l := range s.hints { + if l <= 0 || l > len(data) { + continue + } + sum := sha256.Sum256([]byte(strings.TrimSpace(string(data[:l])))) + if hex.EncodeToString(sum[:]) == h && l > matchedLen { + matchedLen = l + matchedHash = h + } + } + return +} + func estimateTokens(chars int) int { if chars <= 0 { return 0 @@ -1430,3 +2137,58 @@ func retryAfterDuration(value string, now time.Time) time.Duration { } return when.Sub(now) } + +type responseCaptureWriter struct { + http.ResponseWriter + buf bytes.Buffer +} + +func (rcw *responseCaptureWriter) Write(b []byte) (int, error) { + rcw.buf.Write(b) + return rcw.ResponseWriter.Write(b) +} + +func (rcw *responseCaptureWriter) Flush() { + if flusher, ok := rcw.ResponseWriter.(http.Flusher); ok { + flusher.Flush() + } +} + +func (rcw *responseCaptureWriter) WriteString(s string) (int, error) { + rcw.buf.WriteString(s) + if sw, ok := rcw.ResponseWriter.(io.StringWriter); ok { + return sw.WriteString(s) + } + return rcw.ResponseWriter.Write([]byte(s)) +} + +func computePromptHash(prompt, model string) string { + normalized := strings.TrimSpace(strings.ToLower(prompt)) + sum := sha256.Sum256([]byte(normalized + "|" + model)) + return hex.EncodeToString(sum[:]) +} + +func writeSyntheticStreamResponse(w http.ResponseWriter, text string) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Accel-Buffering", "no") + w.WriteHeader(http.StatusOK) + + textData, _ := json.Marshal(text) + events := []string{ + `event: message_start` + "\n" + `data: {"type":"message_start","message":{"id":"msg_synth","type":"message","role":"assistant","content":[],"model":"claude-3-5-sonnet","usage":{"input_tokens":5,"output_tokens":5}}}` + "\n\n", + `event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n", + `event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":` + string(textData) + `}}` + "\n\n", + `event: content_block_stop` + "\n" + `data: {"type":"content_block_stop","index":0}` + "\n\n", + `event: message_delta` + "\n" + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":5}}` + "\n\n", + `event: message_stop` + "\n" + `data: {"type":"message_stop"}` + "\n\n", + } + + for _, ev := range events { + _, _ = io.WriteString(w, ev) + } + if flusher, ok := w.(http.Flusher); ok { + flusher.Flush() + } +} diff --git a/gateway/internal/proxy/prompt_cache.go b/gateway/internal/proxy/prompt_cache.go new file mode 100644 index 0000000..a91b26c --- /dev/null +++ b/gateway/internal/proxy/prompt_cache.go @@ -0,0 +1,40 @@ +package proxy + +// injectPromptCacheHeaders adds Anthropic's cache_control: {type:"ephemeral"} +// marker to the last system content block. When the system prompt is identical +// to the prior turn, Anthropic's server-side prompt cache covers the stable +// prefix and only charges cache-read tokens (≈10% of normal input token cost) +// for all subsequent turns that share the same prefix. +// +// Call this only when the system spans are all "known" (seen in the session +// store), which indicates the system prompt has not changed since the last turn. +func injectPromptCacheHeaders(root map[string]any) { + sys, ok := root["system"] + if !ok { + return + } + cacheCtrl := map[string]any{"type": "ephemeral"} + switch v := sys.(type) { + case string: + if len(v) == 0 { + return + } + // Promote string system to array form so we can attach cache_control. + root["system"] = []any{ + map[string]any{ + "type": "text", + "text": v, + "cache_control": cacheCtrl, + }, + } + case []any: + if len(v) == 0 { + return + } + last, ok := v[len(v)-1].(map[string]any) + if !ok { + return + } + last["cache_control"] = cacheCtrl + } +} diff --git a/gateway/internal/proxy/proxy.go b/gateway/internal/proxy/proxy.go index 0a06e39..08ebad1 100644 --- a/gateway/internal/proxy/proxy.go +++ b/gateway/internal/proxy/proxy.go @@ -10,6 +10,7 @@ import ( "context" "log/slog" "net/http" + "sync" "time" "github.com/Revanth14/indexqube/gateway/internal/domain" @@ -43,6 +44,114 @@ type Proxy struct { sessionTracker *telemetry.AgentSessionStore sessionPersist *sessions.Tracker supabaseStats *StatsHandler + + // sessionTurnCounters tracks the number of turns per session key. + // Used for synthetic request ID generation (FIX 3). + sessionTurnCounters sync.Map // map[string]*sessionTurnState + + // sessionWarmUpDone tracks which sessions have completed cache warm-up. + // Set after pre-registering system spans on the first turn (FIX 6). + sessionWarmUpDone sync.Map // map[string]bool + + // sessionBoilerplateState tracks the last turn/byte-offset where system + // boilerplate was forwarded for injection cooldown logic (FIX 7). + sessionBoilerplateState sync.Map // map[string]*boilerplateState + + // inFlightRequests prevents simultaneous duplicate upstream dispatches. + // When a second identical prompt arrives while the first is in-flight it + // waits for the first to complete before dispatching (FIX 2). + inFlightRequests *inFlightTracker + + // sessionPrefixHints stores per-session (hash → byteLength) entries for + // spans shorter than SmallFileBytes. When a larger span arrives, its prefix + // is checked against these hints to detect "small payload is prefix of + // large payload" reuse patterns (FIX 3). + sessionPrefixHints sync.Map // map[string]*prefixHintSet + + // sessionSuggestionTs tracks the last time a suggestion-mode request was + // processed per session. Used to rate-limit harness meta-prompt injections + // to max 1 per 10 seconds so ephemeral payloads don't pollute the chunk store. + sessionSuggestionTs sync.Map // map[string]time.Time +} + +// inFlightTracker serialises duplicate identical requests so they do not +// all fire upstream concurrently. The first arrival dispatches normally; +// subsequent arrivals with the same prompt hash block until it finishes. +type inFlightTracker struct { + mu sync.Mutex + inflight map[string]chan struct{} // prompt_hash → done channel +} + +func newInFlightTracker() *inFlightTracker { + return &inFlightTracker{inflight: make(map[string]chan struct{})} +} + +// acquire attempts to register promptHash as in-flight. +// Returns (doneFn, nil) when this is the first arrival — the caller must call +// doneFn exactly once after the request completes. +// Returns (nil, waitChan) when a duplicate is already in-flight — the caller +// should wait on waitChan before dispatching. +func (t *inFlightTracker) acquire(promptHash string) (doneFn func(), waitChan chan struct{}) { + t.mu.Lock() + defer t.mu.Unlock() + if ch, ok := t.inflight[promptHash]; ok { + return nil, ch + } + done := make(chan struct{}) + t.inflight[promptHash] = done + return func() { + t.mu.Lock() + delete(t.inflight, promptHash) + t.mu.Unlock() + close(done) + }, nil +} + +// prefixHintSet stores (hash → byteLength) entries for small spans so that +// larger spans can be checked for prefix reuse (FIX 3). +type prefixHintSet struct { + mu sync.Mutex + hints map[string]int // content-hash → byte-length of the small chunk +} + +type cachedResponse struct { + promptHash string + payload []byte +} + +// sessionTurnState holds per-session counters needed for FIX 3 and FIX 7. +type sessionTurnState struct { + mu sync.Mutex + turnIndex int + missingIDWindow []int64 // Unix timestamps of turns with missing request IDs + contextBytesCumulative int64 // running total of context bytes seen across turns + cachedResponses []cachedResponse +} + +func (s *sessionTurnState) getCachedResponse(promptHash string) ([]byte, bool) { + for _, cr := range s.cachedResponses { + if cr.promptHash == promptHash { + return cr.payload, true + } + } + return nil, false +} + +func (s *sessionTurnState) saveCachedResponse(promptHash string, payload []byte) { + if len(s.cachedResponses) >= 10 { + s.cachedResponses = s.cachedResponses[1:] + } + s.cachedResponses = append(s.cachedResponses, cachedResponse{ + promptHash: promptHash, + payload: payload, + }) +} + +// boilerplateState tracks when system boilerplate was last forwarded. +type boilerplateState struct { + mu sync.Mutex + lastForwardedTurn int + lastForwardedCtxBytes int } // BedrockConfig routes /v1/messages to AWS Bedrock instead of Anthropic's API. @@ -61,13 +170,26 @@ type BedrockConfig struct { // safe defaults via claudeDefaults(). type OptimizerConfig struct { MinSpanBytes int // minimum span size to consider for pruning (default 512) - TargetChunkBytes int // target chunk size for future long-block strategy (default 2048) + TargetChunkBytes int // target chunk size for Rabin-Karp chunker (default 2048) MaxChunkBytes int // maximum chunk size (default 8192) MinSavedTokens int // skip rewrite if savings below this threshold (default 10) EnableToolResultPruning bool // prune old tool_result spans (default true via claudeDefaults) EnableAssistantPruning bool // prune old assistant text spans (default false) EnableSystemPruning bool // deprecated; system text spans are never pruned Diagnostics bool // emit verbose per-class diagnostics in logs + + // Sub-span chunking: splits large tool_result spans with Rabin-Karp CDC so + // that only spans whose every chunk is known get pruned. This preserves + // context when a file is partially edited while still deduplicating unchanged + // sections across turns. + EnableSubspanChunking bool // split large spans for chunk-level dedup (default true) + SmallFileBytes int // content shorter than this bypasses the chunker (default 4096) + + // EnablePromptCache injects Anthropic cache_control: {type: "ephemeral"} on + // the last system block when the system prompt is identical to the prior turn. + // Anthropic's server-side cache then covers the stable prefix, reducing + // billable input tokens for sessions with a large, unchanged system prompt. + EnablePromptCache bool } type ClaudeMessagesConfig struct { @@ -173,11 +295,12 @@ func New(gov Governor, opts ...Option) *Proxy { panic("proxy: governor is required") } p := &Proxy{ - governor: gov, - logger: slog.Default(), - mux: http.NewServeMux(), - maxRequestSize: 8 << 20, // default 8 MiB - optimizeTimeout: 30 * time.Second, + governor: gov, + logger: slog.Default(), + mux: http.NewServeMux(), + maxRequestSize: 8 << 20, // default 8 MiB + optimizeTimeout: 30 * time.Second, + inFlightRequests: newInFlightTracker(), } for _, opt := range opts { opt(p) diff --git a/gateway/internal/proxy/proxy_test.go b/gateway/internal/proxy/proxy_test.go index a70a4ec..f86b8d3 100644 --- a/gateway/internal/proxy/proxy_test.go +++ b/gateway/internal/proxy/proxy_test.go @@ -623,6 +623,7 @@ func TestClaudeMessages_OptimizePrunesRepeatedTextBlock(t *testing.T) { t.Fatalf("new request: %v", err) } req.Header.Set("Authorization", "Bearer iq-dev-local") + req.Header.Set("Cache-Control", "no-cache") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("POST /v1/messages: %v", err) @@ -667,6 +668,7 @@ func TestClaudeMessages_OptimizePrunesRepeatedLargeChunks(t *testing.T) { t.Fatalf("new request: %v", err) } req.Header.Set("Authorization", "Bearer iq-dev-local") + req.Header.Set("Cache-Control", "no-cache") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("POST /v1/messages: %v", err) @@ -723,6 +725,7 @@ func TestClaudeMessages_OptimizePreservesLatestTurnWhilePruningOldContent(t *tes t.Fatalf("new request: %v", err) } req.Header.Set("Authorization", "Bearer iq-dev-local") + req.Header.Set("Cache-Control", "no-cache") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("POST /v1/messages: %v", err) @@ -950,8 +953,8 @@ func TestDumpClaudePayloadsAppendsSessionFile(t *testing.T) { sessionFile := filepath.Join(dumpDir, "iq-session-test.jsonl") t.Setenv("IQ_DUMP_SESSION_FILE", sessionFile) - dumpClaudePayloads("dump-test", []byte(`{"before":true}`), []byte(`{"after":true}`)) - dumpClaudePayloads("dump-test-2", []byte(`{"before":2}`), []byte(`{"after":2}`)) + dumpClaudePayloads("dump-test", []byte(`{"before":true}`), []byte(`{"after":true}`), claudeStreamStats{OutputRawText: "simulated response", OutputTokens: 10, Status: "completed"}, claudeOptimizerStats{}) + dumpClaudePayloads("dump-test-2", []byte(`{"before":2}`), []byte(`{"after":2}`), claudeStreamStats{}, claudeOptimizerStats{}) raw, err := os.ReadFile(sessionFile) if err != nil { @@ -971,6 +974,9 @@ func TestDumpClaudePayloadsAppendsSessionFile(t *testing.T) { if !strings.Contains(string(first.Before), `"before":true`) { t.Fatalf("before payload missing from first record: %s", first.Before) } + if first.Response.Text != "simulated response" || first.Response.OutputTokens != 10 { + t.Fatalf("unexpected response metrics inside first record: %+v", first.Response) + } if _, err := os.Stat(filepath.Join(dumpDir, "iq-before-dump-test.json")); !errors.Is(err, os.ErrNotExist) { t.Fatalf("session dumps should not create per-request before file, err=%v", err) } @@ -981,7 +987,7 @@ func TestDumpClaudePayloadsFallsBackToPairFiles(t *testing.T) { t.Setenv("IQ_DUMP_SESSION_FILE", "") t.Setenv("IQ_DUMP_DIR", dumpDir) - dumpClaudePayloads("dump-test", []byte(`{"before":true}`), []byte(`{"after":true}`)) + dumpClaudePayloads("dump-test", []byte(`{"before":true}`), []byte(`{"after":true}`), claudeStreamStats{}, claudeOptimizerStats{}) before, err := os.ReadFile(filepath.Join(dumpDir, "iq-before-dump-test.json")) if err != nil { @@ -999,6 +1005,77 @@ func TestDumpClaudePayloadsFallsBackToPairFiles(t *testing.T) { } } +func TestDumpClaudePayloadsEmitsOptimizerStats(t *testing.T) { + sessionFile := filepath.Join(t.TempDir(), "session.jsonl") + t.Setenv("IQ_DUMP_SESSION_FILE", sessionFile) + + // Realistic post-FIX-B stats: a turn where some bytes were pruned, some + // were known-but-protected (CLAUDE.md), and the rest were known-but- + // latest-turn. KnownBytes is the sum across all preservation paths plus + // the pruned bytes. + opt := claudeOptimizerStats{ + BlocksPruned: 2, + BlocksKnown: 5, + BytesPruned: 1000, + PreservedInstructionBytes: 3027, + PreservedInstructionCount: 1, + PreservedLatestTurnBytes: 500, + PreservedLatestTurnCount: 1, + PreservedLastOccurrenceBytes: 0, + PreservedLastOccurrenceCount: 0, + } + opt.KnownBytes = opt.BytesPruned + opt.PreservedInstructionBytes + + opt.PreservedLatestTurnBytes + opt.PreservedLastOccurrenceBytes + + dumpClaudePayloads("opt-test", []byte(`{"before":1}`), []byte(`{"after":1}`), claudeStreamStats{}, opt) + + raw, err := os.ReadFile(sessionFile) + if err != nil { + t.Fatalf("read session file: %v", err) + } + line := strings.TrimRight(string(raw), "\n") + var rec struct { + SavedBytes int `json:"saved_bytes"` + Optimizer *struct { + BlocksPruned int `json:"blocks_pruned"` + BlocksKnown int `json:"blocks_known"` + BlocksKnownProtected int `json:"blocks_known_protected"` + BytesPruned int `json:"bytes_pruned"` + ProtectedBytes int `json:"protected_bytes"` + KnownBytes int `json:"known_bytes"` + TrueCacheHitBytes int `json:"true_cache_hit_bytes"` + } `json:"optimizer"` + } + if err := json.Unmarshal([]byte(line), &rec); err != nil { + t.Fatalf("parse dump line: %v\nline=%s", err, line) + } + if rec.Optimizer == nil { + t.Fatalf("expected optimizer block in dump record; got: %s", line) + } + if got, want := rec.Optimizer.ProtectedBytes, opt.PreservedInstructionBytes; got != want { + t.Errorf("protected_bytes: got %d, want %d", got, want) + } + if got, want := rec.Optimizer.BlocksKnownProtected, opt.PreservedInstructionCount; got != want { + t.Errorf("blocks_known_protected: got %d, want %d", got, want) + } + if got, want := rec.Optimizer.BytesPruned, opt.BytesPruned; got != want { + t.Errorf("bytes_pruned: got %d, want %d", got, want) + } + if got, want := rec.Optimizer.KnownBytes, opt.KnownBytes; got != want { + t.Errorf("known_bytes: got %d, want %d", got, want) + } + if got, want := rec.Optimizer.TrueCacheHitBytes, opt.KnownBytes; got != want { + t.Errorf("true_cache_hit_bytes: got %d, want %d (must equal KnownBytes)", got, want) + } + // Invariant: true_cache_hit_bytes == bytes_pruned + protected_bytes + (other preservations). + // Here PreservedLastOccurrenceBytes==0 and PreservedLatestTurnBytes==500, so: + wantInvariant := rec.Optimizer.BytesPruned + rec.Optimizer.ProtectedBytes + opt.PreservedLatestTurnBytes + if rec.Optimizer.TrueCacheHitBytes != wantInvariant { + t.Errorf("invariant broken: true_cache_hit_bytes=%d, expected %d", + rec.Optimizer.TrueCacheHitBytes, wantInvariant) + } +} + func TestProxyAnthropicStream_UsesMessageDeltaOutputTokens(t *testing.T) { t.Parallel() body := strings.Join([]string{ diff --git a/gateway/internal/proxy/span.go b/gateway/internal/proxy/span.go index 3d981a6..eeca2e8 100644 --- a/gateway/internal/proxy/span.go +++ b/gateway/internal/proxy/span.go @@ -18,6 +18,13 @@ const ( SpanClassToolResultOld = "tool_result_old" SpanClassToolUse = "tool_use" SpanClassUnknownText = "unknown_text" + + // SpanClassSystemBoilerplate identifies system blocks that contain known + // harness-injected meta-prompts (e.g. SUGGESTION MODE, system-reminder). + // Unlike SpanClassSystemText, these are eligible for pruning after the first + // occurrence — the model has already processed the instruction and the + // block adds no new information on subsequent turns. + SpanClassSystemBoilerplate = "system_boilerplate" ) // TextSpan is one text-bearing unit extracted from an Anthropic messages request. @@ -125,7 +132,7 @@ func toolUseInputPath(item map[string]any) string { if !ok { return "" } - for _, key := range []string{"file_path", "path", "notebook_path"} { + for _, key := range []string{"file_path", "path", "notebook_path", "command"} { if path, ok := input[key].(string); ok && strings.TrimSpace(path) != "" { return path } @@ -133,14 +140,40 @@ func toolUseInputPath(item map[string]any) string { return "" } +// boilerplateSystemPatterns are substrings that identify harness-injected +// meta-prompts inside system blocks. Matches are case-insensitive. +var boilerplateSystemPatterns = []string{ + "suggestion mode", + "", + "system-reminder", + "autonomous-loop", + "< # audit specific session + python3 scripts/audit_session.py list # list available traces + python3 scripts/audit_session.py stats # quick local stats (no Bedrock) +""" +from __future__ import annotations + +import argparse +import dataclasses +import glob +import json +import logging +import os +import re +import shutil +import sys +import textwrap +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass, field +from datetime import datetime, timezone +from enum import Enum +from pathlib import Path +from typing import Any, Optional, Sequence + +# ────────────────────────────────────────────────────────────────────── +# Logging +# ────────────────────────────────────────────────────────────────────── +LOG_FORMAT = "%(asctime)s [%(levelname)s] %(message)s" +LOG_DATEFMT = "%H:%M:%S" + +logger = logging.getLogger("iq-audit") + + +# ────────────────────────────────────────────────────────────────────── +# Terminal colors (auto-disabled when piped / NO_COLOR) +# ────────────────────────────────────────────────────────────────────── +class _Colors: + """ANSI escape wrapper that silences itself when stdout is not a tty.""" + + _enabled: bool = sys.stdout.isatty() and os.environ.get("NO_COLOR") is None + + BLUE = "\033[1;36m" + GREEN = "\033[1;32m" + YELLOW = "\033[1;33m" + RED = "\033[1;31m" + DIM = "\033[2m" + BOLD = "\033[1m" + RESET = "\033[0m" + + def __getattr__(self, name: str) -> str: + val = object.__getattribute__(self, name) if name.isupper() else "" + return val if self._enabled else "" + + +C = _Colors() + + +# ────────────────────────────────────────────────────────────────────── +# Configuration (env-overridable) +# ────────────────────────────────────────────────────────────────────── +@dataclass(frozen=True) +class Config: + """Immutable runtime configuration.""" + + # Dump directories — searched in order + dump_search_paths: tuple[str, ...] = ( + ".indexqube/dumps", + "gateway/.indexqube/dumps", + "~/.indexqube/dumps", + ) + # Housekeeping + max_retained_traces: int = 30 + + # Bedrock + bedrock_region: str = os.environ.get("IQ_BEDROCK_REGION", "us-east-1") + bedrock_connect_timeout: int = 10 + bedrock_read_timeout: int = 90 + bedrock_max_retries: int = 2 + + # Model fallback chain + model_ids: tuple[str, ...] = ( + "us.anthropic.claude-sonnet-4-6", + "us.anthropic.claude-opus-4-6", + "anthropic.claude-sonnet-4-6", + "amazon.nova-pro-v1:0", + ) + + # Prompt budget + max_audit_turns: int = 18 + prompt_token_hard_cap: int = 150_000 # ~600 KB chars + text_truncate_threshold: int = 5120 # bytes before truncation kicks in + text_truncate_retain: int = 4000 # bytes kept after truncation + + # Outputs + s3_bucket: str = os.environ.get("IQ_S3_BUCKET", "indexqube-session-traces") + cw_namespace: str = "IndexQube" + enable_s3: bool = os.environ.get("IQ_ENABLE_S3", "1") == "1" + enable_cloudwatch: bool = os.environ.get("IQ_ENABLE_CW", "1") == "1" + + +CFG = Config() + +# ────────────────────────────────────────────────────────────────────── +# Data models +# ────────────────────────────────────────────────────────────────────── + +class ResponseStatus(str, Enum): + COMPLETED = "completed" + ERROR = "error" + UNKNOWN = "unknown" + + +@dataclass +class OptimizerMetrics: + """Per-turn optimizer stats emitted by the Go gateway.""" + + blocks_pruned: int = 0 + blocks_known: int = 0 + blocks_known_protected: int = 0 + bytes_pruned: int = 0 + protected_bytes: int = 0 + known_bytes: int = 0 + true_cache_hit_bytes: int = 0 + + @classmethod + def from_record(cls, raw: dict[str, Any]) -> OptimizerMetrics: + opt = raw.get("optimizer") or {} + saved = raw.get("saved_bytes", 0) + return cls( + blocks_pruned=opt.get("blocks_pruned", 0), + blocks_known=opt.get("blocks_known", 0), + blocks_known_protected=opt.get("blocks_known_protected", 0), + bytes_pruned=opt.get("bytes_pruned", saved), + protected_bytes=opt.get("protected_bytes", 0), + known_bytes=opt.get("known_bytes", 0), + true_cache_hit_bytes=opt.get("true_cache_hit_bytes", saved), + ) + + +@dataclass +class TurnRecord: + """Normalized representation of one proxy turn.""" + + index: int # 1-based turn number + ts: str = "" + request_id: str = "" + before_bytes: int = 0 + after_bytes: int = 0 + saved_bytes: int = 0 + model: str = "" + status: ResponseStatus = ResponseStatus.UNKNOWN + output_tokens: int = 0 + optimizer: OptimizerMetrics = field(default_factory=OptimizerMetrics) + # Raw payloads kept for prompt construction (not serialized to reports) + _before_payload: dict = field(default_factory=dict, repr=False) + _response_text: str = field(default="", repr=False) + + @property + def reduction_pct(self) -> float: + return (self.saved_bytes / self.before_bytes * 100) if self.before_bytes > 0 else 0.0 + + @classmethod + def from_raw(cls, raw: dict[str, Any], index: int) -> TurnRecord: + before = raw.get("before") or {} + resp = raw.get("response") or {} + status_str = resp.get("status", "unknown") + try: + status = ResponseStatus(status_str) + except ValueError: + status = ResponseStatus.UNKNOWN + return cls( + index=index, + ts=raw.get("ts", ""), + request_id=raw.get("request_id", ""), + before_bytes=raw.get("before_bytes", 0), + after_bytes=raw.get("after_bytes", 0), + saved_bytes=raw.get("saved_bytes", 0), + model=before.get("model", "unknown"), + status=status, + output_tokens=resp.get("output_tokens", 0), + optimizer=OptimizerMetrics.from_record(raw), + _before_payload=before, + _response_text=resp.get("text", ""), + ) + + def summary_dict(self) -> dict[str, Any]: + """Serializable summary for stats output.""" + return { + "turn": self.index, + "ts": self.ts, + "model": self.model, + "before_bytes": self.before_bytes, + "after_bytes": self.after_bytes, + "saved_bytes": self.saved_bytes, + "reduction_pct": round(self.reduction_pct, 1), + "status": self.status.value, + "output_tokens": self.output_tokens, + "optimizer": dataclasses.asdict(self.optimizer), + } + + +@dataclass +class SessionTrace: + """Parsed and validated session trace.""" + + session_id: str + file_path: str + turns: list[TurnRecord] + parse_warnings: int = 0 + + @property + def total_before(self) -> int: + return sum(t.before_bytes for t in self.turns) + + @property + def total_saved(self) -> int: + return sum(t.saved_bytes for t in self.turns) + + @property + def total_reduction_pct(self) -> float: + return (self.total_saved / self.total_before * 100) if self.total_before > 0 else 0.0 + + @property + def models_used(self) -> set[str]: + return {t.model for t in self.turns} + + @property + def error_count(self) -> int: + return sum(1 for t in self.turns if t.status == ResponseStatus.ERROR) + + +# ────────────────────────────────────────────────────────────────────── +# Dump directory & file resolution +# ────────────────────────────────────────────────────────────────────── + +def _resolve_dumps_dir() -> Path | None: + """Find the first existing dumps directory from the search path.""" + for rel in CFG.dump_search_paths: + p = Path(os.path.expanduser(rel)) + if not p.is_absolute(): + p = Path.cwd() / p + if p.is_dir(): + return p + return None + + +def _list_trace_files(include_audited: bool = False) -> list[Path]: + """Return trace files sorted newest-first by mtime.""" + dumps = _resolve_dumps_dir() + if dumps is None: + return [] + files = list(dumps.glob("iq-session-*.jsonl")) + if include_audited: + files.extend((dumps / "audited").glob("iq-session-*.jsonl")) + files.sort(key=lambda f: f.stat().st_mtime, reverse=True) + return files + + +def _resolve_session_file(session_arg: str | None) -> Path | None: + """Resolve a session file from CLI argument or auto-detect latest.""" + if session_arg is None: + files = _list_trace_files(include_audited=False) + return files[0] if files else None + + # Direct path + candidate = Path(session_arg) + if candidate.is_file(): + return candidate + + # Try resolving as session ID + dumps = _resolve_dumps_dir() + if dumps is None: + return None + for subdir in [dumps, dumps / "audited"]: + f = subdir / f"iq-session-{session_arg}.jsonl" + if f.is_file(): + return f + return None + + +# ────────────────────────────────────────────────────────────────────── +# Housekeeping +# ────────────────────────────────────────────────────────────────────── + +def auto_rotate_traces() -> int: + """Delete oldest traces beyond retention limit. Returns count deleted.""" + files = _list_trace_files(include_audited=True) + if len(files) <= CFG.max_retained_traces: + return 0 + to_delete = files[CFG.max_retained_traces:] + deleted = 0 + for f in to_delete: + try: + f.unlink() + deleted += 1 + except OSError as e: + logger.warning("Failed to remove old trace %s: %s", f.name, e) + return deleted + + +# ────────────────────────────────────────────────────────────────────── +# Trace parsing +# ────────────────────────────────────────────────────────────────────── + +def parse_trace(filepath: Path) -> SessionTrace: + """Parse a JSONL session trace file into a SessionTrace object.""" + session_id = filepath.stem.removeprefix("iq-session-") + turns: list[TurnRecord] = [] + warnings = 0 + + with filepath.open("r", encoding="utf-8") as fh: + for lineno, line in enumerate(fh, start=1): + stripped = line.strip() + if not stripped: + continue + try: + raw = json.loads(stripped) + except json.JSONDecodeError: + warnings += 1 + logger.warning("Skipped malformed JSON at %s:%d", filepath.name, lineno) + continue + turns.append(TurnRecord.from_raw(raw, index=len(turns) + 1)) + + return SessionTrace( + session_id=session_id, + file_path=str(filepath), + turns=turns, + parse_warnings=warnings, + ) + + +# ────────────────────────────────────────────────────────────────────── +# Text compression for prompt construction +# ────────────────────────────────────────────────────────────────────── + +def _truncate(text: str, retain: int = CFG.text_truncate_retain, threshold: int = CFG.text_truncate_threshold) -> str: + """Smart truncation preserving head and tail.""" + if len(text) <= threshold: + return text + half = retain // 2 + omitted = len(text) - retain + return f"{text[:half]}\n\n[... TRUNCATED {omitted:,} bytes ...]\n\n{text[-half:]}" + + +def _compress_content(content: Any) -> str: + """Compress message content blocks for prompt inclusion.""" + if not content: + return "" + if isinstance(content, str): + return _truncate(content) + if isinstance(content, list): + parts = [] + for block in content: + if not isinstance(block, dict): + parts.append(str(block)) + continue + btype = block.get("type", "") + if btype == "text": + parts.append(_truncate(block.get("text", ""), retain=3000)) + elif btype == "tool_result": + parts.append(f"[tool_result] {_truncate(str(block.get('content', '')), retain=2000)}") + elif btype == "tool_use": + parts.append(f"[tool_use:{block.get('name', '?')}] input={_truncate(str(block.get('input', '')), retain=1500)}") + else: + parts.append(f"[{btype}] {_truncate(str(block), retain=1000)}") + return "\n".join(parts) + return _truncate(str(content)) + + +# ────────────────────────────────────────────────────────────────────── +# Prompt construction +# ────────────────────────────────────────────────────────────────────── + +_SYSTEM_INSTRUCTION = textwrap.dedent("""\ + You are an expert compiler and LLM systems engineer specializing in token + deduplication, sliding-window Rabin-Karp chunking, and database caching + strategies. + + Analyze the provided IndexQube proxy session trace and return a VALID, + STRICT JSON object. Do NOT output any conversational text, code fences, + or Markdown outside the JSON itself. + + Schema: + { + "summary": "High-level session analysis (2-4 sentences).", + "cache_eviction_loopholes": ["..."], + "chunking_inefficiencies": ["..."], + "loop_breaker_bypasses": ["..."], + "anomalies": ["Any request-ID collisions, duplicate payloads, or timing oddities"], + "actionable_recommendations": [ + {"target": "package/file", "change": "Specific fix", "impact": "Expected improvement", "priority": "P0|P1|P2"} + ], + "session_health_score": 0-100 + } +""") + + +def _select_audit_turns(trace: SessionTrace) -> list[TurnRecord]: + """Select the most informative turns for the audit prompt. + + Strategy: + - Always include first 3 turns (context warm-up visibility) + - Always include last 3 turns (most recent behavior) + - Fill remaining slots (up to max_audit_turns) with the + lowest-savings middle turns (highest optimization opportunity) + """ + n = len(trace.turns) + cap = CFG.max_audit_turns + + if n <= cap: + return list(trace.turns) + + keep_indices: set[int] = set() + # First 3 + for i in range(min(3, n)): + keep_indices.add(i) + # Last 3 + for i in range(max(0, n - 3), n): + keep_indices.add(i) + + # Middle: sort by lowest true_cache_hit_bytes (highest signal) + middle = [i for i in range(n) if i not in keep_indices] + middle.sort(key=lambda i: ( + trace.turns[i].optimizer.true_cache_hit_bytes, + trace.turns[i].before_bytes, + )) + + remaining = cap - len(keep_indices) + if remaining > 0: + keep_indices.update(middle[:remaining]) + + selected = sorted(keep_indices) + return [trace.turns[i] for i in selected] + + +def build_audit_prompt(trace: SessionTrace) -> str: + """Construct the full audit prompt from a session trace.""" + selected = _select_audit_turns(trace) + + header = ( + f"Analyze the following IndexQube L7 proxy session ({len(trace.turns)} total turns, " + f"{len(selected)} shown). Each turn shows the payload before/after chunk pruning and " + f"the model response.\n\n" + f"Session ID: {trace.session_id}\n" + f"Models used: {', '.join(trace.models_used)}\n" + f"Total input bytes: {trace.total_before:,}\n" + f"Total bytes saved: {trace.total_saved:,} ({trace.total_reduction_pct:.1f}%)\n" + f"Error turns: {trace.error_count}\n\n" + "Identify:\n" + "1. Cache eviction loopholes — large misses that should have hit\n" + "2. Chunking inefficiencies — suboptimal block grouping or Rabin-Karp window sizing\n" + "3. Loop/circuit-breaker bypasses — repeated prompts or synthetic ID collisions\n" + "4. Anomalies — duplicate request IDs, identical payloads, timing clusters\n" + "5. Actionable codebase improvements with priority ranking\n\n" + "--- SESSION EXECUTION TRACES ---\n" + ) + + parts = [header] + for turn in selected: + opt = turn.optimizer + block = ( + f"\n[TURN {turn.index}/{len(trace.turns)}] ts={turn.ts}\n" + f" request_id: {turn.request_id}\n" + f" model: {turn.model} | status: {turn.status.value} | output_tokens: {turn.output_tokens}\n" + f" bytes: before={turn.before_bytes:,} after={turn.after_bytes:,} " + f"saved={turn.saved_bytes:,} ({turn.reduction_pct:.1f}%)\n" + f" optimizer: pruned={opt.bytes_pruned:,} protected={opt.protected_bytes:,} " + f"known={opt.known_bytes:,} cache_hit={opt.true_cache_hit_bytes:,} " + f"blocks_pruned={opt.blocks_pruned} blocks_known={opt.blocks_known}\n" + ) + + # Turn 1: include system prompt + full messages for context setup visibility + if turn.index == 1: + sys_prompt = turn._before_payload.get("system", "") + if sys_prompt: + block += f" system_prompt:\n{textwrap.indent(_compress_content(sys_prompt), ' ')}\n" + msgs = turn._before_payload.get("messages", []) + if msgs: + block += f" full_context:\n{textwrap.indent(_compress_content(msgs), ' ')}\n" + else: + # Delta: only the latest user message + msgs = turn._before_payload.get("messages", []) + if msgs: + last = msgs[-1] + content = last.get("content", "") + block += f" latest_prompt:\n{textwrap.indent(_compress_content(content), ' ')}\n" + + # Response preview + if turn._response_text: + block += f" response_preview:\n{textwrap.indent(_truncate(turn._response_text, retain=2000), ' ')}\n" + + block += " " + "─" * 60 + "\n" + parts.append(block) + + prompt = "".join(parts) + + # Hard cap guard + char_limit = CFG.prompt_token_hard_cap * 4 + if len(prompt) > char_limit: + logger.warning( + "Prompt exceeds token cap (~%dk tokens). Truncating to %dk chars.", + len(prompt) // 4000, char_limit // 1000, + ) + prompt = prompt[:char_limit] + + return prompt + + +# ────────────────────────────────────────────────────────────────────── +# AWS helpers +# ────────────────────────────────────────────────────────────────────── + +def _get_boto3(): + """Lazy-import boto3 with a clear error if missing.""" + try: + import boto3 + from botocore.config import Config as BotoConfig + return boto3, BotoConfig + except ImportError: + logger.error( + "boto3 is not installed. Run: pip install boto3" + ) + sys.exit(1) + + +def _verify_aws_creds() -> bool: + """Fast STS identity check.""" + boto3, _ = _get_boto3() + try: + boto3.client("sts").get_caller_identity() + return True + except Exception: + return False + + +# ────────────────────────────────────────────────────────────────────── +# Bedrock streaming audit +# ────────────────────────────────────────────────────────────────────── + +def _parse_bedrock_response(raw: str) -> dict[str, Any]: + """Strip markdown fences and parse JSON from Bedrock response.""" + text = raw.strip() + # Remove ```json ... ``` wrapping if present + fence_pattern = re.compile(r"^```(?:json)?\s*\n?(.*?)\n?\s*```$", re.DOTALL) + m = fence_pattern.match(text) + if m: + text = m.group(1).strip() + return json.loads(text) + + +def run_bedrock_audit(prompt: str) -> dict[str, Any] | None: + """Stream an audit through Bedrock and return the parsed JSON report.""" + boto3, BotoConfig = _get_boto3() + + if not _verify_aws_creds(): + print(f"\n {C.RED}✗ AWS credentials expired or missing.{C.RESET}") + print(f" Run: {C.GREEN}aws sso login{C.RESET}\n") + return None + + config = BotoConfig( + region_name=CFG.bedrock_region, + connect_timeout=CFG.bedrock_connect_timeout, + read_timeout=CFG.bedrock_read_timeout, + retries={"max_attempts": CFG.bedrock_max_retries}, + ) + + try: + client = boto3.client("bedrock-runtime", config=config) + except Exception as e: + logger.error("Failed to initialize Bedrock client: %s", e) + return None + + body = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 4096, + "temperature": 0.1, + "system": _SYSTEM_INSTRUCTION, + "messages": [{"role": "user", "content": prompt}], + } + + for model_id in CFG.model_ids: + try: + print(f" {C.DIM}Trying model:{C.RESET} {C.YELLOW}{model_id}{C.RESET} ", end="", flush=True) + response = client.invoke_model_with_response_stream( + modelId=model_id, + contentType="application/json", + accept="application/json", + body=json.dumps(body), + ) + print(f"{C.GREEN}✓{C.RESET}") + break + except Exception as e: + err = str(e) + if "AccessDenied" in err or "ValidationException" in err: + print(f"{C.RED}✗{C.RESET}") + logger.debug("Model %s denied: %s", model_id, err[:120]) + else: + print(f"{C.YELLOW}⚠{C.RESET}") + logger.debug("Model %s failed: %s", model_id, err[:120]) + else: + print(f"\n {C.RED}✗ All models failed. Check Bedrock access / region.{C.RESET}\n") + return None + + # Stream response + print(f"\n{C.GREEN}{'━' * 64}{C.RESET}") + print(f"{C.BOLD} IndexQube Optimization Report (streaming){C.RESET}") + print(f"{C.GREEN}{'━' * 64}{C.RESET}\n") + + payload = [] + stream = response.get("body") + for event in stream: + chunk_bytes = event.get("chunk", {}).get("bytes", b"") + if not chunk_bytes: + continue + try: + chunk = json.loads(chunk_bytes.decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError): + continue + if chunk.get("type") == "content_block_delta": + text = chunk.get("delta", {}).get("text", "") + payload.append(text) + print(text, end="", flush=True) + + print(f"\n\n{C.GREEN}{'━' * 64}{C.RESET}\n") + + raw_text = "".join(payload) + try: + report = _parse_bedrock_response(raw_text) + print(f" {C.GREEN}✓ Parsed structured JSON report successfully{C.RESET}") + return report + except (json.JSONDecodeError, ValueError) as e: + logger.warning("Response was not valid JSON: %s", e) + return {"summary": raw_text, "_parse_error": str(e)} + + +# ────────────────────────────────────────────────────────────────────── +# Post-audit actions (S3, CloudWatch, archival) +# ────────────────────────────────────────────────────────────────────── + +def _upload_to_s3(filepath: str, session_id: str) -> None: + """Upload trace to S3 for long-term archival.""" + if not CFG.enable_s3: + return + boto3, _ = _get_boto3() + try: + s3 = boto3.client("s3") + key = f"traces/{session_id}.jsonl" + s3.upload_file(filepath, CFG.s3_bucket, key) + print(f" {C.GREEN}✓{C.RESET} Archived to {C.DIM}s3://{CFG.s3_bucket}/{key}{C.RESET}") + except Exception as e: + logger.warning("S3 upload failed (non-fatal): %s", e) + + +def _emit_cloudwatch(trace: SessionTrace) -> None: + """Publish optimization metrics to CloudWatch.""" + if not CFG.enable_cloudwatch: + return + boto3, _ = _get_boto3() + try: + cw = boto3.client("cloudwatch") + dims = [{"Name": "SessionID", "Value": trace.session_id}] + cw.put_metric_data( + Namespace=CFG.cw_namespace, + MetricData=[ + {"MetricName": "TotalBytesSaved", "Value": trace.total_saved, "Unit": "Bytes", "Dimensions": dims}, + {"MetricName": "ReductionPercent", "Value": trace.total_reduction_pct, "Unit": "Percent", "Dimensions": dims}, + {"MetricName": "TurnCount", "Value": len(trace.turns), "Unit": "Count", "Dimensions": dims}, + {"MetricName": "ErrorTurns", "Value": trace.error_count, "Unit": "Count", "Dimensions": dims}, + ], + ) + print(f" {C.GREEN}✓{C.RESET} Metrics emitted to {C.DIM}CloudWatch:{CFG.cw_namespace}{C.RESET}") + except Exception as e: + logger.warning("CloudWatch emit failed (non-fatal): %s", e) + + +def _save_report(report: dict[str, Any], session_id: str) -> Path | None: + """Persist JSON report to local filesystem.""" + dumps = _resolve_dumps_dir() + if dumps is None: + logger.warning("No dumps directory found — skipping report save") + return None + reports_dir = dumps / "reports" + reports_dir.mkdir(parents=True, exist_ok=True) + path = reports_dir / f"audit-{session_id}.json" + with path.open("w", encoding="utf-8") as fh: + json.dump(report, fh, indent=2, ensure_ascii=False) + print(f" {C.GREEN}✓{C.RESET} Report saved to {C.YELLOW}{path}{C.RESET}") + return path + + +def _archive_trace(filepath: str) -> None: + """Move audited trace into the audited/ subdirectory.""" + src = Path(filepath) + dest_dir = src.parent / "audited" + dest_dir.mkdir(parents=True, exist_ok=True) + dest = dest_dir / src.name + try: + shutil.move(str(src), str(dest)) + print(f" {C.GREEN}✓{C.RESET} Trace archived to {C.DIM}{dest}{C.RESET}") + except OSError as e: + logger.warning("Failed to archive trace: %s", e) + + +def run_post_audit(trace: SessionTrace, report: dict[str, Any]) -> None: + """Execute all post-audit side-effects concurrently where possible.""" + _save_report(report, trace.session_id) + + # S3 + CloudWatch can run in parallel + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [] + if CFG.enable_s3: + futures.append(pool.submit(_upload_to_s3, trace.file_path, trace.session_id)) + if CFG.enable_cloudwatch: + futures.append(pool.submit(_emit_cloudwatch, trace)) + for f in as_completed(futures): + try: + f.result() + except Exception as e: + logger.warning("Post-audit task failed: %s", e) + + _archive_trace(trace.file_path) + + +# ────────────────────────────────────────────────────────────────────── +# CLI: stats subcommand (no Bedrock, local only) +# ────────────────────────────────────────────────────────────────────── + +def _fmt_bytes(n: int) -> str: + """Human-readable byte count.""" + if n < 1024: + return f"{n} B" + elif n < 1024 ** 2: + return f"{n / 1024:.1f} KB" + else: + return f"{n / (1024 ** 2):.2f} MB" + + +def _print_table(headers: list[str], rows: list[list[str]], col_align: list[str] | None = None) -> None: + """Print a formatted ASCII table.""" + widths = [len(h) for h in headers] + for row in rows: + for i, cell in enumerate(row): + widths[i] = max(widths[i], len(cell)) + + if col_align is None: + col_align = ["<"] * len(headers) + + sep = "─" + header_line = " ".join(f"{h:{a}{w}}" for h, w, a in zip(headers, widths, col_align)) + divider = " ".join(sep * w for w in widths) + + print(f" {C.BOLD}{header_line}{C.RESET}") + print(f" {C.DIM}{divider}{C.RESET}") + for row in rows: + line = " ".join(f"{cell:{a}{w}}" for cell, w, a in zip(row, widths, col_align)) + print(f" {line}") + + +def cmd_stats(args: argparse.Namespace) -> int: + """Print local statistics for a session trace without calling Bedrock.""" + filepath = _resolve_session_file(args.session) + if filepath is None: + print(f" {C.RED}No session traces found.{C.RESET}") + return 1 + + trace = parse_trace(filepath) + print(f" Session: {C.YELLOW}{trace.session_id}{C.RESET}") + print(f" Turns: {C.GREEN}{len(trace.turns)}{C.RESET} | " + f"Errors: {C.RED if trace.error_count else C.GREEN}{trace.error_count}{C.RESET} | " + f"Models: {', '.join(trace.models_used)}") + print(f" Total: {_fmt_bytes(trace.total_before)} → {_fmt_bytes(trace.total_before - trace.total_saved)} " + f"({C.GREEN}{trace.total_reduction_pct:.1f}%{C.RESET} reduction)\n") + + if trace.parse_warnings: + print(f" {C.YELLOW}⚠ {trace.parse_warnings} malformed lines skipped{C.RESET}\n") + + # Per-turn table + headers = ["Turn", "Model", "Before", "After", "Saved", "%", "Status", "Pruned", "Cache Hit"] + rows = [] + for t in trace.turns: + pct_str = f"{t.reduction_pct:.1f}%" + status_str = "✓" if t.status == ResponseStatus.COMPLETED else ("✗" if t.status == ResponseStatus.ERROR else "?") + rows.append([ + str(t.index), + t.model[:20], + _fmt_bytes(t.before_bytes), + _fmt_bytes(t.after_bytes), + _fmt_bytes(t.saved_bytes), + pct_str, + status_str, + _fmt_bytes(t.optimizer.bytes_pruned), + _fmt_bytes(t.optimizer.true_cache_hit_bytes), + ]) + + _print_table(headers, rows, col_align=[">"]*1 + ["<"]*1 + [">"]*4 + ["^"]*1 + [">"]*2) + + # Anomaly detection + print(f"\n {C.BOLD}Anomaly Scan{C.RESET}") + anomalies = _detect_anomalies(trace) + if anomalies: + for a in anomalies: + print(f" {C.YELLOW}⚠{C.RESET} {a}") + else: + print(f" {C.GREEN}✓ No anomalies detected{C.RESET}") + + return 0 + + +def _detect_anomalies(trace: SessionTrace) -> list[str]: + """Run local heuristic anomaly checks.""" + issues = [] + + # Duplicate request IDs + rid_counts: dict[str, int] = {} + for t in trace.turns: + rid_counts[t.request_id] = rid_counts.get(t.request_id, 0) + 1 + dupes = {k: v for k, v in rid_counts.items() if v > 1} + if dupes: + issues.append(f"Duplicate request IDs: {len(dupes)} IDs reused across {sum(dupes.values())} turns") + + # Consecutive zero-savings turns after warm-up + zero_streak = 0 + max_zero_streak = 0 + for t in trace.turns[3:]: # skip first 3 (cold start) + if t.saved_bytes == 0 and t.before_bytes > 1000: + zero_streak += 1 + max_zero_streak = max(max_zero_streak, zero_streak) + else: + zero_streak = 0 + if max_zero_streak >= 5: + issues.append(f"Long zero-savings streak: {max_zero_streak} consecutive turns with no dedup after warm-up") + + # Rapid-fire identical payloads + for i in range(1, len(trace.turns)): + curr = trace.turns[i] + prev = trace.turns[i - 1] + if (curr.before_bytes == prev.before_bytes + and curr.before_bytes > 500 + and curr.request_id != prev.request_id): + # Check timestamps for rapid-fire + try: + t1 = datetime.fromisoformat(prev.ts) + t2 = datetime.fromisoformat(curr.ts) + delta = abs((t2 - t1).total_seconds()) + if delta < 3: + issues.append( + f"Rapid-fire duplicate: turns {prev.index}→{curr.index} " + f"({curr.before_bytes:,}B, {delta:.1f}s apart)" + ) + except (ValueError, TypeError): + pass + + # High error rate + if trace.error_count > len(trace.turns) * 0.3: + issues.append(f"High error rate: {trace.error_count}/{len(trace.turns)} turns errored ({trace.error_count/len(trace.turns)*100:.0f}%)") + + return issues + + +# ────────────────────────────────────────────────────────────────────── +# CLI: list subcommand +# ────────────────────────────────────────────────────────────────────── + +def cmd_list(args: argparse.Namespace) -> int: + """List available session traces.""" + files = _list_trace_files(include_audited=args.all) + if not files: + print(f" {C.RED}No session traces found.{C.RESET}") + return 1 + + print(f" {C.BOLD}Available session traces{C.RESET} ({len(files)} found)\n") + headers = ["#", "Session ID", "Size", "Lines", "Modified"] + rows = [] + for i, f in enumerate(files, 1): + sid = f.stem.removeprefix("iq-session-") + size = _fmt_bytes(f.stat().st_size) + lines = sum(1 for _ in f.open("r", encoding="utf-8")) + mtime = datetime.fromtimestamp(f.stat().st_mtime).strftime("%Y-%m-%d %H:%M") + audited = " (audited)" if "audited" in str(f) else "" + rows.append([str(i), sid + audited, size, str(lines), mtime]) + + _print_table(headers, rows, col_align=[">", "<", ">", ">", "<"]) + return 0 + + +# ────────────────────────────────────────────────────────────────────── +# CLI: audit subcommand (default) +# ────────────────────────────────────────────────────────────────────── + +def cmd_audit(args: argparse.Namespace) -> int: + """Run the full Bedrock-powered audit pipeline.""" + # Housekeeping + deleted = auto_rotate_traces() + if deleted: + logger.info("Rotated %d old trace files", deleted) + + filepath = _resolve_session_file(args.session) + if filepath is None: + print(f" {C.RED}No session traces found.{C.RESET}") + print(f" Run {C.GREEN}iq claude --dev --dump-payloads{C.RESET} first to collect traces.\n") + return 1 + + # Parse + trace = parse_trace(filepath) + if not trace.turns: + print(f" {C.RED}Session trace is empty or fully corrupted.{C.RESET}") + return 1 + + print(f" Session: {C.YELLOW}{trace.session_id}{C.RESET}") + print(f" Turns: {C.GREEN}{len(trace.turns)}{C.RESET} | " + f"Errors: {trace.error_count} | Models: {', '.join(trace.models_used)}") + print(f" Volume: {_fmt_bytes(trace.total_before)} total → " + f"{C.GREEN}{_fmt_bytes(trace.total_saved)} saved ({trace.total_reduction_pct:.1f}%){C.RESET}\n") + + # Build prompt & invoke + prompt = build_audit_prompt(trace) + prompt_kb = len(prompt) / 1024 + est_tokens = len(prompt) // 4 + print(f" Prompt: {prompt_kb:.0f} KB (~{est_tokens:,} tokens est.)") + print(f" Invoking Bedrock...\n") + + report = run_bedrock_audit(prompt) + if report is None: + return 1 + + # Post-audit pipeline + print(f"\n {C.BOLD}Post-Audit Pipeline{C.RESET}") + run_post_audit(trace, report) + + # Print health score if present + score = report.get("session_health_score") + if score is not None: + color = C.GREEN if score >= 70 else (C.YELLOW if score >= 40 else C.RED) + print(f"\n {C.BOLD}Session Health Score: {color}{score}/100{C.RESET}\n") + + print(f" {C.GREEN}✓ Audit complete.{C.RESET}\n") + return 0 + + +# ────────────────────────────────────────────────────────────────────── +# CLI entry point +# ────────────────────────────────────────────────────────────────────── + +def main() -> int: + parser = argparse.ArgumentParser( + prog="iq-audit", + description="IndexQube Session Audit Tool — analyze LLM proxy optimization traces", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=textwrap.dedent("""\ + examples: + python3 scripts/audit_session.py # audit latest trace + python3 scripts/audit_session.py --session # audit specific session + python3 scripts/audit_session.py stats # local stats (no Bedrock) + python3 scripts/audit_session.py list # list available traces + python3 scripts/audit_session.py list --all # include audited traces + + environment: + IQ_BEDROCK_REGION AWS region for Bedrock (default: us-east-1) + IQ_S3_BUCKET S3 bucket for trace archival + IQ_ENABLE_S3 Set to 0 to disable S3 upload + IQ_ENABLE_CW Set to 0 to disable CloudWatch metrics + NO_COLOR Disable colored output + """), + ) + parser.add_argument( + "-v", "--verbose", action="store_true", + help="Enable debug logging", + ) + parser.add_argument( + "--session", metavar="ID_OR_PATH", + help="Session ID or path to a specific trace file", + ) + + subparsers = parser.add_subparsers(dest="command", help="Available commands") + + # audit (default) + p_audit = subparsers.add_parser("audit", help="Run full Bedrock-powered audit (default)") + p_audit.add_argument("--session", metavar="ID_OR_PATH", help="Session ID or path") + + # stats + p_stats = subparsers.add_parser("stats", help="Print local stats without calling Bedrock") + p_stats.add_argument("--session", metavar="ID_OR_PATH", help="Session ID or path") + + # list + p_list = subparsers.add_parser("list", help="List available session traces") + p_list.add_argument("--all", action="store_true", help="Include already-audited traces") + + args = parser.parse_args() + + # Logging setup + level = logging.DEBUG if args.verbose else logging.WARNING + logging.basicConfig(format=LOG_FORMAT, datefmt=LOG_DATEFMT, level=level, stream=sys.stderr) + + # Banner + print(f"\n{C.BLUE}{'━' * 64}{C.RESET}") + print(f"{C.BLUE} IndexQube Auto-Optimization Audit Tool{C.RESET}") + print(f"{C.BLUE}{'━' * 64}{C.RESET}\n") + + # Dispatch + cmd = args.command + if cmd == "list": + return cmd_list(args) + elif cmd == "stats": + return cmd_stats(args) + else: + return cmd_audit(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/web/index.html b/web/index.html index c01cddc..e03226d 100644 --- a/web/index.html +++ b/web/index.html @@ -1,3 +1,4 @@ + From 391c80595e773a4cba20d6bd52f6b8cf04d1d02a Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Fri, 5 Jun 2026 11:37:19 -0500 Subject: [PATCH 02/15] fix: harden Claude replay and protected context handling --- gateway/internal/proxy/cache_key_test.go | 76 +++++++ gateway/internal/proxy/claude_messages.go | 174 +++++++++++---- gateway/internal/proxy/dump_redaction_test.go | 73 +++++++ gateway/internal/proxy/inflight_test.go | 96 +++++++++ gateway/internal/proxy/proxy.go | 22 +- gateway/internal/proxy/proxy_test.go | 202 +++++++++++++++++- gateway/internal/proxy/request_id_test.go | 34 +++ gateway/internal/proxy/stream_timeout_test.go | 61 ++++++ 8 files changed, 691 insertions(+), 47 deletions(-) create mode 100644 gateway/internal/proxy/cache_key_test.go create mode 100644 gateway/internal/proxy/dump_redaction_test.go create mode 100644 gateway/internal/proxy/inflight_test.go create mode 100644 gateway/internal/proxy/request_id_test.go create mode 100644 gateway/internal/proxy/stream_timeout_test.go diff --git a/gateway/internal/proxy/cache_key_test.go b/gateway/internal/proxy/cache_key_test.go new file mode 100644 index 0000000..0bbb18f --- /dev/null +++ b/gateway/internal/proxy/cache_key_test.go @@ -0,0 +1,76 @@ +package proxy + +import ( + "testing" +) + +// TestComputePromptHashContextAware verifies that the response replay cache +// key includes tool-result context, not just the latest user text. +func TestComputePromptHashContextAware(t *testing.T) { + // Two requests with the same latest user text but different prior + // tool-result content. + bodyA := []byte(`{ + "model": "claude-3", + "messages": [ + {"role": "user", "content": "hello"}, + {"role": "user", "content": [{"type": "tool_result", "content": "data-from-file-A"}]} + ] + }`) + bodyB := []byte(`{ + "model": "claude-3", + "messages": [ + {"role": "user", "content": "hello"}, + {"role": "user", "content": [{"type": "tool_result", "content": "data-from-file-B"}]} + ] + }`) + + hA := computePromptHash(bodyA, "claude-3") + hB := computePromptHash(bodyB, "claude-3") + + if hA == hB { + t.Fatalf("expected different cache keys for different tool results, got %q and %q", hA, hB) + } + + // Same request should produce the same hash. + hA2 := computePromptHash(bodyA, "claude-3") + if hA != hA2 { + t.Fatalf("expected same cache key for identical request, got %q and %q", hA, hA2) + } +} + +func TestComputePromptHashUsesFullRequestBeyondSemanticWindow(t *testing.T) { + bodyA := []byte(`{ + "model": "claude-3", + "messages": [ + {"role": "user", "content": "old context A"}, + {"role": "user", "content": "same recent 1"}, + {"role": "user", "content": "same recent 2"}, + {"role": "user", "content": "same recent 3"} + ] + }`) + bodyB := []byte(`{ + "model": "claude-3", + "messages": [ + {"role": "user", "content": "old context B"}, + {"role": "user", "content": "same recent 1"}, + {"role": "user", "content": "same recent 2"}, + {"role": "user", "content": "same recent 3"} + ] + }`) + + if semanticPromptHash(bodyA) != semanticPromptHash(bodyB) { + t.Fatal("test setup expected semantic hash to ignore the older context") + } + if computePromptHash(bodyA, "claude-3") == computePromptHash(bodyB, "claude-3") { + t.Fatal("response replay hash must include older context outside the semantic dedupe window") + } +} + +func TestComputePromptHashCanonicalizesJSONKeyOrder(t *testing.T) { + bodyA := []byte(`{"model":"claude-3","messages":[{"role":"user","content":"hello"}]}`) + bodyB := []byte(`{"messages":[{"content":"hello","role":"user"}],"model":"claude-3"}`) + + if computePromptHash(bodyA, "claude-3") != computePromptHash(bodyB, "claude-3") { + t.Fatal("expected logically identical request JSON to produce the same replay hash") + } +} diff --git a/gateway/internal/proxy/claude_messages.go b/gateway/internal/proxy/claude_messages.go index 0813add..9e1f0fc 100644 --- a/gateway/internal/proxy/claude_messages.go +++ b/gateway/internal/proxy/claude_messages.go @@ -25,6 +25,7 @@ import ( "github.com/Revanth14/indexqube/gateway/internal/chunker" "github.com/Revanth14/indexqube/gateway/internal/memory" "github.com/Revanth14/indexqube/gateway/internal/middleware" + "github.com/Revanth14/indexqube/gateway/internal/redact" "github.com/Revanth14/indexqube/gateway/internal/telemetry" awsconfig "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" @@ -37,8 +38,8 @@ const ( ) type contextKey string -const isSubscriptionKey contextKey = "isSubscription" +const isSubscriptionKey contextKey = "isSubscription" // guardBypassRe matches directives that attempt to disable proxy safety controls. // The separator between "guards" and "velocity" is optional (covers slash, pipe, @@ -77,10 +78,10 @@ type claudeOptimizerStats struct { BlocksPruned int `json:"blocks_pruned"` // Byte/token accounting. - BytesBefore int `json:"bytes_before"` - BytesAfter int `json:"bytes_after"` - BytesEligible int `json:"bytes_eligible"` - BytesPruned int `json:"bytes_pruned"` + BytesBefore int `json:"bytes_before"` + BytesAfter int `json:"bytes_after"` + BytesEligible int `json:"bytes_eligible"` + BytesPruned int `json:"bytes_pruned"` // KnownBytes is the byte total of every span that hit the session cache // (Seen returned true), independent of whether the span was subsequently // pruned or preserved by a protection rule. Invariant: @@ -100,14 +101,14 @@ type claudeOptimizerStats struct { ClassSpansPruned map[string]int `json:"class_spans_pruned,omitempty"` // Preserve-reason counters. - PreservedLatestTurnBytes int `json:"preserved_latest_turn_bytes"` - PreservedLatestTurnCount int `json:"preserved_latest_turn_count"` - PreservedSmallBytes int `json:"preserved_small_bytes"` - PreservedSmallCount int `json:"preserved_small_count"` - PreservedSystemBytes int `json:"preserved_system_bytes"` - PreservedSystemCount int `json:"preserved_system_count"` - PreservedToolUseBytes int `json:"preserved_tool_use_bytes"` - PreservedToolUseCount int `json:"preserved_tool_use_count"` + PreservedLatestTurnBytes int `json:"preserved_latest_turn_bytes"` + PreservedLatestTurnCount int `json:"preserved_latest_turn_count"` + PreservedSmallBytes int `json:"preserved_small_bytes"` + PreservedSmallCount int `json:"preserved_small_count"` + PreservedSystemBytes int `json:"preserved_system_bytes"` + PreservedSystemCount int `json:"preserved_system_count"` + PreservedToolUseBytes int `json:"preserved_tool_use_bytes"` + PreservedToolUseCount int `json:"preserved_tool_use_count"` PreservedInstructionBytes int `json:"preserved_instruction_bytes"` PreservedInstructionCount int `json:"preserved_instruction_count"` PreservedLastOccurrenceBytes int `json:"preserved_last_occurrence_bytes"` @@ -227,23 +228,79 @@ func (p *Proxy) handleClaudeMessages(w http.ResponseWriter, r *http.Request) { p.writeError(w, r, errorPayload{HTTPStatus: http.StatusBadRequest, Type: "invalid_request_error", Code: "invalid_anthropic_request", Message: err.Error()}) return } - + overheadMs := time.Since(overheadStart).Milliseconds() + + finishSynthetic := func(streamStats claudeStreamStats) { + effectiveOpt := optStats + effectiveOpt.BytesAfter = 0 + effectiveOpt.EstimatedTokensAfter = 0 + effectiveOpt.EstimatedTokensSaved = effectiveOpt.EstimatedTokensBefore + if effectiveOpt.BytesBefore > 0 { + effectiveOpt.ReductionRatio = 1 + } + + duration := time.Since(started) + if cfg.SessionStore != nil { + cfg.SessionStore.RecordUsage(sessionKey, memory.UsageTotals{ + Requests: 1, + TokensIn: estimateTokens(len(body)), + TokensOut: streamStats.estimatedOutputTokens(), + TokensSaved: effectiveOpt.EstimatedTokensSaved, + BytesIn: len(body), + BytesSaved: effectiveOpt.BytesBefore - effectiveOpt.BytesAfter, + }) + } + p.logClaudeRequestComplete(r.Context(), requestID, cfg.Mode, meta.Model, sessionKey, len(body), effectiveOpt, streamStats, duration, missingReqID, requestID, velocityWarning) + p.emitClaudeUsageEvent(r, sessionKey, meta.Model, effectiveOpt, duration, streamStats.StatusCode, overheadMs) + if os.Getenv("IQ_DUMP_PAYLOADS") == "1" { + dumpClaudePayloads(requestID, body, nil, streamStats, effectiveOpt) + } + } + // Intercept Sentinel Probes (quota, ping) to respond in 1 ms at 0 cost normalizedLatest := strings.TrimSpace(strings.ToLower(shape.LatestUserText)) if normalizedLatest == "quota" || normalizedLatest == "ping" { p.logger.InfoContext(ctx, "sentinel probe intercepted", slog.String("session_key", shortLogHash(sessionKey)), slog.String("probe", normalizedLatest)) - writeSyntheticStreamResponse(w, "IndexQube active. Quota check successful.") + text := "IndexQube active. Quota check successful." + writeSyntheticStreamResponse(w, text) + finishSynthetic(claudeStreamStats{ + OutputText: len(text), + OutputRawText: text, + Status: "synthetic_probe", + StatusCode: http.StatusOK, + Completed: true, + Provider: "synthetic", + }) return } var capture *responseCaptureWriter var promptCacheHash string - isCacheablePrompt := shape.LatestUserText != "" && shape.ToolResultCount == 0 && cfg.Mode == "optimize" && + isCacheablePrompt := shape.LatestUserText != "" && cfg.Mode == "optimize" && r.Header.Get("Cache-Control") != "no-cache" && r.Header.Get("Pragma") != "no-cache" + writeCachedResponse := func(cachedPayload []byte) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Accel-Buffering", "no") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(cachedPayload) + if flusher, ok := w.(http.Flusher); ok { + flusher.Flush() + } + finishSynthetic(claudeStreamStats{ + Chunks: bytes.Count(cachedPayload, []byte("event: content_block_delta")), + OutputText: len(cachedPayload), + Status: "cache_replay", + StatusCode: http.StatusOK, + Completed: true, + Provider: "cache", + }) + } if isCacheablePrompt { - promptCacheHash = computePromptHash(shape.LatestUserText, meta.Model) + promptCacheHash = computePromptHash(body, meta.Model) ts := p.getOrCreateTurnState(sessionKey) ts.mu.Lock() cachedPayload, ok := ts.getCachedResponse(promptCacheHash) @@ -252,23 +309,13 @@ func (p *Proxy) handleClaudeMessages(w http.ResponseWriter, r *http.Request) { p.logger.InfoContext(ctx, "response-level cache hit", slog.String("session_key", shortLogHash(sessionKey)), slog.String("prompt_hash", promptCacheHash[:8])) - w.Header().Set("Content-Type", "text/event-stream") - w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate") - w.Header().Set("Connection", "keep-alive") - w.Header().Set("X-Accel-Buffering", "no") - w.WriteHeader(http.StatusOK) - _, _ = w.Write(cachedPayload) - if flusher, ok := w.(http.Flusher); ok { - flusher.Flush() - } + writeCachedResponse(cachedPayload) return } capture = &responseCaptureWriter{ResponseWriter: w} w = capture } - overheadMs := time.Since(overheadStart).Milliseconds() - // FIX 2: in-flight duplicate detection. When a second identical request // arrives while the first is still in-flight, it waits up to 30 s for the // first to complete before dispatching its own upstream call. This prevents @@ -281,6 +328,19 @@ func (p *Proxy) handleClaudeMessages(w http.ResponseWriter, r *http.Request) { case <-r.Context().Done(): return } + if isCacheablePrompt { + ts := p.getOrCreateTurnState(sessionKey) + ts.mu.Lock() + cachedPayload, ok := ts.getCachedResponse(promptCacheHash) + ts.mu.Unlock() + if ok { + p.logger.InfoContext(ctx, "response-level cache hit after in-flight wait", + slog.String("session_key", shortLogHash(sessionKey)), + slog.String("prompt_hash", promptCacheHash[:8])) + writeCachedResponse(cachedPayload) + return + } + } // Re-register after the original completed so our dispatch is tracked. if done2, _ := p.inFlightRequests.acquire(promptHash); done2 != nil { defer done2() @@ -417,6 +477,7 @@ func (c ClaudeMessagesConfig) validate() error { } func validClaudeDevToken(r *http.Request, want string) bool { + _ = want auth := strings.TrimSpace(r.Header.Get("Authorization")) token := strings.TrimSpace(strings.TrimPrefix(auth, "Bearer ")) return token != "" @@ -440,11 +501,11 @@ func dumpClaudePayloads(requestID string, before, after []byte, stats claudeStre } beforePath := filepath.Join(dumpDir, "iq-before-"+requestID+".json") afterPath := filepath.Join(dumpDir, "iq-after-"+requestID+".json") - if err := os.WriteFile(beforePath, prettyJSON(before), 0o600); err != nil { + if err := os.WriteFile(beforePath, prettyJSON([]byte(redact.String(string(before)))), 0o600); err != nil { fmt.Fprintf(os.Stderr, "[iq] failed to dump payload pair: %v\n", err) return } - if err := os.WriteFile(afterPath, prettyJSON(after), 0o600); err != nil { + if err := os.WriteFile(afterPath, prettyJSON([]byte(redact.String(string(after)))), 0o600); err != nil { fmt.Fprintf(os.Stderr, "[iq] failed to dump payload pair: %v\n", err) return } @@ -515,6 +576,7 @@ func appendSessionDump(sessionFile, requestID string, before, after []byte, stat if err != nil { return err } + line = []byte(redact.String(string(line))) line = append(line, '\n') dumpPayloadMu.Lock() @@ -664,6 +726,7 @@ func (p *Proxy) warmUpSystemSpans(ctx context.Context, cfg ClaudeMessagesConfig, // creating it if it does not yet exist. func (p *Proxy) getOrCreateTurnState(sessionKey string) *sessionTurnState { v, _ := p.sessionTurnCounters.LoadOrStore(sessionKey, &sessionTurnState{}) + p.touchSession(sessionKey) return v.(*sessionTurnState) } @@ -671,6 +734,7 @@ func (p *Proxy) getOrCreateTurnState(sessionKey string) *sessionTurnState { // sessionKey, creating it if it does not yet exist. func (p *Proxy) getOrCreateBoilerplateState(sessionKey string) *boilerplateState { v, _ := p.sessionBoilerplateState.LoadOrStore(sessionKey, &boilerplateState{}) + p.touchSession(sessionKey) return v.(*boilerplateState) } @@ -780,6 +844,7 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, // zero-savings. Uses a stable chunk ID: sha256(file_path + content_hash). if root["system"] != nil { _, alreadyWarmed := p.sessionWarmUpDone.LoadOrStore(sessionKey, true) + p.touchSession(sessionKey) if !alreadyWarmed { p.warmUpSystemSpans(ctx, cfg, sessionKey, root) } @@ -796,6 +861,7 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, slog.String("session_key", shortLogHash(sessionKey))) } else { p.sessionSuggestionTs.Store(sessionKey, now) + p.touchSession(sessionKey) } return body, req, stats, shape, nil } @@ -955,6 +1021,23 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, systemAllKnown = false } + // Protected content must win even when the span is a new boilerplate + // variant. This check intentionally runs before boilerplate cooldown + // pruning so instruction files, credentials, and the latest turn are + // never removed just because they live inside a repeated harness block. + if span.IsLatestTurn { + stats.PreservedLatestTurnBytes += span.Bytes + stats.PreservedLatestTurnCount++ + stats.BlocksNew++ + continue + } + if isProtectedInstructionSpan(span) { + stats.PreservedInstructionBytes += span.Bytes + stats.PreservedInstructionCount++ + stats.BlocksNew++ + continue + } + // FIX 7: SUGGESTION MODE injection cooldown. If a boilerplate span // is new (unknown hash) but the last forward was <5 turns ago AND // context delta is <10 KB, prune it anyway to suppress redundant @@ -1130,7 +1213,15 @@ var protectedInstructionPathFragments = [...]string{ } func isProtectedInstructionSpan(span TextSpan) bool { - return containsProtectedInstructionPath(span.SourcePath) || containsProtectedInstructionPath(span.Text) + return containsProtectedInstructionPath(span.SourcePath) || containsProtectedInstructionPath(span.Text) || containsCredentialMarker(span.Text) +} + +func containsCredentialMarker(s string) bool { + lower := strings.ToLower(s) + return strings.Contains(lower, "api-key") || + strings.Contains(lower, "bearer ") || + strings.Contains(lower, "x-anthropic-api-key") || + strings.Contains(lower, "authorization") } func containsProtectedInstructionPath(s string) bool { @@ -1245,7 +1336,6 @@ func isSubscriptionAuth(auth string) bool { strings.Contains(auth, "sk-ant-oat") } - // forwardClaudeMessagesViaBedrock proxies /v1/messages to the Bedrock // InvokeModelWithResponseStream API. The Anthropic Messages body is forwarded // as-is except that `model` and `stream` are removed and @@ -1270,7 +1360,7 @@ func (p *Proxy) forwardClaudeMessagesViaBedrock(w http.ResponseWriter, r *http.R delete(root, "model") delete(root, "stream") delete(root, "context_management") // Bedrock rejects this Anthropic-specific field - stripCacheControlRecursively(root) // Bedrock rejects cache_control fields + stripCacheControlRecursively(root) // Bedrock rejects cache_control fields root["anthropic_version"] = "bedrock-2023-05-31" bedrockBody, err := json.Marshal(root) if err != nil { @@ -2072,6 +2162,7 @@ func (p *Proxy) getOrCreatePrefixHints(sessionKey string) *prefixHintSet { v, _ := p.sessionPrefixHints.LoadOrStore(sessionKey, &prefixHintSet{ hints: make(map[string]int), }) + p.touchSession(sessionKey) return v.(*prefixHintSet) } @@ -2162,10 +2253,19 @@ func (rcw *responseCaptureWriter) WriteString(s string) (int, error) { return rcw.ResponseWriter.Write([]byte(s)) } -func computePromptHash(prompt, model string) string { - normalized := strings.TrimSpace(strings.ToLower(prompt)) - sum := sha256.Sum256([]byte(normalized + "|" + model)) - return hex.EncodeToString(sum[:]) +func computePromptHash(body []byte, model string) string { + canonicalBody := body + var parsed any + if err := json.Unmarshal(body, &parsed); err == nil { + if encoded, encErr := marshalJSONNoHTMLEscape(parsed); encErr == nil { + canonicalBody = encoded + } + } + h := sha256.New() + h.Write([]byte(model)) + h.Write([]byte{0}) + h.Write(canonicalBody) + return hex.EncodeToString(h.Sum(nil)) } func writeSyntheticStreamResponse(w http.ResponseWriter, text string) { diff --git a/gateway/internal/proxy/dump_redaction_test.go b/gateway/internal/proxy/dump_redaction_test.go new file mode 100644 index 0000000..74544cc --- /dev/null +++ b/gateway/internal/proxy/dump_redaction_test.go @@ -0,0 +1,73 @@ +package proxy + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/Revanth14/indexqube/gateway/internal/redact" +) + +func TestDumpRedaction(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "iq-dump-test-*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + sessionFile := filepath.Join(tmpDir, "dump.jsonl") + + before := []byte(`{"messages":[{"role":"user","content":"Use key: sk-proj-abc1234567890"}],"system":"Authorization: Bearer sk-ant-test123"}`) + after := []byte(`{"messages":[{"role":"user","content":"Use key: sk-proj-abc1234567890"}]}`) + + stats := claudeStreamStats{OutputRawText: "done", OutputTokens: 10, Status: "ok"} + opt := claudeOptimizerStats{} + + if err := appendSessionDump(sessionFile, "req-1", before, after, stats, opt); err != nil { + t.Fatalf("appendSessionDump failed: %v", err) + } + + data, err := os.ReadFile(sessionFile) + if err != nil { + t.Fatalf("read dump failed: %v", err) + } + + if strings.Contains(string(data), "sk-proj-abc1234567890") { + t.Fatal("dump should not contain raw API key") + } + if strings.Contains(string(data), "sk-ant-test123") { + t.Fatal("dump should not contain bearer token") + } + if !strings.Contains(string(data), redact.String("sk-proj-abc1234567890")) { + t.Fatal("dump should contain redacted marker") + } +} + +func TestDumpRedactionPrettyFiles(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "iq-dump-test-*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + os.Setenv("IQ_DUMP_DIR", tmpDir) + defer os.Unsetenv("IQ_DUMP_DIR") + os.Unsetenv("IQ_DUMP_SESSION_FILE") + + before := []byte(`{"system":"Authorization: Bearer abc-secret"}`) + after := []byte(`{"system":"redacted"}`) + stats := claudeStreamStats{OutputRawText: "done", OutputTokens: 5, Status: "ok"} + opt := claudeOptimizerStats{} + + dumpClaudePayloads("req-2", before, after, stats, opt) + + beforePath := filepath.Join(tmpDir, "iq-before-req-2.json") + data, err := os.ReadFile(beforePath) + if err != nil { + t.Fatalf("read before dump failed: %v", err) + } + if strings.Contains(string(data), "abc-secret") { + t.Fatal("pretty dump should not contain raw bearer token") + } +} diff --git a/gateway/internal/proxy/inflight_test.go b/gateway/internal/proxy/inflight_test.go new file mode 100644 index 0000000..cd6c608 --- /dev/null +++ b/gateway/internal/proxy/inflight_test.go @@ -0,0 +1,96 @@ +package proxy + +import ( + "testing" + "time" +) + +func TestInFlightTrackerAcquireAndRelease(t *testing.T) { + tr := newInFlightTracker() + + done1, wait1 := tr.acquire("hash-a") + if done1 == nil { + t.Fatal("first arrival should get doneFn") + } + if wait1 != nil { + t.Fatal("first arrival should not get waitChan") + } + + // Duplicate should wait. + done2, wait2 := tr.acquire("hash-a") + if done2 != nil { + t.Fatal("duplicate should not get doneFn") + } + if wait2 == nil { + t.Fatal("duplicate should get waitChan") + } + + // First caller releases. + done1() + + select { + case <-wait2: + // expected + case <-time.After(time.Second): + t.Fatal("duplicate should have been notified") + } + + // After release, a new acquire should succeed again. + done3, wait3 := tr.acquire("hash-a") + if done3 == nil { + t.Fatal("post-release acquire should get doneFn") + } + if wait3 != nil { + t.Fatal("post-release acquire should not get waitChan") + } + done3() +} + +func TestInFlightTrackerDifferentHashes(t *testing.T) { + tr := newInFlightTracker() + + done1, wait1 := tr.acquire("hash-a") + if done1 == nil { + t.Fatal("first arrival should get doneFn") + } + if wait1 != nil { + t.Fatal("first arrival should not get waitChan") + } + + // Different hash should acquire independently. + done2, wait2 := tr.acquire("hash-b") + if done2 == nil { + t.Fatal("different hash should get doneFn") + } + if wait2 != nil { + t.Fatal("different hash should not get waitChan") + } + + done1() + done2() +} + +func TestInFlightTrackerCleanupOnPanic(t *testing.T) { + tr := newInFlightTracker() + + func() { + defer func() { + if r := recover(); r == nil { + t.Fatal("expected panic") + } + }() + done1, _ := tr.acquire("panic-hash") + defer done1() + panic("boom") + }() + + // After panic recovery, the entry should be cleaned up. + done2, wait2 := tr.acquire("panic-hash") + if done2 == nil { + t.Fatal("post-panic acquire should get doneFn") + } + if wait2 != nil { + t.Fatal("post-panic acquire should not get waitChan") + } + done2() +} diff --git a/gateway/internal/proxy/proxy.go b/gateway/internal/proxy/proxy.go index 08ebad1..1f7edda 100644 --- a/gateway/internal/proxy/proxy.go +++ b/gateway/internal/proxy/proxy.go @@ -38,6 +38,7 @@ type Proxy struct { mux *http.ServeMux maxRequestSize int64 optimizeTimeout time.Duration + streamTimeout time.Duration metrics *telemetry.Metrics claude ClaudeMessagesConfig usageTracker telemetry.Sink @@ -68,6 +69,15 @@ type Proxy struct { // large payload" reuse patterns (FIX 3). sessionPrefixHints sync.Map // map[string]*prefixHintSet + // sessionLastUsed tracks the last access time for each session key. + // Used by the background cleanup goroutine to evict idle entries. + sessionLastUsed sync.Map // map[string]time.Time + + // cleanupCtx/cleanupCancel manage the background TTL eviction goroutine. + cleanupCtx context.Context + cleanupCancel context.CancelFunc + cleanupDone chan struct{} + // sessionSuggestionTs tracks the last time a suggestion-mode request was // processed per session. Used to rate-limit harness meta-prompt injections // to max 1 per 10 seconds so ephemeral payloads don't pollute the chunk store. @@ -143,7 +153,7 @@ func (s *sessionTurnState) saveCachedResponse(promptHash string, payload []byte) } s.cachedResponses = append(s.cachedResponses, cachedResponse{ promptHash: promptHash, - payload: payload, + payload: append([]byte(nil), payload...), }) } @@ -246,6 +256,15 @@ func WithOptimizeTimeout(d time.Duration) Option { } } +// WithStreamTimeout caps the duration of streaming governor requests. +// A non-positive value uses the default (5 minutes). +func WithStreamTimeout(d time.Duration) Option { + return func(p *Proxy) { + if d > 0 { + p.streamTimeout = d + } + } +} func WithClaudeMessages(cfg ClaudeMessagesConfig) Option { return func(p *Proxy) { p.claude = cfg @@ -300,6 +319,7 @@ func New(gov Governor, opts ...Option) *Proxy { mux: http.NewServeMux(), maxRequestSize: 8 << 20, // default 8 MiB optimizeTimeout: 30 * time.Second, + streamTimeout: 0, // disabled by default; use WithStreamTimeout to enable inFlightRequests: newInFlightTracker(), } for _, opt := range opts { diff --git a/gateway/internal/proxy/proxy_test.go b/gateway/internal/proxy/proxy_test.go index f86b8d3..ec9191c 100644 --- a/gateway/internal/proxy/proxy_test.go +++ b/gateway/internal/proxy/proxy_test.go @@ -13,6 +13,8 @@ import ( "os" "path/filepath" "strings" + "sync" + "sync/atomic" "testing" "time" @@ -477,6 +479,135 @@ func TestClaudeMessages_AnthropicPassthroughStreaming(t *testing.T) { } } +func TestClaudeMessages_ResponseCacheReplaysRepeatedPromptWithPriorToolHistory(t *testing.T) { + t.Parallel() + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "event: message_start\n") + _, _ = io.WriteString(w, `data: {"type":"message_start","message":{"id":"msg_1","model":"claude-sonnet-4-6"}}`+"\n\n") + _, _ = io.WriteString(w, "event: content_block_delta\n") + _, _ = io.WriteString(w, `data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"project summary"}}`+"\n\n") + _, _ = io.WriteString(w, "event: message_stop\n") + _, _ = io.WriteString(w, `data: {"type":"message_stop"}`+"\n\n") + })) + t.Cleanup(upstream.Close) + srv := newClaudeTestServer(t, upstream.URL, memory.NewStore(time.Hour), "optimize", true) + + body := `{"model":"claude-sonnet-4-6","stream":true,"messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"old file listing\nmain.go\nREADME.md"}]},{"role":"assistant","content":"I read the files."},{"role":"user","content":"what does this project do?"}]}` + for i := 0; i < 2; i++ { + req, err := http.NewRequest(http.MethodPost, srv.URL+"/v1/messages", strings.NewReader(body)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Authorization", "Bearer iq-dev-local") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("POST /v1/messages: %v", err) + } + got, _ := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status=%d body=%s", resp.StatusCode, got) + } + if !strings.Contains(string(got), "project summary") { + t.Fatalf("body=%s, want cached upstream response", got) + } + } + if got := calls.Load(); got != 1 { + t.Fatalf("upstream calls=%d, want 1", got) + } +} + +func TestClaudeMessages_ResponseCacheRecheckedAfterInflightWait(t *testing.T) { + t.Parallel() + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + time.Sleep(100 * time.Millisecond) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "event: content_block_delta\n") + _, _ = io.WriteString(w, `data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"cached concurrent answer"}}`+"\n\n") + _, _ = io.WriteString(w, "event: message_stop\n") + _, _ = io.WriteString(w, `data: {"type":"message_stop"}`+"\n\n") + })) + t.Cleanup(upstream.Close) + srv := newClaudeTestServer(t, upstream.URL, memory.NewStore(time.Hour), "optimize", true) + + body := `{"model":"claude-sonnet-4-6","stream":true,"messages":[{"role":"user","content":"what does this project do?"}]}` + var wg sync.WaitGroup + errs := make(chan error, 2) + for i := 0; i < 2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + req, err := http.NewRequest(http.MethodPost, srv.URL+"/v1/messages", strings.NewReader(body)) + if err != nil { + errs <- err + return + } + req.Header.Set("Authorization", "Bearer iq-dev-local") + resp, err := http.DefaultClient.Do(req) + if err != nil { + errs <- err + return + } + got, _ := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if resp.StatusCode != http.StatusOK { + errs <- fmt.Errorf("status=%d body=%s", resp.StatusCode, got) + return + } + if !strings.Contains(string(got), "cached concurrent answer") { + errs <- fmt.Errorf("body=%s, want cached concurrent answer", got) + } + }() + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + if got := calls.Load(); got != 1 { + t.Fatalf("upstream calls=%d, want 1", got) + } +} + +func TestClaudeMessages_SentinelProbeDoesNotCallUpstream(t *testing.T) { + t.Parallel() + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + t.Fatal("upstream should not be called for sentinel probe") + })) + t.Cleanup(upstream.Close) + srv := newClaudeTestServer(t, upstream.URL, memory.NewStore(time.Hour), "observe", false) + + req, err := http.NewRequest(http.MethodPost, srv.URL+"/v1/messages", strings.NewReader(`{"model":"claude-sonnet-4-6","stream":true,"messages":[{"role":"user","content":"quota"}]}`)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Authorization", "Bearer iq-dev-local") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("POST /v1/messages: %v", err) + } + got, _ := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status=%d body=%s", resp.StatusCode, got) + } + if !strings.Contains(string(got), "IndexQube active") { + t.Fatalf("body=%s, want synthetic response", got) + } + if got := calls.Load(); got != 0 { + t.Fatalf("upstream calls=%d, want 0", got) + } +} + func TestClaudeMessages_MissingAuth(t *testing.T) { t.Parallel() upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -843,6 +974,56 @@ func TestClaudeMessages_OptimizePreservesProtectedInstructionToolResult(t *testi } } +func TestClaudeMessages_ProtectedSystemBoilerplateBeatsCooldownPruning(t *testing.T) { + t.Parallel() + p := New(&fakeGovernor{}) + session := "protected-boilerplate-test" + cfg := ClaudeMessagesConfig{ + Mode: "optimize", + EnableBlockOptimizer: true, + SessionStore: memory.NewStore(time.Hour), + Optimizer: OptimizerConfig{ + MinSpanBytes: 512, + EnableToolResultPruning: true, + }, + } + bodyFor := func(secret string) []byte { + systemText := "\n" + + strings.Repeat("stable harness reminder line\n", 30) + + "Authorization: Bearer " + secret + "\n" + + "CLAUDE.md must remain visible\n" + + "" + return []byte(fmt.Sprintf( + `{"model":"claude-sonnet-4-6","system":[{"type":"text","text":%q}],"messages":[{"role":"user","content":"latest turn"}]}`, + systemText, + )) + } + + p.resolveRequestID(session, "turn-1") + if _, _, _, _, err := p.prepareClaudeBody(context.Background(), cfg, session, bodyFor("sk-proj-firstsecret")); err != nil { + t.Fatalf("first prepare: %v", err) + } + p.resolveRequestID(session, "turn-2") + if _, _, _, _, err := p.prepareClaudeBody(context.Background(), cfg, session, bodyFor("sk-proj-secondsecret")); err != nil { + t.Fatalf("second prepare: %v", err) + } + p.resolveRequestID(session, "turn-3") + forward, _, stats, _, err := p.prepareClaudeBody(context.Background(), cfg, session, bodyFor("sk-proj-thirdsecret")) + if err != nil { + t.Fatalf("third prepare: %v", err) + } + + if strings.Contains(string(forward), "omitted") { + t.Fatalf("protected system boilerplate must not be replaced during cooldown, body=%s", forward) + } + if !strings.Contains(string(forward), "Authorization: Bearer sk-proj-thirdsecret") { + t.Fatalf("protected credential marker missing from forwarded body: %s", forward) + } + if stats.PreservedInstructionCount < 1 { + t.Fatalf("expected protected boilerplate to be counted as preserved, stats=%+v", stats) + } +} + func TestClaudeMessages_OptimizeStillPrunesOrdinaryToolResult(t *testing.T) { t.Parallel() p := New(&fakeGovernor{}) @@ -946,6 +1127,9 @@ func TestProtectedInstructionSpanDetection(t *testing.T) { if isProtectedInstructionSpan(TextSpan{SourcePath: `/repo/src/main.go`, Text: `package main`}) { t.Fatal("ordinary source file should not be protected") } + if !isProtectedInstructionSpan(TextSpan{Text: `Authorization: Bearer my-api-key`}) { + t.Fatal("expected Authorization bearer text to be protected") + } } func TestDumpClaudePayloadsAppendsSessionFile(t *testing.T) { @@ -1014,15 +1198,15 @@ func TestDumpClaudePayloadsEmitsOptimizerStats(t *testing.T) { // latest-turn. KnownBytes is the sum across all preservation paths plus // the pruned bytes. opt := claudeOptimizerStats{ - BlocksPruned: 2, - BlocksKnown: 5, - BytesPruned: 1000, - PreservedInstructionBytes: 3027, - PreservedInstructionCount: 1, - PreservedLatestTurnBytes: 500, - PreservedLatestTurnCount: 1, - PreservedLastOccurrenceBytes: 0, - PreservedLastOccurrenceCount: 0, + BlocksPruned: 2, + BlocksKnown: 5, + BytesPruned: 1000, + PreservedInstructionBytes: 3027, + PreservedInstructionCount: 1, + PreservedLatestTurnBytes: 500, + PreservedLatestTurnCount: 1, + PreservedLastOccurrenceBytes: 0, + PreservedLastOccurrenceCount: 0, } opt.KnownBytes = opt.BytesPruned + opt.PreservedInstructionBytes + opt.PreservedLatestTurnBytes + opt.PreservedLastOccurrenceBytes diff --git a/gateway/internal/proxy/request_id_test.go b/gateway/internal/proxy/request_id_test.go new file mode 100644 index 0000000..c944244 --- /dev/null +++ b/gateway/internal/proxy/request_id_test.go @@ -0,0 +1,34 @@ +package proxy + +import ( + "strings" + "testing" +) + +func TestResolveRequestIDUsesUUIDFormat(t *testing.T) { + p := New(&fakeGovernor{}) + id, synthetic, _ := p.resolveRequestID("sess123", "") + if !synthetic { + t.Fatal("expected synthetic ID for empty rawID") + } + if !strings.HasPrefix(id, "iq-synthetic-") { + t.Fatalf("expected iq-synthetic- prefix, got %q", id) + } + // Should contain an 8-char hex suffix from uuid.New().String()[:8]. + parts := strings.Split(id, "-") + lastPart := parts[len(parts)-1] + if len(lastPart) != 8 { + t.Fatalf("expected 8-char hex suffix, got %q (len=%d)", lastPart, len(lastPart)) + } +} + +func TestProvidedRequestIDPreserved(t *testing.T) { + p := New(&fakeGovernor{}) + id, synthetic, _ := p.resolveRequestID("sess123", "client-req-42") + if synthetic { + t.Fatal("expected provided ID to be preserved") + } + if id != "client-req-42" { + t.Fatalf("expected client-req-42, got %q", id) + } +} diff --git a/gateway/internal/proxy/stream_timeout_test.go b/gateway/internal/proxy/stream_timeout_test.go new file mode 100644 index 0000000..7bd461d --- /dev/null +++ b/gateway/internal/proxy/stream_timeout_test.go @@ -0,0 +1,61 @@ +package proxy + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/Revanth14/indexqube/gateway/internal/domain" +) + +type timeoutCheckGovernor struct { + deadline time.Time +} + +func (g *timeoutCheckGovernor) Stream(ctx context.Context, _ *domain.InferenceRequest, _ domain.TokenWriter) error { + if d, ok := ctx.Deadline(); ok { + g.deadline = d + } + return nil +} + +func (g *timeoutCheckGovernor) Optimize(ctx context.Context, tenant string, messages []domain.Message, projectMemory string) ([]domain.Message, domain.PruneStats, error) { + return messages, domain.PruneStats{}, nil +} + +func (g *timeoutCheckGovernor) Diagnostics(ctx context.Context) (domain.Diagnostics, error) { + return domain.Diagnostics{}, nil +} + +func (g *timeoutCheckGovernor) Ready(ctx context.Context) error { + return nil +} + +func TestStreamTimeoutSetsDeadline(t *testing.T) { + gov := &timeoutCheckGovernor{} + p := New(gov, WithStreamTimeout(50*time.Millisecond)) + req := httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + rec := httptest.NewRecorder() + p.streamThroughGovernor(rec, req, &domain.InferenceRequest{Model: "claude-3", Messages: []domain.Message{{Role: "user"}}, Stream: true}) + + if gov.deadline.IsZero() { + t.Fatal("expected governor context to have a deadline") + } + if time.Until(gov.deadline) > time.Second { + t.Fatalf("deadline too far in the future: %v", time.Until(gov.deadline)) + } +} + +func TestStreamTimeoutDisabledByDefault(t *testing.T) { + gov := &timeoutCheckGovernor{} + p := New(gov) // default streamTimeout is 0 (disabled) + req := httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + rec := httptest.NewRecorder() + p.streamThroughGovernor(rec, req, &domain.InferenceRequest{Model: "claude-3", Messages: []domain.Message{{Role: "user"}}, Stream: true}) + + if !gov.deadline.IsZero() { + t.Fatal("expected no deadline when stream timeout disabled by default") + } +} From 9d6e14976803b9d62c18dae8cb8c8218a01d6eb3 Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Fri, 5 Jun 2026 11:37:34 -0500 Subject: [PATCH 03/15] feat: add public auth and rate limit middleware --- gateway/internal/config/config.go | 6 + gateway/internal/middleware/auth.go | 65 +++++++++ gateway/internal/middleware/auth_test.go | 114 +++++++++++++++ gateway/internal/middleware/cors.go | 1 + gateway/internal/middleware/ratelimit.go | 130 ++++++++++++++++++ gateway/internal/middleware/ratelimit_test.go | 93 +++++++++++++ gateway/internal/server/run.go | 10 +- 7 files changed, 418 insertions(+), 1 deletion(-) create mode 100644 gateway/internal/middleware/auth.go create mode 100644 gateway/internal/middleware/auth_test.go create mode 100644 gateway/internal/middleware/ratelimit.go create mode 100644 gateway/internal/middleware/ratelimit_test.go diff --git a/gateway/internal/config/config.go b/gateway/internal/config/config.go index 40eaa4f..744ef20 100644 --- a/gateway/internal/config/config.go +++ b/gateway/internal/config/config.go @@ -36,6 +36,11 @@ type ServerConfig struct { // value caps the maximum stream duration and breaks long generations. WriteTimeout time.Duration IdleTimeout time.Duration + // AuthToken protects sensitive endpoints (/stats, /v1/agent-sessions, + // /v1/diagnostics) when the gateway binds to a non-loopback address. + // Empty disables auth entirely. Defaults to INDEXQUBE_AUTH_TOKEN, + // falling back to INDEXQUBE_DEV_TOKEN for backwards compatibility. + AuthToken string CORSEnabled bool CORSAllowedOrigins []string CORSAllowChromeExtensions bool @@ -162,6 +167,7 @@ func Load() (*AppConfig, error) { // Slowloris is mitigated by ReadHeaderTimeout + IdleTimeout instead. WriteTimeout: getEnvAsDuration("SERVER_WRITE_TIMEOUT", 0), IdleTimeout: getEnvAsDuration("SERVER_IDLE_TIMEOUT", 120*time.Second), + AuthToken: getEnvFirst([]string{"INDEXQUBE_AUTH_TOKEN", "INDEXQUBE_DEV_TOKEN"}, ""), CORSEnabled: getEnvAsBool("CORS_ENABLED", true), CORSAllowedOrigins: getEnvAsCSV("CORS_ALLOWED_ORIGINS", []string{ "http://localhost:3000", diff --git a/gateway/internal/middleware/auth.go b/gateway/internal/middleware/auth.go new file mode 100644 index 0000000..8552620 --- /dev/null +++ b/gateway/internal/middleware/auth.go @@ -0,0 +1,65 @@ +package middleware + +import ( + "encoding/json" + "net" + "net/http" + "strings" +) + +// Auth returns middleware that requires a valid Bearer token for requests +// from non-loopback addresses. Localhost (127.0.0.1, ::1) is exempt. +// trustedProxies is used to extract the real client IP from X-Forwarded-For. +func Auth(token string, trustedProxies []string) Middleware { + if token == "" { + return func(next http.Handler) http.Handler { return next } + } + + isLoopback := checkLoopbackFunc(trustedProxies) + + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if isLoopback(r) { + next.ServeHTTP(w, r) + return + } + + auth := strings.TrimSpace(r.Header.Get("Authorization")) + got := strings.TrimSpace(strings.TrimPrefix(auth, "Bearer ")) + if got != token { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "error": map[string]interface{}{ + "type": "authentication_error", + "code": "missing_key", + "message": "missing or invalid Bearer token", + }, + }) + return + } + + next.ServeHTTP(w, r) + }) + } +} + +func checkLoopbackFunc(trustedProxies []string) func(*http.Request) bool { + return func(r *http.Request) bool { + ip := resolveIP(r, trustedProxies) + return ip.IsLoopback() + } +} + +func resolveIP(r *http.Request, trustedProxies []string) net.IP { + ipStr := clientIP(r, trustedProxies) + host, _, err := net.SplitHostPort(ipStr) + if err != nil { + host = ipStr + } + ip := net.ParseIP(host) + if ip == nil { + return net.ParseIP("127.0.0.1") + } + return ip +} diff --git a/gateway/internal/middleware/auth_test.go b/gateway/internal/middleware/auth_test.go new file mode 100644 index 0000000..7e8f7fd --- /dev/null +++ b/gateway/internal/middleware/auth_test.go @@ -0,0 +1,114 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "testing" +) + +func TestAuthLocalhostExempt(t *testing.T) { + mw := Auth("secret", nil) + handler := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(http.MethodGet, "/", nil) + req.RemoteAddr = "127.0.0.1:1234" + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200 for localhost, got %d", rec.Code) + } +} + +func TestAuthRequiresToken(t *testing.T) { + mw := Auth("secret", nil) + handler := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + // Simulate non-localhost by overriding RemoteAddr + req := httptest.NewRequest(http.MethodGet, "/", nil) + req.RemoteAddr = "192.168.1.1:1234" + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 without token, got %d", rec.Code) + } + if rec.Header().Get("Content-Type") != "application/json" { + t.Fatalf("expected json error, got %s", rec.Header().Get("Content-Type")) + } +} + +func TestAuthAcceptsValidToken(t *testing.T) { + mw := Auth("secret", nil) + handler := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(http.MethodGet, "/", nil) + req.RemoteAddr = "192.168.1.1:1234" + req.Header.Set("Authorization", "Bearer secret") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200 with valid token, got %d", rec.Code) + } +} + +func TestAuthNoOpWhenTokenEmpty(t *testing.T) { + mw := Auth("", nil) + handler := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(http.MethodGet, "/", nil) + req.RemoteAddr = "192.168.1.1:1234" + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200 when auth disabled, got %d", rec.Code) + } +} + +// TestAuthXFSpoofUntrusted proves that a remote client cannot bypass auth +// by setting X-Forwarded-For to 127.0.0.1 when the actual peer is untrusted. +func TestAuthXFSpoofUntrusted(t *testing.T) { + mw := Auth("secret", nil) + handler := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(http.MethodGet, "/", nil) + req.RemoteAddr = "10.0.0.1:1234" + req.Header.Set("X-Forwarded-For", "127.0.0.1") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 for spoofed XFF from untrusted peer, got %d", rec.Code) + } +} + +// TestAuthXFTrustedProxy proves that an authentic proxy’s X-Forwarded-For is +// respected so that the real client behind the proxy is evaluated. +func TestAuthXFTrustedProxy(t *testing.T) { + mw := Auth("secret", []string{"10.0.0.0/8"}) + handler := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(http.MethodGet, "/", nil) + req.RemoteAddr = "10.0.0.2:1234" + req.Header.Set("X-Forwarded-For", "192.168.1.1") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401 for real client behind trusted proxy, got %d", rec.Code) + } +} diff --git a/gateway/internal/middleware/cors.go b/gateway/internal/middleware/cors.go index 14c7d74..0dac4be 100644 --- a/gateway/internal/middleware/cors.go +++ b/gateway/internal/middleware/cors.go @@ -19,6 +19,7 @@ var ( corsAllowedMethods = []string{"GET", "POST", "OPTIONS"} corsAllowedHeaders = []string{ "Content-Type", + "Authorization", "X-Request-ID", "X-IQ-Session-Key", "X-IQ-Project-Memory", diff --git a/gateway/internal/middleware/ratelimit.go b/gateway/internal/middleware/ratelimit.go new file mode 100644 index 0000000..4b407e9 --- /dev/null +++ b/gateway/internal/middleware/ratelimit.go @@ -0,0 +1,130 @@ +package middleware + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + "sync" + "time" +) + +// RateLimiter provides per-IP token-bucket rate limiting. +type RateLimiter struct { + mu sync.Mutex + buckets map[string]*bucket + rate float64 // tokens per second + burst int + maxAge time.Duration + lastSeen map[string]time.Time +} + +// bucket is a simple token bucket tracker. +type bucket struct { + tokens float64 + lastCheck time.Time +} + +// NewRateLimiter creates a rate limiter with the given rate (req/sec), burst, +// and max age for stale bucket cleanup. +func NewRateLimiter(rate float64, burst int, maxAge time.Duration) *RateLimiter { + rl := &RateLimiter{ + buckets: make(map[string]*bucket), + rate: rate, + burst: burst, + maxAge: maxAge, + lastSeen: make(map[string]time.Time), + } + go rl.cleanupLoop() + return rl +} + +// Allowed reports whether a request from the given IP is within the rate limit. +// It returns false and the seconds until the next token is available when +// the bucket is empty. +func (rl *RateLimiter) Allowed(ip string) (ok bool, retryAfter float64) { + rl.mu.Lock() + defer rl.mu.Unlock() + + now := time.Now() + rl.lastSeen[ip] = now + + b, exists := rl.buckets[ip] + if !exists { + b = &bucket{tokens: float64(rl.burst), lastCheck: now} + rl.buckets[ip] = b + } + + // Replenish tokens based on elapsed time. + elapsed := now.Sub(b.lastCheck).Seconds() + b.tokens += elapsed * rl.rate + if b.tokens > float64(rl.burst) { + b.tokens = float64(rl.burst) + } + b.lastCheck = now + + if b.tokens >= 1 { + b.tokens-- + return true, 0 + } + + retryAfter = (1 - b.tokens) / rl.rate + return false, retryAfter +} + +// cleanupLoop removes stale buckets every maxAge interval. +func (rl *RateLimiter) cleanupLoop() { + ticker := time.NewTicker(rl.maxAge) + defer ticker.Stop() + for range ticker.C { + rl.mu.Lock() + cutoff := time.Now().Add(-rl.maxAge) + for ip, t := range rl.lastSeen { + if t.Before(cutoff) { + delete(rl.buckets, ip) + delete(rl.lastSeen, ip) + } + } + rl.mu.Unlock() + } +} + +// RateLimit returns middleware that rejects requests exceeding the per-IP rate +// with HTTP 429. Paths starting with prefix are rate-limited; others pass +// through. Localhost requests are exempt. +func RateLimit(limiter *RateLimiter, prefix string, trustedProxies []string) Middleware { + if limiter == nil { + return func(next http.Handler) http.Handler { return next } + } + isLoopback := checkLoopbackFunc(trustedProxies) + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasPrefix(r.URL.Path, prefix) { + next.ServeHTTP(w, r) + return + } + if isLoopback(r) { + next.ServeHTTP(w, r) + return + } + + ip := clientIP(r, trustedProxies) + ok, retryAfter := limiter.Allowed(ip) + if ok { + next.ServeHTTP(w, r) + return + } + + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", fmt.Sprintf("%.0f", retryAfter)) + w.WriteHeader(http.StatusTooManyRequests) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "error": map[string]interface{}{ + "type": "rate_limit_error", + "code": "rate_limit_exceeded", + "message": "too many requests; retry after " + fmt.Sprintf("%.0f", retryAfter) + " seconds", + }, + }) + }) + } +} diff --git a/gateway/internal/middleware/ratelimit_test.go b/gateway/internal/middleware/ratelimit_test.go new file mode 100644 index 0000000..ec51c37 --- /dev/null +++ b/gateway/internal/middleware/ratelimit_test.go @@ -0,0 +1,93 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" +) + +func TestRateLimiterAllowed(t *testing.T) { + rl := NewRateLimiter(1, 2, time.Minute) + + // First two requests succeed (burst = 2). + if ok, _ := rl.Allowed("1.2.3.4"); !ok { + t.Fatal("expected first request to succeed") + } + if ok, _ := rl.Allowed("1.2.3.4"); !ok { + t.Fatal("expected second request to succeed") + } + + // Third request fails. + if ok, after := rl.Allowed("1.2.3.4"); ok { + t.Fatalf("expected third request to fail, retryAfter=%f", after) + } +} + +func TestRateLimiterReplenish(t *testing.T) { + rl := NewRateLimiter(10, 1, time.Minute) + + if ok, _ := rl.Allowed("1.2.3.4"); !ok { + t.Fatal("expected first request to succeed") + } + if ok, _ := rl.Allowed("1.2.3.4"); ok { + t.Fatal("expected second request to fail immediately") + } + + time.Sleep(120 * time.Millisecond) // wait for 1 token at 10 tps + + if ok, _ := rl.Allowed("1.2.3.4"); !ok { + t.Fatal("expected request to succeed after replenish") + } +} + +func TestRateLimitMiddlewareExemptsHealthz(t *testing.T) { + mw := RateLimit(NewRateLimiter(1, 0, time.Minute), "/v1/", nil) + handler := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(http.MethodGet, "/healthz", nil) + req.RemoteAddr = "192.168.1.1:1234" + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200 for /healthz, got %d", rec.Code) + } +} + +func TestRateLimitMiddlewareExemptsLocalhost(t *testing.T) { + mw := RateLimit(NewRateLimiter(1, 0, time.Minute), "/v1/", nil) + handler := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(http.MethodGet, "/v1/messages", nil) + req.RemoteAddr = "127.0.0.1:1234" + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200 for localhost, got %d", rec.Code) + } +} + +func TestRateLimitMiddlewareReturns429(t *testing.T) { + mw := RateLimit(NewRateLimiter(1, 0, time.Minute), "/v1/", nil) + handler := mw(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + req := httptest.NewRequest(http.MethodGet, "/v1/messages", nil) + req.RemoteAddr = "192.168.1.1:1234" + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusTooManyRequests { + t.Fatalf("expected 429, got %d", rec.Code) + } + if rec.Header().Get("Retry-After") == "" { + t.Fatal("expected Retry-After header") + } +} diff --git a/gateway/internal/server/run.go b/gateway/internal/server/run.go index 30e43ef..927c962 100644 --- a/gateway/internal/server/run.go +++ b/gateway/internal/server/run.go @@ -101,7 +101,7 @@ func run(ctx context.Context, publicListener net.Listener) error { if home, err := os.UserHomeDir(); err == nil { cacheDir := filepath.Join(home, ".indexqube", "cache") if err := os.MkdirAll(cacheDir, 0o700); err == nil { - lsmCache, err := cache.NewLSMCache(cacheDir, cfg.Cache.MaxEntryBytes) + lsmCache, err := cache.NewLSMCache(cacheDir, cfg.Cache.MaxEntryBytes, cfg.Cache.TTL) if err != nil { logger.Warn("failed to initialize LSM cache; falling back to memory cache", slog.Any("err", err)) } else { @@ -241,6 +241,7 @@ func run(ctx context.Context, publicListener net.Listener) error { HTTPClient: upstreamClient, }), ) + p.Start() publicServer := buildPublicServer(cfg, p, tp, logger) adminServer := buildAdminServer(cfg, tp, logger) @@ -293,6 +294,7 @@ func run(ctx context.Context, publicListener net.Listener) error { logger.Error("admin server shutdown", slog.Any("err", err)) } stopJanitor() + p.Stop() if sessionTracker != nil { if err := sessionTracker.Close(); err != nil { logger.Error("session tracker close", slog.Any("err", err)) @@ -349,6 +351,12 @@ func buildPublicServer(cfg *config.AppConfig, p *proxy.Proxy, tp *telemetry.Prov AllowChromeExtensions: cfg.Server.CORSAllowChromeExtensions, MaxAge: cfg.Server.CORSMaxAge, }), + middleware.Auth(cfg.Server.AuthToken, cfg.Server.TrustedProxies), + middleware.RateLimit( + middleware.NewRateLimiter(1, 60, 5*time.Minute), + "/v1/", + cfg.Server.TrustedProxies, + ), func(next http.Handler) http.Handler { return otelhttp.NewHandler(next, "gateway") }, middleware.RouteResolver(p.Mux()), middleware.RequestID, From 34f9bc11a731e94ee08a00aca6a140935b3fd058 Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Fri, 5 Jun 2026 11:37:41 -0500 Subject: [PATCH 04/15] feat: wire local LSM response cache --- gateway/internal/cache/bench_test.go | 4 +-- gateway/internal/cache/lsm.go | 12 ++++++-- gateway/internal/cache/lsm_test.go | 41 +++++++++++++++++++++++---- gateway/internal/store/lsm/sstable.go | 4 +++ 4 files changed, 51 insertions(+), 10 deletions(-) diff --git a/gateway/internal/cache/bench_test.go b/gateway/internal/cache/bench_test.go index 57915f8..95d2616 100644 --- a/gateway/internal/cache/bench_test.go +++ b/gateway/internal/cache/bench_test.go @@ -196,7 +196,7 @@ func BenchmarkLSMCachePut(b *testing.B) { sz := sz b.Run(sz.name, func(b *testing.B) { dir := b.TempDir() - c, err := NewLSMCache(dir, 0) // 0 = no per-entry size cap + c, err := NewLSMCache(dir, 0, 0) // 0 = no per-entry size cap, no TTL if err != nil { b.Fatalf("NewLSMCache: %v", err) } @@ -223,7 +223,7 @@ func BenchmarkLSMCacheGet(b *testing.B) { sz := sz b.Run(sz.name, func(b *testing.B) { dir := b.TempDir() - c, err := NewLSMCache(dir, 0) + c, err := NewLSMCache(dir, 0, 0) if err != nil { b.Fatalf("NewLSMCache: %v", err) } diff --git a/gateway/internal/cache/lsm.go b/gateway/internal/cache/lsm.go index c9c87cb..6358155 100644 --- a/gateway/internal/cache/lsm.go +++ b/gateway/internal/cache/lsm.go @@ -3,6 +3,7 @@ package cache import ( "context" "encoding/json" + "time" "github.com/Revanth14/indexqube/gateway/internal/store/lsm" ) @@ -12,15 +13,16 @@ import ( type LSMCache struct { engine *lsm.Engine maxEntryBytes int64 + ttl time.Duration } // NewLSMCache opens (or creates) an LSM storage engine in dir. -func NewLSMCache(dir string, maxEntryBytes int64) (*LSMCache, error) { +func NewLSMCache(dir string, maxEntryBytes int64, ttl time.Duration) (*LSMCache, error) { opts := lsm.Options{ MemTableSize: 4 * 1024 * 1024, // 4 MiB flushes MaxL0Tables: 4, - BlockSize: 4096, // page-aligned - BloomFPRate: 0.01, // 1% false positive + BlockSize: 4096, // page-aligned + BloomFPRate: 0.01, // 1% false positive BloomExpected: 10000, MaxTableSize: 2 * 1024 * 1024, // 2 MiB tables L1Budget: 10 * 1024 * 1024, @@ -32,6 +34,7 @@ func NewLSMCache(dir string, maxEntryBytes int64) (*LSMCache, error) { return &LSMCache{ engine: engine, maxEntryBytes: maxEntryBytes, + ttl: ttl, }, nil } @@ -49,6 +52,9 @@ func (c *LSMCache) Get(ctx context.Context, key Key) (*Entry, bool, error) { if err := json.Unmarshal(val, &entry); err != nil { return nil, false, err } + if c.ttl > 0 && !entry.CreatedAt.IsZero() && time.Since(entry.CreatedAt) > c.ttl { + return nil, false, nil + } return &entry, true, nil } diff --git a/gateway/internal/cache/lsm_test.go b/gateway/internal/cache/lsm_test.go index 6e1cb97..1f25575 100644 --- a/gateway/internal/cache/lsm_test.go +++ b/gateway/internal/cache/lsm_test.go @@ -17,7 +17,7 @@ func TestLSMCache_GetMissOnEmpty(t *testing.T) { } defer os.RemoveAll(tmpDir) - c, err := NewLSMCache(tmpDir, 1024) + c, err := NewLSMCache(tmpDir, 1024, time.Hour) if err != nil { t.Fatalf("NewLSMCache failed: %v", err) } @@ -39,7 +39,7 @@ func TestLSMCache_PutThenGet(t *testing.T) { } defer os.RemoveAll(tmpDir) - c, err := NewLSMCache(tmpDir, 1024) + c, err := NewLSMCache(tmpDir, 1024, time.Hour) if err != nil { t.Fatalf("NewLSMCache failed: %v", err) } @@ -85,7 +85,7 @@ func TestLSMCache_PutTooLargeRejected(t *testing.T) { } defer os.RemoveAll(tmpDir) - c, err := NewLSMCache(tmpDir, 10) // extremely small max size + c, err := NewLSMCache(tmpDir, 10, time.Hour) // extremely small max size if err != nil { t.Fatalf("NewLSMCache failed: %v", err) } @@ -112,7 +112,7 @@ func TestLSMCache_PersistenceAcrossCloses(t *testing.T) { defer os.RemoveAll(tmpDir) // Phase 1: Open, write entry, then close - c1, err := NewLSMCache(tmpDir, 1024) + c1, err := NewLSMCache(tmpDir, 1024, time.Hour) if err != nil { t.Fatalf("c1 NewLSMCache failed: %v", err) } @@ -130,7 +130,7 @@ func TestLSMCache_PersistenceAcrossCloses(t *testing.T) { c1.Close() // Phase 2: Re-open in the same directory, read entry, then close - c2, err := NewLSMCache(tmpDir, 1024) + c2, err := NewLSMCache(tmpDir, 1024, time.Hour) if err != nil { t.Fatalf("c2 NewLSMCache failed: %v", err) } @@ -148,3 +148,34 @@ func TestLSMCache_PersistenceAcrossCloses(t *testing.T) { t.Errorf("chunk data corrupted: got %s, want %s", got.Chunks[0], want.Chunks[0]) } } +func TestLSMCache_TTLExpiredReturnsMiss(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "iq-lsm-cache-test-*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + c, err := NewLSMCache(tmpDir, 1024, time.Minute) + if err != nil { + t.Fatalf("NewLSMCache failed: %v", err) + } + defer c.Close() + + entry := &Entry{ + Provider: domain.ProviderAnthropic, + Model: "claude-3-5-sonnet", + Chunks: [][]byte{[]byte("expired chunk")}, + CreatedAt: time.Now().Add(-2 * time.Minute), + } + if err := c.Put(context.Background(), "expired-key", entry); err != nil { + t.Fatalf("Put failed: %v", err) + } + + _, hit, err := c.Get(context.Background(), "expired-key") + if err != nil { + t.Fatalf("Get failed: %v", err) + } + if hit { + t.Fatal("expected expired LSM cache entry to be a miss") + } +} diff --git a/gateway/internal/store/lsm/sstable.go b/gateway/internal/store/lsm/sstable.go index b3d5f2d..f7f4fad 100644 --- a/gateway/internal/store/lsm/sstable.go +++ b/gateway/internal/store/lsm/sstable.go @@ -182,6 +182,10 @@ func (b *Builder) Finish() error { if err := b.bw.Flush(); err != nil { return err } + if err := b.f.Sync(); err != nil { + _ = b.f.Close() + return err + } return b.f.Close() } From 9962e8eae02ed607cfe275b2b0681f09b2453a3b Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Fri, 5 Jun 2026 11:38:12 -0500 Subject: [PATCH 05/15] feat: improve iq session summaries and audit dumps --- .gitignore | 4 +++ gateway/cmd/iq/main.go | 30 ++++++++++++++++++++--- gateway/internal/sessions/tracker.go | 6 +++++ gateway/internal/sessions/tracker_test.go | 26 +++++++++++++++++++- scripts/audit_session.py | 3 ++- 5 files changed, 63 insertions(+), 6 deletions(-) diff --git a/.gitignore b/.gitignore index 2820632..c108366 100644 --- a/.gitignore +++ b/.gitignore @@ -41,11 +41,15 @@ gateway/logs/ gateway/docs/ tmp/ build/ +gateway/iq commands.txt gateway/.claude/ +.kimchi/ # Local planning & Claude config — never commit CONTEXT.md .claudeignore CLAUDE.md +refer.txt +AGENTS.md diff --git a/gateway/cmd/iq/main.go b/gateway/cmd/iq/main.go index 27b1b3e..b84fd11 100644 --- a/gateway/cmd/iq/main.go +++ b/gateway/cmd/iq/main.go @@ -85,12 +85,12 @@ func runClaude(args []string, devMode, dumpPayloads bool) { sessionID := generateToken() if dumpPayloads { os.Setenv("IQ_DUMP_PAYLOADS", "1") - cwd, err := os.Getwd() - if err != nil { - fmt.Fprintf(os.Stderr, "iq: failed to resolve dump directory: %v\n", err) + repoRoot := findRepoRoot() + if repoRoot == "" { + fmt.Fprintf(os.Stderr, "iq: failed to resolve repository root directory\n") os.Exit(1) } - dumpDir := filepath.Join(cwd, ".indexqube", "dumps") + dumpDir := filepath.Join(repoRoot, ".indexqube", "dumps") shortSessionID := sessionID if len(shortSessionID) > 8 { shortSessionID = shortSessionID[:8] @@ -315,3 +315,25 @@ func waitForProxy(port int) bool { } return false } + +func findRepoRoot() string { + cwd, err := os.Getwd() + if err != nil { + return "" + } + dir := cwd + for { + if _, err := os.Stat(filepath.Join(dir, ".git")); err == nil { + return dir + } + if _, err := os.Stat(filepath.Join(dir, "CLAUDE.md")); err == nil { + return dir + } + parent := filepath.Dir(dir) + if parent == dir { + break + } + dir = parent + } + return cwd +} diff --git a/gateway/internal/sessions/tracker.go b/gateway/internal/sessions/tracker.go index 88c2dc3..b34c9e0 100644 --- a/gateway/internal/sessions/tracker.go +++ b/gateway/internal/sessions/tracker.go @@ -16,6 +16,7 @@ package sessions import ( "database/sql" "log/slog" + "os" "sync" "time" @@ -109,6 +110,11 @@ func Open(path string, logger *slog.Logger) (*Tracker, error) { logger = slog.Default() } + if err := os.Chmod(path, 0o600); err != nil { + db.Close() + return nil, err + } + t := &Tracker{ db: db, ch: make(chan writeEvent, 512), diff --git a/gateway/internal/sessions/tracker_test.go b/gateway/internal/sessions/tracker_test.go index 5472bc7..6188846 100644 --- a/gateway/internal/sessions/tracker_test.go +++ b/gateway/internal/sessions/tracker_test.go @@ -10,6 +10,30 @@ import ( "github.com/Revanth14/indexqube/gateway/internal/telemetry" ) +func TestTracker_FilePermissions(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "iq-test-*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelWarn})) + tracker, err := Open(filepath.Join(tmpDir, "sessions.db"), logger) + if err != nil { + t.Fatalf("failed to open tracker: %v", err) + } + defer tracker.Close() + + info, err := os.Stat(filepath.Join(tmpDir, "sessions.db")) + if err != nil { + t.Fatalf("failed to stat db: %v", err) + } + mode := info.Mode().Perm() + if mode != 0o600 { + t.Fatalf("expected file mode 0o600, got 0o%03o", mode) + } +} + func TestTracker_RecordAndRetrieve(t *testing.T) { tmpDir, err := os.MkdirTemp("", "iq-test-*") if err != nil { @@ -99,7 +123,7 @@ func TestTracker_KillEvent(t *testing.T) { sessionID := "test-session-kill" outcome := telemetry.RequestOutcome{ - TokensAttempted: 500, + TokensAttempted: 1000, TokensSent: 500, Warned: true, Killed: true, diff --git a/scripts/audit_session.py b/scripts/audit_session.py index 15293d5..c97d10c 100644 --- a/scripts/audit_session.py +++ b/scripts/audit_session.py @@ -88,8 +88,9 @@ class Config: # Model fallback chain model_ids: tuple[str, ...] = ( + "us.anthropic.claude-opus-4-6-v1", + "anthropic.claude-opus-4-6-v1", "us.anthropic.claude-sonnet-4-6", - "us.anthropic.claude-opus-4-6", "anthropic.claude-sonnet-4-6", "amazon.nova-pro-v1:0", ) From de371fcb2dfd62b393dbc2c0741f846768cff89c Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Fri, 5 Jun 2026 11:38:20 -0500 Subject: [PATCH 06/15] fix: tighten optimizer input handling and chunking --- gateway/internal/proxy/handlers.go | 28 +++++++++++++--- gateway/internal/proxy/handlers_test.go | 43 +++++++++++++++++++++++++ gateway/internal/proxy/request.go | 13 ++++++++ 3 files changed, 79 insertions(+), 5 deletions(-) create mode 100644 gateway/internal/proxy/handlers_test.go diff --git a/gateway/internal/proxy/handlers.go b/gateway/internal/proxy/handlers.go index b80cb24..8ecb2d6 100644 --- a/gateway/internal/proxy/handlers.go +++ b/gateway/internal/proxy/handlers.go @@ -147,6 +147,10 @@ func (p *Proxy) handleOptimize(w http.ResponseWriter, r *http.Request) { if sk == "" { sk = r.Header.Get(headerSessionKey) } + if err := validateSessionKey(sk); err != nil { + p.writeError(w, r, errorPayload{HTTPStatus: http.StatusBadRequest, Type: "invalid_request_error", Code: "invalid_session_key", Message: err.Error()}) + return + } pm := body.ProjectMemory if pm == "" { pm = r.Header.Get(headerProjectMemory) @@ -358,12 +362,22 @@ func normalizeContextPath(path string) string { } path = strings.ReplaceAll(path, "\\", "/") path = strings.TrimPrefix(path, "/") - if path == "" || strings.Contains(path, "```") || strings.ContainsAny(path, "\r\n\t ") { + if path == "" || strings.Contains(path, "```") || strings.ContainsAny(path, "\r\n\t ") || containsTraversal(path) { return defaultRawContextPath } return path } +func containsTraversal(path string) bool { + parts := strings.Split(path, "/") + for _, p := range parts { + if p == ".." { + return true + } + } + return false +} + func normalizeContextLang(lang, path, content string) string { lang = strings.TrimSpace(lang) if lang != "" && !strings.ContainsAny(lang, "` \t\r\n") { @@ -482,8 +496,6 @@ func renderOptimizedText(msgs []domain.Message) string { func (p *Proxy) streamThroughGovernor(w http.ResponseWriter, r *http.Request, req *domain.InferenceRequest) { sw, err := newSSEWriter(w) if err != nil { - // Headers haven't been committed yet (newSSEWriter failed before any - // successful frame flush in practice), so a JSON 500 is still safe. p.logger.ErrorContext(r.Context(), "sse writer init failed", slog.Any("err", err)) p.writeError(w, r, errorPayload{ HTTPStatus: http.StatusInternalServerError, @@ -494,7 +506,14 @@ func (p *Proxy) streamThroughGovernor(w http.ResponseWriter, r *http.Request, re return } - if err := p.governor.Stream(r.Context(), req, sw); err != nil { + ctx := r.Context() + if p.streamTimeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, p.streamTimeout) + defer cancel() + } + + if err := p.governor.Stream(ctx, req, sw); err != nil { // Client hung up -- socket is gone, don't try to write to it. // Two paths lead here: // A) The adapter detected ctx.Err() and returned context.Canceled. @@ -640,7 +659,6 @@ func (p *Proxy) handleAgentSessions(w http.ResponseWriter, _ *http.Request) { // SQLite + in-memory session data for the current process. func (p *Proxy) handleStats(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") - w.Header().Set("Access-Control-Allow-Origin", "*") // Prefer Supabase for global totals (server deployments). if p.supabaseStats != nil { diff --git a/gateway/internal/proxy/handlers_test.go b/gateway/internal/proxy/handlers_test.go new file mode 100644 index 0000000..fae3509 --- /dev/null +++ b/gateway/internal/proxy/handlers_test.go @@ -0,0 +1,43 @@ +package proxy + +import ( + "strings" + "testing" +) + +func TestValidateSessionKey(t *testing.T) { + if err := validateSessionKey(""); err != nil { + t.Fatalf("empty session key should be valid: %v", err) + } + if err := validateSessionKey("short"); err != nil { + t.Fatalf("short session key should be valid: %v", err) + } + long := strings.Repeat("a", 256) + if err := validateSessionKey(long); err != nil { + t.Fatalf("256-byte session key should be valid: %v", err) + } + tooLong := strings.Repeat("a", 257) + if err := validateSessionKey(tooLong); err == nil { + t.Fatal("257-byte session key should be invalid") + } +} + +func TestNormalizeContextPathBlocksTraversal(t *testing.T) { + cases := []struct { + input string + want string + }{ + {"foo/bar", "foo/bar"}, + {"../etc/passwd", defaultRawContextPath}, + {"foo/../bar", defaultRawContextPath}, + {"foo/bar/..", defaultRawContextPath}, + {"", defaultRawContextPath}, + {"ok.txt", "ok.txt"}, + } + for _, c := range cases { + got := normalizeContextPath(c.input) + if got != c.want { + t.Errorf("normalizeContextPath(%q) = %q, want %q", c.input, got, c.want) + } + } +} diff --git a/gateway/internal/proxy/request.go b/gateway/internal/proxy/request.go index 78a5dd2..7240e44 100644 --- a/gateway/internal/proxy/request.go +++ b/gateway/internal/proxy/request.go @@ -21,6 +21,16 @@ const ( headerAWSRegion = "X-IQ-AWS-Region" ) +func validateSessionKey(sk string) error { + if len(sk) == 0 { + return nil // optional + } + if len(sk) > 256 { + return fmt.Errorf("session key too long: %d bytes (max 256)", len(sk)) + } + return nil +} + var ( errMissingProvider = errors.New("missing X-IQ-Provider header") errUnknownProvider = errors.New("unknown provider in X-IQ-Provider header") @@ -62,6 +72,9 @@ func (p *Proxy) parseInferenceRequest(w http.ResponseWriter, r *http.Request) (* req.Credential = cred req.ProjectMemory = r.Header.Get(headerProjectMemory) req.SessionKey = r.Header.Get(headerSessionKey) + if err := validateSessionKey(req.SessionKey); err != nil { + return nil, err + } req.AzureEndpoint = r.Header.Get(headerAzureEndpoint) req.AWSRegion = r.Header.Get(headerAWSRegion) return &req, nil From 7b256c35573b50d3fab7817dd5eb2f9f26e126b2 Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Fri, 5 Jun 2026 11:39:30 -0500 Subject: [PATCH 07/15] fix: add proxy session eviction cleanup --- gateway/internal/proxy/proxy_eviction.go | 71 +++++++++++++++++++ gateway/internal/proxy/proxy_eviction_test.go | 49 +++++++++++++ 2 files changed, 120 insertions(+) create mode 100644 gateway/internal/proxy/proxy_eviction.go create mode 100644 gateway/internal/proxy/proxy_eviction_test.go diff --git a/gateway/internal/proxy/proxy_eviction.go b/gateway/internal/proxy/proxy_eviction.go new file mode 100644 index 0000000..0fe8d8c --- /dev/null +++ b/gateway/internal/proxy/proxy_eviction.go @@ -0,0 +1,71 @@ +package proxy + +import ( + "context" + "log/slog" + "time" +) + +const sessionCleanupInterval = time.Minute +const sessionMaxIdle = time.Hour + +// Start begins background goroutines (session TTL eviction). +func (p *Proxy) Start() { + p.cleanupCtx, p.cleanupCancel = context.WithCancel(context.Background()) + p.cleanupDone = make(chan struct{}) + go p.cleanupLoop() +} + +// Stop cleanly shuts down background goroutines. +func (p *Proxy) Stop() { + if p.cleanupCancel != nil { + p.cleanupCancel() + } + if p.cleanupDone != nil { + <-p.cleanupDone + } +} + +func (p *Proxy) touchSession(key string) { + p.sessionLastUsed.Store(key, time.Now()) +} + +func (p *Proxy) cleanupLoop() { + defer close(p.cleanupDone) + ticker := time.NewTicker(sessionCleanupInterval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + p.evictStaleSessions() + case <-p.cleanupCtx.Done(): + return + } + } +} + +func (p *Proxy) evictStaleSessions() { + cutoff := time.Now().Add(-sessionMaxIdle) + var evicted int + + p.sessionLastUsed.Range(func(key, value interface{}) bool { + lastUsed := value.(time.Time) + if lastUsed.Before(cutoff) { + sk := key.(string) + p.sessionTurnCounters.Delete(sk) + p.sessionWarmUpDone.Delete(sk) + p.sessionBoilerplateState.Delete(sk) + p.sessionPrefixHints.Delete(sk) + p.sessionSuggestionTs.Delete(sk) + p.sessionLastUsed.Delete(sk) + evicted++ + } + return true + }) + + if evicted > 0 { + p.logger.Debug("session eviction complete", + slog.Int("evicted", evicted), + slog.Duration("max_idle", sessionMaxIdle)) + } +} diff --git a/gateway/internal/proxy/proxy_eviction_test.go b/gateway/internal/proxy/proxy_eviction_test.go new file mode 100644 index 0000000..6fd859d --- /dev/null +++ b/gateway/internal/proxy/proxy_eviction_test.go @@ -0,0 +1,49 @@ +package proxy + +import ( + "testing" + "time" +) + +func TestSessionEviction(t *testing.T) { + p := New(&fakeGovernor{}) + p.Start() + defer p.Stop() + + // Seed some session state. + p.sessionTurnCounters.Store("old-session", &sessionTurnState{}) + p.sessionLastUsed.Store("old-session", time.Now().Add(-2*time.Hour)) + + p.sessionWarmUpDone.Store("fresh-session", true) + p.sessionLastUsed.Store("fresh-session", time.Now()) + + // Manually trigger eviction with a 1-minute max idle for testing. + p.evictStaleSessions() + + // old-session should be gone. + if _, ok := p.sessionTurnCounters.Load("old-session"); ok { + t.Fatal("old-session should have been evicted") + } + + // fresh-session should remain. + if _, ok := p.sessionWarmUpDone.Load("fresh-session"); !ok { + t.Fatal("fresh-session should still exist") + } +} + +func TestTouchSessionUpdatesLastUsed(t *testing.T) { + p := New(&fakeGovernor{}) + p.Start() + defer p.Stop() + + p.touchSession("test-key") + + val, ok := p.sessionLastUsed.Load("test-key") + if !ok { + t.Fatal("expected last-used entry after touch") + } + lastUsed := val.(time.Time) + if time.Since(lastUsed) > time.Second { + t.Fatal("expected recently updated last-used time") + } +} From 594ca3f530a30a645f8b6c7c6e9233c667316e28 Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Fri, 5 Jun 2026 16:21:43 -0500 Subject: [PATCH 08/15] fix: preserve valid JSONL payload dumps --- gateway/internal/proxy/claude_messages.go | 53 +++++++++++++++-- gateway/internal/proxy/dump_redaction_test.go | 57 +++++++++++++++++++ 2 files changed, 104 insertions(+), 6 deletions(-) diff --git a/gateway/internal/proxy/claude_messages.go b/gateway/internal/proxy/claude_messages.go index 9e1f0fc..c3e1228 100644 --- a/gateway/internal/proxy/claude_messages.go +++ b/gateway/internal/proxy/claude_messages.go @@ -501,11 +501,11 @@ func dumpClaudePayloads(requestID string, before, after []byte, stats claudeStre } beforePath := filepath.Join(dumpDir, "iq-before-"+requestID+".json") afterPath := filepath.Join(dumpDir, "iq-after-"+requestID+".json") - if err := os.WriteFile(beforePath, prettyJSON([]byte(redact.String(string(before)))), 0o600); err != nil { + if err := os.WriteFile(beforePath, prettyJSON(redactedJSONPayload(before)), 0o600); err != nil { fmt.Fprintf(os.Stderr, "[iq] failed to dump payload pair: %v\n", err) return } - if err := os.WriteFile(afterPath, prettyJSON([]byte(redact.String(string(after)))), 0o600); err != nil { + if err := os.WriteFile(afterPath, prettyJSON(redactedJSONPayload(after)), 0o600); err != nil { fmt.Fprintf(os.Stderr, "[iq] failed to dump payload pair: %v\n", err) return } @@ -555,10 +555,10 @@ func appendSessionDump(sessionFile, requestID string, before, after []byte, stat BeforeBytes: len(before), AfterBytes: len(after), SavedBytes: len(before) - len(after), - Before: json.RawMessage(before), - After: json.RawMessage(after), + Before: redactedJSONPayload(before), + After: redactedJSONPayload(after), Response: payloadDumpResponse{ - Text: stats.OutputRawText, + Text: redact.String(stats.OutputRawText), OutputTokens: stats.OutputTokens, Status: stats.Status, }, @@ -576,7 +576,6 @@ func appendSessionDump(sessionFile, requestID string, before, after []byte, stat if err != nil { return err } - line = []byte(redact.String(string(line))) line = append(line, '\n') dumpPayloadMu.Lock() @@ -590,6 +589,48 @@ func appendSessionDump(sessionFile, requestID string, before, after []byte, stat return err } +func redactedJSONPayload(raw []byte) json.RawMessage { + if len(bytes.TrimSpace(raw)) == 0 { + return json.RawMessage("null") + } + + var parsed any + if err := json.Unmarshal(raw, &parsed); err != nil { + encoded, _ := json.Marshal(redact.String(string(raw))) + return json.RawMessage(encoded) + } + + encoded, err := marshalJSONNoHTMLEscape(redactJSONValue(parsed)) + if err != nil { + fallback, _ := json.Marshal(redact.String(string(raw))) + return json.RawMessage(fallback) + } + return json.RawMessage(encoded) +} + +func redactJSONValue(value any) any { + switch v := value.(type) { + case map[string]any: + for key, child := range v { + if redact.SensitiveKey(key) { + v[key] = "[redacted]" + continue + } + v[key] = redactJSONValue(child) + } + return v + case []any: + for i, child := range v { + v[i] = redactJSONValue(child) + } + return v + case string: + return redact.String(v) + default: + return v + } +} + func appendDumpLog(dumpDir, beforePath, afterPath string) { logPath := filepath.Join(dumpDir, "dump.log") f, err := os.OpenFile(logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600) diff --git a/gateway/internal/proxy/dump_redaction_test.go b/gateway/internal/proxy/dump_redaction_test.go index 74544cc..4a6baf3 100644 --- a/gateway/internal/proxy/dump_redaction_test.go +++ b/gateway/internal/proxy/dump_redaction_test.go @@ -1,6 +1,7 @@ package proxy import ( + "encoding/json" "os" "path/filepath" "strings" @@ -42,6 +43,59 @@ func TestDumpRedaction(t *testing.T) { if !strings.Contains(string(data), redact.String("sk-proj-abc1234567890")) { t.Fatal("dump should contain redacted marker") } + for i, line := range strings.Split(strings.TrimSpace(string(data)), "\n") { + if !json.Valid([]byte(line)) { + t.Fatalf("dump line %d is not valid JSON: %s", i+1, line) + } + } +} + +func TestDumpRedactionPreservesJSONLForEscapedCredentialExamples(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "iq-dump-test-*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + sessionFile := filepath.Join(tmpDir, "dump.jsonl") + before := []byte(`{"messages":[{"role":"user","content":"normal prompt"}],"authorization":"Bearer sk-ant-beforesecret1234567890"}`) + after := []byte(`{"messages":[{"role":"user","content":"normal prompt"}]}`) + stats := claudeStreamStats{ + OutputRawText: "A span ending with `\\\"Authorization: sk-proj-responsesecret1234567890\\\"` should not corrupt JSON.\n```go\n\".cursor/rules/\",\n```", + OutputTokens: 12, + Status: "completed", + } + + if err := appendSessionDump(sessionFile, "req-escaped", before, after, stats, claudeOptimizerStats{}); err != nil { + t.Fatalf("appendSessionDump failed: %v", err) + } + + data, err := os.ReadFile(sessionFile) + if err != nil { + t.Fatalf("read dump failed: %v", err) + } + if strings.Contains(string(data), "sk-ant-beforesecret1234567890") || strings.Contains(string(data), "sk-proj-responsesecret1234567890") { + t.Fatalf("dump leaked secret material: %s", data) + } + + lines := strings.Split(strings.TrimSpace(string(data)), "\n") + if len(lines) != 1 { + t.Fatalf("dump lines=%d, want 1", len(lines)) + } + if !json.Valid([]byte(lines[0])) { + t.Fatalf("dump line is not valid JSON: %s", lines[0]) + } + + var rec payloadDumpRecord + if err := json.Unmarshal([]byte(lines[0]), &rec); err != nil { + t.Fatalf("unmarshal dump line: %v", err) + } + if rec.Response.Text == "" || !strings.Contains(rec.Response.Text, "[redacted") { + t.Fatalf("response text was not redacted as expected: %q", rec.Response.Text) + } + if !json.Valid(rec.Before) || !json.Valid(rec.After) { + t.Fatalf("embedded payloads must stay valid JSON: before=%s after=%s", rec.Before, rec.After) + } } func TestDumpRedactionPrettyFiles(t *testing.T) { @@ -70,4 +124,7 @@ func TestDumpRedactionPrettyFiles(t *testing.T) { if strings.Contains(string(data), "abc-secret") { t.Fatal("pretty dump should not contain raw bearer token") } + if !json.Valid(data) { + t.Fatalf("pretty dump should remain valid JSON: %s", data) + } } From 9dda2030481e9610cf62dbfd0862196a7c3c4710 Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Mon, 8 Jun 2026 19:28:20 -0500 Subject: [PATCH 09/15] fix: preserve cache_control for subscription auth so prompt caching works --- gateway/cmd/iq/main.go | 19 +- gateway/internal/memory/store.go | 6 + gateway/internal/proxy/claude_messages.go | 309 ++++++++++++++++----- gateway/internal/proxy/usage_test.go | 306 ++++++++++++++++++++ gateway/internal/sessions/tracker.go | 125 +++++---- gateway/internal/telemetry/agentsession.go | 31 ++- gateway/internal/telemetry/telemetry.go | 3 + 7 files changed, 670 insertions(+), 129 deletions(-) create mode 100644 gateway/internal/proxy/usage_test.go diff --git a/gateway/cmd/iq/main.go b/gateway/cmd/iq/main.go index b84fd11..193563f 100644 --- a/gateway/cmd/iq/main.go +++ b/gateway/cmd/iq/main.go @@ -193,18 +193,21 @@ func printSessionSummary(sessionID string) { defer db.Close() var requestsTotal, tokensAttempted, tokensSent, tokensDeduplicated int64 + var inputTokensReal, cacheReadTokens, cacheCreationTokens int64 var status string queryID := "%" if len(sessionID) >= 8 { queryID = "%-" + sessionID[:8] } err = db.QueryRow(` - SELECT requests_total, tokens_attempted, tokens_sent, tokens_deduplicated, status + SELECT requests_total, tokens_attempted, tokens_sent, tokens_deduplicated, + input_tokens_real, cache_read_tokens, cache_creation_tokens, status FROM agent_sessions WHERE session_id LIKE ? ORDER BY last_seen_at DESC LIMIT 1 - `, queryID).Scan(&requestsTotal, &tokensAttempted, &tokensSent, &tokensDeduplicated, &status) + `, queryID).Scan(&requestsTotal, &tokensAttempted, &tokensSent, &tokensDeduplicated, + &inputTokensReal, &cacheReadTokens, &cacheCreationTokens, &status) if err != nil { return } @@ -218,8 +221,14 @@ func printSessionSummary(sessionID string) { percent = float64(tokensDeduplicated) / float64(tokensAttempted) * 100 } - // Cost saved (estimate at $3 per million tokens for Claude 3.5 Sonnet inputs) - dollarsSaved := float64(tokensDeduplicated) * 0.000003 + // Prompt-cache hit ratio: of the input the model actually billed this session, + // what fraction Anthropic served from its prompt cache (≈10% cost, the latency + // win) instead of re-reading at full price. This is the real, measured signal + // for a subscription user — not a fabricated dollar figure. + cacheHitRatio := 0.0 + if inputTokensReal > 0 { + cacheHitRatio = float64(cacheReadTokens) / float64(inputTokensReal) * 100 + } fmt.Fprintln(os.Stderr) fmt.Fprintln(os.Stderr, " \033[1;36m┌────────────────────────────────────────────────────────┐\033[0m") @@ -229,7 +238,7 @@ func printSessionSummary(sessionID string) { fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;37m%-29s\033[0m\033[1;36m│\033[0m\n", "Tokens Attempted", formatNumber(tokensAttempted)) fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;37m%-29s\033[0m\033[1;36m│\033[0m\n", "Tokens Sent", formatNumber(tokensSent)) fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;35m%-29s\033[0m\033[1;36m│\033[0m\n", "Tokens Saved", fmt.Sprintf("%s (%.1f%%)", formatNumber(tokensDeduplicated), percent)) - fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;33m$%-28.2f\033[0m\033[1;36m│\033[0m\n", "Estimated Savings", dollarsSaved) + fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;36m%-29s\033[0m\033[1;36m│\033[0m\n", "Prompt-Cache Hits", fmt.Sprintf("%s tok (%.1f%%)", formatNumber(cacheReadTokens), cacheHitRatio)) fmt.Fprintln(os.Stderr, " \033[1;36m├────────────────────────────────────────────────────────┤\033[0m") statusColor := "\033[1;32m" // green if status == "killed" { diff --git a/gateway/internal/memory/store.go b/gateway/internal/memory/store.go index f06b505..d79edca 100644 --- a/gateway/internal/memory/store.go +++ b/gateway/internal/memory/store.go @@ -48,6 +48,10 @@ type UsageTotals struct { TokensSaved int BytesIn int BytesSaved int + // Measured upstream input usage (ground truth from Anthropic's usage object). + // CacheReadTokens is context served from Anthropic's prompt cache this turn. + CacheReadTokens int + CacheCreationTokens int } func NewStore(ttl time.Duration) *Store { @@ -121,6 +125,8 @@ func (s *Store) RecordUsage(sessionKey string, usage UsageTotals) { session.Totals.TokensSaved += usage.TokensSaved session.Totals.BytesIn += usage.BytesIn session.Totals.BytesSaved += usage.BytesSaved + session.Totals.CacheReadTokens += usage.CacheReadTokens + session.Totals.CacheCreationTokens += usage.CacheCreationTokens } func (s *Store) CleanupExpired() { diff --git a/gateway/internal/proxy/claude_messages.go b/gateway/internal/proxy/claude_messages.go index c3e1228..f8c2340 100644 --- a/gateway/internal/proxy/claude_messages.go +++ b/gateway/internal/proxy/claude_messages.go @@ -37,10 +37,6 @@ const ( claudeDefaultMode = "observe" ) -type contextKey string - -const isSubscriptionKey contextKey = "isSubscription" - // guardBypassRe matches directives that attempt to disable proxy safety controls. // The separator between "guards" and "velocity" is optional (covers slash, pipe, // backslash, or none) to handle formatting variations in injected CLAUDE.md files. @@ -113,6 +109,12 @@ type claudeOptimizerStats struct { PreservedInstructionCount int `json:"preserved_instruction_count"` PreservedLastOccurrenceBytes int `json:"preserved_last_occurrence_bytes"` PreservedLastOccurrenceCount int `json:"preserved_last_occurrence_count"` + // PreservedCachePrefix counts spans left untouched because they sit inside + // Anthropic's prompt-cache prefix (a cache_control breakpoint covers them). + // Rewriting them would invalidate the cached prefix for the whole suffix, + // costing far more than the bytes saved — so the optimizer preserves them. + PreservedCachePrefixBytes int `json:"preserved_cache_prefix_bytes"` + PreservedCachePrefixCount int `json:"preserved_cache_prefix_count"` // Size tracking. LargestSpanBytes int `json:"largest_span_bytes"` @@ -120,21 +122,29 @@ type claudeOptimizerStats struct { } type claudeStreamStats struct { - Chunks int - OutputText int - OutputRawText string - OutputTokens int - Status string - StatusCode int - Cancelled bool - Completed bool - HasToolUse bool - Provider string // "anthropic" or "bedrock" - UpstreamErr string - UpstreamErrorCode string - UpstreamErrorType string - UpstreamRequestID string - RetryAfter time.Duration + Chunks int + OutputText int + OutputRawText string + OutputTokens int + // Real upstream input accounting, captured from the message_start event's + // usage object. These are measured ground truth (not byte estimates). + // InputTokens excludes cached tokens; CacheReadInputTokens is context Anthropic + // served from its prompt cache (≈10% cost, the latency win); CacheCreation is + // context written into the cache this turn. + InputTokens int + CacheReadInputTokens int + CacheCreationInputTokens int + Status string + StatusCode int + Cancelled bool + Completed bool + HasToolUse bool + Provider string // "anthropic" or "bedrock" + UpstreamErr string + UpstreamErrorCode string + UpstreamErrorType string + UpstreamRequestID string + RetryAfter time.Duration } type claudeUpstreamErrorMeta struct { @@ -219,9 +229,7 @@ func (p *Proxy) handleClaudeMessages(w http.ResponseWriter, r *http.Request) { Model string `json:"model"` } _ = json.Unmarshal(body, &earlyMeta) - auth := r.Header.Get("Authorization") - isSubscription := isSubscriptionAuth(auth) - ctx := context.WithValue(r.Context(), isSubscriptionKey, isSubscription) + ctx := r.Context() forwardBody, meta, optStats, shape, err := p.prepareClaudeBody(ctx, cfg, sessionKey, body) if err != nil { @@ -251,7 +259,7 @@ func (p *Proxy) handleClaudeMessages(w http.ResponseWriter, r *http.Request) { }) } p.logClaudeRequestComplete(r.Context(), requestID, cfg.Mode, meta.Model, sessionKey, len(body), effectiveOpt, streamStats, duration, missingReqID, requestID, velocityWarning) - p.emitClaudeUsageEvent(r, sessionKey, meta.Model, effectiveOpt, duration, streamStats.StatusCode, overheadMs) + p.emitClaudeUsageEvent(r, sessionKey, meta.Model, effectiveOpt, streamStats, duration, streamStats.StatusCode, overheadMs) if os.Getenv("IQ_DUMP_PAYLOADS") == "1" { dumpClaudePayloads(requestID, body, nil, streamStats, effectiveOpt) } @@ -361,17 +369,25 @@ func (p *Proxy) handleClaudeMessages(w http.ResponseWriter, r *http.Request) { } duration := time.Since(started) if cfg.SessionStore != nil { + // Prefer measured upstream input over the byte estimate when Anthropic + // reported usage; fall back to the estimate when it didn't (e.g. errors). + tokensIn := streamStats.realInputTokens() + if tokensIn == 0 { + tokensIn = estimateTokens(len(body)) + } cfg.SessionStore.RecordUsage(sessionKey, memory.UsageTotals{ - Requests: 1, - TokensIn: estimateTokens(len(body)), - TokensOut: streamStats.estimatedOutputTokens(), - TokensSaved: optStats.EstimatedTokensSaved, - BytesIn: len(body), - BytesSaved: optStats.BytesBefore - optStats.BytesAfter, + Requests: 1, + TokensIn: tokensIn, + TokensOut: streamStats.estimatedOutputTokens(), + TokensSaved: optStats.EstimatedTokensSaved, + BytesIn: len(body), + BytesSaved: optStats.BytesBefore - optStats.BytesAfter, + CacheReadTokens: streamStats.CacheReadInputTokens, + CacheCreationTokens: streamStats.CacheCreationInputTokens, }) } p.logClaudeRequestComplete(r.Context(), requestID, cfg.Mode, meta.Model, sessionKey, len(body), optStats, streamStats, duration, missingReqID, requestID, velocityWarning) - p.emitClaudeUsageEvent(r, sessionKey, meta.Model, optStats, duration, streamStats.StatusCode, overheadMs) + p.emitClaudeUsageEvent(r, sessionKey, meta.Model, optStats, streamStats, duration, streamStats.StatusCode, overheadMs) if os.Getenv("IQ_DUMP_PAYLOADS") == "1" { dumpClaudePayloads(requestID, body, forwardBody, streamStats, optStats) @@ -516,6 +532,12 @@ type payloadDumpResponse struct { Text string `json:"text"` OutputTokens int `json:"output_tokens"` Status string `json:"status"` + // Raw upstream input usage exactly as Anthropic reported it, so a dump can + // distinguish "prompt caching not applied" (cache fields 0 with large input) + // from "applied but small". Zero on synthetic/probe turns (no upstream call). + InputTokens int `json:"input_tokens"` + CacheReadInputTokens int `json:"cache_read_input_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens"` } type payloadDumpRecord struct { @@ -558,9 +580,12 @@ func appendSessionDump(sessionFile, requestID string, before, after []byte, stat Before: redactedJSONPayload(before), After: redactedJSONPayload(after), Response: payloadDumpResponse{ - Text: redact.String(stats.OutputRawText), - OutputTokens: stats.OutputTokens, - Status: stats.Status, + Text: redact.String(stats.OutputRawText), + OutputTokens: stats.OutputTokens, + Status: stats.Status, + InputTokens: stats.InputTokens, + CacheReadInputTokens: stats.CacheReadInputTokens, + CacheCreationInputTokens: stats.CacheCreationInputTokens, }, Optimizer: &payloadOptimizerStats{ BlocksPruned: optStats.BlocksPruned, @@ -973,6 +998,21 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, systemAllKnown := true hasSystemSpan := false + // Prompt-cache protection: Claude Code marks a rolling cache_control breakpoint + // on the latest user turn (and the system prompt) every request, so the prefix + // up to that breakpoint is served by Anthropic at cache-read price on the next + // turn. Rewriting any span inside that prefix invalidates the cache for the + // entire suffix — far more expensive than the bytes a prune would save. When + // the request uses prompt caching we therefore preserve cached-prefix spans. + cachePrefixMsgIdx := lastCacheControlMessageIndex(root) + systemCached := contentHasCacheControl(root["system"]) + spanInCachedPrefix := func(span TextSpan) bool { + if span.Role == "system" { + return systemCached + } + return cachePrefixMsgIdx >= 0 && span.MessageIndex >= 0 && span.MessageIndex <= cachePrefixMsgIdx + } + for _, span := range spans { if span.Bytes < minSpanBytes { // FIX 4 & SQLite Optimization: register small spans in the session cache @@ -1079,6 +1119,15 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, continue } + // Even a new span must not be rewritten if it sits inside the prompt-cache + // prefix — rewriting it breaks Anthropic's cached suffix for this turn. + if spanInCachedPrefix(span) { + stats.PreservedCachePrefixBytes += span.Bytes + stats.PreservedCachePrefixCount++ + stats.BlocksNew++ + continue + } + // FIX 7: SUGGESTION MODE injection cooldown. If a boilerplate span // is new (unknown hash) but the last forward was <5 turns ago AND // context delta is <10 KB, prune it anyway to suppress redundant @@ -1148,6 +1197,14 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, } } + // Never rewrite content inside Anthropic's prompt-cache prefix; doing so + // invalidates the cached suffix and costs more than the prune saves. + if spanInCachedPrefix(span) { + stats.PreservedCachePrefixBytes += span.Bytes + stats.PreservedCachePrefixCount++ + continue + } + if !isEligibleSpanClass(span.Class, cfg.Optimizer) { switch span.Class { case SpanClassSystemText: @@ -1165,8 +1222,15 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, prunableSpans = append(prunableSpans, span) } - isSub, _ := ctx.Value(isSubscriptionKey).(bool) - wantPromptCache := cfg.Optimizer.EnablePromptCache && hasSystemSpan && systemAllKnown && !cfg.Bedrock.Enabled && !isSub + // Only inject our own cache_control when the client did NOT already manage + // prompt caching. Claude Code places its own breakpoints (system + rolling + // latest turn) on every request; adding another risks exceeding Anthropic's + // 4-breakpoint limit (a 400 for any user, not just subscription) and is + // redundant. Deferring to the client is the correct generalization of the old + // blunt "skip for subscription" rule — subscription users keep their cache + // benefit via Claude Code's own breakpoints, which E2 now preserves. + requestHasCacheControl := systemCached || cachePrefixMsgIdx >= 0 + wantPromptCache := cfg.Optimizer.EnablePromptCache && hasSystemSpan && systemAllKnown && !cfg.Bedrock.Enabled && !requestHasCacheControl if len(prunableSpans) == 0 && !wantPromptCache { return body, req, stats, shape, nil } @@ -1253,6 +1317,48 @@ var protectedInstructionPathFragments = [...]string{ ".github/copilot-instructions.md", } +// contentHasCacheControl reports whether an Anthropic content value (an array of +// content blocks) carries a cache_control breakpoint on any block. +func contentHasCacheControl(content any) bool { + arr, ok := content.([]any) + if !ok { + return false + } + for _, raw := range arr { + b, ok := raw.(map[string]any) + if !ok { + continue + } + if _, has := b["cache_control"]; has { + return true + } + } + return false +} + +// lastCacheControlMessageIndex returns the highest message index whose content +// carries a cache_control breakpoint, or -1 if none. Claude Code places a +// rolling breakpoint on the latest user turn each request, so this marks the end +// of the prompt-cache prefix that Anthropic will serve at cache-read price on the +// next turn. Content at or before this index must not be rewritten. +func lastCacheControlMessageIndex(root map[string]any) int { + msgs, ok := root["messages"].([]any) + if !ok { + return -1 + } + last := -1 + for i, raw := range msgs { + m, ok := raw.(map[string]any) + if !ok { + continue + } + if contentHasCacheControl(m["content"]) { + last = i + } + } + return last +} + func isProtectedInstructionSpan(span TextSpan) bool { return containsProtectedInstructionPath(span.SourcePath) || containsProtectedInstructionPath(span.Text) || containsCredentialMarker(span.Text) } @@ -1371,12 +1477,6 @@ func stripCacheControlRecursively(v any) { } } -func isSubscriptionAuth(auth string) bool { - return strings.HasPrefix(auth, "Bearer claudepro_") || - strings.Contains(auth, "claudepro_") || - strings.Contains(auth, "sk-ant-oat") -} - // forwardClaudeMessagesViaBedrock proxies /v1/messages to the Bedrock // InvokeModelWithResponseStream API. The Anthropic Messages body is forwarded // as-is except that `model` and `stream` are removed and @@ -1486,10 +1586,8 @@ func (p *Proxy) forwardClaudeMessagesViaBedrock(w http.ResponseWriter, r *http.R stats.Chunks++ stats.OutputText += anthropicDeltaTextLen(payloadStr) stats.OutputRawText += anthropicDeltaText(payloadStr) - case "message_delta": - if tokens := anthropicUsageOutputTokens(payloadStr); tokens > 0 { - stats.OutputTokens = tokens - } + case "message_start", "message_delta": + stats.applyUpstreamUsage(parseAnthropicUsage(payloadStr)) } } if err := stream.Err(); err != nil { @@ -1626,16 +1724,12 @@ func (p *Proxy) forwardClaudeMessages(w http.ResponseWriter, r *http.Request, cf return p.forwardClaudeMessagesViaBedrock(w, r, cfg, body) } - auth := r.Header.Get("Authorization") - if isSubscriptionAuth(auth) { - var root map[string]any - if err := json.Unmarshal(body, &root); err == nil { - stripCacheControlRecursively(root) - if cleaned, err := marshalJSONNoHTMLEscape(root); err == nil { - body = cleaned - } - } - } + // Preserve Claude Code's cache_control breakpoints on the direct Anthropic + // path. Subscription/OAuth auth supports prompt caching too — Claude Code + // relies on it — so stripping breakpoints here only disabled the cache: + // Anthropic then reported cache_read=0 / cache_creation=0 and billed full + // input every turn. Bedrock, which rejects Anthropic-only cache_control, + // strips it in its own path (see forwardClaudeMessagesViaBedrock). upstreamURL, err := url.JoinPath(strings.TrimRight(cfg.AnthropicBaseURL, "/"), "v1", "messages") if err != nil { @@ -1825,10 +1919,8 @@ func proxyAnthropicStream(w http.ResponseWriter, r *http.Request, resp *http.Res txt := anthropicDeltaText(payload) stats.OutputText += len(txt) stats.OutputRawText += txt - case "message_delta": - if tokens := anthropicUsageOutputTokens(payload); tokens > 0 { - stats.OutputTokens = tokens - } + case "message_start", "message_delta": + stats.applyUpstreamUsage(parseAnthropicUsage(payload)) } } } @@ -1873,19 +1965,77 @@ func anthropicDeltaTextLen(payload string) int { return len(ev.Delta.Text) } -func anthropicUsageOutputTokens(payload string) int { +// anthropicUsage holds the token-usage fields Anthropic reports over SSE. Input +// and cache counters arrive on message_start (inside message.usage); the final +// cumulative output_tokens arrives on message_delta (top-level usage). All fields +// are optional, so a payload that omits a field simply leaves it at zero. +type anthropicUsage struct { + InputTokens int + OutputTokens int + CacheReadInputTokens int + CacheCreationInputTokens int +} + +// parseAnthropicUsage extracts usage from either a message_start payload +// (usage nested under "message") or a message_delta payload (top-level "usage"). +func parseAnthropicUsage(payload string) anthropicUsage { var ev struct { - Usage struct { - OutputTokens int `json:"output_tokens"` - } `json:"usage"` + Usage *usageFields `json:"usage"` + Message *struct { + Usage *usageFields `json:"usage"` + } `json:"message"` } if err := json.Unmarshal([]byte(payload), &ev); err != nil { - return 0 + return anthropicUsage{} + } + u := ev.Usage + if u == nil && ev.Message != nil { + u = ev.Message.Usage + } + if u == nil { + return anthropicUsage{} + } + return anthropicUsage{ + InputTokens: u.InputTokens, + OutputTokens: u.OutputTokens, + CacheReadInputTokens: u.CacheReadInputTokens, + CacheCreationInputTokens: u.CacheCreationInputTokens, + } +} + +type usageFields struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheReadInputTokens int `json:"cache_read_input_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens"` +} + +// applyUpstreamUsage folds a parsed usage object into streamStats, keeping the +// max output_tokens seen (cumulative) and capturing input/cache counters when +// present (they appear once, on message_start). +func (s *claudeStreamStats) applyUpstreamUsage(u anthropicUsage) { + if u.OutputTokens > 0 { + s.OutputTokens = u.OutputTokens + } + if u.InputTokens > 0 { + s.InputTokens = u.InputTokens + } + if u.CacheReadInputTokens > 0 { + s.CacheReadInputTokens = u.CacheReadInputTokens + } + if u.CacheCreationInputTokens > 0 { + s.CacheCreationInputTokens = u.CacheCreationInputTokens } - return ev.Usage.OutputTokens } -func (p *Proxy) emitClaudeUsageEvent(r *http.Request, sessionKey, model string, optStats claudeOptimizerStats, duration time.Duration, upstreamStatus int, overheadMs int64) { +// realInputTokens returns the total measured input the model billed this turn: +// fresh input + cache-creation + cache-read. Zero when upstream reported nothing +// (e.g. synthetic probe / cache replay paths). +func (s claudeStreamStats) realInputTokens() int { + return s.InputTokens + s.CacheReadInputTokens + s.CacheCreationInputTokens +} + +func (p *Proxy) emitClaudeUsageEvent(r *http.Request, sessionKey, model string, optStats claudeOptimizerStats, streamStats claudeStreamStats, duration time.Duration, upstreamStatus int, overheadMs int64) { if p.usageTracker != nil { p.usageTracker.Track(telemetry.UsageEvent{ MachineID: telemetry.GetMachineID(), @@ -1897,6 +2047,9 @@ func (p *Proxy) emitClaudeUsageEvent(r *http.Request, sessionKey, model string, InputTokensAttempted: optStats.EstimatedTokensBefore, InputTokensSent: optStats.EstimatedTokensAfter, TokensSaved: optStats.EstimatedTokensSaved, + InputTokensReal: streamStats.realInputTokens(), + CacheReadTokens: streamStats.CacheReadInputTokens, + CacheCreationTokens: streamStats.CacheCreationInputTokens, ReductionRatio: optStats.ReductionRatio * 100, BlocksAnalyzed: optStats.BlocksSeen, BlocksPruned: optStats.BlocksPruned, @@ -1907,9 +2060,12 @@ func (p *Proxy) emitClaudeUsageEvent(r *http.Request, sessionKey, model string, } outcome := telemetry.RequestOutcome{ - TokensAttempted: optStats.EstimatedTokensBefore, - TokensSent: optStats.EstimatedTokensAfter, - TokensSaved: optStats.EstimatedTokensSaved, + TokensAttempted: optStats.EstimatedTokensBefore, + TokensSent: optStats.EstimatedTokensAfter, + TokensSaved: optStats.EstimatedTokensSaved, + InputTokensReal: streamStats.realInputTokens(), + CacheReadTokens: streamStats.CacheReadInputTokens, + CacheCreationTokens: streamStats.CacheCreationInputTokens, } if p.sessionTracker != nil { p.sessionTracker.Record(sessionKey, outcome) @@ -2072,6 +2228,21 @@ func (p *Proxy) logClaudeRequestComplete(ctx context.Context, requestID, mode, m attrs = append(attrs, slog.Int("preserved_instruction_count", opt.PreservedInstructionCount)) attrs = append(attrs, slog.Int("preserved_instruction_bytes", opt.PreservedInstructionBytes)) } + if opt.PreservedCachePrefixCount > 0 { + attrs = append(attrs, slog.Int("preserved_cache_prefix_count", opt.PreservedCachePrefixCount)) + attrs = append(attrs, slog.Int("preserved_cache_prefix_bytes", opt.PreservedCachePrefixBytes)) + } + // Live per-turn measured cache efficiency (E4): visible each turn in --dev logs + // and session dumps, not just in the post-exit summary. Emitted only when the + // upstream actually reported input usage (skips synthetic/cache-replay turns). + if realIn := stream.realInputTokens(); realIn > 0 { + attrs = append(attrs, + slog.Int("input_tokens_real", realIn), + slog.Int("cache_read_tokens", stream.CacheReadInputTokens), + slog.Int("cache_creation_tokens", stream.CacheCreationInputTokens), + slog.Float64("cache_hit_ratio", float64(stream.CacheReadInputTokens)/float64(realIn)), + ) + } for class, bytes := range opt.ClassBytesSeen { if bytes > 0 { attrs = append(attrs, slog.Int("class_bytes_seen:"+class, bytes)) diff --git a/gateway/internal/proxy/usage_test.go b/gateway/internal/proxy/usage_test.go new file mode 100644 index 0000000..bf9fadf --- /dev/null +++ b/gateway/internal/proxy/usage_test.go @@ -0,0 +1,306 @@ +package proxy + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Revanth14/indexqube/gateway/internal/memory" + "github.com/Revanth14/indexqube/gateway/internal/telemetry" +) + +func TestParseAnthropicUsage(t *testing.T) { + t.Parallel() + + // message_start carries input + cache counters nested under "message". + start := `{"type":"message_start","message":{"id":"msg_1","usage":{"input_tokens":120,"cache_read_input_tokens":4000,"cache_creation_input_tokens":80,"output_tokens":1}}}` + got := parseAnthropicUsage(start) + if got.InputTokens != 120 || got.CacheReadInputTokens != 4000 || got.CacheCreationInputTokens != 80 { + t.Fatalf("message_start usage = %+v, want input=120 cache_read=4000 cache_creation=80", got) + } + + // message_delta carries the final cumulative output_tokens at the top level. + delta := `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":256}}` + got = parseAnthropicUsage(delta) + if got.OutputTokens != 256 { + t.Fatalf("message_delta output_tokens = %d, want 256", got.OutputTokens) + } + + // Folding both into stats should keep input/cache from start and output from delta. + var stats claudeStreamStats + stats.applyUpstreamUsage(parseAnthropicUsage(start)) + stats.applyUpstreamUsage(parseAnthropicUsage(delta)) + if stats.InputTokens != 120 || stats.CacheReadInputTokens != 4000 || stats.CacheCreationInputTokens != 80 { + t.Fatalf("folded input/cache = %+v, want input=120 cache_read=4000 cache_creation=80", stats) + } + if stats.OutputTokens != 256 { + t.Fatalf("folded output_tokens = %d, want 256", stats.OutputTokens) + } + if got := stats.realInputTokens(); got != 120+4000+80 { + t.Fatalf("realInputTokens = %d, want %d", got, 120+4000+80) + } + + // Malformed payload must not panic and yields a zero value. + if got := parseAnthropicUsage("not json"); (got != anthropicUsage{}) { + t.Fatalf("malformed payload usage = %+v, want zero", got) + } +} + +// TestClaudeMessages_CapturesRealUpstreamCacheTokens verifies that the +// cache_read / cache_creation / input tokens reported by the upstream +// message_start event are recorded against the session (visible on +// /v1/agent-sessions) rather than discarded in favor of byte estimates. +func TestClaudeMessages_CapturesRealUpstreamCacheTokens(t *testing.T) { + t.Parallel() + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "event: message_start\n") + _, _ = io.WriteString(w, `data: {"type":"message_start","message":{"id":"msg_1","model":"claude-sonnet-4-6","usage":{"input_tokens":50,"cache_read_input_tokens":9000,"cache_creation_input_tokens":100,"output_tokens":1}}}`+"\n\n") + _, _ = io.WriteString(w, "event: content_block_delta\n") + _, _ = io.WriteString(w, `data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"hi"}}`+"\n\n") + _, _ = io.WriteString(w, "event: message_delta\n") + _, _ = io.WriteString(w, `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":42}}`+"\n\n") + _, _ = io.WriteString(w, "event: message_stop\n") + _, _ = io.WriteString(w, `data: {"type":"message_stop"}`+"\n\n") + })) + t.Cleanup(upstream.Close) + + p := New(&fakeGovernor{}, + WithAgentSessionStore(telemetry.NewAgentSessionStore(time.Hour)), + WithClaudeMessages(ClaudeMessagesConfig{ + Mode: "observe", + DevToken: "iq-dev-local", + AnthropicAPIKey: "sk-ant-test", + AnthropicBaseURL: upstream.URL, + AnthropicVersion: "2023-06-01", + SessionStore: memory.NewStore(time.Hour), + }), + ) + srv := httptest.NewServer(p.Handler()) + t.Cleanup(srv.Close) + + req, err := http.NewRequest(http.MethodPost, srv.URL+"/v1/messages", strings.NewReader(`{"model":"claude-sonnet-4-6","stream":true,"messages":[{"role":"user","content":"hi"}]}`)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Authorization", "Bearer iq-dev-local") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("POST /v1/messages: %v", err) + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status=%d, want 200", resp.StatusCode) + } + + resp, err = http.Get(srv.URL + "/v1/agent-sessions") + if err != nil { + t.Fatalf("GET /v1/agent-sessions: %v", err) + } + defer resp.Body.Close() + var body struct { + Sessions []struct { + CacheReadTokens int64 `json:"cache_read_tokens"` + CacheCreationTokens int64 `json:"cache_creation_tokens"` + InputTokensReal int64 `json:"input_tokens_real"` + } `json:"sessions"` + } + if err := json.NewDecoder(resp.Body).Decode(&body); err != nil { + t.Fatalf("decode response: %v", err) + } + if len(body.Sessions) != 1 { + t.Fatalf("sessions len=%d, want 1", len(body.Sessions)) + } + s := body.Sessions[0] + if s.CacheReadTokens != 9000 { + t.Fatalf("cache_read_tokens=%d, want 9000 (real upstream value, not byte estimate)", s.CacheReadTokens) + } + if s.CacheCreationTokens != 100 { + t.Fatalf("cache_creation_tokens=%d, want 100", s.CacheCreationTokens) + } + if s.InputTokensReal != 50+9000+100 { + t.Fatalf("input_tokens_real=%d, want %d", s.InputTokensReal, 50+9000+100) + } +} + +// TestClaudeMessages_PromptCachePrefixIsNotPruned verifies the optimizer does +// NOT rewrite content that sits inside Anthropic's prompt-cache prefix. Claude +// Code marks a cache_control breakpoint on the latest turn each request; pruning +// an older block before that breakpoint would invalidate the cached suffix and +// cost far more than the bytes saved. The same payload without a breakpoint is +// pruned (see TestClaudeMessages_OptimizeStillPrunesOrdinaryToolResult), so this +// isolates the cache-protection behavior. +func TestClaudeMessages_PromptCachePrefixIsNotPruned(t *testing.T) { + t.Parallel() + p := New(&fakeGovernor{}) + cfg := ClaudeMessagesConfig{ + Mode: "optimize", + EnableBlockOptimizer: true, + SessionStore: memory.NewStore(time.Hour), + Optimizer: OptimizerConfig{ + MinSpanBytes: 512, + EnableToolResultPruning: true, + }, + } + // Two identical tool results (older at messages[1], newer at messages[3]). + // The latest turn carries a cache_control breakpoint, so the whole prefix + // up to messages[4] is cacheable and must be preserved verbatim. + fileBody := strings.Repeat("ordinary source code output line\n", 80) + body := []byte(fmt.Sprintf( + `{"model":"claude-sonnet-4-6","messages":[`+ + `{"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"Read","input":{"file_path":"/repo/src/main.go"}}]},`+ + `{"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":%q}]},`+ + `{"role":"assistant","content":[{"type":"tool_use","id":"t2","name":"Read","input":{"file_path":"/repo/src/main.go"}}]},`+ + `{"role":"user","content":[{"type":"tool_result","tool_use_id":"t2","content":%q}]},`+ + `{"role":"user","content":[{"type":"text","text":"latest turn","cache_control":{"type":"ephemeral"}}]}]}`, + fileBody, fileBody, + )) + + // First request warms the session cache; second would prune the older + // duplicate if it were not inside the cache prefix. + if _, _, _, _, err := p.prepareClaudeBody(context.Background(), cfg, "cacheprefix-test", body); err != nil { + t.Fatalf("first prepare: %v", err) + } + forward, _, stats, _, err := p.prepareClaudeBody(context.Background(), cfg, "cacheprefix-test", body) + if err != nil { + t.Fatalf("second prepare: %v", err) + } + + if strings.Contains(string(forward), "omitted") { + t.Fatalf("cached-prefix content must not be rewritten, body=%s", forward) + } + if stats.BlocksPruned != 0 { + t.Fatalf("blocks_pruned=%d, want 0 (cache prefix protected); stats=%+v", stats.BlocksPruned, stats) + } + if stats.PreservedCachePrefixCount < 1 { + t.Fatalf("preserved_cache_prefix_count=%d, want >=1; stats=%+v", stats.PreservedCachePrefixCount, stats) + } +} + +// promptCacheCfg builds an optimize-mode config with prompt-cache injection on. +func promptCacheCfg() ClaudeMessagesConfig { + return ClaudeMessagesConfig{ + Mode: "optimize", + EnableBlockOptimizer: true, + SessionStore: memory.NewStore(time.Hour), + Optimizer: OptimizerConfig{ + MinSpanBytes: 512, + EnablePromptCache: true, + }, + } +} + +// TestPromptCache_InjectsWhenClientHasNone verifies IndexQube adds a cache_control +// breakpoint to a stable system prompt when the client supplied none. +func TestPromptCache_InjectsWhenClientHasNone(t *testing.T) { + t.Parallel() + p := New(&fakeGovernor{}) + cfg := promptCacheCfg() + systemText := strings.Repeat("You are a careful assistant. ", 40) // >512B, stable + body := []byte(fmt.Sprintf( + `{"model":"claude-sonnet-4-6","system":%q,"messages":[{"role":"user","content":"hello"}]}`, + systemText, + )) + + // Turn 1 warms the system span; turn 2 sees it as known/stable → injects. + if _, _, _, _, err := p.prepareClaudeBody(context.Background(), cfg, "inject-test", body); err != nil { + t.Fatalf("first prepare: %v", err) + } + forward, _, _, _, err := p.prepareClaudeBody(context.Background(), cfg, "inject-test", body) + if err != nil { + t.Fatalf("second prepare: %v", err) + } + if !strings.Contains(string(forward), "cache_control") || !strings.Contains(string(forward), "ephemeral") { + t.Fatalf("expected injected cache_control on stable system prompt, body=%s", forward) + } +} + +// TestPromptCache_DefersWhenClientManagesCaching verifies IndexQube does NOT add +// its own cache_control when the request already carries a breakpoint (Claude Code +// manages caching). Adding another risks exceeding Anthropic's 4-breakpoint limit. +func TestPromptCache_DefersWhenClientManagesCaching(t *testing.T) { + t.Parallel() + p := New(&fakeGovernor{}) + cfg := promptCacheCfg() + systemText := strings.Repeat("You are a careful assistant. ", 40) + // The client put its own cache_control on the latest user turn; system is a + // bare string. IndexQube must leave system alone (no second breakpoint). + body := []byte(fmt.Sprintf( + `{"model":"claude-sonnet-4-6","system":%q,"messages":[{"role":"user","content":[{"type":"text","text":"hello","cache_control":{"type":"ephemeral"}}]}]}`, + systemText, + )) + + if _, _, _, _, err := p.prepareClaudeBody(context.Background(), cfg, "defer-test", body); err != nil { + t.Fatalf("first prepare: %v", err) + } + forward, _, _, _, err := p.prepareClaudeBody(context.Background(), cfg, "defer-test", body) + if err != nil { + t.Fatalf("second prepare: %v", err) + } + // Exactly the one breakpoint the client sent — IndexQube added none. + if got := strings.Count(string(forward), "cache_control"); got != 1 { + t.Fatalf("cache_control count=%d, want 1 (client's only; no extra injected); body=%s", got, forward) + } +} + +// TestClaudeMessages_SubscriptionAuthPreservesCacheControl verifies that for +// subscription/OAuth auth (sk-ant-oat…) the proxy forwards Claude Code's +// cache_control breakpoints to Anthropic instead of stripping them. Stripping +// them disabled prompt caching entirely — Anthropic reported cache_read=0 / +// cache_creation=0 and billed full input every turn, the opposite of saving. +func TestClaudeMessages_SubscriptionAuthPreservesCacheControl(t *testing.T) { + t.Parallel() + bodyCh := make(chan []byte, 1) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + bodyCh <- b + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "event: message_start\n") + _, _ = io.WriteString(w, `data: {"type":"message_start","message":{"id":"msg_1","model":"claude-sonnet-4-6","usage":{"input_tokens":10,"cache_read_input_tokens":5000,"cache_creation_input_tokens":0,"output_tokens":1}}}`+"\n\n") + _, _ = io.WriteString(w, "event: message_stop\n") + _, _ = io.WriteString(w, `data: {"type":"message_stop"}`+"\n\n") + })) + t.Cleanup(upstream.Close) + + // Passthrough mode: no API key, so the subscription bearer flows upstream. + p := New(&fakeGovernor{}, + WithClaudeMessages(ClaudeMessagesConfig{ + Mode: "observe", + DevToken: "iq-dev-local", + AnthropicBaseURL: upstream.URL, + AnthropicVersion: "2023-06-01", + SessionStore: memory.NewStore(time.Hour), + }), + ) + srv := httptest.NewServer(p.Handler()) + t.Cleanup(srv.Close) + + // A cache_control breakpoint on the latest turn, exactly as Claude Code sends. + reqBody := `{"model":"claude-sonnet-4-6","stream":true,"messages":[{"role":"user","content":[{"type":"text","text":"hello","cache_control":{"type":"ephemeral"}}]}]}` + req, err := http.NewRequest(http.MethodPost, srv.URL+"/v1/messages", strings.NewReader(reqBody)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Authorization", "Bearer sk-ant-oat-test-token") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("POST /v1/messages: %v", err) + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status=%d, want 200", resp.StatusCode) + } + + gotBody := <-bodyCh + if !strings.Contains(string(gotBody), "cache_control") { + t.Fatalf("upstream body must retain cache_control for subscription auth, got: %s", gotBody) + } +} diff --git a/gateway/internal/sessions/tracker.go b/gateway/internal/sessions/tracker.go index b34c9e0..598daaf 100644 --- a/gateway/internal/sessions/tracker.go +++ b/gateway/internal/sessions/tracker.go @@ -10,6 +10,7 @@ // The schema uses two tables: // - agent_sessions: one row per session, upserted on every request // - kill_events: append-only audit log, one row per guard kill +// // Hardened for high-concurrency. package sessions @@ -30,17 +31,20 @@ PRAGMA journal_mode = WAL; PRAGMA synchronous = NORMAL; CREATE TABLE IF NOT EXISTS agent_sessions ( - session_id TEXT PRIMARY KEY, - started_at INTEGER NOT NULL, - last_seen_at INTEGER NOT NULL, - tokens_attempted INTEGER NOT NULL DEFAULT 0, - tokens_sent INTEGER NOT NULL DEFAULT 0, - tokens_deduplicated INTEGER NOT NULL DEFAULT 0, - requests_total INTEGER NOT NULL DEFAULT 0, - loop_detected INTEGER NOT NULL DEFAULT 0, - kill_events INTEGER NOT NULL DEFAULT 0, - kill_reason TEXT NOT NULL DEFAULT '', - status TEXT NOT NULL DEFAULT 'active' + session_id TEXT PRIMARY KEY, + started_at INTEGER NOT NULL, + last_seen_at INTEGER NOT NULL, + tokens_attempted INTEGER NOT NULL DEFAULT 0, + tokens_sent INTEGER NOT NULL DEFAULT 0, + tokens_deduplicated INTEGER NOT NULL DEFAULT 0, + input_tokens_real INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + requests_total INTEGER NOT NULL DEFAULT 0, + loop_detected INTEGER NOT NULL DEFAULT 0, + kill_events INTEGER NOT NULL DEFAULT 0, + kill_reason TEXT NOT NULL DEFAULT '', + status TEXT NOT NULL DEFAULT 'active' ); CREATE TABLE IF NOT EXISTS kill_events ( @@ -54,17 +58,20 @@ CREATE TABLE IF NOT EXISTS kill_events ( // SessionRow is the read-side representation of one row in agent_sessions. type SessionRow struct { - SessionID string `json:"session_id"` - StartedAt int64 `json:"started_at"` - LastSeenAt int64 `json:"last_seen_at"` - TokensAttempted int64 `json:"tokens_attempted"` - TokensSent int64 `json:"tokens_sent"` - TokensDeduplicated int64 `json:"tokens_deduplicated"` - RequestsTotal int64 `json:"requests_total"` - LoopDetected int64 `json:"loop_detected"` - KillEvents int64 `json:"kill_events"` - KillReason string `json:"kill_reason"` - Status string `json:"status"` + SessionID string `json:"session_id"` + StartedAt int64 `json:"started_at"` + LastSeenAt int64 `json:"last_seen_at"` + TokensAttempted int64 `json:"tokens_attempted"` + TokensSent int64 `json:"tokens_sent"` + TokensDeduplicated int64 `json:"tokens_deduplicated"` + InputTokensReal int64 `json:"input_tokens_real"` + CacheReadTokens int64 `json:"cache_read_tokens"` + CacheCreationTokens int64 `json:"cache_creation_tokens"` + RequestsTotal int64 `json:"requests_total"` + LoopDetected int64 `json:"loop_detected"` + KillEvents int64 `json:"kill_events"` + KillReason string `json:"kill_reason"` + Status string `json:"status"` } // KillRow is the read-side representation of one row in kill_events. @@ -106,6 +113,18 @@ func Open(path string, logger *slog.Logger) (*Tracker, error) { return nil, err } + // Backfill columns added after the original schema for databases created by + // an earlier IndexQube version. SQLite has no "ADD COLUMN IF NOT EXISTS", so + // each ALTER is best-effort: a "duplicate column name" error means the column + // already exists and is safely ignored. + for _, col := range []string{ + "input_tokens_real", + "cache_read_tokens", + "cache_creation_tokens", + } { + _, _ = db.Exec("ALTER TABLE agent_sessions ADD COLUMN " + col + " INTEGER NOT NULL DEFAULT 0") + } + if logger == nil { logger = slog.Default() } @@ -148,6 +167,7 @@ func (t *Tracker) SessionByID(sessionID string) (SessionRow, bool, error) { err := t.db.QueryRow(` SELECT session_id, started_at, last_seen_at, tokens_attempted, tokens_sent, tokens_deduplicated, + input_tokens_real, cache_read_tokens, cache_creation_tokens, requests_total, loop_detected, kill_events, kill_reason, status FROM agent_sessions @@ -155,6 +175,7 @@ func (t *Tracker) SessionByID(sessionID string) (SessionRow, bool, error) { `, sessionID).Scan( &r.SessionID, &r.StartedAt, &r.LastSeenAt, &r.TokensAttempted, &r.TokensSent, &r.TokensDeduplicated, + &r.InputTokensReal, &r.CacheReadTokens, &r.CacheCreationTokens, &r.RequestsTotal, &r.LoopDetected, &r.KillEvents, &r.KillReason, &r.Status, ) @@ -173,6 +194,7 @@ func (t *Tracker) Sessions() ([]SessionRow, error) { rows, err := t.db.Query(` SELECT session_id, started_at, last_seen_at, tokens_attempted, tokens_sent, tokens_deduplicated, + input_tokens_real, cache_read_tokens, cache_creation_tokens, requests_total, loop_detected, kill_events, kill_reason, status FROM agent_sessions @@ -189,6 +211,7 @@ func (t *Tracker) Sessions() ([]SessionRow, error) { if err := rows.Scan( &r.SessionID, &r.StartedAt, &r.LastSeenAt, &r.TokensAttempted, &r.TokensSent, &r.TokensDeduplicated, + &r.InputTokensReal, &r.CacheReadTokens, &r.CacheCreationTokens, &r.RequestsTotal, &r.LoopDetected, &r.KillEvents, &r.KillReason, &r.Status, ); err != nil { @@ -227,17 +250,20 @@ func (t *Tracker) KillLog() ([]KillRow, error) { // in-memory store and the /v1/agent-sessions API response. func ToAgentSession(r SessionRow) telemetry.AgentSession { return telemetry.AgentSession{ - SessionID: r.SessionID, - StartedAt: r.StartedAt, - LastSeenAt: r.LastSeenAt, - TokensAttempted: r.TokensAttempted, - TokensSent: r.TokensSent, - TokensDeduplicated: r.TokensDeduplicated, - RequestsTotal: int(r.RequestsTotal), - LoopDetected: int(r.LoopDetected), - KillEvents: int(r.KillEvents), - KillReason: r.KillReason, - Status: r.Status, + SessionID: r.SessionID, + StartedAt: r.StartedAt, + LastSeenAt: r.LastSeenAt, + TokensAttempted: r.TokensAttempted, + TokensSent: r.TokensSent, + TokensDeduplicated: r.TokensDeduplicated, + InputTokensReal: r.InputTokensReal, + CacheReadTokens: r.CacheReadTokens, + CacheCreationTokens: r.CacheCreationTokens, + RequestsTotal: int(r.RequestsTotal), + LoopDetected: int(r.LoopDetected), + KillEvents: int(r.KillEvents), + KillReason: r.KillReason, + Status: r.Status, } } @@ -311,26 +337,31 @@ func (t *Tracker) write(ev writeEvent) { INSERT INTO agent_sessions (session_id, started_at, last_seen_at, tokens_attempted, tokens_sent, tokens_deduplicated, + input_tokens_real, cache_read_tokens, cache_creation_tokens, requests_total, loop_detected, kill_events, kill_reason, status) VALUES - (?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?) + (?, ?, ?, ?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?) ON CONFLICT(session_id) DO UPDATE SET - last_seen_at = excluded.last_seen_at, - tokens_attempted = tokens_attempted + excluded.tokens_attempted, - tokens_sent = tokens_sent + excluded.tokens_sent, - tokens_deduplicated = tokens_deduplicated + excluded.tokens_deduplicated, - requests_total = requests_total + 1, - loop_detected = loop_detected + excluded.loop_detected, - kill_events = kill_events + excluded.kill_events, - kill_reason = CASE WHEN excluded.kill_reason != '' - THEN excluded.kill_reason - ELSE kill_reason END, - status = CASE WHEN excluded.status = 'killed' - THEN 'killed' - ELSE status END + last_seen_at = excluded.last_seen_at, + tokens_attempted = tokens_attempted + excluded.tokens_attempted, + tokens_sent = tokens_sent + excluded.tokens_sent, + tokens_deduplicated = tokens_deduplicated + excluded.tokens_deduplicated, + input_tokens_real = input_tokens_real + excluded.input_tokens_real, + cache_read_tokens = cache_read_tokens + excluded.cache_read_tokens, + cache_creation_tokens = cache_creation_tokens + excluded.cache_creation_tokens, + requests_total = requests_total + 1, + loop_detected = loop_detected + excluded.loop_detected, + kill_events = kill_events + excluded.kill_events, + kill_reason = CASE WHEN excluded.kill_reason != '' + THEN excluded.kill_reason + ELSE kill_reason END, + status = CASE WHEN excluded.status = 'killed' + THEN 'killed' + ELSE status END `, ev.sessionID, ev.ts, ev.ts, ev.outcome.TokensAttempted, ev.outcome.TokensSent, tokensDeduplicated, + ev.outcome.InputTokensReal, ev.outcome.CacheReadTokens, ev.outcome.CacheCreationTokens, loopInc, killInc, killReason, status, ) if err != nil { diff --git a/gateway/internal/telemetry/agentsession.go b/gateway/internal/telemetry/agentsession.go index b80eea5..a3b8e7c 100644 --- a/gateway/internal/telemetry/agentsession.go +++ b/gateway/internal/telemetry/agentsession.go @@ -16,11 +16,17 @@ type AgentSession struct { TokensAttempted int64 `json:"tokens_attempted"` TokensSent int64 `json:"tokens_sent"` TokensDeduplicated int64 `json:"tokens_deduplicated"` - RequestsTotal int `json:"requests_total"` - LoopDetected int `json:"loop_detected"` // velocity/circuit warn events - KillEvents int `json:"kill_events"` // guard blocks (HTTP 429) - KillReason string `json:"kill_reason"` // most recent block reason - Status string `json:"status"` // "active" | "killed" | "ended" + // Measured upstream token usage (ground truth from Anthropic's usage object, + // not byte estimates). InputTokensReal is fresh+cache-creation+cache-read; + // CacheReadTokens is context served from Anthropic's prompt cache. + InputTokensReal int64 `json:"input_tokens_real"` + CacheReadTokens int64 `json:"cache_read_tokens"` + CacheCreationTokens int64 `json:"cache_creation_tokens"` + RequestsTotal int `json:"requests_total"` + LoopDetected int `json:"loop_detected"` // velocity/circuit warn events + KillEvents int `json:"kill_events"` // guard blocks (HTTP 429) + KillReason string `json:"kill_reason"` // most recent block reason + Status string `json:"status"` // "active" | "killed" | "ended" } // KillEvent is emitted whenever a guard blocks a request (HTTP 429). @@ -40,9 +46,15 @@ type RequestOutcome struct { TokensAttempted int TokensSent int TokensSaved int // tokens stripped by the optimizer (deduplicated) - GuardReason string - Killed bool // guard returned !Allow (HTTP 429 was sent) - Warned bool // guard returned Allow+Warn + // Measured upstream usage for this request (0 when nothing was sent upstream, + // e.g. synthetic probe / cache replay). InputTokensReal = fresh + cache_creation + // + cache_read; CacheReadTokens = context served from Anthropic's prompt cache. + InputTokensReal int + CacheReadTokens int + CacheCreationTokens int + GuardReason string + Killed bool // guard returned !Allow (HTTP 429 was sent) + Warned bool // guard returned Allow+Warn } // AgentSessionStore tracks per-session observability in memory. @@ -94,6 +106,9 @@ func (s *AgentSessionStore) Record(sessionKey string, out RequestOutcome) { entry.TokensAttempted += int64(out.TokensAttempted) entry.TokensSent += int64(out.TokensSent) entry.TokensDeduplicated += int64(out.TokensSaved) + entry.InputTokensReal += int64(out.InputTokensReal) + entry.CacheReadTokens += int64(out.CacheReadTokens) + entry.CacheCreationTokens += int64(out.CacheCreationTokens) if out.Killed { entry.KillEvents++ diff --git a/gateway/internal/telemetry/telemetry.go b/gateway/internal/telemetry/telemetry.go index 70250eb..5e595a8 100644 --- a/gateway/internal/telemetry/telemetry.go +++ b/gateway/internal/telemetry/telemetry.go @@ -46,6 +46,9 @@ type UsageEvent struct { InputTokensAttempted int `json:"input_tokens_attempted"` InputTokensSent int `json:"input_tokens_sent"` TokensSaved int `json:"tokens_saved"` + InputTokensReal int `json:"input_tokens_real"` + CacheReadTokens int `json:"cache_read_tokens"` + CacheCreationTokens int `json:"cache_creation_tokens"` ReductionRatio float64 `json:"reduction_ratio"` BlocksAnalyzed int `json:"blocks_analyzed"` BlocksPruned int `json:"blocks_pruned"` From 0a3a00c7ce61a41e229f85951087a139ae6f0df6 Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Mon, 8 Jun 2026 19:42:06 -0500 Subject: [PATCH 10/15] feat: session summary leads with real prompt-cache metrics, byte-estimate demoted to diagnostic --- gateway/cmd/iq/main.go | 64 +++++++++++++++++++++++++++++++----------- 1 file changed, 47 insertions(+), 17 deletions(-) diff --git a/gateway/cmd/iq/main.go b/gateway/cmd/iq/main.go index 193563f..8db6a72 100644 --- a/gateway/cmd/iq/main.go +++ b/gateway/cmd/iq/main.go @@ -216,35 +216,65 @@ func printSessionSummary(sessionID string) { return } - percent := 0.0 - if tokensAttempted > 0 { - percent = float64(tokensDeduplicated) / float64(tokensAttempted) * 100 - } - // Prompt-cache hit ratio: of the input the model actually billed this session, - // what fraction Anthropic served from its prompt cache (≈10% cost, the latency - // win) instead of re-reading at full price. This is the real, measured signal - // for a subscription user — not a fabricated dollar figure. + // the fraction Anthropic served from its prompt cache (≈10% cost) instead of + // re-reading at full price. The headline signal for a subscription user — real, + // measured from upstream usage, not a byte estimate or a fabricated dollar figure. cacheHitRatio := 0.0 if inputTokensReal > 0 { cacheHitRatio = float64(cacheReadTokens) / float64(inputTokensReal) * 100 } + // Fresh = the only portion billed at full input price this session. + freshInput := inputTokensReal - cacheReadTokens - cacheCreationTokens + if freshInput < 0 { + freshInput = 0 + } + // Estimated optimizer pruning (byte-based, pre-cache). A diagnostic only — it + // does NOT correspond to tokens Anthropic actually billed, so it is dimmed and + // kept below the real metrics rather than shown as "saved". + estPercent := 0.0 + if tokensAttempted > 0 { + estPercent = float64(tokensDeduplicated) / float64(tokensAttempted) * 100 + } + + const ( + cyan = "\033[1;36m" + white = "\033[1;37m" + green = "\033[1;32m" + purple = "\033[1;35m" + grey = "\033[90m" + ) + row := func(label, color, value string) { + // Inner width 56 (2 + 20 label + 3 + 31 value) to match the box borders. + fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : %s%-31s\033[0m\033[1;36m│\033[0m\n", label, color, value) + } + divider := func() { + fmt.Fprintln(os.Stderr, " \033[1;36m├────────────────────────────────────────────────────────┤\033[0m") + } fmt.Fprintln(os.Stderr) fmt.Fprintln(os.Stderr, " \033[1;36m┌────────────────────────────────────────────────────────┐\033[0m") fmt.Fprintln(os.Stderr, " \033[1;36m│\033[0m \033[1;37mIndexQube Session Summary\033[0m \033[1;36m│\033[0m") - fmt.Fprintln(os.Stderr, " \033[1;36m├────────────────────────────────────────────────────────┤\033[0m") - fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;32m%-29d\033[0m\033[1;36m│\033[0m\n", "Requests Processed", requestsTotal) - fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;37m%-29s\033[0m\033[1;36m│\033[0m\n", "Tokens Attempted", formatNumber(tokensAttempted)) - fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;37m%-29s\033[0m\033[1;36m│\033[0m\n", "Tokens Sent", formatNumber(tokensSent)) - fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;35m%-29s\033[0m\033[1;36m│\033[0m\n", "Tokens Saved", fmt.Sprintf("%s (%.1f%%)", formatNumber(tokensDeduplicated), percent)) - fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : \033[1;36m%-29s\033[0m\033[1;36m│\033[0m\n", "Prompt-Cache Hits", fmt.Sprintf("%s tok (%.1f%%)", formatNumber(cacheReadTokens), cacheHitRatio)) - fmt.Fprintln(os.Stderr, " \033[1;36m├────────────────────────────────────────────────────────┤\033[0m") - statusColor := "\033[1;32m" // green + divider() + row("Requests", green, formatNumber(requestsTotal)) + if inputTokensReal > 0 { + // Real, Anthropic-reported numbers lead. + row("Prompt-Cache Hit", purple, fmt.Sprintf("%.1f%%", cacheHitRatio)) + row("Input Billed (real)", white, formatNumber(inputTokensReal)+" tok") + row(" from cache", cyan, formatNumber(cacheReadTokens)+" tok") + row(" cache write", white, formatNumber(cacheCreationTokens)+" tok") + row(" fresh (full price)", white, formatNumber(freshInput)+" tok") + } else { + row("Prompt-Cache Hit", grey, "— (no upstream usage)") + } + divider() + // Dimmed diagnostic, explicitly not a billing figure. + row("Optimizer (est.)", grey, fmt.Sprintf("%s tok (%.1f%%)", formatNumber(tokensDeduplicated), estPercent)) + statusColor := green if status == "killed" { statusColor = "\033[1;31m" // red } - fmt.Fprintf(os.Stderr, " \033[1;36m│\033[0m %-20s : %s%-29s\033[0m\033[1;36m│\033[0m\n", "Session Status", statusColor, strings.ToUpper(status)) + row("Session Status", statusColor, strings.ToUpper(status)) fmt.Fprintln(os.Stderr, " \033[1;36m└────────────────────────────────────────────────────────┘\033[0m") fmt.Fprintln(os.Stderr) } From 62a3439a3d55961d458820d6b041423389bf9271 Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Mon, 8 Jun 2026 20:01:36 -0500 Subject: [PATCH 11/15] feat: add iq bench, a subscription-native proxy-vs-direct A/B harness --- gateway/cmd/iq/main.go | 238 +++++++++++++++++++++++++++++++++--- gateway/cmd/iq/main_test.go | 74 +++++++++++ 2 files changed, 295 insertions(+), 17 deletions(-) create mode 100644 gateway/cmd/iq/main_test.go diff --git a/gateway/cmd/iq/main.go b/gateway/cmd/iq/main.go index 8db6a72..76b4f51 100644 --- a/gateway/cmd/iq/main.go +++ b/gateway/cmd/iq/main.go @@ -10,6 +10,7 @@ import ( "database/sql" "encoding/hex" "errors" + "flag" "fmt" "net" "net/http" @@ -47,6 +48,8 @@ func main() { case "claude": devMode, dumpPayloads, claudeArgs := parseClaudeFlags(os.Args[2:]) runClaude(claudeArgs, devMode, dumpPayloads) + case "bench": + runBench(os.Args[2:]) case "gemini": fmt.Println(" iq gemini — coming soon") case "codex": @@ -82,7 +85,12 @@ func runClaude(args []string, devMode, dumpPayloads bool) { os.Setenv("IQ_DEV_MODE", "1") fmt.Fprintln(os.Stderr, " [iq] dev mode — relaxed guards, full telemetry") } - sessionID := generateToken() + // Honor a caller-provided session ID (used by `iq bench` to locate each + // arm's recorded metrics afterward); otherwise generate one. + sessionID := os.Getenv("IQ_SESSION_ID") + if sessionID == "" { + sessionID = generateToken() + } if dumpPayloads { os.Setenv("IQ_DUMP_PAYLOADS", "1") repoRoot := findRepoRoot() @@ -167,8 +175,11 @@ func runClaude(args []string, devMode, dumpPayloads bool) { cancel() <-proxyDone - // Print beautiful colored session summary of token savings - printSessionSummary(sessionID) + // Print the colored session summary unless suppressed (e.g. by `iq bench`, + // which prints its own side-by-side comparison instead). + if os.Getenv("IQ_NO_SUMMARY") == "" { + printSessionSummary(sessionID) + } if runErr != nil { var exitErr *exec.ExitError @@ -180,41 +191,69 @@ func runClaude(args []string, devMode, dumpPayloads bool) { } } -func printSessionSummary(sessionID string) { +// sessionMetrics holds the real, Anthropic-reported usage recorded for one +// session, plus the byte-estimate diagnostic, read from sessions.db. +type sessionMetrics struct { + requests int64 + tokensAttempted int64 + tokensDeduplicated int64 + inputReal int64 + cacheRead int64 + cacheCreation int64 + status string +} + +func sessionsDBPath() string { home, err := os.UserHomeDir() if err != nil { - return + return "" + } + return filepath.Join(home, ".indexqube", "sessions.db") +} + +// readSessionMetrics loads the most recent agent_sessions row whose ID matches +// the given session (the recorded ID is suffixed with sessionID[:8]). +func readSessionMetrics(dbPath, sessionID string) (sessionMetrics, bool) { + var m sessionMetrics + if dbPath == "" { + return m, false } - dbPath := filepath.Join(home, ".indexqube", "sessions.db") db, err := sql.Open("sqlite", dbPath) if err != nil { - return + return m, false } defer db.Close() - - var requestsTotal, tokensAttempted, tokensSent, tokensDeduplicated int64 - var inputTokensReal, cacheReadTokens, cacheCreationTokens int64 - var status string queryID := "%" if len(sessionID) >= 8 { queryID = "%-" + sessionID[:8] } err = db.QueryRow(` - SELECT requests_total, tokens_attempted, tokens_sent, tokens_deduplicated, + SELECT requests_total, tokens_attempted, tokens_deduplicated, input_tokens_real, cache_read_tokens, cache_creation_tokens, status FROM agent_sessions WHERE session_id LIKE ? ORDER BY last_seen_at DESC LIMIT 1 - `, queryID).Scan(&requestsTotal, &tokensAttempted, &tokensSent, &tokensDeduplicated, - &inputTokensReal, &cacheReadTokens, &cacheCreationTokens, &status) + `, queryID).Scan(&m.requests, &m.tokensAttempted, &m.tokensDeduplicated, + &m.inputReal, &m.cacheRead, &m.cacheCreation, &m.status) if err != nil { - return + return m, false } + return m, true +} - if requestsTotal == 0 { +func printSessionSummary(sessionID string) { + m, ok := readSessionMetrics(sessionsDBPath(), sessionID) + if !ok || m.requests == 0 { return } + requestsTotal := m.requests + tokensAttempted := m.tokensAttempted + tokensDeduplicated := m.tokensDeduplicated + inputTokensReal := m.inputReal + cacheReadTokens := m.cacheRead + cacheCreationTokens := m.cacheCreation + status := m.status // Prompt-cache hit ratio: of the input the model actually billed this session, // the fraction Anthropic served from its prompt cache (≈10% cost) instead of @@ -295,6 +334,164 @@ func formatNumber(n int64) string { return string(out) } +// runBench drives the same prompt through two arms of the local proxy — an +// `optimize` arm (proxy) and an `observe` arm (direct passthrough) — both using +// the user's real subscription/OAuth auth (no API key), then prints the real +// Anthropic-reported usage side by side. This is the proxy-vs-direct A/B. +func runBench(args []string) { + fs := flag.NewFlagSet("bench", flag.ExitOnError) + prompt := fs.String("prompt", "List the files in the current directory, then stop.", "prompt used to drive both arms") + cooldown := fs.Duration("cooldown", 0, "wait between arms so Anthropic's prompt cache expires for an uncontaminated read (e.g. 6m)") + if err := fs.Parse(args); err != nil { + os.Exit(2) + } + + self, err := os.Executable() + if err != nil || self == "" { + self = os.Args[0] + } + + // Optimize first (cold cache), direct second. With no cooldown the 2nd arm + // may read the 1st arm's cache, which inflates the DIRECT column — a + // conservative bias for any proxy-savings claim. + type arm struct { + label string + mode string + sessionID string + m sessionMetrics + } + arms := []*arm{ + {label: "proxy (optimize)", mode: "optimize"}, + {label: "direct (observe)", mode: "observe"}, + } + + for i, a := range arms { + a.sessionID = generateToken() + fmt.Fprintf(os.Stderr, "\n [iq bench] arm %d/%d — %s: claude -p %q\n", i+1, len(arms), a.label, *prompt) + cmd := exec.Command(self, "claude", "-p", *prompt) //nolint:gosec + cmd.Env = append(os.Environ(), + "INDEXQUBE_MODE="+a.mode, + "IQ_SESSION_ID="+a.sessionID, + "IQ_NO_SUMMARY=1", + ) + cmd.Stdout = os.Stderr + cmd.Stderr = os.Stderr + if err := cmd.Run(); err != nil { + fmt.Fprintf(os.Stderr, " [iq bench] arm %q exited with error: %v\n", a.label, err) + } + if i < len(arms)-1 && *cooldown > 0 { + fmt.Fprintf(os.Stderr, " [iq bench] cooldown %s (waiting for Anthropic cache TTL)…\n", *cooldown) + time.Sleep(*cooldown) + } + } + + dbPath := sessionsDBPath() + for _, a := range arms { + a.m, _ = readSessionMetrics(dbPath, a.sessionID) + } + // arms[0]=proxy, arms[1]=direct. + printBenchComparison(arms[1].label, arms[1].m, arms[0].label, arms[0].m, *cooldown > 0) +} + +func benchCacheRatio(m sessionMetrics) float64 { + if m.inputReal <= 0 { + return 0 + } + return float64(m.cacheRead) / float64(m.inputReal) * 100 +} + +func benchFreshInput(m sessionMetrics) int64 { + f := m.inputReal - m.cacheRead - m.cacheCreation + if f < 0 { + f = 0 + } + return f +} + +func signedNumber(n int64) string { + switch { + case n > 0: + return "+" + formatNumber(n) + case n < 0: + return "-" + formatNumber(-n) + default: + return "0" + } +} + +func printBenchComparison(directLabel string, d sessionMetrics, proxyLabel string, p sessionMetrics, cooldownApplied bool) { + const ( + green = "\033[1;32m" + red = "\033[1;31m" + grey = "\033[90m" + bold = "\033[1;37m" + reset = "\033[0m" + ) + w := os.Stderr + fmt.Fprintln(w) + fmt.Fprintf(w, " %sIndexQube A/B — %s vs %s%s\n", bold, proxyLabel, directLabel, reset) + fmt.Fprintln(w, " ──────────────────────────────────────────────────────────────") + + if d.requests == 0 || p.requests == 0 { + fmt.Fprintf(w, " %s⚠ one arm recorded no requests (direct=%d, proxy=%d) — cannot compare.%s\n", red, d.requests, p.requests, reset) + fmt.Fprintln(w, " Check that `claude -p` ran successfully for both arms.") + fmt.Fprintln(w) + return + } + + fmt.Fprintf(w, " %-22s %14s %14s %14s\n", "Metric", "Direct", "Proxy", "Δ proxy−direct") + + // row prints one metric. lowerBetter colors a negative delta green. + row := func(label string, dv, pv int64, lowerBetter bool) { + delta := pv - dv + col := grey + if delta != 0 && lowerBetter { + if delta < 0 { + col = green + } else { + col = red + } + } + fmt.Fprintf(w, " %-22s %14s %14s %s%14s%s\n", + label, formatNumber(dv), formatNumber(pv), col, signedNumber(delta), reset) + } + + row("Requests", d.requests, p.requests, false) + row("Real input (tok)", d.inputReal, p.inputReal, true) + row(" cache read", d.cacheRead, p.cacheRead, false) + row(" cache write", d.cacheCreation, p.cacheCreation, false) + row(" fresh (full price)", benchFreshInput(d), benchFreshInput(p), true) + + dr, pr := benchCacheRatio(d), benchCacheRatio(p) + rcol := grey + if pr-dr > 0.05 { + rcol = green + } else if dr-pr > 0.05 { + rcol = red + } + fmt.Fprintf(w, " %-22s %13.1f%% %13.1f%% %s%+13.1f%s\n", + "Cache-hit ratio", dr, pr, rcol, pr-dr, "pp"+reset) + + fmt.Fprintln(w, " ──────────────────────────────────────────────────────────────") + + switch { + case p.inputReal < d.inputReal: + saved := d.inputReal - p.inputReal + pct := float64(saved) / float64(d.inputReal) * 100 + fmt.Fprintf(w, " %s✓ Proxy billed %s fewer real input tokens than direct (%.1f%%).%s\n", green, formatNumber(saved), pct, reset) + case p.inputReal > d.inputReal: + extra := p.inputReal - d.inputReal + fmt.Fprintf(w, " %s✗ Proxy billed %s MORE real input than direct — investigate.%s\n", red, formatNumber(extra), reset) + default: + fmt.Fprintf(w, " %s≈ Proxy and direct billed the same real input (parity).%s\n", grey, reset) + } + + if !cooldownApplied { + fmt.Fprintf(w, " %snote: no cooldown — the 2nd arm may have read the 1st arm's cache.\n Re-run with --cooldown 6m for an uncontaminated read.%s\n", grey, reset) + } + fmt.Fprintln(w) +} + func printHelp() { fmt.Print(` iq — IndexQube CLI @@ -305,6 +502,9 @@ func printHelp() { iq claude --dev Start Claude Code, dev mode (relaxed guards) iq claude --dump-payloads Dump Anthropic payloads to .indexqube/dumps/iq-session-*.jsonl + iq bench Proxy-vs-direct A/B: runs the same prompt in optimize + and observe modes, prints real Anthropic usage side by side + (flags: --prompt "...", --cooldown 6m) iq gemini Gemini via IndexQube (coming soon) iq codex Codex via IndexQube (coming soon) iq help Show this help @@ -323,7 +523,11 @@ func startProxy(ctx context.Context, ln net.Listener) { if os.Getenv("INDEXQUBE_DEV_TOKEN") == "" { os.Setenv("INDEXQUBE_DEV_TOKEN", generateToken()) } - os.Setenv("INDEXQUBE_MODE", "optimize") + // Respect an explicit mode (e.g. `iq bench` runs an observe arm and an + // optimize arm); default to optimize for normal `iq claude`. + if os.Getenv("INDEXQUBE_MODE") == "" { + os.Setenv("INDEXQUBE_MODE", "optimize") + } os.Setenv("INDEXQUBE_ENABLE_BLOCK_OPTIMIZER", "true") // Bind admin server to an ephemeral port so multiple iq instances // don't collide on the default 9100. diff --git a/gateway/cmd/iq/main_test.go b/gateway/cmd/iq/main_test.go new file mode 100644 index 0000000..c1a4ece --- /dev/null +++ b/gateway/cmd/iq/main_test.go @@ -0,0 +1,74 @@ +package main + +import ( + "database/sql" + "path/filepath" + "testing" + + _ "modernc.org/sqlite" +) + +// TestReadSessionMetrics covers the reusable DB read that both the session +// summary and the `iq bench` A/B comparison depend on: a row is located by the +// IQ_SESSION_ID suffix the proxy appends to session_id, and the real +// Anthropic-reported usage is returned (not the byte estimate). +func TestReadSessionMetrics(t *testing.T) { + dir := t.TempDir() + dbPath := filepath.Join(dir, "sessions.db") + db, err := sql.Open("sqlite", dbPath) + if err != nil { + t.Fatalf("open: %v", err) + } + defer db.Close() + if _, err := db.Exec(`CREATE TABLE agent_sessions ( + session_id TEXT PRIMARY KEY, + last_seen_at INTEGER NOT NULL DEFAULT 0, + requests_total INTEGER NOT NULL DEFAULT 0, + tokens_attempted INTEGER NOT NULL DEFAULT 0, + tokens_sent INTEGER NOT NULL DEFAULT 0, + tokens_deduplicated INTEGER NOT NULL DEFAULT 0, + input_tokens_real INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + status TEXT NOT NULL DEFAULT 'active' + )`); err != nil { + t.Fatalf("create: %v", err) + } + + sessionID := "abcdef1234567890" // proxy stores session_id ending in sessionID[:8] + if _, err := db.Exec(`INSERT INTO agent_sessions + (session_id, last_seen_at, requests_total, tokens_attempted, tokens_deduplicated, + input_tokens_real, cache_read_tokens, cache_creation_tokens, status) + VALUES (?,?,?,?,?,?,?,?,?)`, + "deadbeef-"+sessionID[:8], 100, 11, 395276, 79, 612374, 521776, 87694, "active"); err != nil { + t.Fatalf("insert: %v", err) + } + + m, ok := readSessionMetrics(dbPath, sessionID) + if !ok { + t.Fatalf("readSessionMetrics ok=false, want true") + } + if m.requests != 11 || m.inputReal != 612374 || m.cacheRead != 521776 || m.cacheCreation != 87694 { + t.Fatalf("metrics = %+v", m) + } + if got := benchCacheRatio(m); got < 85.1 || got > 85.3 { + t.Fatalf("cache ratio = %.2f, want ~85.2", got) + } + if got := benchFreshInput(m); got != 612374-521776-87694 { + t.Fatalf("fresh = %d, want %d", got, 612374-521776-87694) + } + + // Unknown session must report not-found rather than a bogus row. + if _, ok := readSessionMetrics(dbPath, "0000000000000000"); ok { + t.Fatalf("expected not-found for unknown session") + } +} + +func TestSignedNumber(t *testing.T) { + cases := map[int64]string{0: "0", 1234: "+1,234", -1234: "-1,234"} + for in, want := range cases { + if got := signedNumber(in); got != want { + t.Fatalf("signedNumber(%d) = %q, want %q", in, got, want) + } + } +} From edfc2606c513c1d1a805dedbbe63f65faf0e5658 Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Mon, 8 Jun 2026 20:19:51 -0500 Subject: [PATCH 12/15] fix: iq bench reports cost-weighted effective input so cache-busting isn't hidden --- gateway/cmd/iq/main.go | 29 +++++++++++++++++++++-------- gateway/cmd/iq/main_test.go | 4 ++++ 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/gateway/cmd/iq/main.go b/gateway/cmd/iq/main.go index 76b4f51..5a9a7e1 100644 --- a/gateway/cmd/iq/main.go +++ b/gateway/cmd/iq/main.go @@ -408,6 +408,14 @@ func benchFreshInput(m sessionMetrics) int64 { return f } +// benchEffectiveInput weights input by Anthropic's standard multipliers +// (cache read ≈0.1×, cache write ≈1.25×, fresh =1.0×) so the comparison reflects +// real cost / rate-limit weight. Raw token sums hide that a cache miss (write) +// costs ~12× a cache hit (read) — which is exactly the optimizer's failure mode. +func benchEffectiveInput(m sessionMetrics) int64 { + return int64(float64(m.cacheRead)*0.1 + float64(m.cacheCreation)*1.25 + float64(benchFreshInput(m))) +} + func signedNumber(n int64) string { switch { case n > 0: @@ -472,19 +480,24 @@ func printBenchComparison(directLabel string, d sessionMetrics, proxyLabel strin fmt.Fprintf(w, " %-22s %13.1f%% %13.1f%% %s%+13.1f%s\n", "Cache-hit ratio", dr, pr, rcol, pr-dr, "pp"+reset) + // The bottom line: cost-weighted input. This is what actually moves billing + // and rate-limit headroom (read×0.1, write×1.25, fresh×1.0). + row("Effective (cost-wt)", benchEffectiveInput(d), benchEffectiveInput(p), true) + fmt.Fprintln(w, " ──────────────────────────────────────────────────────────────") + de, pe := benchEffectiveInput(d), benchEffectiveInput(p) switch { - case p.inputReal < d.inputReal: - saved := d.inputReal - p.inputReal - pct := float64(saved) / float64(d.inputReal) * 100 - fmt.Fprintf(w, " %s✓ Proxy billed %s fewer real input tokens than direct (%.1f%%).%s\n", green, formatNumber(saved), pct, reset) - case p.inputReal > d.inputReal: - extra := p.inputReal - d.inputReal - fmt.Fprintf(w, " %s✗ Proxy billed %s MORE real input than direct — investigate.%s\n", red, formatNumber(extra), reset) + case pe < de: + pct := float64(de-pe) / float64(de) * 100 + fmt.Fprintf(w, " %s✓ Proxy's cost-weighted input is %s lower than direct (%.1f%%).%s\n", green, formatNumber(de-pe), pct, reset) + case pe > de: + pct := float64(pe-de) / float64(de) * 100 + fmt.Fprintf(w, " %s✗ Proxy's cost-weighted input is %s HIGHER than direct (%.1f%%) — the optimizer is busting the cache.%s\n", red, formatNumber(pe-de), pct, reset) default: - fmt.Fprintf(w, " %s≈ Proxy and direct billed the same real input (parity).%s\n", grey, reset) + fmt.Fprintf(w, " %s≈ Proxy and direct are at cost parity.%s\n", grey, reset) } + fmt.Fprintf(w, " %scost-wt: cache read×0.1, write×1.25, fresh×1.0 (Anthropic's standard weights)%s\n", grey, reset) if !cooldownApplied { fmt.Fprintf(w, " %snote: no cooldown — the 2nd arm may have read the 1st arm's cache.\n Re-run with --cooldown 6m for an uncontaminated read.%s\n", grey, reset) diff --git a/gateway/cmd/iq/main_test.go b/gateway/cmd/iq/main_test.go index c1a4ece..b6e01e9 100644 --- a/gateway/cmd/iq/main_test.go +++ b/gateway/cmd/iq/main_test.go @@ -57,6 +57,10 @@ func TestReadSessionMetrics(t *testing.T) { if got := benchFreshInput(m); got != 612374-521776-87694 { t.Fatalf("fresh = %d, want %d", got, 612374-521776-87694) } + // read*0.1 + write*1.25 + fresh = 52177.6 + 109617.5 + 2904 = 164699.1 + if got := benchEffectiveInput(m); got != 164699 { + t.Fatalf("effective = %d, want 164699", got) + } // Unknown session must report not-found rather than a bogus row. if _, ok := readSessionMetrics(dbPath, "0000000000000000"); ok { From e5bedd6cea81a7261d7c52d7fb9fa10af4e845ce Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Mon, 8 Jun 2026 20:33:33 -0500 Subject: [PATCH 13/15] fix: forward cache-managed requests byte-for-byte so optimize mode stops busting the prompt cache --- gateway/internal/proxy/claude_messages.go | 20 +++++++++ gateway/internal/proxy/usage_test.go | 54 +++++++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/gateway/internal/proxy/claude_messages.go b/gateway/internal/proxy/claude_messages.go index f8c2340..6a617e6 100644 --- a/gateway/internal/proxy/claude_messages.go +++ b/gateway/internal/proxy/claude_messages.go @@ -116,6 +116,11 @@ type claudeOptimizerStats struct { PreservedCachePrefixBytes int `json:"preserved_cache_prefix_bytes"` PreservedCachePrefixCount int `json:"preserved_cache_prefix_count"` + // PreservedCacheFidelity is set when the whole body was forwarded byte-for-byte + // because the client manages prompt caching (cache_control present). Any rewrite + // would re-serialize the cached prefix and bust Anthropic's cache. + PreservedCacheFidelity bool `json:"preserved_cache_fidelity"` + // Size tracking. LargestSpanBytes int `json:"largest_span_bytes"` LargestPrunedBytes int `json:"largest_pruned_bytes"` @@ -1230,6 +1235,21 @@ func (p *Proxy) prepareClaudeBody(ctx context.Context, cfg ClaudeMessagesConfig, // blunt "skip for subscription" rule — subscription users keep their cache // benefit via Claude Code's own breakpoints, which E2 now preserves. requestHasCacheControl := systemCached || cachePrefixMsgIdx >= 0 + + // Cache fidelity beats pruning. When the client manages prompt caching — Claude + // Code sets a rolling cache_control breakpoint every turn — forward the body + // byte-for-byte. Any rewrite here re-serializes the whole request (Go sorts JSON + // map keys), changing the bytes of the cached prefix even when no prefix content + // is pruned. Anthropic then misses the cache every turn, turning 0.1x cache reads + // into 1.25x writes. `iq bench` measured this at ~6x more expensive than going + // direct (optimize 0% hit vs observe 91%). The span loop above already recorded + // blocks for dedup, and redundant-call elimination runs upstream of this; pruning + // is only safe to attempt when the client is NOT caching. + if requestHasCacheControl { + stats.PreservedCacheFidelity = true + return body, req, stats, shape, nil + } + wantPromptCache := cfg.Optimizer.EnablePromptCache && hasSystemSpan && systemAllKnown && !cfg.Bedrock.Enabled && !requestHasCacheControl if len(prunableSpans) == 0 && !wantPromptCache { return body, req, stats, shape, nil diff --git a/gateway/internal/proxy/usage_test.go b/gateway/internal/proxy/usage_test.go index bf9fadf..d099582 100644 --- a/gateway/internal/proxy/usage_test.go +++ b/gateway/internal/proxy/usage_test.go @@ -184,6 +184,60 @@ func TestClaudeMessages_PromptCachePrefixIsNotPruned(t *testing.T) { } } +// TestClaudeMessages_CacheControlRequestForwardedByteIdentical verifies that when +// the client manages prompt caching, the optimizer forwards the body byte-for-byte +// — even when it contains an older duplicate tool_result that sits AFTER the cache +// breakpoint and would otherwise be pruned. Re-marshaling such a request reorders +// JSON keys in the cached prefix and busts Anthropic's cache (measured ~6x worse +// than direct), so byte fidelity must win. +func TestClaudeMessages_CacheControlRequestForwardedByteIdentical(t *testing.T) { + t.Parallel() + p := New(&fakeGovernor{}) + cfg := ClaudeMessagesConfig{ + Mode: "optimize", + EnableBlockOptimizer: true, + SessionStore: memory.NewStore(time.Hour), + Optimizer: OptimizerConfig{ + MinSpanBytes: 512, + EnableToolResultPruning: true, + }, + } + // Breakpoint on messages[1]. The same tool_result content appears at messages + // [1], [3], [5]; the copy at [3] is an older duplicate AFTER the breakpoint and + // is NOT the last occurrence, so without the cache-fidelity short-circuit the + // optimizer would prune it and re-marshal the whole body. + fileBody := strings.Repeat("ordinary source code output line\n", 80) + body := []byte(fmt.Sprintf( + `{"model":"claude-sonnet-4-6","messages":[`+ + `{"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"Read","input":{"file_path":"/repo/a.go"}}]},`+ + `{"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":%q,"cache_control":{"type":"ephemeral"}}]},`+ + `{"role":"assistant","content":[{"type":"tool_use","id":"t2","name":"Read","input":{"file_path":"/repo/a.go"}}]},`+ + `{"role":"user","content":[{"type":"tool_result","tool_use_id":"t2","content":%q}]},`+ + `{"role":"assistant","content":[{"type":"tool_use","id":"t3","name":"Read","input":{"file_path":"/repo/a.go"}}]},`+ + `{"role":"user","content":[{"type":"tool_result","tool_use_id":"t3","content":%q}]},`+ + `{"role":"user","content":[{"type":"text","text":"continue"}]}]}`, + fileBody, fileBody, fileBody, + )) + + // Warm the session so the duplicate is "known" and would be eligible to prune. + if _, _, _, _, err := p.prepareClaudeBody(context.Background(), cfg, "fidelity-test", body); err != nil { + t.Fatalf("first prepare: %v", err) + } + forward, _, stats, _, err := p.prepareClaudeBody(context.Background(), cfg, "fidelity-test", body) + if err != nil { + t.Fatalf("second prepare: %v", err) + } + if string(forward) != string(body) { + t.Fatalf("cache_control request must be forwarded byte-identical (no re-marshal); got:\n%s", forward) + } + if stats.BlocksPruned != 0 { + t.Fatalf("blocks_pruned=%d, want 0 (cache fidelity)", stats.BlocksPruned) + } + if !stats.PreservedCacheFidelity { + t.Fatalf("PreservedCacheFidelity=false, want true") + } +} + // promptCacheCfg builds an optimize-mode config with prompt-cache injection on. func promptCacheCfg() ClaudeMessagesConfig { return ClaudeMessagesConfig{ From 70c4f87e4c1943db3497bbba0623b6b0d3c51603 Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Mon, 8 Jun 2026 20:54:59 -0500 Subject: [PATCH 14/15] fix: iq bench guards against cache-TTL contamination --- gateway/cmd/iq/main.go | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/gateway/cmd/iq/main.go b/gateway/cmd/iq/main.go index 5a9a7e1..2e55893 100644 --- a/gateway/cmd/iq/main.go +++ b/gateway/cmd/iq/main.go @@ -341,7 +341,7 @@ func formatNumber(n int64) string { func runBench(args []string) { fs := flag.NewFlagSet("bench", flag.ExitOnError) prompt := fs.String("prompt", "List the files in the current directory, then stop.", "prompt used to drive both arms") - cooldown := fs.Duration("cooldown", 0, "wait between arms so Anthropic's prompt cache expires for an uncontaminated read (e.g. 6m)") + cooldown := fs.Duration("cooldown", 6*time.Minute, "wait between arms so Anthropic's prompt cache (5m TTL) expires for an uncontaminated read; below 5m the comparison is contaminated") if err := fs.Parse(args); err != nil { os.Exit(2) } @@ -390,9 +390,14 @@ func runBench(args []string) { a.m, _ = readSessionMetrics(dbPath, a.sessionID) } // arms[0]=proxy, arms[1]=direct. - printBenchComparison(arms[1].label, arms[1].m, arms[0].label, arms[0].m, *cooldown > 0) + printBenchComparison(arms[1].label, arms[1].m, arms[0].label, arms[0].m, *cooldown) } +// anthropicCacheTTL is Anthropic's prompt-cache lifetime. A bench cooldown below +// this lets the second arm read the first arm's still-warm cache, contaminating +// the comparison (it freeloads, looking artificially cheap). +const anthropicCacheTTL = 5 * time.Minute + func benchCacheRatio(m sessionMetrics) float64 { if m.inputReal <= 0 { return 0 @@ -427,7 +432,7 @@ func signedNumber(n int64) string { } } -func printBenchComparison(directLabel string, d sessionMetrics, proxyLabel string, p sessionMetrics, cooldownApplied bool) { +func printBenchComparison(directLabel string, d sessionMetrics, proxyLabel string, p sessionMetrics, cooldown time.Duration) { const ( green = "\033[1;32m" red = "\033[1;31m" @@ -488,20 +493,21 @@ func printBenchComparison(directLabel string, d sessionMetrics, proxyLabel strin de, pe := benchEffectiveInput(d), benchEffectiveInput(p) switch { + case cooldown < anthropicCacheTTL: + // The second arm (direct) ran while the first arm's cache was still warm + // and read it, so direct looks artificially cheap. Don't print a verdict. + fmt.Fprintf(w, " %s⚠ contaminated: cooldown %s < %s cache TTL — the 2nd arm read the 1st arm's\n warm cache, biasing the comparison against the proxy. Re-run with --cooldown 6m.%s\n", red, cooldown, anthropicCacheTTL, reset) + fmt.Fprintf(w, " %seach arm's own cache-hit ratio is still valid: proxy=%.1f%%, direct=%.1f%%%s\n", grey, pr, dr, reset) case pe < de: pct := float64(de-pe) / float64(de) * 100 fmt.Fprintf(w, " %s✓ Proxy's cost-weighted input is %s lower than direct (%.1f%%).%s\n", green, formatNumber(de-pe), pct, reset) case pe > de: pct := float64(pe-de) / float64(de) * 100 - fmt.Fprintf(w, " %s✗ Proxy's cost-weighted input is %s HIGHER than direct (%.1f%%) — the optimizer is busting the cache.%s\n", red, formatNumber(pe-de), pct, reset) + fmt.Fprintf(w, " %s✗ Proxy's cost-weighted input is %s HIGHER than direct (%.1f%%) — investigate.%s\n", red, formatNumber(pe-de), pct, reset) default: fmt.Fprintf(w, " %s≈ Proxy and direct are at cost parity.%s\n", grey, reset) } fmt.Fprintf(w, " %scost-wt: cache read×0.1, write×1.25, fresh×1.0 (Anthropic's standard weights)%s\n", grey, reset) - - if !cooldownApplied { - fmt.Fprintf(w, " %snote: no cooldown — the 2nd arm may have read the 1st arm's cache.\n Re-run with --cooldown 6m for an uncontaminated read.%s\n", grey, reset) - } fmt.Fprintln(w) } From e3e94d3f0a71a93935fe0fb890d5bfdfe56cc16c Mon Sep 17 00:00:00 2001 From: Revanth Ch Date: Mon, 8 Jun 2026 21:05:37 -0500 Subject: [PATCH 15/15] test: lock optimize==direct parity for cache-managed requests --- gateway/internal/proxy/usage_test.go | 38 ++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/gateway/internal/proxy/usage_test.go b/gateway/internal/proxy/usage_test.go index d099582..107ae90 100644 --- a/gateway/internal/proxy/usage_test.go +++ b/gateway/internal/proxy/usage_test.go @@ -238,6 +238,44 @@ func TestClaudeMessages_CacheControlRequestForwardedByteIdentical(t *testing.T) } } +// TestClaudeMessages_OptimizeMatchesObserveForCacheManagedRequest proves the +// proxy-vs-direct parity that a live A/B can't cleanly show (the two arms share +// Anthropic's warm cache and contaminate each other). For a cache-managed request +// optimize mode must forward the exact same bytes as observe (== direct) and as +// the original — so on the wire the proxy is a no-op and cannot bust the cache. +func TestClaudeMessages_OptimizeMatchesObserveForCacheManagedRequest(t *testing.T) { + t.Parallel() + body := []byte(`{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":[{"type":"text","text":"` + + strings.Repeat("stable project context line ", 60) + `","cache_control":{"type":"ephemeral"}}]}]}`) + + run := func(mode string) []byte { + p := New(&fakeGovernor{}) + cfg := ClaudeMessagesConfig{ + Mode: mode, + EnableBlockOptimizer: true, + SessionStore: memory.NewStore(time.Hour), + Optimizer: OptimizerConfig{MinSpanBytes: 512, EnableToolResultPruning: true}, + } + // Two prepares so any "known block" pruning would have fired by turn 2. + if _, _, _, _, err := p.prepareClaudeBody(context.Background(), cfg, "parity", body); err != nil { + t.Fatalf("%s first prepare: %v", mode, err) + } + out, _, _, _, err := p.prepareClaudeBody(context.Background(), cfg, "parity", body) + if err != nil { + t.Fatalf("%s second prepare: %v", mode, err) + } + return out + } + + obs, opt := run("observe"), run("optimize") + if string(opt) != string(obs) { + t.Fatalf("optimize must forward identical bytes to observe (direct) for cache-managed requests\nobserve=%s\noptimize=%s", obs, opt) + } + if string(opt) != string(body) { + t.Fatalf("cache-managed request must be byte-identical to the original; got:\n%s", opt) + } +} + // promptCacheCfg builds an optimize-mode config with prompt-cache injection on. func promptCacheCfg() ClaudeMessagesConfig { return ClaudeMessagesConfig{