From 6cdd51fce1d2e8f866840f18f0a502b387494604 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Thu, 23 Jul 2026 07:29:22 -0700 Subject: [PATCH] perf(codex): make replay prefix fingerprints incremental codexReasoningReplayTurnAnchorIndex re-hashed the entire input prefix for every probed index while scanning anchors downward, turning the anchor search into O(n^2) SHA-256 over the conversation. On long-context (1M) sessions the gateway spent minutes hashing per request before contacting upstream, which surfaced as client-side 'operation timed out' errors and ~11 cores of steady CPU on the gorillaclaw candidate. Replace the per-probe recomputation with a lazily extended incremental digest that snapshots each prefix sum once, keeping every probe O(1) after a single O(n) pass. Output is bit-identical to codexReplayInputPrefixFingerprint; equivalence covered by a new test. --- internal/runtime/executor/codex_executor.go | 43 +++++++++++++++++-- ...ex_executor_reasoning_replay_cache_test.go | 22 ++++++++++ 2 files changed, 62 insertions(+), 3 deletions(-) diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index 1295d1b43..c6d633592 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -7,6 +7,7 @@ import ( "crypto/sha256" "encoding/hex" "fmt" + "hash" "io" "net/http" "sort" @@ -502,6 +503,7 @@ func insertCodexReasoningReplayTurns(body []byte, replayItems [][]byte) ([]byte, turns := splitCodexReasoningReplayTurns(replayItems) insertions := make(map[int][][]byte) usedAnchorIndexes := make(map[int]bool) + prefixFingerprints := newCodexReplayPrefixFingerprints(inputItems) fallbackAnchorEnd := len(inputItems) - 1 inserted := false for turnIndex := len(turns) - 1; turnIndex >= 0; turnIndex-- { @@ -521,7 +523,7 @@ func insertCodexReasoningReplayTurns(body []byte, replayItems [][]byte) ([]byte, continue } - anchorIndex, matched := codexReasoningReplayTurnAnchorIndex(inputItems, turn, fallbackAnchorEnd, usedAnchorIndexes) + anchorIndex, matched := codexReasoningReplayTurnAnchorIndex(inputItems, turn, fallbackAnchorEnd, usedAnchorIndexes, prefixFingerprints) if !matched { continue } @@ -590,7 +592,7 @@ func splitCodexReasoningReplayTurns(items [][]byte) []codexReasoningReplayTurn { return turns } -func codexReasoningReplayTurnAnchorIndex(inputItems []gjson.Result, turn codexReasoningReplayTurn, fallbackEnd int, used map[int]bool) (int, bool) { +func codexReasoningReplayTurnAnchorIndex(inputItems []gjson.Result, turn codexReasoningReplayTurn, fallbackEnd int, used map[int]bool, prefixFingerprints *codexReplayPrefixFingerprints) (int, bool) { searchEnd := fallbackEnd if turn.requestFingerprint != "" { searchEnd = len(inputItems) - 1 @@ -599,7 +601,7 @@ func codexReasoningReplayTurnAnchorIndex(inputItems []gjson.Result, turn codexRe searchEnd = len(inputItems) - 1 } matchesRequestPrefix := func(index int) bool { - return turn.requestFingerprint == "" || codexReplayInputPrefixFingerprint(inputItems, index) == turn.requestFingerprint + return turn.requestFingerprint == "" || prefixFingerprints.at(index) == turn.requestFingerprint } if len(turn.callIDs) > 0 { callIDs := make(map[string]bool) @@ -740,6 +742,41 @@ func codexReplayInputPrefixFingerprint(inputItems []gjson.Result, end int) strin return hex.EncodeToString(hasher.Sum(nil)) } +// codexReplayPrefixFingerprints answers codexReplayInputPrefixFingerprint queries +// from one incremental hashing pass. The anchor search probes many prefixes per +// turn; recomputing each prefix from scratch is O(n^2) hashing and stalled large +// long-context requests for minutes before anything was sent upstream. +type codexReplayPrefixFingerprints struct { + items []gjson.Result + hasher hash.Hash + // sums[end] is the fingerprint of items[0:end]; extended lazily. + sums []string +} + +func newCodexReplayPrefixFingerprints(items []gjson.Result) *codexReplayPrefixFingerprints { + hasher := sha256.New() + return &codexReplayPrefixFingerprints{ + items: items, + hasher: hasher, + sums: []string{hex.EncodeToString(hasher.Sum(nil))}, + } +} + +func (f *codexReplayPrefixFingerprints) at(end int) string { + if end < 0 || end > len(f.items) { + return "" + } + // Sum copies the running digest state, so absorbing one item and + // snapshotting per step reproduces every prefix fingerprint exactly. + for len(f.sums) <= end { + next := len(f.sums) - 1 + _, _ = f.hasher.Write([]byte("\x00item\x00")) + _, _ = io.WriteString(f.hasher, f.items[next].Raw) + f.sums = append(f.sums, hex.EncodeToString(f.hasher.Sum(nil))) + } + return f.sums[end] +} + func filterCodexReasoningReplayItemsForInput(body []byte, items [][]byte) [][]byte { input := gjson.GetBytes(body, "input") if !input.IsArray() { diff --git a/internal/runtime/executor/codex_executor_reasoning_replay_cache_test.go b/internal/runtime/executor/codex_executor_reasoning_replay_cache_test.go index cd1b6a785..e2704c017 100644 --- a/internal/runtime/executor/codex_executor_reasoning_replay_cache_test.go +++ b/internal/runtime/executor/codex_executor_reasoning_replay_cache_test.go @@ -1090,3 +1090,25 @@ func TestCodexExecutorReasoningReplayCacheMatchesShortenedClaudeToolResultCallID t.Fatalf("input.3.call_id = %q, want shortened call_id %q; body=%s", got, shortCallID, string(secondBody)) } } + +func TestCodexReplayPrefixFingerprintsMatchesDirectComputation(t *testing.T) { + items := []gjson.Result{ + gjson.Parse(`{"type":"message","role":"user","content":"a"}`), + gjson.Parse(`{"type":"reasoning","encrypted_content":"abc"}`), + gjson.Parse(`{"type":"function_call","call_id":"call_1"}`), + gjson.Parse(`{"type":"function_call_output","call_id":"call_1","output":"ok"}`), + } + cache := newCodexReplayPrefixFingerprints(items) + // Out-of-order and repeated probes mirror the downward anchor scan. + for _, end := range []int{4, 2, 0, 3, 1, 4, 2} { + want := codexReplayInputPrefixFingerprint(items, end) + if got := cache.at(end); got != want { + t.Fatalf("cache.at(%d) = %q, want %q", end, got, want) + } + } + for _, end := range []int{-1, 5} { + if got := cache.at(end); got != "" { + t.Fatalf("cache.at(%d) = %q, want empty for out-of-range", end, got) + } + } +}