Skip to content

perf(memory): run the query rewriter once per turn - #416

Closed
plombeer31 wants to merge 1 commit into
fix/memory-subcall-timeoutsfrom
perf/rewriter-once-per-turn
Closed

plombeer31 wants to merge 1 commit into
fix/memory-subcall-timeoutsfrom
perf/rewriter-once-per-turn

Conversation

@plombeer31

Copy link
Copy Markdown
Collaborator

Stacked on #415 (fix/memory-subcall-timeouts) — it extends that PR's config v65 migration step. Review/merge that one first; this PR will be retargeted to main when it lands.

What

Before: the agent loop refreshes memory context before the first step and again after every step, each time with the same user message, and the query rewriter ran on every refresh. Against a slow provider every step paid the rewriter's timeout (3 s) on the hot path, for the same answer.

After:

  • Once per turn. rewriter-aware-recall-provider.ts remembers the rewrite per session, keyed by a digest of exactly what the runner is asked (the user message plus the history slice sent), and reuses it for every later refresh with the same key — including a timed-out or failed attempt, whose outcome is "use the raw message". An attempt whose caller's signal aborted is not remembered, so a cancelled turn never decides the next turn's identical retry. Only the latest key per session is kept, for at most 256 sessions (least recently used dropped).
  • Rewriter timeout 3 000 → 10 000 ms, folded into the same v65 migration step (only a value equal to the old default migrates). It stays well below the background sub-calls' 60 s because the rewriter blocks the turn — now once per turn, not once per step.

Why

Measured on v0.6.0 in live 8-turn sessions: google/gemini-3.8-flash (rewriter median 4.2 s) timed out on 16 of 18 rewriter calls; moonshotai/kimi-k2.6 (median 23.5 s) on 16 of 16. Each of those timeouts was a 3 s stall before a step, repeated for every step of the turn, for a rewrite that never landed.

Known edge, deliberately accepted: MemoryContextProviderInput carries no turn identity, so the memo cannot tell "the next step of this turn" from "the operator re-sent the identical message after a turn that ended without a reply". In that case the retry reuses the previous rewrite outcome (the raw message) for recall; the turn itself runs again normally. Keying on a turn id would mean widening the agent-loop interface, which this PR does not do.

How it was verified

  • npm run lint clean
  • npx vitest run src/config src/memory/retrieve src/memory/reflection src/memory/links — 29 files / 567 tests green
  • new rewriter-aware-recall-provider-memo.test.ts (real runner): 6 refreshes → 1 LLM call; a timeout or failure is not retried within the key; a new message or a changed history triggers a new call; an aborted attempt is not cached and the next identical request reaches the LLM; sessions are isolated; the LRU cap evicts; the migration test covers the rewriter field
  • vacuity: disabling the memo lookup turns the memo tests red; removing the abort guard turns the abort test red
  • live, this branch's build, google/gemini-3.8-flash, the same scripted 8-turn session, started from a v64 config holding 10 000 / 8 000 / 3 000 (rewritten on disk to v65 with 60 000 / 60 000 / 10 000):
v0.6.0 this branch
steps / rewriter rows 24 / 24 23 / 8 (one per turn)
rewriter outcomes 16 timeout, 2 ok, 6 skipped 0 timeout, 5 ok (2.7–4.7 s), 3 skipped
reflection outcomes 2 ok, 5 none, 1 timeout 3 ok, 5 none, 0 timeout
sum of turn durations 118.6 s 84.0 s (−29 %)

Docs: AGENTS.md rewriter lines, a "once per turn" paragraph and a new locked invariant.

The agent loop refreshes memory context before the first step and
again after every step, each time with the same user message. The
rewriter decorator asked its runner on every refresh, so one turn sent
the same rewrite request once per step. It sits on the hot path: every
call blocks the step until it answers or times out. Against a hosted
model that is slower than the cap, each step paid the full timeout and
still recalled with the raw message.

Measured on OpenRouter (median, byte-faithful rewriter requests):
gemini-3.8-flash 4.2 s, glm-5.3-flash 4.2 s, kimi-k2.6 23.5 s. In live
8-turn sessions gemini-3.8-flash's rewriter timed out on 16 of 18 calls
at the 3 s default and kimi-k2.6's on 16 of 16.

The decorator now remembers the outcome per session, keyed by the user
message and a digest of the history slice it sent, and reuses it for
every later refresh with the same key, a timed-out or failed attempt
(outcome "use the raw message") included. A slow provider costs a turn
one timeout, not one per step. Only the latest key per session is kept,
for at most 256 sessions. An attempt whose caller's signal aborted is
not remembered: a cancelled turn says nothing about the provider and
must not decide the next turn's identical retry.

With one call per turn the rewriter can afford a realistic cap, so its
default goes 3 s -> 10 s inside the same v65 migration as the reflection
and link-generator timeouts: a pre-v65 file carrying the old 3000 takes
10000, any other value is kept as a pin. It stays well below the
background sub-calls because it still blocks the turn.
@plombeer31

Copy link
Copy Markdown
Collaborator Author

Shipped in v0.6.1 via #422 — its commits are on main.

@plombeer31 plombeer31 closed this Sep 13, 2026
Bartok9 pushed a commit to Bartok9/atomic-agent that referenced this pull request Sep 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant