From 37c025edffd3bcc7d0d2b1b7735027094edb2bb0 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Sat, 25 Jul 2026 16:21:21 +0200 Subject: [PATCH] fix(memory): run extended-memory predictor concurrently with rerank The per-turn recall has a hard 5s budget, and queryAtomsWithPrediction ran two LLM calls sequentially: the paid rerank of literal-query results followed by predictor.Predict. A slow memory LLM (by default the main reasoning model) let the rerank consume most of the window, so Predict inherited a nearly-expired context and failed immediately: extended memory: predictor LLM failed: context deadline exceeded extended memory: predicted-intent generation failed: predictor: context deadline exceeded Run the predictor concurrently with the literal query instead, so both LLM calls share the full window in parallel. The predictor only touches the LLM (no shared state) and the store/index are mutex-guarded; the result channel is buffered so the early-error path cannot leak the goroutine. Adds TestPredictionRunsConcurrentWithRerank, which deadlocks under the old sequential behavior. --- internal/memory/extended/recall.go | 26 ++++++++- internal/memory/extended/recall_test.go | 75 +++++++++++++++++++++++++ 2 files changed, 98 insertions(+), 3 deletions(-) diff --git a/internal/memory/extended/recall.go b/internal/memory/extended/recall.go index bad20bc..c6ac2a0 100644 --- a/internal/memory/extended/recall.go +++ b/internal/memory/extended/recall.go @@ -112,6 +112,26 @@ func (r *Recall) trackErr(err error) { // blended ranking with a pure retention ordering. func (r *Recall) queryAtomsWithPrediction(ctx context.Context, query string, recent []string, state UserState) ([]MemoryAtom, error) { all := make(map[string]scoredAtomMeta) + + // Run intent prediction concurrently with the literal query (which may + // include a paid LLM rerank). Sequentially, the rerank can consume most of + // the recall context deadline, starving the predictor of budget and making + // it fail with context-deadline-exceeded before its LLM call even starts. + predicting := r.predictor != nil && r.cfg.PredictiveIntents > 0 && + r.cfg.FollowUpAnticipationEnabled != nil && *r.cfg.FollowUpAnticipationEnabled + type predictResult struct { + intents []PredictedIntent + err error + } + var predCh chan predictResult + if predicting { + predCh = make(chan predictResult, 1) + go func() { + intents, err := r.predictor.Predict(ctx, query, recent, state) + predCh <- predictResult{intents, err} + }() + } + literal, err := r.queryAtomsScored(ctx, query, false) if err != nil { r.trackErr(err) @@ -121,9 +141,9 @@ func (r *Recall) queryAtomsWithPrediction(ctx context.Context, query string, rec all[s.atom.ID] = s } - if r.predictor != nil && r.cfg.PredictiveIntents > 0 && - r.cfg.FollowUpAnticipationEnabled != nil && *r.cfg.FollowUpAnticipationEnabled { - intents, err := r.predictor.Predict(ctx, query, recent, state) + if predicting { + res := <-predCh + intents, err := res.intents, res.err if err != nil { r.trackErr(err) log.Printf("extended memory: predicted-intent generation failed: %v", err) diff --git a/internal/memory/extended/recall_test.go b/internal/memory/extended/recall_test.go index 1980276..79b5e32 100644 --- a/internal/memory/extended/recall_test.go +++ b/internal/memory/extended/recall_test.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strings" + "sync" "testing" "time" @@ -309,6 +310,80 @@ func TestQueryAtomsWithPredictionCompositeOrdering(t *testing.T) { } } +// rerankGateLLM blocks its SimpleCall until the predictor has started, +// simulating a slow rerank that would otherwise consume the recall deadline. +type rerankGateLLM struct { + wait <-chan struct{} +} + +func (m *rerankGateLLM) SimpleCall(ctx context.Context, _, _ string) (string, error) { + select { + case <-m.wait: + return "", nil + case <-ctx.Done(): + return "", ctx.Err() + } +} + +// predictorSignalLLM closes the channel on its first call so the test can +// observe that prediction started while the rerank was still blocked. +type predictorSignalLLM struct { + ch chan struct{} + once sync.Once + resp string +} + +func (m *predictorSignalLLM) SimpleCall(_ context.Context, _, _ string) (string, error) { + m.once.Do(func() { close(m.ch) }) + return m.resp, nil +} + +// TestPredictionRunsConcurrentWithRerank is a regression test for the +// "predictor: context deadline exceeded" failure: the predictor must run +// concurrently with the literal query (including its paid rerank) so a slow +// rerank cannot consume the shared recall deadline before prediction starts. +// Under the old sequential behavior the rerank gate never opens and the +// recall deadlocks, failing this test via the outer timeout. +func TestPredictionRunsConcurrentWithRerank(t *testing.T) { + dir := t.TempDir() + cfg := DefaultConfig() + cfg.Enabled = boolPtr(true) + cfg.SemanticSearchRerank = boolPtr(true) + cfg.SemanticSearchMinScore = 0.01 + cfg.PredictiveIntents = 1 + cfg.FollowUpAnticipationEnabled = boolPtr(true) + // The mock embedder rates these two English sentences as near-duplicates; + // disable semantic dedup so both atoms reach the rerank path. + cfg.SemanticDedupThreshold = floatPtr(0) + + predictorStarted := make(chan struct{}) + em := New(dir, &rerankGateLLM{wait: predictorStarted}, cfg) + em.index.newEmb = func() embedding.TextEmbedder { return newMockEmbedder(vectorDim) } + em.index.emb = newMockEmbedder(vectorDim) + em.recall.predictor = NewPredictor(&predictorSignalLLM{ + ch: predictorStarted, + resp: `[{"text":"follow-up question","confidence":0.9}]`, + }, cfg) + defer em.Close() + + _ = em.AddAtom(context.Background(), makeSearchableAtom("User prefers Go for backend services")) + _ = em.AddAtom(context.Background(), makeSearchableAtom("Run go test ./... to verify")) + + done := make(chan error, 1) + go func() { + _, err := em.recall.queryAtomsWithPrediction(context.Background(), "Go backend", nil, UserState{}) + done <- err + }() + select { + case err := <-done: + if err != nil { + t.Fatalf("queryAtomsWithPrediction failed: %v", err) + } + case <-time.After(10 * time.Second): + t.Fatal("recall deadlocked: predictor did not run concurrently with the rerank") + } +} + // newIntentWeightedRecallEM builds an enabled ExtendedMemory whose recall // uses the mock embedder, no rerank, and prediction enabled. The corpus is // two atoms with (almost) disjoint character sets so the literal query