diff --git a/internal/runtime/executor/codex_executor_reasoning.go b/internal/runtime/executor/codex_executor_reasoning.go index 25093b5ae..fc26f2dec 100644 --- a/internal/runtime/executor/codex_executor_reasoning.go +++ b/internal/runtime/executor/codex_executor_reasoning.go @@ -4,6 +4,8 @@ import ( "context" "crypto/sha256" "encoding/hex" + "hash" + "io" "net/http" "strings" @@ -221,6 +223,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-- { @@ -240,7 +243,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 } @@ -309,7 +312,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 @@ -318,7 +321,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) @@ -459,6 +462,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) + } + } +}