Fix retrieval fallback under Docker state pressure

Co-authored-by: Codex <noreply@openai.com>

Author: ericjuta

src/functions/retrieval-blocks.ts

@@ -875,6 +875,71 @@ export async function collectRetrievalBlocksFromState(
   return [...blocks.values()];
 }
 
+export async function collectLightweightRetrievalBlocksFromState(
+  kv: StateKV,
+  options?: {
+    project?: string;
+    sessionId?: string;
+  },
+): Promise<RetrievalBlock[]> {
+  const memories = await kv.list<Memory>(KV.memories).catch(() => []);
+  const semantic = await kv.list<SemanticMemory>(KV.semantic).catch(() => []);
+  const procedural = await kv.list<ProceduralMemory>(KV.procedural).catch(() => []);
+  const summaries = await kv.list<SessionSummary>(KV.summaries).catch(() => []);
+  const handoffs = await kv.list<HandoffPacket>(KV.handoffPackets).catch(() => []);
+  const branchOverlays = await kv.list<BranchOverlay>(KV.branchOverlays).catch(() => []);
+  const guardrails = await kv.list<GuardrailMemory>(KV.guardrails).catch(() => []);
+  const decisions = await kv.list<DecisionMemory>(KV.decisions).catch(() => []);
+  const dossiers = await kv.list<ComponentDossier>(KV.componentDossiers).catch(() => []);
+  const profiles = await kv.list<ProjectProfile>(KV.profiles).catch(() => []);
+  const observations = options?.sessionId
+    ? await kv.list<CompressedObservation>(KV.observations(options.sessionId)).catch(() => [])
+    : [];
+
+  const blocks = new Map<string, RetrievalBlock>();
+  const put = (block: RetrievalBlock | null) => {
+    if (!block) return;
+    if (
+      options?.project &&
+      block.project !== options.project &&
+      block.project !== "global"
+    ) {
+      return;
+    }
+    blocks.set(block.id, block);
+  };
+
+  for (const memory of memories.filter((item) => item.isLatest)) {
+    put(buildMemoryRetrievalBlock(memory));
+  }
+  for (const item of semantic) put(buildSemanticRetrievalBlock(item));
+  for (const item of procedural) put(buildProceduralRetrievalBlock(item));
+  for (const summary of summaries) put(buildSessionSummaryRetrievalBlock(summary));
+  for (const packet of handoffs) put(buildHandoffRetrievalBlock(packet));
+  for (const overlay of branchOverlays.filter((item) => item.status === "active")) {
+    put(buildBranchOverlayRetrievalBlock(overlay));
+  }
+  for (const guardrail of guardrails.filter((item) => item.status === "active")) {
+    put(buildGuardrailRetrievalBlock(guardrail));
+  }
+  for (const decision of decisions.filter((item) => item.status === "active")) {
+    put(buildDecisionRetrievalBlock(decision));
+  }
+  for (const dossier of dossiers) put(buildDossierRetrievalBlock(dossier));
+  for (const profile of profiles) put(buildProfileRetrievalBlock(profile));
+  for (const observation of observations) {
+    if (
+      observation.importance >= 6 ||
+      observation.type === "error" ||
+      observation.type === "decision"
+    ) {
+      put(buildObservationRetrievalBlock(observation, options?.project || "global"));
+    }
+  }
+
+  return [...blocks.values()];
+}
+
 export async function refreshRetrievalBlocksFromState(
   kv: StateKV,
 ): Promise<number> {

src/functions/retrieval-engine.ts

@@ -19,19 +19,22 @@ import { SearchIndex } from "../state/search-index.js";
 import { GraphRetrieval } from "./graph-retrieval.js";
 import { extractEntitiesFromQuery } from "./query-expansion.js";
 import {
+  collectLightweightRetrievalBlocksFromState,
   collectRetrievalBlocksFromState,
 } from "./retrieval-blocks.js";
 
 const QUERY_EMBEDDING_CACHE_MAX_ENTRIES = 128;
 const QUERY_EMBEDDING_CACHE_TTL_MS = 5 * 60_000;
 const QUERY_EMBEDDING_TIMEOUT_MS = 2500;
+const RETRIEVAL_BLOCK_SCOPE_COOLDOWN_MS = 60_000;
 
 type CachedQueryEmbedding = {
   embedding: Float32Array;
   cachedAt: number;
 };
 
 const queryEmbeddingCache = new Map<string, CachedQueryEmbedding>();
+let retrievalBlockScopeUnavailableUntil = 0;
 
 function estimateTokens(text: string): number {
   return Math.ceil(text.length / 3);
@@ -301,12 +304,32 @@ export async function retrieveRelevantBlocks(
     blocks.some(
       (block) => block.project === query.project || block.project === "global",
     );
-  let allBlocks = await kv.list<RetrievalBlock>(KV.retrievalBlocks).catch(() => []);
+  let allBlocks: RetrievalBlock[] = [];
+  let usingStateFallbackBlocks = false;
+  const canReadStoredBlocks = Date.now() >= retrievalBlockScopeUnavailableUntil;
+  if (canReadStoredBlocks) {
+    try {
+      allBlocks = await kv.list<RetrievalBlock>(KV.retrievalBlocks);
+      retrievalBlockScopeUnavailableUntil = 0;
+    } catch {
+      retrievalBlockScopeUnavailableUntil = Date.now() + RETRIEVAL_BLOCK_SCOPE_COOLDOWN_MS;
+    }
+  }
   const shouldRefreshBlocks =
     allBlocks.length === 0 ||
     (Boolean(query.project) && !hasProjectCoverage(allBlocks));
   if (shouldRefreshBlocks) {
-    allBlocks = await collectRetrievalBlocksFromState(kv).catch(() => []);
+    const lightweightBlocks = await collectLightweightRetrievalBlocksFromState(kv, {
+      project: query.project,
+      sessionId: query.sessionId,
+    }).catch(() => []);
+    if (lightweightBlocks.length > 0) {
+      allBlocks = lightweightBlocks;
+      usingStateFallbackBlocks = true;
+    } else {
+      allBlocks = await collectRetrievalBlocksFromState(kv).catch(() => []);
+      usingStateFallbackBlocks = allBlocks.length > 0;
+    }
   }
   const blocks = allBlocks
     .filter((block) =>
@@ -344,8 +367,10 @@ export async function retrieveRelevantBlocks(
   ]).join(" ");
   const lexicalScores = new Map<string, number>();
   if (lexicalQuery.trim()) {
-    let lexicalResults = getRetrievalSearchIndex().searchDocuments(lexicalQuery, 120);
-    if (blocks.length > 0 && lexicalResults.length === 0) {
+    let lexicalResults = usingStateFallbackBlocks
+      ? []
+      : getRetrievalSearchIndex().searchDocuments(lexicalQuery, 120);
+    if (blocks.length > 0 && (usingStateFallbackBlocks || lexicalResults.length === 0)) {
       const fallbackIndex = new SearchIndex();
       for (const block of blocks) {
         fallbackIndex.addDocument(

src/index.ts

@@ -161,6 +161,12 @@ async function main() {
   });
 
   const kv = new StateKV(sdk);
+  const persistenceKv = new StateKV(sdk, {
+    timeoutMs: Math.max(
+      Number.parseInt(getEnvVar("STATE_KV_TIMEOUT_MS") || "5000", 10) || 5000,
+      20000,
+    ),
+  });
   const secret = getEnvVar("AGENTMEMORY_SECRET");
   const metricsStore = new MetricsStore(kv);
   const dedupMap = new DedupMap();
@@ -341,9 +347,9 @@ async function main() {
     };
   });
 
-  indexPersistence = new IndexPersistence(kv, bm25Index, vectorIndex);
+  indexPersistence = new IndexPersistence(persistenceKv, bm25Index, vectorIndex);
   retrievalIndexPersistence = new IndexPersistence(
-    kv,
+    persistenceKv,
     getRetrievalSearchIndex(),
     retrievalVectorIndex,
     KV.retrievalBlockIndex,

test/retrieval-engine.test.ts

@@ -1,7 +1,11 @@
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
-const { collectRetrievalBlocksFromStateMock } = vi.hoisted(() => ({
+const {
+  collectRetrievalBlocksFromStateMock,
+  collectLightweightRetrievalBlocksFromStateMock,
+} = vi.hoisted(() => ({
   collectRetrievalBlocksFromStateMock: vi.fn(async () => []),
+  collectLightweightRetrievalBlocksFromStateMock: vi.fn(async () => []),
 }));
 
 vi.mock("../src/functions/retrieval-blocks.js", async () => {
@@ -11,6 +15,8 @@ vi.mock("../src/functions/retrieval-blocks.js", async () => {
   return {
     ...actual,
     collectRetrievalBlocksFromState: collectRetrievalBlocksFromStateMock,
+    collectLightweightRetrievalBlocksFromState:
+      collectLightweightRetrievalBlocksFromStateMock,
   };
 });
 
@@ -27,6 +33,7 @@ import type { RetrievalBlock } from "../src/types.js";
 describe("retrieveRelevantBlocks", () => {
   beforeEach(() => {
     collectRetrievalBlocksFromStateMock.mockClear();
+    collectLightweightRetrievalBlocksFromStateMock.mockClear();
     getRetrievalSearchIndex().clear();
     configureRetrievalBlockIndexingRuntime({
       embeddingProvider: null,
@@ -77,4 +84,78 @@ describe("retrieveRelevantBlocks", () => {
     expect(result.searchResults).toHaveLength(1);
     expect(result.searchResults[0]?.block.id).toBe(block.id);
   });
+
+  it("falls back to lightweight state collection when the retrieval block scope is unavailable", async () => {
+    const kv = mockKV();
+    const staleBlock: RetrievalBlock = {
+      id: "rblk_stale",
+      sourceType: "semantic_memory",
+      sourceId: "sem_stale",
+      project: "global",
+      scope: "global",
+      freshnessLane: "cold",
+      canonicalText: "Unrelated stale retrieval memory",
+      title: "Stale memory",
+      files: [],
+      concepts: ["stale"],
+      entities: ["stale"],
+      sourceObservationIds: [],
+      hadFailure: false,
+      hadDecision: false,
+      hadAssistantConclusion: true,
+      isResumeArtifact: false,
+      importance: 7,
+      createdAt: "2026-01-01T00:00:00Z",
+      updatedAt: "2026-01-01T00:00:00Z",
+      eventAt: "2026-01-01T00:00:00Z",
+    };
+    const block: RetrievalBlock = {
+      id: "rblk_theta",
+      sourceType: "memory",
+      sourceId: "mem_theta",
+      project: "global",
+      scope: "global",
+      freshnessLane: "cold",
+      canonicalText: "Codex durable retrieval probe theta retrieval memory",
+      title: "Theta memory",
+      files: ["/tmp/codex-theta.txt"],
+      concepts: ["theta sentinel"],
+      entities: ["theta", "sentinel"],
+      sourceObservationIds: [],
+      hadFailure: false,
+      hadDecision: false,
+      hadAssistantConclusion: true,
+      isResumeArtifact: false,
+      importance: 7,
+      createdAt: "2026-01-01T00:00:00Z",
+      updatedAt: "2026-01-01T00:00:00Z",
+      eventAt: "2026-01-01T00:00:00Z",
+    };
+
+    getRetrievalSearchIndex().addDocument(
+      staleBlock.id,
+      staleBlock.project,
+      buildRetrievalBlockLexicalText(staleBlock),
+    );
+
+    const listError = new Error("retrieval block scope timeout");
+    const rawList = kv.list.bind(kv);
+    kv.list = (async <T>(scope: string): Promise<T[]> => {
+      if (scope === KV.retrievalBlocks) throw listError;
+      return rawList(scope);
+    }) as typeof kv.list;
+
+    collectLightweightRetrievalBlocksFromStateMock.mockResolvedValue([block]);
+
+    const result = await retrieveRelevantBlocks(kv as never, {
+      query: "theta sentinel",
+      budget: 300,
+      purpose: "smart-search",
+    });
+
+    expect(collectLightweightRetrievalBlocksFromStateMock).toHaveBeenCalled();
+    expect(collectRetrievalBlocksFromStateMock).not.toHaveBeenCalled();
+    expect(result.searchResults).toHaveLength(1);
+    expect(result.searchResults[0]?.block.id).toBe(block.id);
+  });
 });