diff --git a/benchmark/auto.go b/benchmark/auto.go index 541e45a..2e2e989 100644 --- a/benchmark/auto.go +++ b/benchmark/auto.go @@ -91,6 +91,51 @@ type AutoBenchmarkConfig struct { // remains the normal budget C (--concurrency). ReplayStopAtLowConcurrency bool + // UUID-based cache-coherency validation (--replay-inject-uuids). ROUTER + // PATH ONLY (cfg.RouterReplayFile != ""); the CLI rejects this combined + // with --from-dataset (see cli/benchmark_commands.go). One deterministic + // UUID is injected per user turn, spread through the conversation + // (rather than clustered at a session boundary) — see the package doc + // in replay_router_uuid.go for the full design. Each request recites a + // bounded WINDOW (first turn + up to 3 most-recent turns, excluding the + // current turn, capped at 4), keeping the recite cost/response budget + // constant regardless of session length. + ReplayInjectUUIDs bool + // ReplayUUIDSeed seeds the UUID generator (see newUUIDGenerator); 0 = crypto/rand + // (non-deterministic across runs). + ReplayUUIDSeed int64 + // ReplayReciteEveryRequest: ask the model to recite the first-line UUID + // window on EVERY request (default true), not just each instance's + // final request. + ReplayReciteEveryRequest bool + // replaySessionTurnHashes is the precomputed per-session ordered list of + // qualifying user-turn hashes, populated once by RunAutoBenchmark before + // any per-model goroutine spawns — see computeSessionTurnHashes. + // replaySessionTurnHashes[i][t] = session i's turn-t message hash. + // Shared, read-only, across every model in a multi-model run (same + // sharing rationale as replayConversations below). + replaySessionTurnHashes [][]string + // replayUUIDSets is the precomputed per-session-per-turn UUID + // assignment, populated once alongside replaySessionTurnHashes — see + // buildSessionTurnUUIDs. Index i = session i's ordered turn-UUID list + // (index t = turn t's UUID); i corresponds to seriesNum-1, the order + // sessions are dispatched in (see the sizing note at the router-replay + // precompute call site). Shared, read-only, across every model in a + // multi-model run so every model sees the identical assignment. + replayUUIDSets [][]string + // replayBlockSessionCounts maps a replay-v3 block hash to the number of + // DISTINCT SESSIONS that reference it (see computeBlockSessionCounts) — + // a hash referenced by more than one session is never eligible for + // turn-stamping (isQualifyingUserTurn requires count==1). Populated once + // alongside replayUUIDSets. + replayBlockSessionCounts map[string]int + // replayUUIDOwner is the precomputed reverse uuid -> owning-session-index + // map, populated once alongside replayUUIDSets — see + // buildSessionTurnUUIDs. Read-only after precompute; used by + // findLeakedUUIDsByOwner (replay_uuid.go) to flag cross-session + // contamination in O(response) time. + replayUUIDOwner map[string]int + // RunID is populated internally by RunAutoBenchmark at the start of each // run. It's the UUID injected into every conversation's system prompt // (when ReplayNoStamp is false). Per-run scope — conversations that share @@ -203,6 +248,18 @@ type requestDataRecord struct { Question string `json:"question,omitempty"` ResponseText string `json:"response_text,omitempty"` RawResponseTail string `json:"raw_response_tail,omitempty"` + + // UUID validation (router-replay --replay-inject-uuids only). The three + // counts are always populated (0 when the feature is off); the raw detail + // lists are populated ONLY on a miss or a leak (mirrors the + // failed-request-only policy above — avoid bloating every row). + UUIDExpected int `json:"uuid_expected"` + UUIDFound int `json:"uuid_found"` + UUIDLeaked int `json:"uuid_leaked"` + UUIDExactMatch bool `json:"uuid_exact_match"` + ExpectedUUIDsRaw []string `json:"expected_uuids_raw,omitempty"` + FoundMask []bool `json:"found_mask,omitempty"` + LeakedUUIDsRaw []string `json:"leaked_uuids_raw,omitempty"` } // requestDataWriter writes requestDataRecord entries as JSONL, safe for concurrent use. @@ -896,6 +953,18 @@ type autoState struct { coldStartTTFTCount atomic.Int64 // count of cold-start samples (used for series-scaling gate) ttftDegradedCount atomic.Int64 // requests disqualified from cache-hit by TTFT degradation + // UUID validation (replay --replay-inject-uuids only). All zero when the + // feature is off — recordReplayRequest only touches these when + // metrics.ExpectedUUIDs is non-empty. + valReqs atomic.Int64 // requests that carried >=1 expected UUID (i.e. validation ran) + valUUIDChecks atomic.Int64 // total per-UUID presence checks made + valUUIDFound atomic.Int64 // per-UUID presence checks that found the UUID + valExactMatchReqs atomic.Int64 // requests whose first line was the exact ordered UUID list (output conformity) + valPresenceMissUUIDs atomic.Int64 // per-UUID PRESENCE_MISS count (expected UUID absent) + valCrossContamUUIDs atomic.Int64 // per-UUID CROSS_CONTAMINATION count (other-conversation UUID present) + valPresenceMissReqs atomic.Int64 // requests with >=1 PRESENCE_MISS + valCrossContamReqs atomic.Int64 // requests with >=1 CROSS_CONTAMINATION + // Persistent early-sample buffers — never trimmed, survive stream eviction. printMu sync.Mutex // serialises --print-responses output across concurrent series @@ -989,6 +1058,17 @@ type autoBenchmarkResult struct { totalInputWarm int64 // input tokens from warm requests (prefix was cached) totalOutput int64 // output tokens across all requests totalCachedTokens int64 // server-reported cached prompt tokens + + // UUID validation (replay --replay-inject-uuids only); all zero when the + // feature is off. See autoState's val* atomics for field meanings. + valReqs int64 + valUUIDChecks int64 + valUUIDFound int64 + valExactMatchReqs int64 + valPresenceMissUUIDs int64 + valCrossContamUUIDs int64 + valPresenceMissReqs int64 + valCrossContamReqs int64 } // displaySnapshot is an atomic snapshot of state for the display goroutine. @@ -1140,6 +1220,19 @@ func printAutoSummary(res autoBenchmarkResult, cfg AutoBenchmarkConfig) { fmt.Println(strings.Repeat("-", 62)) fmt.Printf(" Total completed : %d\n", res.totalCompleted) fmt.Printf(" Total errors : %d\n", res.totalErrors) + if cfg.ReplayInjectUUIDs { + fmt.Println(strings.Repeat("-", 62)) + fmt.Println(" UUID validation (replay)") + fmt.Printf(" Requests validated : %d\n", res.valReqs) + // Two tests, mirroring the cache-coherency eval CLI's layout: UUID + // correctness (per-stamp presence, Contains anywhere in the response) + // and output conformity (first line is exactly the ordered, + // comma-joined UUID list — see firstLineConformity). + fmt.Printf(" UUID correctness (presence) : %d/%d\n", res.valUUIDFound, res.valUUIDChecks) + fmt.Printf(" Output conformity (first-line exact) : %d/%d\n", res.valExactMatchReqs, res.valReqs) + fmt.Printf(" PRESENCE_MISS (expected UUID absent) : %d across %d requests\n", res.valPresenceMissUUIDs, res.valPresenceMissReqs) + fmt.Printf(" CROSS_CONTAMINATION (other-conv) : %d across %d requests\n", res.valCrossContamUUIDs, res.valCrossContamReqs) + } fmt.Println(strings.Repeat("=", 62)) } @@ -1560,6 +1653,23 @@ func runSingleModelBenchmark( } pp.outputRatio = cfg.ReplayOutputRatio pp.forceOutput = cfg.ReplayForceOutput + // UUID cache-coherency injection (--replay-inject-uuids). Same + // global/read-only refs set on the per-instance poster in + // replay_router.go — every poster in the run, per-instance or + // pooled-per-endpoint, shares these identical slices/maps, so + // this poster (potentially shared across many concurrent + // sessions once picked by the router) can safely serve any + // session: do()/dryDo() derive the per-request sessionIdx from + // seriesNum and buildInjection looks up that session's data from + // these globals, with no mutable per-session state on the + // poster itself. + if cfg.ReplayInjectUUIDs { + pp.uuidEnabled = true + pp.allUUIDSets = cfg.replayUUIDSets + pp.sessionTurnHashes = cfg.replaySessionTurnHashes + pp.owner = cfg.replayUUIDOwner + pp.reciteEveryRequest = cfg.ReplayReciteEveryRequest + } posters[i] = pp } if ok { @@ -2323,6 +2433,14 @@ func runSingleModelBenchmark( res.totalInputWarm = tt.inputWarm res.totalOutput = tt.output res.totalCachedTokens = tt.cached + res.valReqs = st.valReqs.Load() + res.valUUIDChecks = st.valUUIDChecks.Load() + res.valUUIDFound = st.valUUIDFound.Load() + res.valExactMatchReqs = st.valExactMatchReqs.Load() + res.valPresenceMissUUIDs = st.valPresenceMissUUIDs.Load() + res.valCrossContamUUIDs = st.valCrossContamUUIDs.Load() + res.valPresenceMissReqs = st.valPresenceMissReqs.Load() + res.valCrossContamReqs = st.valCrossContamReqs.Load() // Send a final snapshot with termReason set so multi-model display shows DONE. { @@ -2515,6 +2633,82 @@ func RunAutoBenchmark(ctx context.Context, cfg AutoBenchmarkConfig) error { fmt.Printf("Header: %d sessions / %d instances / %d requests / %d fan-out turns / max fan-out %d\n\n", hdr.Summary.Sessions, hdr.Summary.Instances, hdr.Summary.Requests, hdr.Summary.FanOutTurns, hdr.Summary.MaxFanOutInOneTurn) + + // Precompute UUID cache-coherency injection ONCE here, before any + // per-model goroutine spawns below, so every model's + // runSingleModelBenchmark sees the identical per-session assignment + // (same sharing rationale as replayConversations above). + if cfg.ReplayInjectUUIDs { + // Effective session count: the SAME formula + // openRouterReplayStream uses to compute its own `total` (see + // replay_router.go), so the array is sized to exactly the + // number of sessions that will ever be dispatched (dispatch + // order == array index, since both walk the file in the same + // filtered, sequential order). This is what makes true + // lazy-growth unnecessary: the count is fully determined by + // the header + --replay-series + --replay-series-indices/ + // --replay-series-range before any session is ever pulled. + effectiveSessions := hdr.Summary.Sessions + if cfg.ReplaySeries > 0 && cfg.ReplaySeries < effectiveSessions { + effectiveSessions = cfg.ReplaySeries + } + if len(cfg.RouterReplaySeriesIndices) > 0 { + effectiveSessions = len(cfg.RouterReplaySeriesIndices) + if cfg.ReplaySeries > 0 && cfg.ReplaySeries < effectiveSessions { + effectiveSessions = cfg.ReplaySeries + } + } + if effectiveSessions <= 0 { + fmt.Fprintf(os.Stderr, + "[router-replay] WARNING: --replay-inject-uuids could not determine an effective session count (header sessions=%d) — UUID injection disabled for this run\n", + hdr.Summary.Sessions) + } else { + counts, cerr := computeBlockSessionCounts(cfg.RouterReplayFile, cfg.RouterReplaySeriesIndices, cfg.ReplaySeries) + if cerr != nil { + return fmt.Errorf("compute block session counts for --replay-inject-uuids: %w", cerr) + } + cfg.replayBlockSessionCounts = counts + + // Per-session ordered list of qualifying user-turn hashes + // (role=="user", has a text block, hash referenced by + // exactly this one session — see isQualifyingUserTurn), + // each turn getting its own UUID (buildSessionTurnUUIDs) + // rather than N stamps clustered at a session boundary. + turnHashes, terr := computeSessionTurnHashes(cfg.RouterReplayFile, cfg.RouterReplaySeriesIndices, cfg.ReplaySeries, counts) + if terr != nil { + return fmt.Errorf("compute session turn hashes for --replay-inject-uuids: %w", terr) + } + cfg.replaySessionTurnHashes = turnHashes + + turnCounts := make([]int, len(turnHashes)) + for i, h := range turnHashes { + turnCounts[i] = len(h) + } + cfg.replayUUIDSets, cfg.replayUUIDOwner = buildSessionTurnUUIDs(turnCounts, cfg.ReplayUUIDSeed) + + sharedHashes := 0 + for _, n := range counts { + if n > 1 { + sharedHashes++ + } + } + totalTurns, minTurns, maxTurns := 0, 0, 0 + if len(turnCounts) > 0 { + minTurns, maxTurns = turnCounts[0], turnCounts[0] + for _, n := range turnCounts { + totalTurns += n + if n < minTurns { + minTurns = n + } + if n > maxTurns { + maxTurns = n + } + } + } + fmt.Printf("UUID validation enabled: %d session(s) prepared, %d turn(s) total, %d-%d turns/session, %d cross-session-shared block hash(es), seed=%d, recite-every-request=%v\n", + effectiveSessions, totalTurns, minTurns, maxTurns, sharedHashes, cfg.ReplayUUIDSeed, cfg.ReplayReciteEveryRequest) + } + } } // Create per-run subdirectory for request data if configured. diff --git a/benchmark/replay.go b/benchmark/replay.go index 94d0020..dd91d96 100644 --- a/benchmark/replay.go +++ b/benchmark/replay.go @@ -101,7 +101,7 @@ func runReplaySeriesLoop( seriesNum := convIdx + 1 seriesGUID := conv.ID - fullyWalked := runReplayConversation(benchCtx, cfg, st, rdw, conv, seriesNum, seriesGUID, endpointOverride, reqTimeout, gate) + fullyWalked := runReplayConversation(benchCtx, cfg, st, rdw, conv, convIdx, seriesNum, seriesGUID, endpointOverride, reqTimeout, gate) // The conversation's slot is retired — its context no longer counts // toward the active dataset. st.datasetTracker.Reset(seriesNum) @@ -116,6 +116,9 @@ func runReplaySeriesLoop( // benchmark request per gpt turn. Errors on individual requests are recorded // but don't abort the series — the next turn still runs. // +// convIdx is this conversation's index into cfg.replayConversations +// (== seriesNum-1, passed explicitly rather than re-derived). +// // Returns true if the whole conversation was walked; false if a stop signal // (--total reached, context cancel) cut it short mid-walk. func runReplayConversation( @@ -124,6 +127,7 @@ func runReplayConversation( st *autoState, rdw *requestDataWriter, conv Conversation, + convIdx int, seriesNum int, seriesGUID string, endpointOverride string, @@ -200,7 +204,7 @@ func runReplayConversation( var pending strings.Builder turnNum := 0 - flush := func() bool { + flush := func(gptIdx int) bool { userContent := strings.TrimSpace(pending.String()) pending.Reset() if userContent == "" { @@ -232,7 +236,11 @@ func runReplayConversation( requestNum := int(st.totalCompleted.Load()) + 1 // Observe the accumulated history (including this turn's user message) - // with the content-level estimator BEFORE the server responds. + // with the content-level estimator BEFORE the server responds. The + // recite-instruction tail (below) is deliberately NOT part of what the + // estimator/history see — it's per-request boilerplate, not + // conversation content, and would otherwise skew the cache-ratio + // estimate every single turn (recite-every-turn is the default). history.WriteString(userContent) ratio := st.estimator.Observe(history.String()) @@ -281,6 +289,7 @@ func runReplayConversation( } metrics.LocalCacheRatio = ratio + recordReplayRequest(cfg, st, rdw, metrics, isFirstRequest, &coldStartTTFT) isFirstRequest = false return true @@ -289,7 +298,7 @@ func runReplayConversation( for i := firstIdx; i < len(conv.Turns); i++ { t := conv.Turns[i] if t.From == "gpt" { - if !flush() { + if !flush(i) { return false } continue @@ -347,6 +356,34 @@ func recordReplayRequest( isErr := metrics.Error != nil explicitCache := metrics.UsageData.CachedTokens.Count > 0 + // UUID validation tallies (--replay-inject-uuids only). metrics.ExpectedUUIDs + // is nil/empty for every request when the feature is off (default), so this + // block — and the val* counters it touches — is fully inert then. + uuidExpectedCount := len(metrics.ExpectedUUIDs) + uuidFoundCount := 0 + for _, found := range metrics.UUIDFound { + if found { + uuidFoundCount++ + } + } + uuidLeakedCount := len(metrics.LeakedUUIDs) + if uuidExpectedCount > 0 { + st.valReqs.Add(1) + st.valUUIDChecks.Add(int64(uuidExpectedCount)) + st.valUUIDFound.Add(int64(uuidFoundCount)) + if metrics.ExactMatch { + st.valExactMatchReqs.Add(1) + } + if missCount := uuidExpectedCount - uuidFoundCount; missCount > 0 { + st.valPresenceMissUUIDs.Add(int64(missCount)) + st.valPresenceMissReqs.Add(1) + } + if uuidLeakedCount > 0 { + st.valCrossContamUUIDs.Add(int64(uuidLeakedCount)) + st.valCrossContamReqs.Add(1) + } + } + earlyColdBaseline := st.earlyColdStartTTFT() st.mu.Lock() @@ -390,7 +427,7 @@ func recordReplayRequest( if metrics.Error != nil { errMsg = metrics.Error.Error() } - if writeErr := rdw.write(requestDataRecord{ + rec := requestDataRecord{ StartTime: reqStart, EndTime: reqEnd, TTFT: float64(metrics.TimeToFirstToken.Milliseconds()), @@ -409,7 +446,21 @@ func recordReplayRequest( ErrorMessage: errMsg, IsEmpty: metrics.IsEmpty, LocalCacheRatio: metrics.LocalCacheRatio, - }); writeErr != nil { + UUIDExpected: uuidExpectedCount, + UUIDFound: uuidFoundCount, + UUIDLeaked: uuidLeakedCount, + UUIDExactMatch: metrics.ExactMatch, + } + // Raw detail lists only on a miss or a leak — mirrors the + // failed-request-only policy on PromptText/ResponseText/RawResponseTail + // above (avoid bloating every row with data that matters only when + // something's wrong). + if uuidFoundCount < uuidExpectedCount || uuidLeakedCount > 0 { + rec.ExpectedUUIDsRaw = metrics.ExpectedUUIDs + rec.FoundMask = metrics.UUIDFound + rec.LeakedUUIDsRaw = metrics.LeakedUUIDs + } + if writeErr := rdw.write(rec); writeErr != nil { fmt.Fprintf(os.Stderr, "warning: failed to write request data: %v\n", writeErr) } } diff --git a/benchmark/replay_double_count_test.go b/benchmark/replay_double_count_test.go index 22e78e6..bcabfff 100644 --- a/benchmark/replay_double_count_test.go +++ b/benchmark/replay_double_count_test.go @@ -125,7 +125,7 @@ func TestReplayConsumerInputTokensAreNetOfCache_Streaming(t *testing.T) { defer ts.Close() st := &autoState{stream: newCompletionStream(200)} - metrics := p.do(context.Background(), minReplayReq(true), strings.Repeat("x", 300), 1, "s1", "i1", 1, st) + metrics := p.do(context.Background(), minReplayReq(true), strings.Repeat("x", 300), 1, "s1", "i1", 1, st, true) if metrics.Error != nil { t.Fatalf("unexpected error: %v", metrics.Error) } @@ -180,7 +180,7 @@ func TestReplayConsumerInputTokensAreNetOfCache_Plain(t *testing.T) { defer ts.Close() st := &autoState{stream: newCompletionStream(200)} - metrics := p.do(context.Background(), minReplayReq(false), strings.Repeat("x", 300), 1, "s1", "i1", 1, st) + metrics := p.do(context.Background(), minReplayReq(false), strings.Repeat("x", 300), 1, "s1", "i1", 1, st, true) if metrics.Error != nil { t.Fatalf("unexpected error: %v", metrics.Error) } @@ -247,7 +247,7 @@ func TestReplayConsumerInputTokensAreNetOfCache_AnthropicPlain(t *testing.T) { {Role: "user", Hash: "h1", Bytes: 50, BlockTypes: []string{"text"}}, }, } - metrics := p.do(context.Background(), req, strings.Repeat("x", 300), 1, "s1", "i1", 1, st) + metrics := p.do(context.Background(), req, strings.Repeat("x", 300), 1, "s1", "i1", 1, st, true) if metrics.Error != nil { t.Fatalf("unexpected error: %v", metrics.Error) } diff --git a/benchmark/replay_router.go b/benchmark/replay_router.go index 6e82adc..80340a0 100644 --- a/benchmark/replay_router.go +++ b/benchmark/replay_router.go @@ -683,6 +683,21 @@ func runRouterReplayInstance( if err == nil { poster.outputRatio = cfg.ReplayOutputRatio poster.forceOutput = cfg.ReplayForceOutput + // UUID cache-coherency injection (--replay-inject-uuids, router + // path). All UUID state set here is GLOBAL and READ-ONLY — the same + // cfg.replayUUIDSets/replaySessionTurnHashes/replayUUIDOwner refs get + // set on every poster in the run (see also the picker pool in + // auto.go), so a poster shared across sessions under multi-endpoint + // routing is safe: buildInjection derives the per-session view + // (which sessionIdx, which turn hashes) from these globals per call + // rather than from any state cached on the poster itself. + if cfg.ReplayInjectUUIDs { + poster.uuidEnabled = true + poster.allUUIDSets = cfg.replayUUIDSets + poster.sessionTurnHashes = cfg.replaySessionTurnHashes + poster.owner = cfg.replayUUIDOwner + poster.reciteEveryRequest = cfg.ReplayReciteEveryRequest + } } if err != nil { // Configuration error — record one error per request in this @@ -739,6 +754,8 @@ func runRouterReplayInstance( } } + isLastRequest := ti == len(inst.Requests)-1 + reqCtx, reqCancel := context.WithTimeout(ctx, reqTimeout) var metrics RequestMetrics // Per-request endpoint selection: prefer this series' home endpoint, @@ -749,10 +766,17 @@ func runRouterReplayInstance( if reqPoster == nil { reqPoster = poster } + // reqPoster may come from the picker's shared per-endpoint pool + // (built once in auto.go and reused by every series that lands on + // that endpoint) rather than this instance's own `poster`. No + // per-session copy-across is needed: every poster's UUID fields are + // global/read-only (set identically on construction — see auto.go + // and above), and do()/dryDo() derive this call's sessionIdx from + // seriesNum directly, so a shared poster safely serves any session. if reqPoster.dryRun { - metrics = reqPoster.dryDo(reqCtx, req, docs, ti+1, sessionID, inst.InstanceID, seriesNum, st) + metrics = reqPoster.dryDo(reqCtx, req, docs, ti+1, sessionID, inst.InstanceID, seriesNum, st, isLastRequest) } else { - metrics = reqPoster.do(reqCtx, req, docs, ti+1, sessionID, inst.InstanceID, seriesNum, st) + metrics = reqPoster.do(reqCtx, req, docs, ti+1, sessionID, inst.InstanceID, seriesNum, st, isLastRequest) } picker.release(epIdx) reqCancel() diff --git a/benchmark/replay_router_post.go b/benchmark/replay_router_post.go index 7322d8f..b4c9837 100644 --- a/benchmark/replay_router_post.go +++ b/benchmark/replay_router_post.go @@ -61,6 +61,140 @@ type replayPoster struct { // newReplayPoster, to avoid touching its many existing call sites. outputRatio float64 forceOutput bool + + // UUID cache-coherency injection (--replay-inject-uuids, router path — + // see replay_router_uuid.go). Set directly on the poster after + // construction, same rationale as outputRatio/forceOutput above. + // uuidEnabled gates everything: false leaves do()/dryDo() byte-for-byte + // identical to before this feature existed. + // + // EVERY field below this point is GLOBAL and READ-ONLY for the lifetime + // of the poster: with multi-endpoint routing a single replayPoster is + // SHARED across every session whose requests currently land on that + // endpoint (see endpointPicker in auto.go), so there is no such thing as + // "this poster's session" — a poster serves whichever sessionIdx the + // caller passes into buildInjection per-request. All per-session state is + // therefore looked up from these global slices/maps by an explicit + // sessionIdx argument rather than cached mutable fields on the poster — + // see buildInjection. + uuidEnabled bool + // allUUIDSets is cfg.replayUUIDSets: the full per-session-per-turn UUID + // assignment (index i = session i's ordered turn-UUID list, index t = + // turn t's UUID), shared read-only across every poster in the run. + allUUIDSets [][]string + // sessionTurnHashes is cfg.replaySessionTurnHashes: the full per-session + // ordered turn-hash list (index i = session i's ordered turn-hash list; + // sessionTurnHashes[i][t] is the hash of session i's turn t), shared + // read-only across every poster in the run. buildInjection indexes this + // by the caller-supplied sessionIdx and builds that session's + // hash-to-turn map locally, per call — see buildInjection. + sessionTurnHashes [][]string + // owner is cfg.replayUUIDOwner: the reverse uuid -> owning-session-index + // map (see buildSessionTurnUUIDs), shared read-only across every poster + // in the run — used by findLeakedUUIDsByOwner to flag cross-session + // contamination in O(response) time. + owner map[string]int + // reciteEveryRequest mirrors --replay-recite-every-request: true asks + // for the recite line on every request; false only on each instance's + // final request (see the isLastRequest parameter to do()/dryDo()). + reciteEveryRequest bool +} + +// buildInjection returns this call's *uuidInjection (nil when UUID +// injection is disabled, sessionIdx has no assigned turns — e.g. it fell +// outside the precomputed array, or that session had zero qualifying turns +// — or this particular request has no qualifying turn visible in its +// message history yet). sessionIdx identifies which session's global UUID +// assignment to use (== seriesNum-1 — see callers); it is passed explicitly +// rather than cached on the poster because a single poster may be SHARED +// across sessions under multi-endpoint routing (see the struct doc). +// isLastRequest is whether req is the final request of the CURRENT +// instance's request list (see runRouterReplayInstance) — with +// --replay-recite-every-request=false, only that final request carries the +// recite ask. +// +// Every VISIBLE qualifying turn in req.Messages gets stamped into +// StampByHash (keeping every turn's marker warm in KV as later requests +// repeat it) — see uuidInjection's doc. The recite WINDOW is separate and +// bounded: the first (visible) turn plus up to 3 most-recent turns +// EXCLUDING the current turn (the highest-index turn visible in THIS +// request), deduplicated and capped at 4 (see the package doc in +// replay_router_uuid.go for the design rationale and edge cases at turns +// 1-3). +func (p *replayPoster) buildInjection(req RouterReplayRequest, sessionIdx int, isLastRequest bool) *uuidInjection { + if !p.uuidEnabled || sessionIdx < 0 || sessionIdx >= len(p.allUUIDSets) || sessionIdx >= len(p.sessionTurnHashes) { + return nil + } + uuids := p.allUUIDSets[sessionIdx] + turnHashes := p.sessionTurnHashes[sessionIdx] + if len(uuids) == 0 || len(turnHashes) == 0 { + return nil + } + // hashToTurn is built fresh per call rather than cached: it's a small, + // cheap map (one entry per turn) with no shared mutable state, which is + // exactly the point — a shared poster can serve any sessionIdx + // concurrently with no data race. + hashToTurn := make(map[string]int, len(turnHashes)) + for i, h := range turnHashes { + if h != "" { + hashToTurn[h] = i + } + } + + stampByHash := map[string]turnStamp{} + var visible []int // turn indices visible in this request, in first-appearance order + seenTurn := map[int]bool{} + for _, m := range req.Messages { + t, ok := hashToTurn[m.Hash] + if !ok || t < 0 || t >= len(uuids) { + continue + } + stampByHash[m.Hash] = turnStamp{Idx: t, UUID: uuids[t], Label: fmt.Sprintf("turn-%d", t+1)} + if !seenTurn[t] { + seenTurn[t] = true + visible = append(visible, t) + } + } + if len(visible) == 0 { + return nil + } + + // Window: first visible turn, plus up to 3 most-recent turns EXCLUDING + // the current turn (visible's last entry — the highest turn index + // present, since turns only ever get appended to a growing history). + first := visible[0] + var recentCandidates []int + if len(visible) > 1 { + recentCandidates = visible[:len(visible)-1] + } + recent := recentCandidates + if len(recent) > 3 { + recent = recent[len(recent)-3:] + } + window := []int{first} + for _, t := range recent { + if t == first { + continue + } + window = append(window, t) + } + if len(window) > 4 { + window = window[:4] + } + + labels := make([]string, len(window)) + reciteUUIDs := make([]string, len(window)) + for i, t := range window { + labels[i] = fmt.Sprintf("turn-%d", t+1) + reciteUUIDs[i] = uuids[t] + } + + return &uuidInjection{ + StampByHash: stampByHash, + Recite: p.reciteEveryRequest || isLastRequest, + ReciteLabels: labels, + ReciteUUIDs: reciteUUIDs, + } } func newReplayPoster(modelSpec string, keys llm.APIKeys, endpointOverride string, runID string, dryRun bool, coldTPS, warmTPS, outputTPS int, estimator *cacheEstimator) (*replayPoster, error) { @@ -286,7 +420,9 @@ func (p *replayPoster) sendOnce(ctx context.Context, url string, bodyBytes []byt // do issues one request and returns its metrics. Honors ctx for // cancellation / deadline. Dispatches to the appropriate body builder -// and response parser based on p.apiType. +// and response parser based on p.apiType. isLastRequest is whether req is +// the final request of the calling instance's request list — see +// buildInjection. func (p *replayPoster) do( ctx context.Context, req RouterReplayRequest, @@ -296,17 +432,21 @@ func (p *replayPoster) do( instanceID string, seriesNum int, st *autoState, + isLastRequest bool, ) RequestMetrics { startTime := time.Now() + sessionIdx := seriesNum - 1 + inj := p.buildInjection(req, sessionIdx, isLastRequest) + var bodyBytes []byte var canonical string var err error switch p.apiType { case "openai", "openai_vllm": - bodyBytes, canonical, err = buildOpenAIChatCompletionsBody(req, docs, p.model, p.runID, p.outputRatio, p.forceOutput) + bodyBytes, canonical, err = buildOpenAIChatCompletionsBody(req, docs, p.model, p.runID, p.outputRatio, p.forceOutput, inj) default: - bodyBytes, canonical, err = buildAnthropicMessagesBody(req, docs, p.model, p.runID, p.outputRatio, p.forceOutput) + bodyBytes, canonical, err = buildAnthropicMessagesBody(req, docs, p.model, p.runID, p.outputRatio, p.forceOutput, inj) } if err != nil { return RequestMetrics{ @@ -389,6 +529,43 @@ func (p *replayPoster) do( m.Error = fmt.Errorf("empty response from model") } m.LocalCacheRatio = localCacheRatio + + // UUID cache-coherency validation (--replay-inject-uuids, router path). + // consumeOpenAISSE/consumeOpenAIPlain/consumeSSE/consumePlain already + // merge reasoning/thinking into m.Response (see their doc comments), so + // a single Contains-scan of m.Response covers both — thinking is passed + // as "" here, mirroring the dataset path's own call shape. + // + // Two independent checks, mirroring the cache-coherency eval's two + // reported tests: per-UUID PRESENCE (Contains anywhere in the response, + // scored against inj.ReciteUUIDs — this request's recite WINDOW, not the + // session's full turn history) and output CONFORMITY (the FIRST LINE of + // the response is exactly the ordered, comma-joined window UUID list — + // see firstLineConformity/matchesExpectedUUIDList). Cross-contamination + // uses findLeakedUUIDsByOwner (replay_uuid.go), an O(response) reverse- + // map scan — NOT FindLeakedUUIDs, whose O(population) iteration over + // every session's UUID set no longer fits once turns (not sessions) are + // the stamping unit. + // + // Gated on inj.Recite, NOT just inj != nil: buildInjection returns a + // non-nil *uuidInjection whenever ANY qualifying turn is visible in this + // request (it always stamps every visible turn inline, so those stamps + // stay warm in KV across the session — see the package doc in + // replay_router_uuid.go), but only inj.Recite (reciteEveryRequest || + // isLastRequest) says the model was actually ASKED to recite this turn. + // Scoring a non-recite turn would count "the model didn't volunteer the + // UUID list" as a PRESENCE_MISS/conformity failure even though nothing + // asked it to. + if inj != nil && inj.Recite && m.Error == nil && !m.IsEmpty { + m.ConvIdx = sessionIdx + m.ExpectedUUIDs = append([]string(nil), inj.ReciteUUIDs...) + m.UUIDFound = make([]bool, len(inj.ReciteUUIDs)) + for i, u := range inj.ReciteUUIDs { + m.UUIDFound[i] = strings.Contains(m.Response, u) + } + m.LeakedUUIDs = findLeakedUUIDsByOwner(m.Response, "", sessionIdx, p.owner) + m.ExactMatch = firstLineConformity(m.Response, inj.ReciteUUIDs) + } return m } @@ -601,6 +778,12 @@ func consumeOpenAISSE(body io.Reader, startTime time.Time, m *RequestMetrics) { } // consumeOpenAIPlain reads a non-streaming OpenAI chat/completions response. +// Like consumeOpenAISSE, this merges message.reasoning_content (or +// message.reasoning, the vLLM field name) into m.Response ahead of +// message.content — a non-streaming reasoning-model response carries its +// full reasoning text on the message object rather than as incremental +// deltas, and skipping it here would make a UUID recited/leaked only in the +// reasoning channel invisible to the presence/leak scan. func consumeOpenAIPlain(body io.Reader, startTime time.Time, m *RequestMetrics) { b, err := io.ReadAll(body) if err != nil { @@ -611,7 +794,9 @@ func consumeOpenAIPlain(body io.Reader, startTime time.Time, m *RequestMetrics) var resp struct { Choices []struct { Message struct { - Content string `json:"content"` + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + Reasoning string `json:"reasoning"` // vLLM uses "reasoning" } `json:"message"` } `json:"choices"` Usage struct { @@ -628,7 +813,12 @@ func consumeOpenAIPlain(body io.Reader, startTime time.Time, m *RequestMetrics) return } if len(resp.Choices) > 0 { - m.Response = resp.Choices[0].Message.Content + msg := resp.Choices[0].Message + reasoning := msg.ReasoningContent + if reasoning == "" { + reasoning = msg.Reasoning + } + m.Response = reasoning + msg.Content } cached := 0 if resp.Usage.PromptTokensDetails != nil { @@ -637,7 +827,12 @@ func consumeOpenAIPlain(body io.Reader, startTime time.Time, m *RequestMetrics) m.UsageData = buildReplayUsage(resp.Usage.PromptTokens, cached, resp.Usage.CompletionTokens) } -// consumePlain reads a non-streaming Anthropic response. +// consumePlain reads a non-streaming Anthropic response. Like consumeSSE, +// this merges "thinking" content blocks into m.Response alongside "text" +// blocks — a non-streaming extended-thinking response carries its thinking +// block(s) as ordinary entries in the content array (field "thinking", not +// "text"), and skipping them here would make a UUID recited/leaked only in +// the thinking channel invisible to the presence/leak scan. func consumePlain(body io.Reader, startTime time.Time, m *RequestMetrics) { b, err := io.ReadAll(body) if err != nil { @@ -647,8 +842,9 @@ func consumePlain(body io.Reader, startTime time.Time, m *RequestMetrics) { m.TimeToFirstToken = time.Since(startTime) var resp struct { Content []struct { - Type string `json:"type"` - Text string `json:"text"` + Type string `json:"type"` + Text string `json:"text"` + Thinking string `json:"thinking"` } `json:"content"` StopReason string `json:"stop_reason"` Usage struct { @@ -664,8 +860,11 @@ func consumePlain(body io.Reader, startTime time.Time, m *RequestMetrics) { } var sb strings.Builder for _, c := range resp.Content { - if c.Type == "text" { + switch c.Type { + case "text": sb.WriteString(c.Text) + case "thinking": + sb.WriteString(c.Thinking) } } m.Response = sb.String() @@ -704,14 +903,21 @@ func (p *replayPoster) dryDo( instanceID string, seriesNum int, st *autoState, + isLastRequest bool, ) RequestMetrics { - // Build canonical string for estimator and compute ratio. + // Build canonical string for estimator and compute ratio. Injection is + // threaded through purely so the canonical text (and therefore the + // cache-ratio estimate) stays consistent with a real do() call for the + // same request — dry-run never makes a real request, so there's no + // response to validate. + sessionIdx := seriesNum - 1 + inj := p.buildInjection(req, sessionIdx, isLastRequest) var canonical string switch p.apiType { case "openai", "openai_vllm": - _, canonical, _ = buildOpenAIChatCompletionsBody(req, docs, p.model, p.runID, p.outputRatio, p.forceOutput) + _, canonical, _ = buildOpenAIChatCompletionsBody(req, docs, p.model, p.runID, p.outputRatio, p.forceOutput, inj) default: - _, canonical, _ = buildAnthropicMessagesBody(req, docs, p.model, p.runID, p.outputRatio, p.forceOutput) + _, canonical, _ = buildAnthropicMessagesBody(req, docs, p.model, p.runID, p.outputRatio, p.forceOutput, inj) } var ratio float64 if p.estimator != nil { diff --git a/benchmark/replay_router_post_test.go b/benchmark/replay_router_post_test.go index 44187a9..c82cdca 100644 --- a/benchmark/replay_router_post_test.go +++ b/benchmark/replay_router_post_test.go @@ -64,7 +64,7 @@ func TestReplayEndpointResolution(t *testing.T) { ts, seen := newVLLMStyleServer() defer ts.Close() p := mustPoster(t, fmt.Sprintf("dynamic/%s,type=openai_vllm,model=m", ts.URL)) - if m := p.do(context.Background(), minimalReq, docs, 1, "s", "i", 1, newState()); m.Error != nil { + if m := p.do(context.Background(), minimalReq, docs, 1, "s", "i", 1, newState(), true); m.Error != nil { t.Fatalf("first request: %v", m.Error) } got := seen() @@ -75,7 +75,7 @@ func TestReplayEndpointResolution(t *testing.T) { t.Errorf("latch = %q fellBack=%v, want fallback latched", p.epResolved, p.epFellBack) } // Latched: the second request goes straight to /v1, one wire call. - if m := p.do(context.Background(), minimalReq, docs, 2, "s", "i", 1, newState()); m.Error != nil { + if m := p.do(context.Background(), minimalReq, docs, 2, "s", "i", 1, newState(), true); m.Error != nil { t.Fatalf("second request: %v", m.Error) } if got = seen(); len(got) != 3 || got[2] != "/v1/chat/completions" { @@ -87,7 +87,7 @@ func TestReplayEndpointResolution(t *testing.T) { ts, seen := newVLLMStyleServer() defer ts.Close() p := mustPoster(t, fmt.Sprintf("dynamic/%s/v1,type=openai_vllm,model=m", ts.URL)) - if m := p.do(context.Background(), minimalReq, docs, 1, "s", "i", 1, newState()); m.Error != nil { + if m := p.do(context.Background(), minimalReq, docs, 1, "s", "i", 1, newState(), true); m.Error != nil { t.Fatalf("request: %v", m.Error) } got := seen() @@ -110,7 +110,7 @@ func TestReplayEndpointResolution(t *testing.T) { })) defer ts.Close() p := mustPoster(t, fmt.Sprintf("dynamic/%s,type=openai_vllm,model=m", ts.URL)) - m := p.do(context.Background(), minimalReq, docs, 1, "s", "i", 1, newState()) + m := p.do(context.Background(), minimalReq, docs, 1, "s", "i", 1, newState(), true) if m.Error == nil || !strings.Contains(m.Error.Error(), "status 500") { t.Fatalf("expected status-500 error, got %v", m.Error) } @@ -136,7 +136,7 @@ func TestReplayEndpointResolution(t *testing.T) { wg.Add(1) go func(i int) { defer wg.Done() - m := p.do(context.Background(), minimalReq, docs, 1, "s", fmt.Sprintf("i%d", i), 1, newState()) + m := p.do(context.Background(), minimalReq, docs, 1, "s", fmt.Sprintf("i%d", i), 1, newState(), true) errs[i] = m.Error }(i) } @@ -152,7 +152,7 @@ func TestReplayEndpointResolution(t *testing.T) { // Duplicate probes during the race are allowed; once latched, a new // request adds exactly one wire call. before := len(seen()) - if m := p.do(context.Background(), minimalReq, docs, 2, "s", "i", 1, newState()); m.Error != nil { + if m := p.do(context.Background(), minimalReq, docs, 2, "s", "i", 1, newState(), true); m.Error != nil { t.Fatalf("post-latch request: %v", m.Error) } after := seen() @@ -378,7 +378,7 @@ func TestOpenAIReplayEndToEnd(t *testing.T) { } ctx := context.Background() - metrics := p.do(ctx, req, docs, 1, "session-1", "instance-1", 1, st) + metrics := p.do(ctx, req, docs, 1, "session-1", "instance-1", 1, st, true) // Verify no error. if metrics.Error != nil { @@ -459,7 +459,7 @@ func TestOpenAIReplayToolTranslation(t *testing.T) { }, } - body, _, err := buildOpenAIChatCompletionsBody(req, docs, "test-model", "", 0, false) + body, _, err := buildOpenAIChatCompletionsBody(req, docs, "test-model", "", 0, false, nil) if err != nil { t.Fatalf("buildOpenAIChatCompletionsBody: %v", err) } @@ -519,7 +519,7 @@ func TestOpenAIBodyBuilderExtra(t *testing.T) { {Role: "user", Hash: "h1", Bytes: 50, BlockTypes: []string{"text"}}, }, } - body, _, err := buildOpenAIChatCompletionsBody(req, docs, "test-model", "run-42", 0, false) + body, _, err := buildOpenAIChatCompletionsBody(req, docs, "test-model", "run-42", 0, false, nil) if err != nil { t.Fatalf("build: %v", err) } @@ -548,7 +548,7 @@ func TestOpenAIBodyBuilderExtra(t *testing.T) { {Role: "user", Hash: "h1", Bytes: 50, BlockTypes: []string{"text"}}, }, } - body, _, err := buildOpenAIChatCompletionsBody(req, docs, "test-model", "", 0, false) + body, _, err := buildOpenAIChatCompletionsBody(req, docs, "test-model", "", 0, false, nil) if err != nil { t.Fatalf("build: %v", err) } @@ -576,7 +576,7 @@ func TestOpenAIBodyBuilderExtra(t *testing.T) { {Role: "user", Hash: "h1", Bytes: 50, BlockTypes: []string{"text"}}, }, } - body, _, err := buildOpenAIChatCompletionsBody(req, docs, "test-model", "", 0, false) + body, _, err := buildOpenAIChatCompletionsBody(req, docs, "test-model", "", 0, false, nil) if err != nil { t.Fatalf("build: %v", err) } @@ -602,7 +602,7 @@ func TestOpenAIBodyBuilderExtra(t *testing.T) { Stream: true, OutputTokens: 100, } - body, _, err := buildOpenAIChatCompletionsBody(req, docs, "test-model", "", 0, false) + body, _, err := buildOpenAIChatCompletionsBody(req, docs, "test-model", "", 0, false, nil) if err != nil { t.Fatalf("build: %v", err) } @@ -652,7 +652,7 @@ func TestOpenAINonStreamingEndToEnd(t *testing.T) { } st := &autoState{stream: newCompletionStream(200)} - metrics := p.do(context.Background(), req, docs, 1, "s1", "i1", 1, st) + metrics := p.do(context.Background(), req, docs, 1, "s1", "i1", 1, st, true) if metrics.Error != nil { t.Fatalf("unexpected error: %v", metrics.Error) @@ -698,7 +698,7 @@ func TestOpenAIErrorResponse(t *testing.T) { } st := &autoState{stream: newCompletionStream(200)} - metrics := p.do(context.Background(), req, docs, 1, "s1", "i1", 1, st) + metrics := p.do(context.Background(), req, docs, 1, "s1", "i1", 1, st, true) if metrics.Error == nil { t.Fatal("expected error for 500 response, got nil") @@ -740,7 +740,7 @@ func TestOpenAISSEWithoutUsage(t *testing.T) { } st := &autoState{stream: newCompletionStream(200)} - metrics := p.do(context.Background(), req, docs, 1, "s1", "i1", 1, st) + metrics := p.do(context.Background(), req, docs, 1, "s1", "i1", 1, st, true) if metrics.Error != nil { // "empty response from model" is acceptable for no-usage streams. @@ -750,3 +750,151 @@ func TestOpenAISSEWithoutUsage(t *testing.T) { t.Errorf("response = %q, want %q", metrics.Response, "ok") } } + +// TestUUIDScoringGatedOnRecite covers the H1 fix: buildInjection returns a +// non-nil *uuidInjection on EVERY request once UUID injection is on (see +// its doc comment), but only inj.Recite says the model was actually ASKED +// to recite this turn. With --replay-recite-every-request=false, a +// non-final request must still get the UUID block injected (so it stays +// warm in KV — see replay_router_uuid.go's package doc) but must NOT be +// scored: scoring it would count "the model didn't volunteer the UUID +// list" as a false PRESENCE_MISS/conformity failure even though nothing +// asked it to. This drives the poster through p.do() end-to-end (real HTTP +// round trip against an httptest server) rather than calling buildInjection +// directly, so the assertion covers the actual gate in do(). +func TestUUIDScoringGatedOnRecite(t *testing.T) { + docs := strings.Repeat("uuid-gate-docs ", 100) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(200) + // Response deliberately omits the UUID list entirely — if this + // non-recite request were (wrongly) scored, it would register as a + // presence-miss on every expected UUID. + fmt.Fprintf(w, `{ + "id": "chatcmpl-1", + "object": "chat.completion", + "model": "test-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "just a normal answer, no uuids here"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 30, "completion_tokens": 5, "total_tokens": 35} + }`) + })) + defer ts.Close() + + modelSpec := fmt.Sprintf("dynamic/%s,type=openai,model=test-model", ts.URL) + keys := llm.APIKeys{OpenAI: "sk-test"} + p, err := newReplayPoster(modelSpec, keys, "", "", false, 0, 0, 0, nil) + if err != nil { + t.Fatalf("newReplayPoster: %v", err) + } + // Wire up UUID injection exactly as runRouterReplayInstance does (see + // replay_router.go): global/read-only slices indexed by sessionIdx + // (0, matching seriesNum=1 below via sessionIdx = seriesNum-1), with + // reciteEveryRequest=false so only the FINAL request of an instance's + // list carries the recite ask. + p.uuidEnabled = true + p.allUUIDSets = [][]string{{"uuid-alpha", "uuid-beta"}} + p.sessionTurnHashes = [][]string{{"h1", "h2"}} + p.reciteEveryRequest = false + + req := RouterReplayRequest{ + Stream: false, + OutputTokens: 100, + Messages: []RouterReplayMessage{ + {Role: "user", Hash: "h1", Bytes: 50, BlockTypes: []string{"text"}}, + }, + } + st := &autoState{stream: newCompletionStream(200)} + + // isLastRequest=false -> inj.Recite=false (reciteEveryRequest is also + // false) -> the H1 gate must skip scoring entirely. + metrics := p.do(context.Background(), req, docs, 1, "s1", "i1", 1, st, false) + + if metrics.Error != nil { + t.Fatalf("unexpected error: %v", metrics.Error) + } + if len(metrics.ExpectedUUIDs) != 0 { + t.Errorf("ExpectedUUIDs = %v, want empty (non-recite request must not be scored)", metrics.ExpectedUUIDs) + } + if len(metrics.UUIDFound) != 0 { + t.Errorf("UUIDFound = %v, want empty (non-recite request must not be scored)", metrics.UUIDFound) + } + if len(metrics.LeakedUUIDs) != 0 { + t.Errorf("LeakedUUIDs = %v, want empty (non-recite request must not be scored)", metrics.LeakedUUIDs) + } + if metrics.ExactMatch { + t.Error("ExactMatch = true, want false (non-recite request must not be scored)") + } + + // Sanity check the OTHER half of the invariant: the SAME request, with + // isLastRequest=true (recite asked), DOES get scored — otherwise this + // test could be vacuously passing because scoring never fires at all. + metrics2 := p.do(context.Background(), req, docs, 1, "s1", "i2", 1, st, true) + if metrics2.Error != nil { + t.Fatalf("unexpected error on recite request: %v", metrics2.Error) + } + if len(metrics2.ExpectedUUIDs) == 0 { + t.Fatal("ExpectedUUIDs empty on a recite=true request — scoring gate is broken (never scores) rather than fixed") + } +} + +// TestConsumeOpenAIPlainMergesReasoning covers the M2 fix: a non-streaming +// OpenAI response's message.reasoning_content must be merged into +// m.Response alongside message.content, matching consumeOpenAISSE's +// streaming behavior — otherwise a UUID recited/leaked only in the +// reasoning channel would be invisible to the presence/leak scan. +func TestConsumeOpenAIPlainMergesReasoning(t *testing.T) { + body := strings.NewReader(`{ + "choices": [{"index": 0, "message": {"role": "assistant", "content": "the answer is 42", "reasoning_content": "let me think about uuid-in-reasoning first"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5} + }`) + var m RequestMetrics + consumeOpenAIPlain(body, time.Now(), &m) + + if !strings.Contains(m.Response, "uuid-in-reasoning") { + t.Errorf("m.Response = %q, want it to contain the reasoning_content text", m.Response) + } + if !strings.Contains(m.Response, "the answer is 42") { + t.Errorf("m.Response = %q, want it to also contain message.content", m.Response) + } +} + +// TestConsumeOpenAIPlainMergesReasoningVLLMField covers vLLM's alternate +// "reasoning" field name (used when reasoning_content is absent). +func TestConsumeOpenAIPlainMergesReasoningVLLMField(t *testing.T) { + body := strings.NewReader(`{ + "choices": [{"index": 0, "message": {"role": "assistant", "content": "final text", "reasoning": "uuid-in-vllm-reasoning-field"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5} + }`) + var m RequestMetrics + consumeOpenAIPlain(body, time.Now(), &m) + + if !strings.Contains(m.Response, "uuid-in-vllm-reasoning-field") { + t.Errorf("m.Response = %q, want it to contain the reasoning field text", m.Response) + } +} + +// TestConsumePlainMergesThinking covers the M2 fix: a non-streaming +// Anthropic response's "thinking" content block must be merged into +// m.Response alongside "text" blocks, matching consumeSSE's streaming +// behavior — otherwise a UUID recited/leaked only in the thinking channel +// would be invisible to the presence/leak scan. +func TestConsumePlainMergesThinking(t *testing.T) { + body := strings.NewReader(`{ + "content": [ + {"type": "thinking", "thinking": "pondering uuid-in-thinking-block"}, + {"type": "text", "text": "here is my answer"} + ], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5} + }`) + var m RequestMetrics + consumePlain(body, time.Now(), &m) + + if !strings.Contains(m.Response, "uuid-in-thinking-block") { + t.Errorf("m.Response = %q, want it to contain the thinking block text", m.Response) + } + if !strings.Contains(m.Response, "here is my answer") { + t.Errorf("m.Response = %q, want it to also contain the text block", m.Response) + } +} diff --git a/benchmark/replay_router_uuid.go b/benchmark/replay_router_uuid.go new file mode 100644 index 0000000..75d835c --- /dev/null +++ b/benchmark/replay_router_uuid.go @@ -0,0 +1,396 @@ +package benchmark + +// UUID-based cache-coherency validation for the ROUTER-REPLAY path +// (--router-replay-file, --replay-inject-uuids). +// +// Strategy: one deterministic UUID is injected inline into EVERY qualifying +// user turn — a role=="user" message with >=1 text block whose hash is +// referenced by exactly one session (see isQualifyingUserTurn, which reuses +// computeBlockSessionCounts' distinct-session-count map to exclude blocks +// shared across sessions) — spread through the conversation rather than +// clustered at a session boundary. Turn i's marker +// ("\n\n[turn-N id: ]") is appended to that turn's own synthesized +// text content (see buildMessageContent's stampByHash param in +// replay_router_wire.go), keyed by the turn's session-global +// first-appearance index (see computeSessionTurnHashes/ +// buildSessionTurnUUIDs below). Content synthesis is already deterministic +// in the block hash (synthText, seeded ":block:"), so every request +// that repeats a given turn in its growing history re-emits byte-identical +// content for it — the fidelity invariant that makes cache-hit reproduction +// against the original capture possible: two sessions sharing a leading +// block still collide on the server's prefix cache exactly as in the +// original traffic, because only count==1 (genuinely per-session) turns +// ever carry a stamp. +// +// Each request asks the model to recite a WINDOW of turns rather than the +// whole history: the first (visible) turn plus up to 3 most-recent turns, +// EXCLUDING the current turn, deduplicated and capped at 4 (see +// replayPoster.buildInjection in replay_router_post.go and +// replayReciteWindowInstruction below). This keeps the recite cost/response +// budget constant regardless of session length while still spreading +// coverage across the whole conversation over time — turn N's stamp gets +// asked about at turns N, N+1, N+2, N+3, then ages out of the window (but +// stays warm in KV via the always-embedded StampByHash markers above, which +// cover every visible turn, not just the recited window). + +import ( + "bufio" + "encoding/json" + "fmt" + "os" + "strings" + "sync" +) + +// turnStamp is the per-user-turn UUID marker injected inline into that +// turn's own message content (see buildMessageContent's stampByHash param +// in replay_router_wire.go). +type turnStamp struct { + Idx int // this session's global turn index (0-based, first-appearance order) + UUID string // this turn's deterministic UUID + Label string // "turn-N" (Idx+1 — 1-based, for human-readable instructions) +} + +// uuidInjection describes the per-request UUID injection to apply when +// building a router-replay wire body. A nil *uuidInjection means "no +// injection" — buildAnthropicMessagesBody / buildOpenAIChatCompletionsBody +// must behave identically to before this feature existed. +type uuidInjection struct { + // StampByHash carries one turnStamp per user-turn message VISIBLE in + // this request (keyed by RouterReplayMessage.Hash) — every qualifying + // turn gets its UUID marker embedded inline in its own synthesized + // content (see buildMessageContent), keeping every turn's stamp warm in + // KV as later requests repeat that history in full, not just the turns + // named in the recite window below. + StampByHash map[string]turnStamp + // Recite asks the model to output, on the FIRST line of its response, + // the ReciteUUIDs values (identified to the model by ReciteLabels' + // inline tags — see replayReciteWindowInstruction), then continue + // normally. + Recite bool + // ReciteLabels is the ordered window of turn labels ("turn-N") the + // instruction names — first (visible) turn, then up to 3 most-recent + // turns EXCLUDING the current turn, deduplicated, capped at 4. + ReciteLabels []string + // ReciteUUIDs is ReciteLabels' matching ordered UUID values — what + // firstLineConformity/matchesExpectedUUIDList checks the response's + // first line against. + ReciteUUIDs []string +} + +// firstLineConformity implements this feature's output-conformity check: +// whether the FIRST LINE of resp (up to the first '\n', trimmed) is EXACTLY +// the comma-separated expected UUID list, in order (see +// matchesExpectedUUIDList) — the router-replay analogue of the coherency +// eval's ExactMatch, adapted for a recite-FIRST instruction that lets the +// model keep generating after line 1 (so forced-output/ignore_eos still +// fills the remainder of the output budget with a normal continuation). +// expected is inj.ReciteUUIDs — the current request's recite WINDOW, not +// every UUID ever stamped in the session. +func firstLineConformity(resp string, expected []string) bool { + line := resp + if idx := strings.Index(resp, "\n"); idx >= 0 { + line = resp[:idx] + } + return matchesExpectedUUIDList(strings.TrimSpace(line), expected) +} + +// computeBlockSessionCounts streams a replay-v3 file once and returns, for +// every distinct block hash seen, the number of DISTINCT SESSIONS that +// reference it at least once (not the number of requests — a hash reused +// many times within one session still counts once for that session). +// +// A hash with count > 1 is shared across sessions — the router's own +// leading system prompt(s), repeated verbatim capture after capture, are +// the common case — and is never eligible for UUID stamping: stamping it +// would perturb content this session shares with others, poisoning their +// cross-session prefix-cache key too. computeSessionTurnHashes uses this +// map (via isQualifyingUserTurn) to restrict turn-stamping to hashes with +// count == 1: genuinely per-session, per-turn content. +// +// allowed and sessionLimit mirror openRouterReplayStream's filtering +// exactly (nil allowed = every session; sessionLimit <= 0 = no cap) so the +// set of sessions counted here matches the set the real run will dispatch. +// Only hashes are read here — synthText/synthesized content is never +// generated, so this pass is cheap even over large files. +func computeBlockSessionCounts(path string, allowed map[int]bool, sessionLimit int) (map[string]int, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + br := bufio.NewReaderSize(f, 1<<20) + if _, err := br.ReadBytes('\n'); err != nil { + return nil, fmt.Errorf("read header line: %w", err) + } + + counts := map[string]int{} + lineIdx := 0 + produced := 0 + for { + if sessionLimit > 0 && produced >= sessionLimit { + break + } + if allowed != nil && produced >= len(allowed) { + break + } + line, rerr := br.ReadBytes('\n') + if len(line) > 0 { + line = trimNL(line) + if len(line) > 0 { + currentIdx := lineIdx + lineIdx++ + if allowed != nil && !allowed[currentIdx] { + if rerr != nil { + break + } + continue + } + var sess RouterReplaySession + if jerr := json.Unmarshal(line, &sess); jerr == nil { + seen := map[string]bool{} + for _, inst := range sess.Instances { + for _, req := range inst.Requests { + hashes, _ := BuildReplayRequestPrefix(req) + for _, h := range hashes { + if h != "" { + seen[h] = true + } + } + } + } + for h := range seen { + counts[h]++ + } + produced++ + } + } + } + if rerr != nil { + break + } + } + return counts, nil +} + +// isQualifyingUserTurn reports whether m is a "user-input turn" for the +// purposes of UUID stamping: role=="user", at least one "text" content +// block, and a hash referenced by exactly one session (counts[m.Hash]==1 — +// see computeBlockSessionCounts). tool_result-only messages, assistant/ +// system messages, and any message whose hash is shared across sessions +// never qualify. +func isQualifyingUserTurn(m RouterReplayMessage, counts map[string]int) bool { + if m.Role != "user" { + return false + } + if m.Hash == "" || counts[m.Hash] != 1 { + return false + } + for _, t := range m.BlockTypes { + if t == "text" { + return true + } + } + return false +} + +// computeSessionTurnHashes makes a second streaming pass (same filtering as +// computeBlockSessionCounts — nil allowed = every session; sessionLimit <= 0 +// = no cap) and returns, in dispatch order (index i = the i-th session +// encountered, the SAME order buildSessionTurnUUIDs' returned slice is +// indexed by and cfg.replayUUIDSets ends up in), that session's ordered list +// of DISTINCT qualifying user-turn hashes (see isQualifyingUserTurn), in +// first-appearance order across the session's instances/requests/messages +// (walked in file order — Instances, then each instance's Requests, then +// each request's Messages). A request's Messages carries the FULL growing +// conversation history, so the same turn hash reappears in every later +// request of the same instance; only its FIRST appearance contributes an +// entry here — turnHashes[i][t] is session i's turn-t hash, and len(...) is +// that session's total turn count. +func computeSessionTurnHashes(path string, allowed map[int]bool, sessionLimit int, counts map[string]int) ([][]string, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + br := bufio.NewReaderSize(f, 1<<20) + if _, err := br.ReadBytes('\n'); err != nil { + return nil, fmt.Errorf("read header line: %w", err) + } + + var turnHashes [][]string + lineIdx := 0 + produced := 0 + for { + if sessionLimit > 0 && produced >= sessionLimit { + break + } + if allowed != nil && produced >= len(allowed) { + break + } + line, rerr := br.ReadBytes('\n') + if len(line) > 0 { + line = trimNL(line) + if len(line) > 0 { + currentIdx := lineIdx + lineIdx++ + if allowed != nil && !allowed[currentIdx] { + if rerr != nil { + break + } + continue + } + var sess RouterReplaySession + if jerr := json.Unmarshal(line, &sess); jerr == nil { + seen := map[string]bool{} + var hashes []string + for _, inst := range sess.Instances { + for _, req := range inst.Requests { + for _, m := range req.Messages { + if !isQualifyingUserTurn(m, counts) { + continue + } + if seen[m.Hash] { + continue + } + seen[m.Hash] = true + hashes = append(hashes, m.Hash) + } + } + } + turnHashes = append(turnHashes, hashes) + produced++ + } + } + } + if rerr != nil { + break + } + } + return turnHashes, nil +} + +// buildSessionTurnUUIDs returns, for each session, one UUID per turn +// (sets[i][t] = session i's turn-t UUID), drawn session-major/turn-minor +// from a single seeded generator (see newUUIDGenerator) — same determinism/ +// disjointness rationale as the dataset path's buildReplayUUIDSets: same +// seed -> same per-session-per-turn UUID assignment across runs and across +// every model in a multi-model run (see the precompute call site in +// RunAutoBenchmark, which populates cfg.replayUUIDSets once, before any +// per-model goroutine spawns, so every model sees the identical +// assignment). +// +// owner is the reverse mapping (uuid -> the owning session's index i) used +// by findLeakedUUIDsByOwner (replay_uuid.go) to flag cross-session +// contamination in O(response) time — a substring scan of the response +// plus map lookups, rather than iterating every session's UUID set. +func buildSessionTurnUUIDs(turnCounts []int, seed int64) (sets [][]string, owner map[string]int) { + owner = map[string]int{} + if len(turnCounts) == 0 { + return nil, owner + } + newUUID := newUUIDGenerator(seed) + sets = make([][]string, len(turnCounts)) + for i, n := range turnCounts { + if n <= 0 { + continue + } + uuids := make([]string, n) + for j := range uuids { + u := newUUID() + uuids[j] = u + owner[u] = i + } + sets[i] = uuids + } + return sets, owner +} + +// replayReciteWindowInstruction returns the tail instruction asking the +// model to output, as the FIRST line of its response, the id VALUES for the +// given ordered window of inline "[turn-N id: ...]" tags — in the same +// order, comma-separated — then continue normally. Unlike the retired +// per-session replayReciteFirstLineInstruction (which asked for "every UUID +// shown above" verbatim), this instruction names the exact turns to recite +// by LABEL, not position or value, so the model must locate each tagged +// turn in its own (possibly long) context rather than simply copying +// whatever happens to be nearby. Mirrors the cache-coherency eval's +// userMessage ("List every UUID shown in the request...") but (a) +// front-loads the ask to line 1 specifically, so a forced-output/ +// ignore_eos budget can still fill the remainder with the model's normal +// continuation, and (b) references labels instead of embedding the UUIDs +// themselves, so first-line conformity in the response reflects genuine +// recall from cached KV context, not an echo of the ask. +func replayReciteWindowInstruction(labels []string) string { + if len(labels) == 0 { + return "" + } + tagged := make([]string, len(labels)) + for i, l := range labels { + tagged[i] = "[" + l + "]" + } + return "\n\nSomewhere above, several ids are tagged like [turn-N id: ...]. On the FIRST line output ONLY the id values for these tags, " + + "in this exact order, comma-separated and nothing else: " + strings.Join(tagged, ", ") + ". Then continue normally." +} + +// ---- max_tokens recite floor ---- + +// replayReciteFloorMultiplier mirrors the cache-coherency eval's default +// --max-output-multiplier (see computeMaxOutputTokens): the recite floor is +// sized at multiplier x the expected N-UUID first-line list size, giving +// the model headroom to emit the full list without truncation. +const replayReciteFloorMultiplier = 3.0 + +// replayReciteFloorTokens returns the minimum max_tokens budget to enforce +// on a request that carries the recite ask (see uuidInjection.Recite), sized +// to fit the FIRST-LINE numUUIDs-UUID comma-joined list this feature asks +// for (reuses the cache-coherency eval's computeMaxOutputTokens sizing: +// numUUIDs*36 chars + separating commas, /4 for an approximate token count, +// x replayReciteFloorMultiplier). numUUIDs is now len(inj.ReciteUUIDs) — the +// current request's recite WINDOW, capped at 4 (see uuidInjection) — so the +// floor itself is now bounded and constant regardless of session length, +// unlike the retired per-session-N scheme where a long session's floor grew +// without bound. The tradeoff: for a request whose recorded output budget +// is tiny (a handful of tokens, e.g. a pure tool-call turn), this constant +// floor is a much LARGER ratio of the original budget than an N-scaled +// floor would have been at N=2 — i.e. the recite ask now perturbs a small +// turn's output-size profile proportionally more. Accepted: correctness +// (not truncating the recite line into a false PRESENCE_MISS) takes +// priority over preserving the exact captured output-size ratio. +// +// A router-replay request's max_tokens is normally sized off the ORIGINAL +// capture's output_tokens (see pickMaxTokens) — which for a tool-call-only +// turn can be a handful of tokens, nowhere near enough to also emit the +// first-line UUID list. Without this floor, a tiny budget truncates that +// first line, which would misread as PRESENCE_MISS/NOT_EXACT (coherency +// failure) when it's actually just an output-size artifact. +func replayReciteFloorTokens(numUUIDs int) int { + return computeMaxOutputTokens(numUUIDs, replayReciteFloorMultiplier) +} + +var reciteFloorWarnOnce sync.Once + +// applyReciteFloor raises maxTokens to replayReciteFloorTokens(numUUIDs) +// when recite is requested and the original budget falls short, warning +// once per process (mirrors the dataset path's reciteTruncWarned one-shot +// pattern, but this is a single global warning rather than per-conversation +// since the router path computes one floor value per call, not a per- +// conversation truncation computation). +func applyReciteFloor(maxTokens int, recite bool, numUUIDs int) int { + if !recite { + return maxTokens + } + floor := replayReciteFloorTokens(numUUIDs) + if maxTokens >= floor { + return maxTokens + } + reciteFloorWarnOnce.Do(func() { + fmt.Fprintf(os.Stderr, + "[router-replay] WARNING: max_tokens raised to the UUID-recite floor (%d) for one or more requests — "+ + "a tiny captured output budget would otherwise truncate the first-line recite list into a false PRESENCE_MISS, not real corruption\n", + floor) + }) + return floor +} diff --git a/benchmark/replay_router_uuid_test.go b/benchmark/replay_router_uuid_test.go new file mode 100644 index 0000000..52bc297 --- /dev/null +++ b/benchmark/replay_router_uuid_test.go @@ -0,0 +1,844 @@ +package benchmark + +// Pure, offline unit tests for router-replay UUID cache-coherency injection +// (replay_router_uuid.go, replay_uuid.go, and the buildInjection/wire-body +// plumbing in replay_router_post.go/replay_router_wire.go). No mocked LLM/ +// Chat — these exercise data transforms and the wire builders directly, per +// repo testing policy. + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "testing" +) + +// writeReplayV3File marshals a header + sessions into a replay-v3 JSONL +// file under t.TempDir() and returns its path. +func writeReplayV3File(t *testing.T, sessions []RouterReplaySession) string { + t.Helper() + dir := t.TempDir() + path := filepath.Join(dir, "replay.jsonl") + f, err := os.Create(path) + if err != nil { + t.Fatalf("create temp replay file: %v", err) + } + defer f.Close() + + hdr := RouterReplayHeader{ + Schema: "replay-v3", + Summary: RouterReplaySummary{ + Sessions: len(sessions), + }, + } + enc := json.NewEncoder(f) + if err := enc.Encode(hdr); err != nil { + t.Fatalf("encode header: %v", err) + } + for _, s := range sessions { + if err := enc.Encode(s); err != nil { + t.Fatalf("encode session: %v", err) + } + } + return path +} + +// userText is a shorthand for a qualifying user-turn message (role=="user", +// one "text" block). +func userText(hash string, bytes int) RouterReplayMessage { + return RouterReplayMessage{Role: "user", Hash: hash, BlockTypes: []string{"text"}, Bytes: bytes} +} + +// assistantText is a shorthand for an assistant message — never a +// qualifying turn (wrong role), included in fixtures purely to exercise +// that non-user messages are skipped. +func assistantText(hash string, bytes int) RouterReplayMessage { + return RouterReplayMessage{Role: "assistant", Hash: hash, BlockTypes: []string{"text"}, Bytes: bytes} +} + +// toolResultOnly is a shorthand for a role=="user" message carrying ONLY a +// tool_result block (no text) — the "exclude tool_result-only" case from +// isQualifyingUserTurn. +func toolResultOnly(hash string, bytes int) RouterReplayMessage { + return RouterReplayMessage{Role: "user", Hash: hash, BlockTypes: []string{"tool_result"}, Bytes: bytes, ToolResultIDs: []string{"tr1"}} +} + +// turnFixtureSessions builds two sessions exercising every +// isQualifyingUserTurn exclusion plus multi-instance turn ordering: +// +// - both sessions' first request opens with the SAME message hash +// ("shared-msg", a role=="user" text message) — shared across sessions +// (count==2) so it must NEVER qualify as a turn despite being +// role==user with a text block. +// - each session's main instance accumulates 3 genuine user turns +// (u1,u2,u3) interleaved with assistant replies (never turns) and, for +// session 0 only, a tool_result-only message (never a turn, even +// though role=="user"). +// - session 0 has a SECOND instance ("s0-sub") appearing AFTER the main +// instance in the Instances slice, contributing one more turn +// (s0-sub-u1) — verifies turn indices are session-global and ordered +// by instance/request/message file order, not per-instance. +func turnFixtureSessions() []RouterReplaySession { + mkMainInstance := func(id, prefix string, includeToolResultOnly bool) RouterReplayInstance { + req1 := RouterReplayRequest{ + RequestID: 1, + Messages: []RouterReplayMessage{userText("shared-msg", 60)}, + } + req2Msgs := []RouterReplayMessage{ + userText("shared-msg", 60), + assistantText(prefix+"-a1", 80), + userText(prefix+"-u1", 100), + } + req2 := RouterReplayRequest{RequestID: 2, Messages: req2Msgs} + + req3Msgs := append(append([]RouterReplayMessage{}, req2Msgs...), + assistantText(prefix+"-a2", 80), + userText(prefix+"-u2", 100), + ) + req3 := RouterReplayRequest{RequestID: 3, Messages: req3Msgs} + + req4Msgs := append([]RouterReplayMessage{}, req3Msgs...) + if includeToolResultOnly { + req4Msgs = append(req4Msgs, toolResultOnly(prefix+"-tool1", 40)) + } + req4 := RouterReplayRequest{RequestID: 4, Messages: req4Msgs} + + req5Msgs := append(append([]RouterReplayMessage{}, req4Msgs...), + assistantText(prefix+"-a3", 80), + userText(prefix+"-u3", 100), + ) + req5 := RouterReplayRequest{RequestID: 5, Messages: req5Msgs} + + return RouterReplayInstance{ + InstanceID: id, + Role: "main", + Requests: []RouterReplayRequest{req1, req2, req3, req4, req5}, + } + } + + subInstance := RouterReplayInstance{ + InstanceID: "s0-sub", + Role: "sub-agent", + Requests: []RouterReplayRequest{ + {RequestID: 6, Messages: []RouterReplayMessage{userText("s0-sub-u1", 90)}}, + }, + } + + s0 := RouterReplaySession{ + SessionID: "s0", + Instances: []RouterReplayInstance{ + mkMainInstance("s0-main", "s0", true), + subInstance, + }, + } + s1 := RouterReplaySession{ + SessionID: "s1", + Instances: []RouterReplayInstance{ + mkMainInstance("s1-main", "s1", false), + }, + } + return []RouterReplaySession{s0, s1} +} + +func TestComputeSessionTurnHashesIdentifiesQualifyingTurns(t *testing.T) { + path := writeReplayV3File(t, turnFixtureSessions()) + counts, err := computeBlockSessionCounts(path, nil, 0) + if err != nil { + t.Fatalf("computeBlockSessionCounts: %v", err) + } + if got := counts["shared-msg"]; got != 2 { + t.Fatalf("counts[shared-msg] = %d, want 2 (referenced by both sessions)", got) + } + + turnHashes, err := computeSessionTurnHashes(path, nil, 0, counts) + if err != nil { + t.Fatalf("computeSessionTurnHashes: %v", err) + } + if len(turnHashes) != 2 { + t.Fatalf("len(turnHashes) = %d, want 2 sessions", len(turnHashes)) + } + + wantS0 := []string{"s0-u1", "s0-u2", "s0-u3", "s0-sub-u1"} + if !equalStrSlices(turnHashes[0], wantS0) { + t.Errorf("session 0 turnHashes = %v, want %v", turnHashes[0], wantS0) + } + wantS1 := []string{"s1-u1", "s1-u2", "s1-u3"} + if !equalStrSlices(turnHashes[1], wantS1) { + t.Errorf("session 1 turnHashes = %v, want %v", turnHashes[1], wantS1) + } + + for _, session := range turnHashes { + for _, h := range session { + if h == "shared-msg" { + t.Error("shared-msg (count==2) qualified as a turn — cross-session-shared content must never be stamped") + } + if strings.Contains(h, "tool1") { + t.Errorf("tool_result-only hash %q qualified as a turn", h) + } + if strings.Contains(h, "-a") { + t.Errorf("assistant hash %q qualified as a turn", h) + } + } + } +} + +func TestComputeSessionTurnHashesRespectsFilters(t *testing.T) { + path := writeReplayV3File(t, turnFixtureSessions()) + counts, err := computeBlockSessionCounts(path, nil, 0) + if err != nil { + t.Fatalf("computeBlockSessionCounts: %v", err) + } + + t.Run("sessionLimit caps to the first N sessions", func(t *testing.T) { + turnHashes, err := computeSessionTurnHashes(path, nil, 1, counts) + if err != nil { + t.Fatalf("computeSessionTurnHashes: %v", err) + } + if len(turnHashes) != 1 { + t.Fatalf("len(turnHashes) = %d, want 1 (sessionLimit=1)", len(turnHashes)) + } + }) + + t.Run("allowed index set restricts to those sessions only", func(t *testing.T) { + // counts computed globally (both sessions), but the turn-hash pass + // only walks session index 1 (s1). + allowed := map[int]bool{1: true} + turnHashes, err := computeSessionTurnHashes(path, allowed, 0, counts) + if err != nil { + t.Fatalf("computeSessionTurnHashes: %v", err) + } + if len(turnHashes) != 1 { + t.Fatalf("len(turnHashes) = %d, want 1 (only index 1 allowed)", len(turnHashes)) + } + want := []string{"s1-u1", "s1-u2", "s1-u3"} + if !equalStrSlices(turnHashes[0], want) { + t.Errorf("turnHashes[0] = %v, want %v (session s1)", turnHashes[0], want) + } + }) +} + +func equalStrSlices(a, b []string) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +// TestBuildSessionTurnUUIDsDeterminism verifies buildSessionTurnUUIDs +// matches the dataset path's determinism contract: same seed -> same +// per-session-per-turn UUID assignment; different seed -> different +// assignment; every UUID unique across the whole run; owner correctly +// reverse-maps each UUID to its issuing session. +func TestBuildSessionTurnUUIDsDeterminism(t *testing.T) { + turnCounts := []int{3, 2, 4} + setsA, ownerA := buildSessionTurnUUIDs(turnCounts, 42) + setsB, ownerB := buildSessionTurnUUIDs(turnCounts, 42) + if len(setsA) != 3 || len(setsB) != 3 { + t.Fatalf("expected 3 sets each, got %d and %d", len(setsA), len(setsB)) + } + for i := range setsA { + if len(setsA[i]) != turnCounts[i] || len(setsB[i]) != turnCounts[i] { + t.Fatalf("session %d: expected %d UUIDs, got %v / %v", i, turnCounts[i], setsA[i], setsB[i]) + } + for j := range setsA[i] { + if setsA[i][j] != setsB[i][j] { + t.Errorf("session %d turn %d: same seed produced different UUIDs: %q vs %q", i, j, setsA[i][j], setsB[i][j]) + } + } + } + if len(ownerA) != len(ownerB) { + t.Errorf("owner map sizes differ: %d vs %d", len(ownerA), len(ownerB)) + } + + setsC, _ := buildSessionTurnUUIDs(turnCounts, 43) + same := true + for i := range setsA { + if len(setsA[i]) > 0 && len(setsC[i]) > 0 && setsA[i][0] != setsC[i][0] { + same = false + } + } + if same { + t.Error("different seeds produced an identical UUID assignment") + } + + seen := map[string]bool{} + for i, set := range setsA { + for _, u := range set { + if seen[u] { + t.Errorf("uuid %q assigned to more than one turn", u) + } + seen[u] = true + if got := ownerA[u]; got != i { + t.Errorf("owner[%q] = %d, want %d", u, got, i) + } + } + } + + if sets, owner := buildSessionTurnUUIDs(nil, 42); sets != nil || len(owner) != 0 { + t.Errorf("buildSessionTurnUUIDs(nil, ...) = (%v, %v), want (nil, empty)", sets, owner) + } +} + +// TestBuildSessionTurnUUIDsScalesWithTurnCount verifies each session gets +// EXACTLY turnCounts[i] UUIDs (one per turn, no floor/multiplier — unlike +// the retired per-session-N scheme) and a zero-turn session gets none. +func TestBuildSessionTurnUUIDsScalesWithTurnCount(t *testing.T) { + turnCounts := []int{0, 1, 5} + sets, owner := buildSessionTurnUUIDs(turnCounts, 7) + if len(sets[0]) != 0 { + t.Errorf("session 0 (0 turns): len = %d, want 0", len(sets[0])) + } + if len(sets[1]) != 1 { + t.Errorf("session 1 (1 turn): len = %d, want 1", len(sets[1])) + } + if len(sets[2]) != 5 { + t.Errorf("session 2 (5 turns): len = %d, want 5", len(sets[2])) + } + if len(owner) != 6 { + t.Errorf("len(owner) = %d, want 6 (1+5 issued UUIDs)", len(owner)) + } +} + +// newFixturePoster builds a replayPoster wired up exactly as +// runRouterReplayInstance does (see replay_router.go), for a single session +// (sessionIdx 0) with nTurns turns, without touching HTTP/newReplayPoster. +// All UUID state is global/read-only on the poster (see replay_router_post.go); +// the fixture's session lives at index 0 in the global slices, and callers +// pass sessionIdx=0 into buildInjection. +func newFixturePoster(nTurns int, seed int64, reciteEveryRequest bool) *replayPoster { + turnHashes := make([]string, nTurns) + for i := range turnHashes { + turnHashes[i] = fmt.Sprintf("h%d", i) + } + sets, owner := buildSessionTurnUUIDs([]int{nTurns}, seed) + return &replayPoster{ + uuidEnabled: true, + allUUIDSets: sets, + sessionTurnHashes: [][]string{turnHashes}, + owner: owner, + reciteEveryRequest: reciteEveryRequest, + } +} + +// visibleTurnsRequest builds a RouterReplayRequest whose Messages carry the +// qualifying-turn hashes h0..h(n-1) (in order), simulating a growing +// conversation history at turn n. +func visibleTurnsRequest(n int) RouterReplayRequest { + msgs := make([]RouterReplayMessage, 0, n) + for i := 0; i < n; i++ { + msgs = append(msgs, userText(fmt.Sprintf("h%d", i), 50)) + } + return RouterReplayRequest{Messages: msgs} +} + +// TestBuildInjectionWindowSelection verifies buildInjection's recite window: +// first (visible) turn + up to 3 most-recent turns EXCLUDING the current +// (highest-index visible) turn, deduplicated, capped at 4 — and that +// StampByHash always covers EVERY visible turn, not just the window. +// Exercises the edge cases at turns 1, 2, 3 explicitly (D1/D2/D3 in the +// plan) plus the steady-state 4-cap and window-sliding behavior beyond it. +func TestBuildInjectionWindowSelection(t *testing.T) { + p := newFixturePoster(6, 1, true) + + cases := []struct { + turn int // 1-based "current turn" being requested + wantLabels []string // expected inj.ReciteLabels + }{ + {1, []string{"turn-1"}}, + {2, []string{"turn-1"}}, + {3, []string{"turn-1", "turn-2"}}, + {4, []string{"turn-1", "turn-2", "turn-3"}}, + {5, []string{"turn-1", "turn-2", "turn-3", "turn-4"}}, + {6, []string{"turn-1", "turn-3", "turn-4", "turn-5"}}, + } + for _, c := range cases { + t.Run(fmt.Sprintf("turn-%d", c.turn), func(t *testing.T) { + req := visibleTurnsRequest(c.turn) + inj := p.buildInjection(req, 0, false) + if inj == nil { + t.Fatal("buildInjection returned nil, want a non-nil injection") + } + if !equalStrSlices(inj.ReciteLabels, c.wantLabels) { + t.Errorf("ReciteLabels = %v, want %v", inj.ReciteLabels, c.wantLabels) + } + if len(inj.ReciteLabels) > 4 { + t.Errorf("ReciteLabels len = %d, want <= 4", len(inj.ReciteLabels)) + } + if len(inj.ReciteUUIDs) != len(inj.ReciteLabels) { + t.Fatalf("ReciteUUIDs len = %d, want %d (matching ReciteLabels)", len(inj.ReciteUUIDs), len(inj.ReciteLabels)) + } + // Every UUID in the window must be this session's own (index 0). + for i, u := range inj.ReciteUUIDs { + if p.owner[u] != 0 { + t.Errorf("ReciteUUIDs[%d] = %q owned by session %d, want session 0", i, u, p.owner[u]) + } + } + // StampByHash must cover EVERY visible turn (h0..h(turn-1)), not + // just the recite window — this is what keeps every turn's + // marker warm in KV regardless of whether it's being recited + // this request. + if len(inj.StampByHash) != c.turn { + t.Errorf("StampByHash covers %d turns, want %d (every visible turn)", len(inj.StampByHash), c.turn) + } + for i := 0; i < c.turn; i++ { + h := fmt.Sprintf("h%d", i) + stamp, ok := inj.StampByHash[h] + if !ok { + t.Errorf("StampByHash missing visible turn hash %q", h) + continue + } + if stamp.Idx != i { + t.Errorf("StampByHash[%q].Idx = %d, want %d", h, stamp.Idx, i) + } + if stamp.Label != fmt.Sprintf("turn-%d", i+1) { + t.Errorf("StampByHash[%q].Label = %q, want turn-%d", h, stamp.Label, i+1) + } + } + }) + } +} + +// TestBuildInjectionNilCases verifies buildInjection degrades to "no +// injection" (nil) rather than panicking: disabled, session index out of +// range, session with zero turns, and a request with no visible qualifying +// turn at all. +func TestBuildInjectionNilCases(t *testing.T) { + t.Run("uuidEnabled false", func(t *testing.T) { + p := newFixturePoster(3, 1, true) + p.uuidEnabled = false + if got := p.buildInjection(visibleTurnsRequest(2), 0, false); got != nil { + t.Errorf("buildInjection = %v, want nil (disabled)", got) + } + }) + t.Run("sessionIdx out of range", func(t *testing.T) { + p := newFixturePoster(3, 1, true) + if got := p.buildInjection(visibleTurnsRequest(2), 5, false); got != nil { + t.Errorf("buildInjection = %v, want nil (sessionIdx out of range)", got) + } + }) + t.Run("zero-turn session", func(t *testing.T) { + p := newFixturePoster(0, 1, true) + if got := p.buildInjection(visibleTurnsRequest(0), 0, false); got != nil { + t.Errorf("buildInjection = %v, want nil (no turns)", got) + } + }) + t.Run("no qualifying turn visible in this request", func(t *testing.T) { + p := newFixturePoster(3, 1, true) + req := RouterReplayRequest{Messages: []RouterReplayMessage{assistantText("unrelated", 30)}} + if got := p.buildInjection(req, 0, false); got != nil { + t.Errorf("buildInjection = %v, want nil (nothing visible)", got) + } + }) +} + +// TestBuildInjectionRecite verifies Recite = reciteEveryRequest || +// isLastRequest, independent of window selection. +func TestBuildInjectionRecite(t *testing.T) { + req := visibleTurnsRequest(2) + + pAlways := newFixturePoster(3, 1, true) + if inj := pAlways.buildInjection(req, 0, false); inj == nil || !inj.Recite { + t.Error("reciteEveryRequest=true, isLastRequest=false: expected Recite=true") + } + + pFinalOnly := newFixturePoster(3, 1, false) + if inj := pFinalOnly.buildInjection(req, 0, false); inj == nil || inj.Recite { + t.Error("reciteEveryRequest=false, isLastRequest=false: expected Recite=false") + } + if inj := pFinalOnly.buildInjection(req, 0, true); inj == nil || !inj.Recite { + t.Error("reciteEveryRequest=false, isLastRequest=true: expected Recite=true") + } +} + +// TestWireInjectionDeterminism verifies that, for a fixed set of turn +// stamps, buildOpenAIChatCompletionsBody / buildAnthropicMessagesBody +// produce byte-identical bodies across repeated calls (same request + same +// injection in -> same bytes out) and across two DIFFERENT requests that +// both carry the SAME turn message (the within-session cache-reuse +// property), while two DIFFERENT sessions' turn stamps diverge the body. +func TestWireInjectionDeterminism(t *testing.T) { + docs := strings.Repeat("wire-injection-docs ", 100) + stampA := map[string]turnStamp{"msg1": {Idx: 0, UUID: "uuid-session-A", Label: "turn-1"}} + stampB := map[string]turnStamp{"msg1": {Idx: 0, UUID: "uuid-session-B", Label: "turn-1"}} + req := RouterReplayRequest{ + InputTokens: 500, + Messages: []RouterReplayMessage{userText("msg1", 100)}, + } + + for _, kind := range []string{"openai", "anthropic"} { + build := func(r RouterReplayRequest, inj *uuidInjection) []byte { + var body []byte + var err error + if kind == "openai" { + body, _, err = buildOpenAIChatCompletionsBody(r, docs, "model", "", 0, false, inj) + } else { + body, _, err = buildAnthropicMessagesBody(r, docs, "model", "", 0, false, inj) + } + if err != nil { + t.Fatalf("%s build: %v", kind, err) + } + return body + } + injA := &uuidInjection{StampByHash: stampA} + injA2 := &uuidInjection{StampByHash: stampA} + injB := &uuidInjection{StampByHash: stampB} + + bodyA1 := build(req, injA) + bodyA2 := build(req, injA2) + bodyB := build(req, injB) + + if string(bodyA1) != string(bodyA2) { + t.Errorf("%s: identical injection produced different bytes", kind) + } + if string(bodyA1) == string(bodyB) { + t.Errorf("%s: different sessions' turn stamps produced identical bytes", kind) + } + if !strings.Contains(string(bodyA1), "uuid-session-A") { + t.Errorf("%s: body missing session A's own turn UUID", kind) + } + if strings.Contains(string(bodyA1), "uuid-session-B") { + t.Errorf("%s: body A leaked session B's turn UUID into the wire body", kind) + } + + // A SECOND, different request that repeats the SAME turn message + // (msg1) must emit the byte-identical stamped content for it. + req2 := RouterReplayRequest{ + InputTokens: 500, + Messages: []RouterReplayMessage{ + userText("msg1", 100), + assistantText("msg2", 60), + }, + } + bodyA3 := build(req2, injA) + var p1, p3 map[string]interface{} + if err := json.Unmarshal(bodyA1, &p1); err != nil { + t.Fatalf("unmarshal bodyA1: %v", err) + } + if err := json.Unmarshal(bodyA3, &p3); err != nil { + t.Fatalf("unmarshal bodyA3: %v", err) + } + msgs1, _ := p1["messages"].([]interface{}) + msgs3, _ := p3["messages"].([]interface{}) + if len(msgs1) == 0 || len(msgs3) == 0 { + t.Fatalf("%s: expected non-empty messages in both bodies", kind) + } + first1, _ := json.Marshal(msgs1[0]) + first3, _ := json.Marshal(msgs3[0]) + if string(first1) != string(first3) { + t.Errorf("%s: turn msg1's stamped content diverged across two different requests:\n%s\nvs\n%s", kind, first1, first3) + } + } +} + +// TestInlineMarkerAppendedToTurnMessage verifies the injected marker's exact +// shape ("\n\n[turn-N id: ]") lands inside the stamped message's OWN +// text content (both wire body and canonical text), and that an unrelated +// message in the same request is left untouched. +func TestInlineMarkerAppendedToTurnMessage(t *testing.T) { + docs := strings.Repeat("inline-marker-docs ", 100) + stampByHash := map[string]turnStamp{"stamped-msg": {Idx: 3, UUID: "abc-uuid", Label: "turn-4"}} + inj := &uuidInjection{StampByHash: stampByHash} + req := RouterReplayRequest{ + InputTokens: 500, + Messages: []RouterReplayMessage{ + userText("stamped-msg", 100), + assistantText("other-msg", 80), + }, + } + wantMarker := "\n\n[turn-4 id: abc-uuid]" + + body, canonical, err := buildAnthropicMessagesBody(req, docs, "model", "", 0, false, inj) + if err != nil { + t.Fatalf("build: %v", err) + } + if !strings.Contains(canonical, wantMarker) { + t.Errorf("canonical text missing marker %q", wantMarker) + } + var parsed map[string]interface{} + if err := json.Unmarshal(body, &parsed); err != nil { + t.Fatalf("unmarshal: %v", err) + } + msgs, _ := parsed["messages"].([]interface{}) + if len(msgs) != 2 { + t.Fatalf("expected 2 messages, got %d", len(msgs)) + } + // Compare against the UNMARSHALED content (fmt.Sprintf, not + // json.Marshal) so a "\n" in the marker is a real newline byte, not the + // two-character JSON escape sequence "\\n". + first, _ := msgs[0].(map[string]interface{}) + firstContent := fmt.Sprintf("%v", first["content"]) + if !strings.Contains(firstContent, wantMarker) { + t.Errorf("stamped message content missing marker: %s", firstContent) + } + second, _ := msgs[1].(map[string]interface{}) + secondContent := fmt.Sprintf("%v", second["content"]) + if strings.Contains(secondContent, "abc-uuid") { + t.Errorf("marker leaked into the unrelated (unstamped) message: %s", secondContent) + } +} + +// TestCacheFidelitySharedBlockUnaffectedByInjection verifies the core +// fidelity invariant: a message hash that is NOT a key in StampByHash +// (standing in for a cross-session-shared, count>1 block — see +// isQualifyingUserTurn, which excludes those from ever being stamped) emits +// byte-identical content whether or not UUID injection is active elsewhere +// in the SAME request. +func TestCacheFidelitySharedBlockUnaffectedByInjection(t *testing.T) { + docs := strings.Repeat("fidelity-docs ", 100) + req := RouterReplayRequest{ + InputTokens: 500, + Messages: []RouterReplayMessage{ + userText("shared-unstamped-msg", 100), // stands in for a count>1 block + userText("own-msg", 80), + }, + } + + bodyNoInj, _, err := buildAnthropicMessagesBody(req, docs, "model", "", 0, false, nil) + if err != nil { + t.Fatalf("build (no injection): %v", err) + } + inj := &uuidInjection{StampByHash: map[string]turnStamp{"own-msg": {Idx: 0, UUID: "own-uuid", Label: "turn-1"}}} + bodyWithInj, _, err := buildAnthropicMessagesBody(req, docs, "model", "", 0, false, inj) + if err != nil { + t.Fatalf("build (with injection): %v", err) + } + + var parsedNoInj, parsedWithInj map[string]interface{} + if err := json.Unmarshal(bodyNoInj, &parsedNoInj); err != nil { + t.Fatalf("unmarshal (no injection): %v", err) + } + if err := json.Unmarshal(bodyWithInj, &parsedWithInj); err != nil { + t.Fatalf("unmarshal (with injection): %v", err) + } + msgsNoInj, _ := parsedNoInj["messages"].([]interface{}) + msgsWithInj, _ := parsedWithInj["messages"].([]interface{}) + if len(msgsNoInj) != 2 || len(msgsWithInj) != 2 { + t.Fatalf("expected 2 messages in both bodies, got %d and %d", len(msgsNoInj), len(msgsWithInj)) + } + + sharedNoInj, _ := json.Marshal(msgsNoInj[0]) + sharedWithInj, _ := json.Marshal(msgsWithInj[0]) + if string(sharedNoInj) != string(sharedWithInj) { + t.Errorf("unstamped (shared) message diverged with injection active elsewhere:\nwithout: %s\nwith: %s", sharedNoInj, sharedWithInj) + } + + ownWithInj, _ := json.Marshal(msgsWithInj[1]) + if !strings.Contains(string(ownWithInj), "own-uuid") { + t.Errorf("stamped message missing its own marker: %s", ownWithInj) + } +} + +// TestFindLeakedUUIDsByOwner exercises the router path's O(response) +// contamination scanner: a known OTHER-session UUID is flagged with the +// correct series index; the caller's OWN uuid is never flagged; a +// UUID-shaped string with no entry in owner is silently ignored; and +// multiple leaks are returned in deterministic (scan-order) sequence. +// Mirrors validateReplayResponse/FindLeakedUUIDs' semantics for the dataset +// path, but via the reverse uuid->owner map instead of iterating every +// session's UUID set. +func TestFindLeakedUUIDsByOwner(t *testing.T) { + owner := map[string]int{ + "11111111-1111-1111-1111-111111111111": 0, + "22222222-2222-2222-2222-222222222222": 0, + "33333333-3333-3333-3333-333333333333": 1, + "44444444-4444-4444-4444-444444444444": 2, + } + + t.Run("other-session uuid flagged with correct series", func(t *testing.T) { + resp := "here it is: 33333333-3333-3333-3333-333333333333" + got := findLeakedUUIDsByOwner(resp, "", 0, owner) + if len(got) != 1 || !strings.Contains(got[0], "33333333-3333-3333-3333-333333333333") || !strings.Contains(got[0], "series=1") { + t.Errorf("got %v, want one entry naming series=1", got) + } + }) + + t.Run("own uuid never flagged", func(t *testing.T) { + resp := "own ids: 11111111-1111-1111-1111-111111111111, 22222222-2222-2222-2222-222222222222" + got := findLeakedUUIDsByOwner(resp, "", 0, owner) + if len(got) != 0 { + t.Errorf("got %v, want none (both are the caller's own uuids)", got) + } + }) + + t.Run("unowned uuid-shaped string ignored", func(t *testing.T) { + resp := "random: 99999999-9999-9999-9999-999999999999" + got := findLeakedUUIDsByOwner(resp, "", 0, owner) + if len(got) != 0 { + t.Errorf("got %v, want none (not a real stamp)", got) + } + }) + + t.Run("thinking channel scanned too", func(t *testing.T) { + got := findLeakedUUIDsByOwner("no leak here", "but here: 44444444-4444-4444-4444-444444444444", 0, owner) + if len(got) != 1 || !strings.Contains(got[0], "series=2") { + t.Errorf("got %v, want one entry naming series=2 (found in thinking)", got) + } + }) + + t.Run("deterministic scan order for multiple leaks", func(t *testing.T) { + resp := "first 44444444-4444-4444-4444-444444444444 then 33333333-3333-3333-3333-333333333333" + got1 := findLeakedUUIDsByOwner(resp, "", 0, owner) + got2 := findLeakedUUIDsByOwner(resp, "", 0, owner) + if len(got1) != 2 { + t.Fatalf("got %v, want 2 leaked entries", got1) + } + for i := range got1 { + if got1[i] != got2[i] { + t.Errorf("non-deterministic order: %v vs %v", got1, got2) + } + } + if !strings.Contains(got1[0], "series=2") || !strings.Contains(got1[1], "series=1") { + t.Errorf("got %v, want series=2 before series=1 (scan order)", got1) + } + }) + + t.Run("empty owner map returns nil", func(t *testing.T) { + if got := findLeakedUUIDsByOwner("11111111-1111-1111-1111-111111111111", "", 0, map[string]int{}); got != nil { + t.Errorf("got %v, want nil", got) + } + }) +} + +// TestFirstLineConformity exercises firstLineConformity (the output- +// conformity check --replay-inject-uuids scores, mirroring the coherency +// eval's matchesExpectedUUIDList/ExactMatch): pass on an exact ordered +// comma-joined first line; fail on missing, reordered, or chatty first +// lines; pass when line 1 is exact even though LATER lines contain filler +// (the whole point of front-loading the ask to line 1 while forced-output +// keeps generating). expected here plays the role of inj.ReciteUUIDs — the +// request's recite WINDOW, not the session's full turn history. +func TestFirstLineConformity(t *testing.T) { + expected := []string{"uuid-a", "uuid-b", "uuid-c"} + + cases := []struct { + name string + resp string + want bool + }{ + {"exact single line", "uuid-a, uuid-b, uuid-c", true}, + {"exact with surrounding whitespace tolerated", " uuid-a, uuid-b, uuid-c ", true}, + {"exact first line, filler after", "uuid-a, uuid-b, uuid-c\nHere is more detail about your request...", true}, + {"exact first line, multiple filler lines after", "uuid-a, uuid-b, uuid-c\nline2\nline3 with more text", true}, + {"missing a uuid", "uuid-a, uuid-c", false}, + {"reordered", "uuid-b, uuid-a, uuid-c", false}, + {"chatty first line", "Sure! The UUIDs are uuid-a, uuid-b, uuid-c", false}, + {"uuids only on line 2, not line 1", "Sure, here you go:\nuuid-a, uuid-b, uuid-c", false}, + {"empty response", "", false}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if got := firstLineConformity(c.resp, expected); got != c.want { + t.Errorf("firstLineConformity(%q) = %v, want %v", c.resp, got, c.want) + } + }) + } +} + +// TestApplyReciteFloor verifies the max_tokens recite-floor helper: raises +// a too-small budget to replayReciteFloorTokens(numUUIDs) only when recite +// is requested; leaves larger budgets and non-recite calls untouched. +func TestApplyReciteFloor(t *testing.T) { + cases := []struct { + name string + tokens int + recite bool + numUUIDs int + }{ + {"below floor, recite -> raised", 5, true, 2}, + {"at floor, recite -> unchanged", replayReciteFloorTokens(2), true, 2}, + {"above floor, recite -> unchanged", 100000, true, 2}, + {"below floor, no recite -> unchanged", 5, false, 2}, + {"below floor, more uuids (capped at 4), recite -> raised higher", 5, true, 4}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + floor := replayReciteFloorTokens(c.numUUIDs) + want := c.tokens + if c.recite && c.tokens < floor { + want = floor + } + if got := applyReciteFloor(c.tokens, c.recite, c.numUUIDs); got != want { + t.Errorf("applyReciteFloor(%d, %v, %d) = %d, want %d", c.tokens, c.recite, c.numUUIDs, got, want) + } + }) + } +} + +// TestReciteFloorScalesWithN verifies the recite floor grows with numUUIDs +// up to the window cap (4) — more UUIDs to recite on the first line needs a +// bigger budget — but is now BOUNDED, unlike the retired per-session-N +// scheme where a long session's floor grew without bound. +func TestReciteFloorScalesWithN(t *testing.T) { + small := replayReciteFloorTokens(1) + large := replayReciteFloorTokens(4) + if large <= small { + t.Errorf("replayReciteFloorTokens(4) = %d, want > replayReciteFloorTokens(1) = %d", large, small) + } +} + +// TestMaxTokensFloorAppliedInWireBuilders verifies the floor is actually +// wired into both body builders' emitted max_tokens when a recite +// injection is present and the original/recorded budget is tiny — the +// scenario a real tool-call-only turn would hit. numUUIDs is now +// len(inj.ReciteUUIDs) (the window, capped at 4), not the old per-session N. +func TestMaxTokensFloorAppliedInWireBuilders(t *testing.T) { + docs := strings.Repeat("floor-docs ", 100) + req := RouterReplayRequest{ + InputTokens: 500, + OutputTokens: 5, // tiny recorded budget -- would truncate the recite line + Messages: []RouterReplayMessage{userText("msg1", 100)}, + } + reciteUUIDs := []string{"uuid-floor-0", "uuid-floor-1"} + inj := &uuidInjection{ + StampByHash: map[string]turnStamp{"msg1": {Idx: 0, UUID: reciteUUIDs[0], Label: "turn-1"}}, + Recite: true, + ReciteLabels: []string{"turn-1"}, + ReciteUUIDs: reciteUUIDs, + } + wantFloor := float64(replayReciteFloorTokens(len(reciteUUIDs))) + + anthBody, _, err := buildAnthropicMessagesBody(req, docs, "model", "", 0, false, inj) + if err != nil { + t.Fatalf("anthropic build: %v", err) + } + var anthParsed map[string]interface{} + if err := json.Unmarshal(anthBody, &anthParsed); err != nil { + t.Fatalf("anthropic unmarshal: %v", err) + } + if got := anthParsed["max_tokens"].(float64); got != wantFloor { + t.Errorf("anthropic max_tokens = %v, want %v (floor)", got, wantFloor) + } + + openaiBody, _, err := buildOpenAIChatCompletionsBody(req, docs, "model", "", 0, false, inj) + if err != nil { + t.Fatalf("openai build: %v", err) + } + var openaiParsed map[string]interface{} + if err := json.Unmarshal(openaiBody, &openaiParsed); err != nil { + t.Fatalf("openai unmarshal: %v", err) + } + if got := openaiParsed["max_tokens"].(float64); got != wantFloor { + t.Errorf("openai max_tokens = %v, want %v (floor)", got, wantFloor) + } + + // Without injection (nil), the tiny recorded output_tokens is honored + // as before -- the floor must never apply when there's no recite ask. + plainBody, _, err := buildAnthropicMessagesBody(req, docs, "model", "", 0, false, nil) + if err != nil { + t.Fatalf("anthropic build (no injection): %v", err) + } + var plainParsed map[string]interface{} + if err := json.Unmarshal(plainBody, &plainParsed); err != nil { + t.Fatalf("anthropic unmarshal (no injection): %v", err) + } + if got, want := plainParsed["max_tokens"].(float64), float64(5); got != want { + t.Errorf("anthropic max_tokens (no injection) = %v, want %v (unfloored)", got, want) + } +} diff --git a/benchmark/replay_router_wire.go b/benchmark/replay_router_wire.go index 272badd..fae40c9 100644 --- a/benchmark/replay_router_wire.go +++ b/benchmark/replay_router_wire.go @@ -41,10 +41,20 @@ const verboseOutputInstruction = "Provide a thorough, detailed response and keep // as possible. Returns the marshaled JSON bytes and the canonical text // string (all synthesized content concatenated in generation order) for // feeding into the content-level cache estimator. -func buildAnthropicMessagesBody(req RouterReplayRequest, docs string, modelName string, runID string, outputRatio float64, forceOutput bool) ([]byte, string, error) { +// +// inj carries the UUID cache-coherency injection (--replay-inject-uuids, +// router path — see replay_router_uuid.go); nil means "no injection", +// leaving the body byte-for-byte identical to before this feature existed. +func buildAnthropicMessagesBody(req RouterReplayRequest, docs string, modelName string, runID string, outputRatio float64, forceOutput bool, inj *uuidInjection) ([]byte, string, error) { + var stampByHash map[string]turnStamp + injNumUUIDs := 0 + if inj != nil { + stampByHash = inj.StampByHash + injNumUUIDs = len(inj.ReciteUUIDs) + } body := map[string]interface{}{ "model": modelName, - "max_tokens": pickMaxTokens(req, outputRatio), + "max_tokens": applyReciteFloor(pickMaxTokens(req, outputRatio), inj != nil && inj.Recite, injNumUUIDs), "stream": req.Stream, } if req.Temperature != nil { @@ -71,6 +81,7 @@ func buildAnthropicMessagesBody(req RouterReplayRequest, docs string, modelName } systemArr = append([]map[string]interface{}{stamp}, systemArr...) } + if forceOutput { systemArr = append(systemArr, map[string]interface{}{ "type": "text", @@ -83,11 +94,21 @@ func buildAnthropicMessagesBody(req RouterReplayRequest, docs string, modelName if req.Tools != nil && req.Tools.Count > 0 { body["tools"] = buildTools(req.Tools, docs) } + var msgs []map[string]interface{} if len(req.Messages) > 0 { - body["messages"] = buildMessages(req.Messages, docs) + msgs = buildMessages(req.Messages, docs, stampByHash) + } + if inj != nil && inj.Recite { + msgs = appendTailMessageAnthropic(msgs, replayReciteWindowInstruction(inj.ReciteLabels)) + } + if len(msgs) > 0 { + body["messages"] = msgs } - // Collect canonical text for the cache estimator. + // Collect canonical text for the cache estimator. Uses the SAME + // buildMessageContent(m, docs, stampByHash) call as buildMessages above + // so the cache-estimate canonical matches the wire body exactly + // (including any inline turn-stamp markers). var canonical strings.Builder if runID != "" { canonical.WriteString(fmt.Sprintf("RUN_GUID: %s", runID)) @@ -112,7 +133,7 @@ func buildAnthropicMessagesBody(req RouterReplayRequest, docs string, modelName } } for _, m := range req.Messages { - blocks := buildMessageContent(m, docs) + blocks := buildMessageContent(m, docs, stampByHash) for _, blk := range blocks { if t, ok := blk["text"]; ok { canonical.WriteString(fmt.Sprintf("%v", t)) @@ -132,6 +153,9 @@ func buildAnthropicMessagesBody(req RouterReplayRequest, docs string, modelName } } } + if inj != nil && inj.Recite { + canonical.WriteString(replayReciteWindowInstruction(inj.ReciteLabels)) + } bodyBytes, err := json.Marshal(body) return bodyBytes, canonical.String(), err @@ -235,10 +259,13 @@ func buildTools(spec *RouterReplayToolsSpec, docs string) []map[string]interface // share of the message's total Bytes. For tool_use and tool_result blocks // we preserve the original ids verbatim so the conversation maintains a // valid reference graph that matches what the original capture sent. -func buildMessages(msgs []RouterReplayMessage, docs string) []map[string]interface{} { +// stampByHash is inj.StampByHash (nil when --replay-inject-uuids is off) — +// threaded down to buildMessageContent, which appends each qualifying +// turn's inline UUID marker to its own synthesized content. +func buildMessages(msgs []RouterReplayMessage, docs string, stampByHash map[string]turnStamp) []map[string]interface{} { out := make([]map[string]interface{}, 0, len(msgs)) for _, m := range msgs { - content := buildMessageContent(m, docs) + content := buildMessageContent(m, docs, stampByHash) entry := map[string]interface{}{ "role": roleOrUser(m.Role), "content": content, @@ -264,11 +291,57 @@ func roleOrUser(role string) string { return "user" } +// appendTailMessageAnthropic appends text (the windowed recite ask — see +// replayReciteWindowInstruction) to the end of an Anthropic messages array, +// preserving strict user/ +// assistant role alternation (Anthropic rejects consecutive same-role +// messages, and "system" is not a valid role inside `messages` at all): +// - if the array is empty, or the last message is role "assistant", a +// NEW role="user" message carrying the text is appended (valid +// alternation; a fresh assistant turn should never have text +// injected into it after the fact). +// - if the last message is role "user" (the common case — router-replay +// requests carry history up to, but not including, the response being +// generated), the text is appended as an additional content block on +// THAT message instead of a new one. +func appendTailMessageAnthropic(msgs []map[string]interface{}, text string) []map[string]interface{} { + if text == "" { + return msgs + } + newUserMsg := map[string]interface{}{ + "role": "user", + "content": []map[string]interface{}{{"type": "text", "text": text}}, + } + if len(msgs) == 0 { + return append(msgs, newUserMsg) + } + last := msgs[len(msgs)-1] + if last["role"] != "user" { + return append(msgs, newUserMsg) + } + content, _ := last["content"].([]map[string]interface{}) + content = append(content, map[string]interface{}{"type": "text", "text": text}) + last["content"] = content + msgs[len(msgs)-1] = last + return msgs +} + // buildMessageContent materializes per-block content for one message. // block_types lists the kinds in order; we use tool_use_ids and // tool_result_ids to populate ids on the matching blocks (in order of // appearance in block_types). -func buildMessageContent(m RouterReplayMessage, docs string) []map[string]interface{} { +// +// stampByHash is inj.StampByHash (nil when --replay-inject-uuids is off, or +// when this message isn't a qualifying turn). When m.Hash is a key in +// stampByHash, the labeled marker "\n\n[turn-N id: ]" is appended to +// the LAST "text"-type block's synthesized content, at a position wholly +// determined by m.Hash (same seed synthText already keys on) — so two +// requests carrying the same turn emit byte-identical content for it, +// preserving the within-session cache-hit property. Only count==1 +// (genuinely per-session) messages are ever present in stampByHash — see +// isQualifyingUserTurn — so a cross-session-shared message is never +// perturbed. +func buildMessageContent(m RouterReplayMessage, docs string, stampByHash map[string]turnStamp) []map[string]interface{} { nText := 0 for _, t := range m.BlockTypes { if t == "text" { @@ -347,6 +420,15 @@ func buildMessageContent(m RouterReplayMessage, docs string) []map[string]interf }) } } + if stamp, ok := stampByHash[m.Hash]; ok { + marker := "\n\n[" + stamp.Label + " id: " + stamp.UUID + "]" + for i := len(out) - 1; i >= 0; i-- { + if out[i]["type"] == "text" { + out[i]["text"] = fmt.Sprintf("%v", out[i]["text"]) + marker + break + } + } + } return out } @@ -460,10 +542,19 @@ func buildOpenAITools(spec *RouterReplayToolsSpec, docs string) []map[string]int // - Anthropic-specific fields (top_k, thinking) are dropped. // - Stream options with include_usage are set so we get token counts in // the final SSE chunk. -func buildOpenAIChatCompletionsBody(req RouterReplayRequest, docs string, modelName string, runID string, outputRatio float64, forceOutput bool) ([]byte, string, error) { +// +// inj carries the UUID cache-coherency injection (--replay-inject-uuids, +// router path — see replay_router_uuid.go); nil means "no injection". +func buildOpenAIChatCompletionsBody(req RouterReplayRequest, docs string, modelName string, runID string, outputRatio float64, forceOutput bool, inj *uuidInjection) ([]byte, string, error) { + var stampByHash map[string]turnStamp + injNumUUIDs := 0 + if inj != nil { + stampByHash = inj.StampByHash + injNumUUIDs = len(inj.ReciteUUIDs) + } body := map[string]interface{}{ "model": modelName, - "max_tokens": pickMaxTokens(req, outputRatio), + "max_tokens": applyReciteFloor(pickMaxTokens(req, outputRatio), inj != nil && inj.Recite, injNumUUIDs), "stream": req.Stream, } if req.Temperature != nil { @@ -502,6 +593,7 @@ func buildOpenAIChatCompletionsBody(req RouterReplayRequest, docs string, modelN } messages = append([]map[string]interface{}{stamp}, messages...) } + if forceOutput { // Force the model to generate up to max_tokens: append a short // continue-generating instruction as a system message AND set vLLM's @@ -524,16 +616,24 @@ func buildOpenAIChatCompletionsBody(req RouterReplayRequest, docs string, modelN // each block and join them with newlines, and log a one-time warning so // the operator knows tool fidelity is lost. if len(req.Messages) > 0 { - openaiMsgs := buildOpenAIMessages(req.Messages, docs) + openaiMsgs := buildOpenAIMessages(req.Messages, docs, stampByHash) messages = append(messages, openaiMsgs...) } + if inj != nil && inj.Recite { + messages = appendTailMessageOpenAI(messages, replayReciteWindowInstruction(inj.ReciteLabels)) + } + body["messages"] = messages if req.Tools != nil && req.Tools.Count > 0 { body["tools"] = buildOpenAITools(req.Tools, docs) } - // Collect canonical text for the cache estimator (system blocks + messages). + // Collect canonical text for the cache estimator (system blocks + + // messages). Uses the SAME buildMessageContent(m, docs, stampByHash) + // call as buildOpenAIMessages above so the cache-estimate canonical + // matches the wire body exactly (including any inline turn-stamp + // markers). var canonical strings.Builder if runID != "" { canonical.WriteString(fmt.Sprintf("RUN_GUID: %s", runID)) @@ -542,7 +642,7 @@ func buildOpenAIChatCompletionsBody(req RouterReplayRequest, docs string, modelN canonical.WriteString(synthText(b.Hash, b.Bytes, docs)) } for _, m := range req.Messages { - blocks := buildMessageContent(m, docs) + blocks := buildMessageContent(m, docs, stampByHash) for _, blk := range blocks { if t, ok := blk["text"]; ok { canonical.WriteString(fmt.Sprintf("%v", t)) @@ -555,21 +655,57 @@ func buildOpenAIChatCompletionsBody(req RouterReplayRequest, docs string, modelN } } } + if inj != nil && inj.Recite { + canonical.WriteString(replayReciteWindowInstruction(inj.ReciteLabels)) + } bodyBytes, err := json.Marshal(body) return bodyBytes, canonical.String(), err } +// appendTailMessageOpenAI appends text (the windowed recite ask — see +// replayReciteWindowInstruction) to the end of an OpenAI messages array. +// Unlike Anthropic, OpenAI has no +// strict role-alternation requirement, but we still fold the text into the +// last message's content when that message is one the model would read as +// its own turn's input (user/system/tool) rather than always creating a +// new message — keeping the shape close to a real client's behavior. If +// the last message is "assistant" (or there are no messages at all), a new +// role="user" message carrying the text is appended instead. +func appendTailMessageOpenAI(msgs []map[string]interface{}, text string) []map[string]interface{} { + if text == "" { + return msgs + } + newUserMsg := map[string]interface{}{"role": "user", "content": text} + if len(msgs) == 0 { + return append(msgs, newUserMsg) + } + last := msgs[len(msgs)-1] + role, _ := last["role"].(string) + if role == "assistant" { + return append(msgs, newUserMsg) + } + existing, _ := last["content"].(string) + if existing != "" { + last["content"] = existing + "\n\n" + text + } else { + last["content"] = text + } + msgs[len(msgs)-1] = last + return msgs +} + // buildOpenAIMessages converts router-replay messages into OpenAI chat // messages. tool_use blocks become tool_calls on the assistant message; // tool_result blocks become separate role="tool" messages. Orphaned // tool_result blocks (no matching prior tool_call) are folded into user text. -func buildOpenAIMessages(msgs []RouterReplayMessage, docs string) []map[string]interface{} { +// stampByHash is threaded through to buildMessageContent (see its doc). +func buildOpenAIMessages(msgs []RouterReplayMessage, docs string, stampByHash map[string]turnStamp) []map[string]interface{} { out := make([]map[string]interface{}, 0, len(msgs)) seenToolCallIDs := map[string]bool{} for _, m := range msgs { - blocks := buildMessageContent(m, docs) + blocks := buildMessageContent(m, docs, stampByHash) role := roleOrUser(m.Role) if role == "assistant" { diff --git a/benchmark/replay_router_wire_test.go b/benchmark/replay_router_wire_test.go index 460f214..e461b52 100644 --- a/benchmark/replay_router_wire_test.go +++ b/benchmark/replay_router_wire_test.go @@ -26,7 +26,7 @@ func TestBuildAnthropicMessagesBodyRunGUIDStamp(t *testing.T) { // --- runID set --- runID := "test-run-id" - body, _, err := buildAnthropicMessagesBody(req, docs, modelName, runID, 0, false) + body, _, err := buildAnthropicMessagesBody(req, docs, modelName, runID, 0, false, nil) if err != nil { t.Fatalf("buildAnthropicMessagesBody with runID: %v", err) } @@ -53,7 +53,7 @@ func TestBuildAnthropicMessagesBodyRunGUIDStamp(t *testing.T) { } // --- runID empty --- - body, _, err = buildAnthropicMessagesBody(req, docs, modelName, "", 0, false) + body, _, err = buildAnthropicMessagesBody(req, docs, modelName, "", 0, false, nil) if err != nil { t.Fatalf("buildAnthropicMessagesBody empty runID: %v", err) } @@ -74,7 +74,7 @@ func TestBuildAnthropicMessagesBodyRunGUIDStamp(t *testing.T) { // --- runID set, 0 system blocks --- reqNoSys := RouterReplayRequest{} // no SystemBlocks - body, _, err = buildAnthropicMessagesBody(reqNoSys, docs, modelName, runID, 0, false) + body, _, err = buildAnthropicMessagesBody(reqNoSys, docs, modelName, runID, 0, false, nil) if err != nil { t.Fatalf("buildAnthropicMessagesBody no system blocks: %v", err) } @@ -107,8 +107,8 @@ func TestBuildAnthropicMessagesBodyCanonicalDeterminism(t *testing.T) { {Hash: "m1", Role: "user", BlockTypes: []string{"text"}, Bytes: 30}, }, } - _, c1, err1 := buildAnthropicMessagesBody(req, docs, "model", "run1", 0, false) - _, c2, err2 := buildAnthropicMessagesBody(req, docs, "model", "run1", 0, false) + _, c1, err1 := buildAnthropicMessagesBody(req, docs, "model", "run1", 0, false, nil) + _, c2, err2 := buildAnthropicMessagesBody(req, docs, "model", "run1", 0, false, nil) if err1 != nil || err2 != nil { t.Fatalf("errors: %v %v", err1, err2) } @@ -131,7 +131,7 @@ func TestBuildAnthropicMessagesBodyCanonicalContainsAllBlocks(t *testing.T) { {Hash: "msg1", Role: "user", BlockTypes: []string{"text"}, Bytes: 100}, }, } - _, canonical, err := buildAnthropicMessagesBody(req, docs, "model", "runX", 0, false) + _, canonical, err := buildAnthropicMessagesBody(req, docs, "model", "runX", 0, false, nil) if err != nil { t.Fatalf("error: %v", err) } @@ -174,12 +174,12 @@ func TestEffectiveSystemBlocksSkipsHeader(t *testing.T) { hdrText := synthText("uniq-header-per-req", 106, "") for _, builder := range []struct { name string - fn func(RouterReplayRequest, string, string, string, float64, bool) ([]byte, string, error) + fn func(RouterReplayRequest, string, string, string, float64, bool, *uuidInjection) ([]byte, string, error) }{ {"openai", buildOpenAIChatCompletionsBody}, {"anthropic", buildAnthropicMessagesBody}, } { - body, canonical, err := builder.fn(req, "", "m", "", 0, false) + body, canonical, err := builder.fn(req, "", "m", "", 0, false, nil) if err != nil { t.Fatalf("%s build: %v", builder.name, err) } @@ -274,7 +274,7 @@ func TestBuildAnthropicMessagesBodyForceOutput(t *testing.T) { } // --- force-output off (--replay-natural-output) --- - body, _, err := buildAnthropicMessagesBody(req, docs, "model", "", 0, false) + body, _, err := buildAnthropicMessagesBody(req, docs, "model", "", 0, false, nil) if err != nil { t.Fatalf("build (force-output off): %v", err) } @@ -283,7 +283,7 @@ func TestBuildAnthropicMessagesBodyForceOutput(t *testing.T) { } // --- force-output on (default) --- - body, _, err = buildAnthropicMessagesBody(req, docs, "model", "", 0, true) + body, _, err = buildAnthropicMessagesBody(req, docs, "model", "", 0, true, nil) if err != nil { t.Fatalf("build (force-output on): %v", err) } @@ -305,7 +305,7 @@ func TestBuildAnthropicMessagesBodyForceOutput(t *testing.T) { // --- force-output on, no system blocks at all: instruction must still be injected --- reqNoSys := RouterReplayRequest{InputTokens: 100} - body, _, err = buildAnthropicMessagesBody(reqNoSys, docs, "model", "", 0, true) + body, _, err = buildAnthropicMessagesBody(reqNoSys, docs, "model", "", 0, true, nil) if err != nil { t.Fatalf("build (force-output on, no system blocks): %v", err) } @@ -328,7 +328,7 @@ func TestBuildOpenAIChatCompletionsBodyForceOutput(t *testing.T) { } // force-output off: neither ignore_eos nor the instruction in the body. - body, _, err := buildOpenAIChatCompletionsBody(req, docs, "model", "", 0, false) + body, _, err := buildOpenAIChatCompletionsBody(req, docs, "model", "", 0, false, nil) if err != nil { t.Fatalf("build (force-output off): %v", err) } @@ -344,7 +344,7 @@ func TestBuildOpenAIChatCompletionsBodyForceOutput(t *testing.T) { } // force-output on (default): ignore_eos=true AND the instruction is present. - body, _, err = buildOpenAIChatCompletionsBody(req, docs, "model", "", 0, true) + body, _, err = buildOpenAIChatCompletionsBody(req, docs, "model", "", 0, true, nil) if err != nil { t.Fatalf("build (force-output on): %v", err) } @@ -367,7 +367,7 @@ func TestBuildAnthropicMessagesBodyOutputRatioMaxTokens(t *testing.T) { docs := strings.Repeat("doc content ", 50) req := RouterReplayRequest{InputTokens: 2000, OutputTokens: 10} - anthBody, _, err := buildAnthropicMessagesBody(req, docs, "model", "", 0.25, false) + anthBody, _, err := buildAnthropicMessagesBody(req, docs, "model", "", 0.25, false, nil) if err != nil { t.Fatalf("anthropic build: %v", err) } @@ -379,7 +379,7 @@ func TestBuildAnthropicMessagesBodyOutputRatioMaxTokens(t *testing.T) { t.Errorf("anthropic max_tokens = %v, want %v", got, want) } - openaiBody, _, err := buildOpenAIChatCompletionsBody(req, docs, "model", "", 0.25, false) + openaiBody, _, err := buildOpenAIChatCompletionsBody(req, docs, "model", "", 0.25, false, nil) if err != nil { t.Fatalf("openai build: %v", err) } @@ -392,7 +392,7 @@ func TestBuildAnthropicMessagesBodyOutputRatioMaxTokens(t *testing.T) { } // Without a ratio, max_tokens falls back to the original output_tokens. - anthBodyNoRatio, _, err := buildAnthropicMessagesBody(req, docs, "model", "", 0, false) + anthBodyNoRatio, _, err := buildAnthropicMessagesBody(req, docs, "model", "", 0, false, nil) if err != nil { t.Fatalf("anthropic build (no ratio): %v", err) } diff --git a/benchmark/replay_router_wire_tools_test.go b/benchmark/replay_router_wire_tools_test.go index 412ae55..2c60109 100644 --- a/benchmark/replay_router_wire_tools_test.go +++ b/benchmark/replay_router_wire_tools_test.go @@ -43,11 +43,11 @@ func TestOpenAIVsAnthropicBodySize(t *testing.T) { Messages: msgs, } - anthBody, _, err := buildAnthropicMessagesBody(req, docs, "model-a", "", 0, false) + anthBody, _, err := buildAnthropicMessagesBody(req, docs, "model-a", "", 0, false, nil) if err != nil { t.Fatalf("buildAnthropicMessagesBody: %v", err) } - openaiBody, _, err := buildOpenAIChatCompletionsBody(req, docs, "model-a", "", 0, false) + openaiBody, _, err := buildOpenAIChatCompletionsBody(req, docs, "model-a", "", 0, false, nil) if err != nil { t.Fatalf("buildOpenAIChatCompletionsBody: %v", err) } @@ -134,11 +134,11 @@ func TestOpenAIVsAnthropicBodySize(t *testing.T) { sessCount++ for _, inst := range sess.Instances { for _, r := range inst.Requests { - ab, _, err := buildAnthropicMessagesBody(r, docs, "model-a", "", 0, false) + ab, _, err := buildAnthropicMessagesBody(r, docs, "model-a", "", 0, false, nil) if err != nil { continue } - ob, _, err := buildOpenAIChatCompletionsBody(r, docs, "model-a", "", 0, false) + ob, _, err := buildOpenAIChatCompletionsBody(r, docs, "model-a", "", 0, false, nil) if err != nil { continue } @@ -186,7 +186,7 @@ func TestOpenAIToolUseConversion(t *testing.T) { }, } - body, _, err := buildOpenAIChatCompletionsBody(req, docs, "model-x", "", 0, false) + body, _, err := buildOpenAIChatCompletionsBody(req, docs, "model-x", "", 0, false, nil) if err != nil { t.Fatalf("buildOpenAIChatCompletionsBody: %v", err) } @@ -279,7 +279,7 @@ func TestOpenAIToolUseConversion(t *testing.T) { }, }, } - orphanBody, _, err := buildOpenAIChatCompletionsBody(reqOrphan, docs, "model-x", "", 0, false) + orphanBody, _, err := buildOpenAIChatCompletionsBody(reqOrphan, docs, "model-x", "", 0, false, nil) if err != nil { t.Fatalf("buildOpenAIChatCompletionsBody orphan: %v", err) } diff --git a/benchmark/replay_uuid.go b/benchmark/replay_uuid.go new file mode 100644 index 0000000..b3d39ea --- /dev/null +++ b/benchmark/replay_uuid.go @@ -0,0 +1,128 @@ +package benchmark + +// Path-agnostic UUID-based response validation primitives, shared by the +// router-replay UUID cache-coherency check (replay_router_uuid.go, +// --replay-inject-uuids + --router-replay-file) and the cache-coherency eval +// CLI (cli/eval_commands.go). Distinct UUIDs are stamped somewhere in a +// request's content; the model is later asked (directly or implicitly) to +// recall them, and the response is scored for presence (did it recall its +// own UUID?) and cross-contamination (did it leak a UUID belonging to a +// DIFFERENT session/series, i.e. a KV/scheduling leak?). +// +// This is a presence-based check (Contains), not the cache-coherency eval's +// exact-conformity check (matchesExpectedUUIDList) — replay responses are +// real multi-turn chat turns with real prose, not a bare comma-joined list. + +import ( + "fmt" + "regexp" + "strings" +) + +// uuidRe matches a canonical hyphenated UUID string (8-4-4-4-12 hex digits). +// Used by findLeakedUUIDsByOwner to pull UUID-shaped substrings out of a +// response/thinking blob without needing to know the candidate set up +// front. +var uuidRe = regexp.MustCompile(`[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}`) + +// findLeakedUUIDsByOwner scans resp and thinking for UUID-shaped substrings +// and flags any that are a KNOWN turn UUID (present in owner, the +// uuid -> owning-session-index reverse map — see buildSessionTurnUUIDs) +// belonging to a session OTHER than ownIdx. Unlike FindLeakedUUIDs (which +// iterates every UUID ever assigned across the whole population — O(total +// UUIDs)), this scans only the response/thinking text once via uuidRe and +// does an O(1) map lookup per match — O(response), not O(population) — +// which matters once turns (not sessions) are the stamping unit: the total +// UUID population grows with conversation length instead of session count. +// A UUID-shaped string with no entry in owner (or that happens to match a +// hex pattern but isn't a real stamp) is silently ignored, not flagged. +// Returns "uuid(series=N)" entries, one per distinct leaked UUID found, in +// first-appearance (scan) order — deterministic for a given response. +func findLeakedUUIDsByOwner(resp, thinking string, ownIdx int, owner map[string]int) []string { + if len(owner) == 0 { + return nil + } + combined := resp + if thinking != "" { + combined = resp + "\n" + thinking + } + var leaked []string + seen := map[string]bool{} + for _, u := range uuidRe.FindAllString(combined, -1) { + if seen[u] { + continue + } + seen[u] = true + si, ok := owner[u] + if !ok || si == ownIdx { + continue + } + leaked = append(leaked, fmt.Sprintf("%s(series=%d)", u, si)) + } + return leaked +} + +// injectUUIDMarker appends one visible marker per uuid to turnValue. Unlike +// the cache-coherency eval's ... filler, the model MUST see +// and be able to repeat this text — it's asked to recite every ref-id later +// — so the marker is plain, readable text. Detection downstream is a raw +// substring match on the UUID itself, so the wrapper text ("[ref-id: ...]") +// is purely cosmetic and not load-bearing for scoring. +func injectUUIDMarker(turnValue string, uuids []string) string { + if len(uuids) == 0 { + return turnValue + } + var b strings.Builder + b.WriteString(turnValue) + for _, u := range uuids { + b.WriteString("\n\n[ref-id: ") + b.WriteString(u) + b.WriteString("]") + } + return b.String() +} + +// validateReplayResponse scores one replay response/thinking pair: +// - found[i] reports whether expected[i] (this conversation's in-scope +// ref-ids at this point) appears in resp or thinking (Contains, mirroring +// the cache-coherency eval's presence check). +// - leaked reports any OTHER conversation's UUID found in resp/thinking — +// cross-contamination — via the shared FindLeakedUUIDs helper (moved +// here from cli/eval_commands.go so both the coherency CLI and replay +// validation share one implementation). +func validateReplayResponse(resp, thinking string, expected []string, convIdx int, allSets [][]string) (found []bool, leaked []string) { + found = make([]bool, len(expected)) + for i, u := range expected { + found[i] = strings.Contains(resp, u) || strings.Contains(thinking, u) + } + leaked = FindLeakedUUIDs(resp, thinking, convIdx, allSets) + return found, leaked +} + +// FindLeakedUUIDs scans resp and thinking for UUIDs belonging to a series/ +// conversation OTHER than ownIdx, per the ordered allSets list (allSets[i] = +// full UUID list "owned" by index i — this doubles as the uuid -> owner +// mapping without needing an actual map, keeping iteration order — and +// therefore leak-report order — deterministic for a given seed). Returns +// "uuid(series=N)" entries, one per leaked UUID found. +// +// Exported (moved from cli/eval_commands.go) so both the cache-coherency +// eval CLI (cli/eval_commands.go, where ownIdx is a coherency series index +// and allSets is CacheCoherencyResult.SeriesUUIDs) and replay UUID +// validation above (where ownIdx is a conversation/session index and +// allSets is AutoBenchmarkConfig.replayUUIDSets) share one implementation, +// not two. +func FindLeakedUUIDs(resp, thinking string, ownIdx int, allSets [][]string) []string { + var leaked []string + for si, uuids := range allSets { + if si == ownIdx { + continue + } + for _, u := range uuids { + if strings.Contains(resp, u) || strings.Contains(thinking, u) { + leaked = append(leaked, fmt.Sprintf("%s(series=%d)", u, si)) + } + } + } + return leaked +} diff --git a/benchmark/replay_uuid_test.go b/benchmark/replay_uuid_test.go new file mode 100644 index 0000000..68f80bf --- /dev/null +++ b/benchmark/replay_uuid_test.go @@ -0,0 +1,131 @@ +package benchmark + +import ( + "strings" + "testing" +) + +func TestInjectUUIDMarker(t *testing.T) { + original := "what is the weather today?" + uuids := []string{"uuid-aaa", "uuid-bbb"} + + got := injectUUIDMarker(original, uuids) + + if !strings.Contains(got, original) { + t.Errorf("injectUUIDMarker() dropped the original turn text: %q", got) + } + for _, u := range uuids { + if !strings.Contains(got, u) { + t.Errorf("injectUUIDMarker() result missing uuid %q: %q", u, got) + } + } + + // No uuids -> turnValue passed through unchanged. + if got := injectUUIDMarker(original, nil); got != original { + t.Errorf("injectUUIDMarker(_, nil) = %q, want unchanged %q", got, original) + } +} + +func TestValidateReplayResponse(t *testing.T) { + allSets := [][]string{ + {"own-0", "own-1"}, // conversation 0 (this response's own conversation) + {"other-0", "other-1"}, // conversation 1 + } + + t.Run("presence in content only", func(t *testing.T) { + found, leaked := validateReplayResponse("SEEN_REFS: own-0,own-1", "", []string{"own-0", "own-1"}, 0, allSets) + if len(found) != 2 || !found[0] || !found[1] { + t.Errorf("found = %v, want both true", found) + } + if len(leaked) != 0 { + t.Errorf("leaked = %v, want none", leaked) + } + }) + + t.Run("presence in thinking only", func(t *testing.T) { + found, _ := validateReplayResponse("some unrelated prose", "I recall own-0 and own-1 from earlier", []string{"own-0", "own-1"}, 0, allSets) + if !found[0] || !found[1] { + t.Errorf("found = %v, want both true (present in thinking)", found) + } + }) + + t.Run("presence in general prose (not the SEEN_REFS line)", func(t *testing.T) { + found, _ := validateReplayResponse("Sure! Earlier you mentioned own-0, and also own-1 came up.", "", []string{"own-0", "own-1"}, 0, allSets) + if !found[0] || !found[1] { + t.Errorf("found = %v, want both true (Contains-based, no format requirement)", found) + } + }) + + t.Run("PRESENCE_MISS: expected uuid absent from both", func(t *testing.T) { + found, leaked := validateReplayResponse("I don't remember any ref ids.", "", []string{"own-0", "own-1"}, 0, allSets) + if found[0] || found[1] { + t.Errorf("found = %v, want both false", found) + } + if len(leaked) != 0 { + t.Errorf("leaked = %v, want none", leaked) + } + }) + + t.Run("CROSS_CONTAMINATION: another conversation's uuid appears", func(t *testing.T) { + found, leaked := validateReplayResponse("SEEN_REFS: own-0,own-1,other-0", "", []string{"own-0", "own-1"}, 0, allSets) + if !found[0] || !found[1] { + t.Errorf("found = %v, want both true", found) + } + if len(leaked) != 1 || !strings.Contains(leaked[0], "other-0") || !strings.Contains(leaked[0], "series=1") { + t.Errorf("leaked = %v, want one entry naming other-0 from series=1", leaked) + } + }) +} + +func TestFindLeakedUUIDs(t *testing.T) { + allSets := [][]string{ + {"uuid-0-a", "uuid-0-b"}, // 0 + {"uuid-1-a", "uuid-1-b"}, // 1 + {"uuid-2-a", "uuid-2-b"}, // 2 + } + + t.Run("no leak: response only contains own uuids", func(t *testing.T) { + got := FindLeakedUUIDs("uuid-0-a,uuid-0-b", "", 0, allSets) + if len(got) != 0 { + t.Errorf("expected no leaks, got %v", got) + } + }) + + t.Run("single leak from another set", func(t *testing.T) { + got := FindLeakedUUIDs("uuid-0-a,uuid-1-a", "", 0, allSets) + if len(got) != 1 { + t.Fatalf("expected 1 leak, got %v", got) + } + if !strings.Contains(got[0], "uuid-1-a") || !strings.Contains(got[0], "series=1") { + t.Errorf("leak entry %q missing uuid or owning index", got[0]) + } + }) + + t.Run("multiple leaks, deterministic order", func(t *testing.T) { + resp := "uuid-0-a,uuid-1-a,uuid-2-b" + got1 := FindLeakedUUIDs(resp, "", 0, allSets) + got2 := FindLeakedUUIDs(resp, "", 0, allSets) + if len(got1) != 2 { + t.Fatalf("expected 2 leaks, got %v", got1) + } + for i := range got1 { + if got1[i] != got2[i] { + t.Errorf("non-deterministic leak order: %v vs %v", got1, got2) + } + } + }) + + t.Run("leak detected in thinking, not just response", func(t *testing.T) { + got := FindLeakedUUIDs("uuid-0-a", "I recall uuid-2-a from earlier", 0, allSets) + if len(got) != 1 || !strings.Contains(got[0], "uuid-2-a") { + t.Errorf("expected leak of uuid-2-a via thinking, got %v", got) + } + }) + + t.Run("own index never reported as a leak", func(t *testing.T) { + got := FindLeakedUUIDs("uuid-1-a,uuid-1-b", "", 1, allSets) + if len(got) != 0 { + t.Errorf("expected no leaks for own index, got %v", got) + } + }) +} diff --git a/benchmark/types.go b/benchmark/types.go index 8dbc3b2..b4583e3 100644 --- a/benchmark/types.go +++ b/benchmark/types.go @@ -23,5 +23,13 @@ type RequestMetrics struct { CachedPrompt string // Full system prompt text (only populated on error/empty for diagnostics) Question string // The user question sent with the request (only populated on error/empty) RawResponseTail string // raw SSE tail (last bytes); only populated on error/empty for diagnostics -} + // UUID validation (router-replay --replay-inject-uuids only). All nil/zero + // when the feature is off (default) or on the synthetic path, which never + // populates these. + ConvIdx int // session index within cfg.replayUUIDSets (== seriesNum-1) + ExpectedUUIDs []string // this session's own N-UUID list, in order + UUIDFound []bool // parallel to ExpectedUUIDs: whether each was found in Response or thinking + LeakedUUIDs []string // "uuid(series=N)" entries for any OTHER session's UUID found here + ExactMatch bool // first line of Response is exactly the ordered, comma-joined ExpectedUUIDs list (output conformity) +} diff --git a/cli/benchmark_commands.go b/cli/benchmark_commands.go index 21979a7..720a7e9 100644 --- a/cli/benchmark_commands.go +++ b/cli/benchmark_commands.go @@ -270,6 +270,9 @@ func (c *BenchmarkAutoCommand) Execute(args []string) error { ReplayNoStamp: c.ReplayNoStamp, AbortOnCollapse: c.AbortOnCollapse, ReplayStopAtLowConcurrency: c.ReplayStopAtLowConcurrency, + ReplayInjectUUIDs: c.ReplayInjectUUIDs, + ReplayUUIDSeed: c.ReplayUUIDSeed, + ReplayReciteEveryRequest: c.ReplayReciteEveryRequest != "false", RouterReplayFile: c.RouterReplayFile, RouterReplayRoles: c.RouterReplayRoles, ReplayOutputRatio: c.ReplayOutputRatio, @@ -287,6 +290,20 @@ func (c *BenchmarkAutoCommand) Execute(args []string) error { if c.FromDataset != "" && c.RouterReplayFile != "" { return fmt.Errorf("--from-dataset and --router-replay-file are mutually exclusive") } + + // --replay-inject-uuids is ROUTER-REPLAY PATH ONLY: it splices a + // per-session UUID marker at the boundary between cross-session-shared + // prefix blocks (computed from the replay-v3 block-hash schema) and + // per-session content — the dataset-replay path has no such block-hash + // schema to compute that boundary from. + if c.ReplayInjectUUIDs { + if c.RouterReplayFile == "" { + return fmt.Errorf("--replay-inject-uuids requires --router-replay-file") + } + if c.FromDataset != "" { + return fmt.Errorf("--replay-inject-uuids and --from-dataset are mutually exclusive") + } + } if c.DryRun && c.RouterReplayFile == "" { return fmt.Errorf("--dry-run requires --router-replay-file") } diff --git a/cli/benchmark_options.go b/cli/benchmark_options.go index 47558c0..4d5432d 100644 --- a/cli/benchmark_options.go +++ b/cli/benchmark_options.go @@ -64,18 +64,26 @@ type BenchmarkAutoOptions struct { ReplayNoStamp bool `long:"replay-no-stamp" description:"Disable per-run RUN_GUID stamping in replay mode. By default each replay run prepends a fresh UUID to every request's system prompt (both --from-dataset and --router-replay-file paths) so server prefix caches from prior runs can't be reused — pristine per-run cache state." env:"BENCHMARK_REPLAY_NO_STAMP"` AbortOnCollapse bool `long:"abort-on-collapse" description:"Abort the benchmark if the windowed cache hit rate stays below 50% for 2 minutes. Off by default — this heuristic fires on legitimate workloads with low cache reuse (e.g. replay across many distinct conversations)." env:"BENCHMARK_ABORT_ON_COLLAPSE"` ReplayStopAtLowConcurrency bool `long:"replay-stop-at-low-concurrency" description:"Terminate the replay run once the queue is drained AND the number of active worker goroutines has dropped below --concurrency. Avoids long-tail measurements where only a handful of long conversations remain and the gate is underutilized." env:"BENCHMARK_REPLAY_STOP_AT_LOW_CONCURRENCY"` - RouterReplayFile string `long:"router-replay-file" description:"Path to a tree-aware replay file produced by 'wekai router replay-prepare'. Each series = one CLI session; within the session sub-agents fan out concurrently, honoring parent->child sequencing baked into the file. Mutually exclusive with --from-dataset." env:"BENCHMARK_ROUTER_REPLAY_FILE"` - ReplayOutputRatio float64 `long:"replay-output-ratio" description:"Router-replay only: retarget each request's max_tokens to InputTokens * this ratio, overriding the recorded original output_tokens (which otherwise pins max_tokens to what the model produced in the original capture, so the model stops almost immediately on replay). 0 = off (default: use original output_tokens/max_tokens)." default:"0" env:"BENCHMARK_REPLAY_OUTPUT_RATIO"` - ReplayNaturalOutput bool `long:"replay-natural-output" description:"Let the model stop generation naturally instead of forcing it to fill max_tokens (disables the continue-generating instruction and vLLM ignore_eos). Router-replay only." env:"BENCHMARK_REPLAY_NATURAL_OUTPUT"` - RouterReplayRoles string `long:"router-replay-roles" description:"Comma-separated list of instance roles to replay (default: all). E.g. 'main,sub-agent' excludes the CLI's background helper:title / helper:summarize / ephemeral (haiku side-calls) instances and keeps only the agentic workload — gives a cleaner 'in_flight ~= series' steady state. Other values: 'helper-or-isolated', 'ephemeral (no system)', 'other'." env:"BENCHMARK_ROUTER_REPLAY_ROLES"` - RouterReplaySeriesIndices string `long:"replay-series-indices" description:"Comma-separated list of 0-based session indices to replay from --router-replay-file (e.g. '3,7,42'). Only sessions at those line positions (0 = first session after the header) are dispatched; others are skipped. Mutually exclusive with --replay-series-range. Overrides --replay-series." env:"BENCHMARK_REPLAY_SERIES_INDICES"` - RouterReplaySeriesRange string `long:"replay-series-range" description:"Inclusive range of 0-based session indices to replay from --router-replay-file (e.g. '0-50' or '100-199'). Mutually exclusive with --replay-series-indices. Overrides --replay-series." env:"BENCHMARK_REPLAY_SERIES_RANGE"` - DryRun bool `long:"dry-run" description:"Dry run: skip remote HTTP requests; drive the router-replay pipeline with synthetic timing so gcache evolution can be observed offline. Requires --router-replay-file." env:"BENCHMARK_DRY_RUN"` - DryRunColdTPS int `long:"dry-run-cold-tps" default:"1000000" description:"Dry-run: cold (uncached) input tokens processed per second." env:"BENCHMARK_DRY_RUN_COLD_TPS"` - DryRunWarmTPS int `long:"dry-run-warm-tps" default:"10000000" description:"Dry-run: warm (cached) input tokens processed per second." env:"BENCHMARK_DRY_RUN_WARM_TPS"` - DryRunOutputTPS int `long:"dry-run-output-tps" default:"100000" description:"Dry-run: output tokens generated per second." env:"BENCHMARK_DRY_RUN_OUTPUT_TPS"` - CacheSimChunkBytes int `long:"cache-sim-chunk-bytes" description:"Chunk size in bytes for the content-level cache estimator (0 = default 1024)." default:"0" env:"BENCHMARK_CACHE_SIM_CHUNK_BYTES"` - RandomGateOrder string `long:"random-gate-order" choice:"true" choice:"false" default:"true" optional:"yes" optional-value:"true" description:"Wake the concurrency gate's waiting series in uniformly random order when oversubscribed (the DEFAULT). Strict FIFO forces every series to wait behind all other waiting series before its next turn -- the adversarial worst case for GPU prefix-cache LRU. Pass --random-gate-order=false for the legacy exact-FIFO order. Cold-start waiters are unaffected (always served first, FIFO)." env:"BENCHMARK_RANDOM_GATE_ORDER"` + ReplayInjectUUIDs bool `long:"replay-inject-uuids" description:"Inject one deterministic UUID per user turn into --router-replay-file replay sessions (spread through the conversation), and validate a windowed recite of first + up to 3 most recent turns (excluding the current turn) in later responses -- a coherency check (PRESENCE_MISS / CROSS_CONTAMINATION) for the KV-offload path under realistic multi-turn agentic traffic. ROUTER-REPLAY PATH ONLY: requires --router-replay-file; rejected together with --from-dataset." env:"BENCHMARK_REPLAY_INJECT_UUIDS"` + ReplayUUIDSeed int64 `long:"replay-uuid-seed" description:"PRNG seed for --replay-inject-uuids' per-session UUID generation (0 = crypto/rand, non-deterministic across runs)." default:"0" env:"BENCHMARK_REPLAY_UUID_SEED"` + // ReplayReciteEveryRequest is a string (not bool) choice, same workaround + // as RandomGateOrder above: a plain bool flag defaulting true can never + // be turned OFF via the CLI with this parser (go-flags only accepts + // `--flag=false` for bool options when the parser is built with + // AllowBoolValues, which main.go's flags.Default does not set). + ReplayReciteEveryRequest string `long:"replay-recite-every-request" choice:"true" choice:"false" default:"true" optional:"yes" optional-value:"true" description:"With --replay-inject-uuids, ask the model to recite the ref-id marker on EVERY request (the DEFAULT) rather than only on each agent instance's final request. Pass --replay-recite-every-request=false to recite only on each instance's last request." env:"BENCHMARK_REPLAY_RECITE_EVERY_REQUEST"` + RouterReplayFile string `long:"router-replay-file" description:"Path to a tree-aware replay file produced by 'wekai router replay-prepare'. Each series = one CLI session; within the session sub-agents fan out concurrently, honoring parent->child sequencing baked into the file. Mutually exclusive with --from-dataset." env:"BENCHMARK_ROUTER_REPLAY_FILE"` + ReplayOutputRatio float64 `long:"replay-output-ratio" description:"Router-replay only: retarget each request's max_tokens to InputTokens * this ratio, overriding the recorded original output_tokens (which otherwise pins max_tokens to what the model produced in the original capture, so the model stops almost immediately on replay). 0 = off (default: use original output_tokens/max_tokens)." default:"0" env:"BENCHMARK_REPLAY_OUTPUT_RATIO"` + ReplayNaturalOutput bool `long:"replay-natural-output" description:"Let the model stop generation naturally instead of forcing it to fill max_tokens (disables the continue-generating instruction and vLLM ignore_eos). Router-replay only." env:"BENCHMARK_REPLAY_NATURAL_OUTPUT"` + RouterReplayRoles string `long:"router-replay-roles" description:"Comma-separated list of instance roles to replay (default: all). E.g. 'main,sub-agent' excludes the CLI's background helper:title / helper:summarize / ephemeral (haiku side-calls) instances and keeps only the agentic workload — gives a cleaner 'in_flight ~= series' steady state. Other values: 'helper-or-isolated', 'ephemeral (no system)', 'other'." env:"BENCHMARK_ROUTER_REPLAY_ROLES"` + RouterReplaySeriesIndices string `long:"replay-series-indices" description:"Comma-separated list of 0-based session indices to replay from --router-replay-file (e.g. '3,7,42'). Only sessions at those line positions (0 = first session after the header) are dispatched; others are skipped. Mutually exclusive with --replay-series-range. Overrides --replay-series." env:"BENCHMARK_REPLAY_SERIES_INDICES"` + RouterReplaySeriesRange string `long:"replay-series-range" description:"Inclusive range of 0-based session indices to replay from --router-replay-file (e.g. '0-50' or '100-199'). Mutually exclusive with --replay-series-indices. Overrides --replay-series." env:"BENCHMARK_REPLAY_SERIES_RANGE"` + DryRun bool `long:"dry-run" description:"Dry run: skip remote HTTP requests; drive the router-replay pipeline with synthetic timing so gcache evolution can be observed offline. Requires --router-replay-file." env:"BENCHMARK_DRY_RUN"` + DryRunColdTPS int `long:"dry-run-cold-tps" default:"1000000" description:"Dry-run: cold (uncached) input tokens processed per second." env:"BENCHMARK_DRY_RUN_COLD_TPS"` + DryRunWarmTPS int `long:"dry-run-warm-tps" default:"10000000" description:"Dry-run: warm (cached) input tokens processed per second." env:"BENCHMARK_DRY_RUN_WARM_TPS"` + DryRunOutputTPS int `long:"dry-run-output-tps" default:"100000" description:"Dry-run: output tokens generated per second." env:"BENCHMARK_DRY_RUN_OUTPUT_TPS"` + CacheSimChunkBytes int `long:"cache-sim-chunk-bytes" description:"Chunk size in bytes for the content-level cache estimator (0 = default 1024)." default:"0" env:"BENCHMARK_CACHE_SIM_CHUNK_BYTES"` + RandomGateOrder string `long:"random-gate-order" choice:"true" choice:"false" default:"true" optional:"yes" optional-value:"true" description:"Wake the concurrency gate's waiting series in uniformly random order when oversubscribed (the DEFAULT). Strict FIFO forces every series to wait behind all other waiting series before its next turn -- the adversarial worst case for GPU prefix-cache LRU. Pass --random-gate-order=false for the legacy exact-FIFO order. Cold-start waiters are unaffected (always served first, FIFO)." env:"BENCHMARK_RANDOM_GATE_ORDER"` // Positional arguments Args struct { diff --git a/cli/command_misc_cache_coherency_test.go b/cli/command_misc_cache_coherency_test.go index 9d28027..3875829 100644 --- a/cli/command_misc_cache_coherency_test.go +++ b/cli/command_misc_cache_coherency_test.go @@ -62,61 +62,10 @@ func TestResolveGarbageChars(t *testing.T) { } } -func TestFindLeakedUUIDs(t *testing.T) { - seriesUUIDs := [][]string{ - {"uuid-0-a", "uuid-0-b"}, // series 0 - {"uuid-1-a", "uuid-1-b"}, // series 1 - {"uuid-2-a", "uuid-2-b"}, // series 2 - } - - t.Run("no leak: response only contains own series uuids", func(t *testing.T) { - resp := "uuid-0-a,uuid-0-b" - got := findLeakedUUIDs(resp, "", 0, seriesUUIDs) - if len(got) != 0 { - t.Errorf("expected no leaks, got %v", got) - } - }) - - t.Run("single leak from another series in response", func(t *testing.T) { - resp := "uuid-0-a,uuid-1-a" - got := findLeakedUUIDs(resp, "", 0, seriesUUIDs) - if len(got) != 1 { - t.Fatalf("expected 1 leak, got %v", got) - } - if !strings.Contains(got[0], "uuid-1-a") || !strings.Contains(got[0], "series=1") { - t.Errorf("leak entry %q missing uuid or owning series", got[0]) - } - }) - - t.Run("multiple leaks from multiple series, deterministic order", func(t *testing.T) { - resp := "uuid-0-a,uuid-1-a,uuid-2-b" - got1 := findLeakedUUIDs(resp, "", 0, seriesUUIDs) - got2 := findLeakedUUIDs(resp, "", 0, seriesUUIDs) - if len(got1) != 2 { - t.Fatalf("expected 2 leaks, got %v", got1) - } - // Order must be deterministic (series-index order) across repeated calls. - for i := range got1 { - if got1[i] != got2[i] { - t.Errorf("non-deterministic leak order: %v vs %v", got1, got2) - } - } - }) - - t.Run("leak detected in thinking, not just response", func(t *testing.T) { - got := findLeakedUUIDs("uuid-0-a", "I recall uuid-2-a from earlier", 0, seriesUUIDs) - if len(got) != 1 || !strings.Contains(got[0], "uuid-2-a") { - t.Errorf("expected leak of uuid-2-a via thinking, got %v", got) - } - }) - - t.Run("own series never reported as a leak", func(t *testing.T) { - got := findLeakedUUIDs("uuid-1-a,uuid-1-b", "", 1, seriesUUIDs) - if len(got) != 0 { - t.Errorf("expected no leaks for own series, got %v", got) - } - }) -} +// TestFindLeakedUUIDs moved to benchmark.FindLeakedUUIDs's own test +// (benchmark/replay_uuid_test.go) — the function itself moved to the +// benchmark package (benchmark/replay_uuid.go) so both this CLI and +// dataset-replay UUID validation share one implementation. func TestColdWarmCycleCounts(t *testing.T) { tests := []struct { diff --git a/cli/eval_commands.go b/cli/eval_commands.go index d448979..7619cdb 100644 --- a/cli/eval_commands.go +++ b/cli/eval_commands.go @@ -5,7 +5,6 @@ import ( "fmt" "io" "os" - "strings" "github.com/weka/wekai/benchmark" "github.com/weka/wekai/config" @@ -204,7 +203,7 @@ func (c *EvalCacheCoherencyCommand) Execute(args []string) error { } // Any uuid belonging to a DIFFERENT series found in this response/thinking // is cross-contamination (KV/scheduling leak across series). - leaked = findLeakedUUIDs(r.Response, r.Thinking, r.SeriesIdx, result.SeriesUUIDs) + leaked = benchmark.FindLeakedUUIDs(r.Response, r.Thinking, r.SeriesIdx, result.SeriesUUIDs) } uuidMissingFlakyCount += len(missing) crossContamCount += len(leaked) @@ -353,25 +352,9 @@ func resolveGarbageChars(garbageCharacters, garbageTokens int, w io.Writer) int } } -// findLeakedUUIDs scans resp and thinking for UUIDs belonging to a series OTHER than -// ownSeries, per the ordered seriesUUIDs list (seriesUUIDs[i] = full UUID stamp list of -// series i — this doubles as the uuid -> owning-series mapping without needing an actual -// map, keeping iteration order — and therefore leak-report order — deterministic for a -// given seed). Returns "uuid(series=N)" entries, one per leaked UUID found. -func findLeakedUUIDs(resp, thinking string, ownSeries int, seriesUUIDs [][]string) []string { - var leaked []string - for si, uuids := range seriesUUIDs { - if si == ownSeries { - continue - } - for _, u := range uuids { - if strings.Contains(resp, u) || strings.Contains(thinking, u) { - leaked = append(leaked, fmt.Sprintf("%s(series=%d)", u, si)) - } - } - } - return leaked -} +// findLeakedUUIDs moved to benchmark.FindLeakedUUIDs (benchmark/replay_uuid.go) +// so both this CLI and the dataset-replay UUID validation path share one +// implementation. // coldWarmCycleCounts returns the number of distinct COLD (cycle 1) and WARM (cycle >= 2) // cycle numbers present in cyclesSeen — cycle-number CARDINALITY, not a request count.