From 5d96bcb4323aaeb1d001111f18297acdd83ee671 Mon Sep 17 00:00:00 2001 From: Fengzdadi <453788063@qq.com> Date: Sun, 25 Oct 2026 01:28:17 -0400 Subject: [PATCH 1/2] fix(context): bound a preparation's query embedding waits by one budget recall_embedding_timeout_seconds applied to each search, so a preparation that searched its Scope, context references and Topic Memory waited a full budget for each of them when the provider was slow. Measured against a stalled provider, a three-Scope preparation took 1.06 s at the 0.25 s default, past the 1 s prepare request of Kimi Code, ZCode and DSH. The budget now covers the whole preparation. Each search reports how long it waited for its query embedding and only that wait is charged, so database reads, Experience recall and reranking between two searches do not take a later search's share. A search that finds the budget spent skips its embedding and matches lexically; the timeout that spent it has already recorded the preparation's degradation cause. That bound lets the default rise to 0.5 s, in RuntimeConfig and in BuiltinRuntime. With text-embedding-3-small through OpenRouter (p50 0.24 s), 0.25 s lost the semantic channel on 20 of 60 measured preparations of the #2039 Hook prompt, and 0.5 s on 1 of 60. --- docs/en/docs/operate/configuration.md | 2 +- .../en/docs/workflows/prepare-context-text.md | 9 +- docs/zh/docs/operate/configuration.md | 2 +- .../zh/docs/workflows/prepare-context-text.md | 6 +- .../builtin/artifacts/topic_memory/models.py | 2 + .../builtin/runtime/application.py | 39 ++++++-- .../builtin/runtime/atomic_memory.py | 8 +- src/powercontext/builtin/runtime/config.py | 11 +-- .../builtin/runtime/topic_memory_search.py | 7 +- .../test_recall_embedding_timeout_config.py | 2 +- tests/e2e/test_recall_embedding_timeout.py | 56 ++++++++++++ tests/e2e/test_recall_sufficiency_gate.py | 90 +++++++++++++++++-- 12 files changed, 205 insertions(+), 29 deletions(-) diff --git a/docs/en/docs/operate/configuration.md b/docs/en/docs/operate/configuration.md index ac49d8a06c..f3325224e8 100644 --- a/docs/en/docs/operate/configuration.md +++ b/docs/en/docs/operate/configuration.md @@ -170,7 +170,7 @@ still occupies an HTTP slot until it finishes; reserve other slots for concurren | `POWERCONTEXT_SERVER_RUNTIME_RECALL_GATE_ALLOW_WITH_RERANK` | `false` | Whether expansion may still run when `MEMORY_RERANK_ENABLED` is set; by default reranking ends the expansion | | `POWERCONTEXT_SERVER_RUNTIME_RECALL_FTS_WEIGHT` | `1.0` | Relative FTS contribution in hybrid Atomic Memory and Topic Memory RRF; `0` excludes FTS-only matches from hybrid results | | `POWERCONTEXT_SERVER_RUNTIME_RECALL_VECTOR_WEIGHT` | `1.0` | Relative vector contribution in hybrid Atomic Memory and Topic Memory RRF; `0` excludes vector-only matches from hybrid results | -| `POWERCONTEXT_SERVER_RUNTIME_RECALL_EMBEDDING_TIMEOUT_SECONDS` | `0.25` | Seconds each Atomic Memory or Topic Memory search inside a context preparation waits for the query embedding, at most `60`. A preparation searches the Scope and each of its context references in turn, so the waits add up. When a budget expires, that search continues with lexical matching only and logs `memory.search.embedding_fallback` or `topic_memory.search.embedding_fallback`. A remote embedding provider needs more than the default; keep the summed waits inside the host's Hook request timeout, 3 seconds on most hosts | +| `POWERCONTEXT_SERVER_RUNTIME_RECALL_EMBEDDING_TIMEOUT_SECONDS` | `0.5` | Seconds one context preparation waits for query embeddings in total, at most `60`. The Atomic Memory searches of the Scope and each of its context references and the Topic Memory search share this budget, and only the time spent waiting for an embedding is charged to it: each search waits only for what the earlier ones left, so the preparation's wait stays within it however many Scopes are searched, while database reads and reranking between the searches do not use it up. A search that finds the budget spent continues with lexical matching only and logs `memory.search.embedding_fallback` or `topic_memory.search.embedding_fallback`. The default covers a typical remote embedding round trip; keep it below the prepare request timeout of the hosts in use, the shortest default being 1 second on Kimi Code, ZCode, and DSH | | `POWERCONTEXT_SERVER_RUNTIME_MEMORY_SCHEDULE_SECONDS` | unset | Memory automatic admission interval; `SCHEDULE_SECONDS` remains a compatibility alias | | `POWERCONTEXT_SERVER_RUNTIME_TOPIC_MEMORY_SCHEDULE_SECONDS` | unset | Topic Memory automatic admission interval; unset disables new automatic admission | | `POWERCONTEXT_SERVER_RUNTIME_TOPIC_MEMORY_SOURCE_WINDOW_LIMIT` | `10` | Maximum Sources per Topic Memory Window, capped at 100; one Scope invocation can finish several Windows | diff --git a/docs/en/docs/workflows/prepare-context-text.md b/docs/en/docs/workflows/prepare-context-text.md index 721bc87502..485581f0a3 100644 --- a/docs/en/docs/workflows/prepare-context-text.md +++ b/docs/en/docs/workflows/prepare-context-text.md @@ -102,9 +102,12 @@ asyncio.run(export_context()) Default prepare requests without `assembly` also recall Topic Memory from the current Scope when available. With explicit `assembly`, select `topic-memory` to include it. `assembly: {}` still selects only Memory and Experience. -Topic recall skips query Embedding when the Scope has no current Topics. During context preparation, each Topic -query Embedding attempt has a 250 ms budget and falls back to FTS on timeout. The same prepare request reuses that FTS -choice in recall-gate expansion rounds; a subsequent request can try Embedding again. A slow optional provider does not +Topic recall skips query Embedding when the Scope has no current Topics. During context preparation, Topic recall +reuses the Memory query vector when both use the same embedding model and the query fits Topic search (64 distinct +terms and 8,192 characters); otherwise it embeds its own query. That attempt waits only for what Memory recall left of +the preparation's `recall_embedding_timeout_seconds` budget (500 ms by default), so a second embedding round trip has +to fit in the same budget, and falls back to FTS when the budget runs out. The same prepare request reuses that +FTS choice in recall-gate expansion rounds; a subsequent request can try Embedding again. A slow optional provider does not hold up already retrieved Memory. Direct Topic search keeps the configured provider timeout. Entries retain retrieval order within a family. Existing Memory reranking remains authoritative. Rank can have gaps diff --git a/docs/zh/docs/operate/configuration.md b/docs/zh/docs/operate/configuration.md index 5e96b0a468..495ed11c17 100644 --- a/docs/zh/docs/operate/configuration.md +++ b/docs/zh/docs/operate/configuration.md @@ -157,7 +157,7 @@ MCP 使用 stateless Streamable HTTP，每次工具执行都采用当前请求 | `POWERCONTEXT_SERVER_RUNTIME_RECALL_GATE_ALLOW_WITH_RERANK` | `false` | 已启用 `MEMORY_RERANK_ENABLED` 时是否仍允许追加搜索；默认在 rerank 之后不再追加 | | `POWERCONTEXT_SERVER_RUNTIME_RECALL_FTS_WEIGHT` | `1.0` | hybrid Atomic Memory 和 Topic Memory RRF 中 FTS 的相对贡献；`0` 会从 hybrid 结果中排除仅命中 FTS 的候选 | | `POWERCONTEXT_SERVER_RUNTIME_RECALL_VECTOR_WEIGHT` | `1.0` | hybrid Atomic Memory 和 Topic Memory RRF 中 vector 的相对贡献；`0` 会从 hybrid 结果中排除仅命中 vector 的候选 | -| `POWERCONTEXT_SERVER_RUNTIME_RECALL_EMBEDDING_TIMEOUT_SECONDS` | `0.25` | 一次上下文准备里每次 Atomic Memory 或 Topic Memory 检索等待查询 embedding 的秒数，上限 `60`。准备会依次检索该 Scope 及其每个 context reference，等待时间会累加。某次等待超时后，该次检索只用词法匹配继续，并记录 `memory.search.embedding_fallback` 或 `topic_memory.search.embedding_fallback`。远程 embedding 提供商需要比默认值更长的时间；累计等待应控制在宿主 Hook 请求超时之内，多数宿主为 3 秒 | +| `POWERCONTEXT_SERVER_RUNTIME_RECALL_EMBEDDING_TIMEOUT_SECONDS` | `0.5` | 一次上下文准备等待查询 embedding 的总秒数，上限 `60`。该 Scope 及其每个 context reference 的 Atomic Memory 检索与 Topic Memory 检索共用这一预算，且只有等待 embedding 的时间计入：每次检索只等待前面检索剩下的时间，所以无论检索多少个 Scope，准备的等待都不超过它，而检索之间的数据库读取和重排不会占用它。发现预算已用完的检索只用词法匹配继续，并记录 `memory.search.embedding_fallback` 或 `topic_memory.search.embedding_fallback`。默认值覆盖一次典型的远程 embedding 往返；应小于所用宿主的 prepare 请求超时，其中默认值最短的是 Kimi Code、ZCode 和 DSH 的 1 秒 | | `POWERCONTEXT_SERVER_RUNTIME_MEMORY_SCHEDULE_SECONDS` | 未设置 | Memory 自动准入间隔；`SCHEDULE_SECONDS` 保留为兼容别名 | | `POWERCONTEXT_SERVER_RUNTIME_TOPIC_MEMORY_SCHEDULE_SECONDS` | 未设置 | Topic Memory 自动准入间隔；未设置时不接纳新的自动调用 | | `POWERCONTEXT_SERVER_RUNTIME_TOPIC_MEMORY_SOURCE_WINDOW_LIMIT` | `10` | 每个 Topic Memory Window 的 Source 数量上限，硬上限为 100；一次 Scope 调用可完成多个 Window | diff --git a/docs/zh/docs/workflows/prepare-context-text.md b/docs/zh/docs/workflows/prepare-context-text.md index dc5c4fa733..e051ae24b4 100644 --- a/docs/zh/docs/workflows/prepare-context-text.md +++ b/docs/zh/docs/workflows/prepare-context-text.md @@ -99,8 +99,10 @@ asyncio.run(export_context()) 省略 `assembly` 的默认 prepare 请求还会召回当前 Scope 中可用的 Topic Memory。 显式 `assembly` 通过选择 `topic-memory` 包含主题记忆；`assembly: {}` 仍只选择 Memory 和 Experience。 -当前 Scope 没有 Topic 时，主题召回会跳过查询 Embedding。组装上下文时，每次 Topic 查询 Embedding 最多等待 -250 毫秒，超时后回退到 FTS。同一次 prepare 的 recall gate 扩展轮次复用这个选择，后续请求可以重新尝试 Embedding， +当前 Scope 没有 Topic 时，主题召回会跳过查询 Embedding。组装上下文时，若 Topic 与 Memory 使用同一个 Embedding 模型， +且查询在 Topic 检索的限制之内（64 个不同词项、8,192 个字符），主题召回复用 Memory 的查询向量；否则单独嵌入查询。 +这次嵌入只等待 Memory 召回用剩的本次准备 `recall_embedding_timeout_seconds` 预算（默认 500 毫秒），所以第二次 +Embedding 往返也要落在同一预算内，预算用完后回退到 FTS。同一次 prepare 的 recall gate 扩展轮次复用这个选择，后续请求可以重新尝试 Embedding， 避免缓慢的可选模型重复阻塞已召回的 Memory。直接检索 Topic 仍使用配置的模型超时。 同一类别内保留召回顺序，包括已有 Memory reranker 的排序。前面的候选因预算无法装入时，rank 可能不连续。 diff --git a/src/powercontext/builtin/artifacts/topic_memory/models.py b/src/powercontext/builtin/artifacts/topic_memory/models.py index 810d06fa03..2d4805717d 100644 --- a/src/powercontext/builtin/artifacts/topic_memory/models.py +++ b/src/powercontext/builtin/artifacts/topic_memory/models.py @@ -251,6 +251,8 @@ class TopicMemorySearchResult(BaseModel): artifacts: tuple[TopicMemory, ...] = Field(default=(), exclude=True) #: Set only when a wanted query embedding was abandoned, never when FTS was the plan. embedding_fallback: str | None = Field(default=None, exclude=True) + #: Seconds spent waiting for the query embedding, so a context preparation can share one budget across searches. + embedding_wait_seconds: float = Field(default=0.0, exclude=True) class PublishedTopicMemory(BaseModel): diff --git a/src/powercontext/builtin/runtime/application.py b/src/powercontext/builtin/runtime/application.py index c2c030ecda..9154637a45 100644 --- a/src/powercontext/builtin/runtime/application.py +++ b/src/powercontext/builtin/runtime/application.py @@ -469,6 +469,21 @@ def _degradation(current: str | None, observed: str | None) -> str | None: return _TIMEOUT_REASON if observed == _TIMEOUT_REASON else (current or observed) +@dataclass +class _EmbeddingBudget: + """What is left of one context preparation's budget for waiting on query embeddings. + + Only the waits are charged, so database reads, Experience recall and reranking between two searches do not take a + later search's share. A spent budget is reached only through an embedding that timed out, which already recorded + the preparation's degradation cause, so a later search can skip its embedding without losing that cause. + """ + + remaining: float + + def spend(self, seconds: float) -> None: + self.remaining = max(self.remaining - seconds, 0.0) + + @dataclass(frozen=True) class _RecallRoundOutcome: """All observable results produced by one bounded recall pass.""" @@ -1140,6 +1155,9 @@ async def _prepare_build( request.bootstrap_receipt_id, ) + # One budget covers every query embedding this preparation waits for. Each search gets what the earlier ones + # left, so the waits cannot add up past it however many Scopes and families are searched. + embedding_budget = _EmbeddingBudget(self._runtime.recall_embedding_timeout_seconds) round_zero = await self._recall_round( request, scope_ids, @@ -1148,6 +1166,7 @@ async def _prepare_build( admission=None, reuse=reuse, topic_reuse=topic_reuse, + embedding_budget=embedding_budget, excluded_memory=injected_memory_keys, ) memory_candidates = list(round_zero.memory) @@ -1182,6 +1201,7 @@ async def _prepare_build( profile_candidates=profile_candidates, reuse=reuse, topic_reuse=topic_reuse, + embedding_budget=embedding_budget, round_zero=round_zero, excluded_memory=injected_memory_keys, ) @@ -1310,6 +1330,7 @@ async def _gated_recall_effort( # noqa: C901 - the bounded expansion loop is in profile_candidates: Sequence[PreparedProfileCandidate], reuse: dict[str, MemoryQueryEmbedding | None], topic_reuse: dict[str, MemoryQueryEmbedding | None], + embedding_budget: _EmbeddingBudget, round_zero: _RecallRoundOutcome, excluded_memory: frozenset[tuple[object, ...]], ) -> tuple[ @@ -1397,6 +1418,7 @@ async def _gated_recall_effort( # noqa: C901 - the bounded expansion loop is in admission=plan.admission, reuse=reuse, topic_reuse=topic_reuse, + embedding_budget=embedding_budget, excluded_memory=excluded_memory, ) memory_ranks: dict[Any, int] = {} @@ -1543,6 +1565,7 @@ async def _recall_round( admission: AdmissionFloor | None, reuse: dict[str, MemoryQueryEmbedding | None], topic_reuse: dict[str, MemoryQueryEmbedding | None], + embedding_budget: _EmbeddingBudget, excluded_memory: frozenset[tuple[object, ...]], ) -> _RecallRoundOutcome: memory_candidates: list[PreparedMemoryCandidates] = [] @@ -1567,6 +1590,7 @@ async def _recall_round( admission=admission, reuse=reuse.get(scope_id, shared_embedding), allow_embedding=scope_id not in reuse or reuse[scope_id] is not None, + embedding_budget=embedding_budget, ) memory_recoverable |= outcome.memory_recoverable degradation = _degradation(degradation, outcome.memory_embedding_fallback) @@ -1612,6 +1636,7 @@ async def _recall_round( admission=admission, reuse=topic_embedding, allow_embedding=self.scope_id not in topic_reuse or topic_reuse[self.scope_id] is not None, + embedding_budget=embedding_budget, ) if TOPIC_MEMORY_FAMILY in families else TopicMemoryRecallOutcome() @@ -1639,6 +1664,7 @@ async def _recall_scope( admission: AdmissionFloor | None, reuse: MemoryQueryEmbedding | None, allow_embedding: bool = True, + embedding_budget: _EmbeddingBudget, ) -> _ScopeRecallOutcome: async with self._runtime._scoped_operation(scope_id, embedding_purpose=ModelUsagePurpose.MEMORY_RECALL): with self._runtime._stage( @@ -1661,12 +1687,13 @@ async def _recall_scope( mode="auto", admission=admission, query_embedding=reuse, - embedding_timeout_seconds=self._runtime.recall_embedding_timeout_seconds, - allow_embedding=allow_embedding, + embedding_timeout_seconds=embedding_budget.remaining, + allow_embedding=allow_embedding and embedding_budget.remaining > 0, recovery_admission=None if policy is None else policy.round2_admission, context=_PREPARE_ATOMIC_CONTEXT.get(), _trace=False, ) + embedding_budget.spend(result.embedding_wait_seconds) memory_hits = result.hits memory_admission = result.admission memory_recoverable = result.recoverable @@ -1734,6 +1761,7 @@ async def _topic_memory_hits( admission: AdmissionFloor | None, reuse: MemoryQueryEmbedding | None, allow_embedding: bool, + embedding_budget: _EmbeddingBudget, ) -> TopicMemoryRecallOutcome: configured = self._runtime._topic_memory_search is not None bounded_query = bounded_topic_memory_query(query) @@ -1752,11 +1780,12 @@ async def _topic_memory_hits( SearchTopicMemoryRequest(query=bounded_query, limit=limit), admission=admission, query_embedding=reuse, - embedding_timeout_seconds=self._runtime.recall_embedding_timeout_seconds, - allow_embedding=allow_embedding, + embedding_timeout_seconds=embedding_budget.remaining, + allow_embedding=allow_embedding and embedding_budget.remaining > 0, ) ) ) + embedding_budget.spend(result.embedding_wait_seconds) if span is not None: span.set_attributes({"powercontext.topic_memory.search.result_count": len(result.hits)}) return TopicMemoryRecallOutcome( @@ -3830,7 +3859,7 @@ def __init__( source_window_limit: int = 100, context_assembly_max_entries: int = 8, recall_sufficiency_policy: RecallSufficiencyPolicy | None = None, - recall_embedding_timeout_seconds: float = 0.25, + recall_embedding_timeout_seconds: float = 0.5, scope_cache_size: int = DEFAULT_SCOPE_CACHE_SIZE, scope_evictor: ScopeEvictor | None = None, scope_cache_observer: ScopeCacheObserver | None = None, diff --git a/src/powercontext/builtin/runtime/atomic_memory.py b/src/powercontext/builtin/runtime/atomic_memory.py index 0ab5824c3b..36e07776b0 100644 --- a/src/powercontext/builtin/runtime/atomic_memory.py +++ b/src/powercontext/builtin/runtime/atomic_memory.py @@ -20,7 +20,7 @@ import logging from collections.abc import Sequence from contextlib import asynccontextmanager, nullcontext -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from hashlib import sha256 from time import perf_counter from typing import cast @@ -237,6 +237,8 @@ class AtomicMemorySearchPage: #: Set only when the query embedding was abandoned for lexical retrieval, so a #: caller can tell "no vector channel was ever wanted" apart from "it was lost". embedding_fallback: str | None = None + #: Seconds spent waiting for the query embedding, so a context preparation can share one budget across searches. + embedding_wait_seconds: float = field(default=0.0, compare=False) class AtomicMemoryApplication: @@ -541,6 +543,7 @@ async def search( # noqa: C901 raise InvalidBaseAccessRequestError("mode", "must be auto, text, vector, or hybrid") if kind is not None: AtomicMemoryContent(kind=kind, text="validation") + started = asyncio.get_running_loop().time() mode, vector, query_embedding, embedding_calls, embedding_fallback = await self._resolve_query_embedding( query, mode, @@ -549,6 +552,7 @@ async def search( # noqa: C901 allow_embedding, public_search=artifact_request is not None, ) + embedding_wait_seconds = asyncio.get_running_loop().time() - started channel_weights = application.recall_channel_weights if mode == "hybrid" else None profile = application.index.capabilities.embedding_profile if vector is not None else None # Inference finishes before the business snapshot opens. Builtin policy @@ -641,7 +645,7 @@ async def search( # noqa: C901 ) ) page = replace(page, matches=tuple(matches)) - return replace(page, recoverable=recoverable) + return replace(page, recoverable=recoverable, embedding_wait_seconds=embedding_wait_seconds) async def _validate_search_hits(self, connection, hits, *, verify_vectors) -> tuple[AtomicMemory, ...]: """Verify only returned candidates against one bounded authority snapshot. diff --git a/src/powercontext/builtin/runtime/config.py b/src/powercontext/builtin/runtime/config.py index 64e3ba3230..1fbacc9e48 100644 --- a/src/powercontext/builtin/runtime/config.py +++ b/src/powercontext/builtin/runtime/config.py @@ -180,11 +180,12 @@ def validate_memory_capacity_order(self) -> RuntimeConfig: recall_gate_allow_with_rerank: bool = False recall_fts_weight: float = Field(default=1.0, ge=0.0, allow_inf_nan=False) recall_vector_weight: float = Field(default=1.0, ge=0.0, allow_inf_nan=False) - # Budget for embedding the query in each Atomic Memory or Topic Memory search of a context preparation. A - # preparation searches the Scope and each context reference in turn, so the waits add up inside the Hook - # request; when a budget expires, that search continues with lexical matching only. A remote embedding - # provider needs more than the default. - recall_embedding_timeout_seconds: float = Field(default=0.25, gt=0.0, le=60.0, allow_inf_nan=False) + # Budget for waiting on query embeddings in one context preparation. The Atomic Memory searches of the Scope and + # each context reference and the Topic Memory search share it, and only their waits are charged, so database reads + # and reranking in between do not take a later search's share; a search that finds it spent continues with lexical + # matching only. The default covers a typical remote embedding round trip and leaves the rest of a one-second Hook + # request to the other work. + recall_embedding_timeout_seconds: float = Field(default=0.5, gt=0.0, le=60.0, allow_inf_nan=False) profile_schedule_enabled: bool = False profile_cron: str = "0 2 * * *" profile_timezone: str = "Asia/Shanghai" diff --git a/src/powercontext/builtin/runtime/topic_memory_search.py b/src/powercontext/builtin/runtime/topic_memory_search.py index 1d1d40a8f7..ab88d0afbf 100644 --- a/src/powercontext/builtin/runtime/topic_memory_search.py +++ b/src/powercontext/builtin/runtime/topic_memory_search.py @@ -175,16 +175,20 @@ async def _execute( **keywords, ) return result.model_copy(update={"query_embedding": query_embedding, "embedding_calls": 0}), None + started = asyncio.get_running_loop().time() embedded, fallback = await self._embed_for_search( embedding, query, timeout_seconds=embedding_timeout_seconds, allow_fallback=plan is None or plan.fallback_allowed, ) + waited = asyncio.get_running_loop().time() - started if embedded is None: result = await self._search(scope_id, query, mode="fts", **keywords) return ( - result.model_copy(update={"embedding_calls": 1, "embedding_fallback": fallback}), + result.model_copy( + update={"embedding_calls": 1, "embedding_fallback": fallback, "embedding_wait_seconds": waited} + ), fallback, ) result = await self._search( @@ -201,6 +205,7 @@ async def _execute( query_vector=tuple(embedded.vectors[0]), embedding_profile=embedding.profile ), "embedding_calls": 1, + "embedding_wait_seconds": waited, } ), None diff --git a/tests/builtin/runtime/test_recall_embedding_timeout_config.py b/tests/builtin/runtime/test_recall_embedding_timeout_config.py index 482e3720d7..4767366e00 100644 --- a/tests/builtin/runtime/test_recall_embedding_timeout_config.py +++ b/tests/builtin/runtime/test_recall_embedding_timeout_config.py @@ -26,7 +26,7 @@ def test_recall_embedding_timeout_defaults_to_the_hook_sized_budget() -> None: - assert RuntimeConfig().recall_embedding_timeout_seconds == 0.25 + assert RuntimeConfig().recall_embedding_timeout_seconds == 0.5 @pytest.mark.parametrize("value", [0, -1, 61, float("nan"), True]) diff --git a/tests/e2e/test_recall_embedding_timeout.py b/tests/e2e/test_recall_embedding_timeout.py index d8d0a3583c..a1ae8862f4 100644 --- a/tests/e2e/test_recall_embedding_timeout.py +++ b/tests/e2e/test_recall_embedding_timeout.py @@ -79,6 +79,25 @@ async def embed(self, texts: tuple[str, ...], /) -> EmbeddingResult: return EmbeddingResult(vectors=tuple((value, math.sqrt(1.0 - value**2)) for value in similarities)) +class _TimedQueryEmbedding: + """Embed entries at once and record how long each query embedding is waited for before it is abandoned.""" + + profile = _PROFILE + + def __init__(self, delay: float) -> None: + self.delay = delay + self.waits: list[float] = [] + + async def embed(self, texts: tuple[str, ...], /) -> EmbeddingResult: + if any(text == _QUERY for text in texts): + started = asyncio.get_running_loop().time() + try: + await asyncio.sleep(self.delay) + finally: + self.waits.append(asyncio.get_running_loop().time() - started) + return EmbeddingResult(vectors=tuple((0.3, math.sqrt(1.0 - 0.3**2)) for _ in texts)) + + @asynccontextmanager async def _runtime(database: Path, budget: float, embedding: Any) -> AsyncIterator[BuiltinRuntime]: config = BuiltinConfig( @@ -186,6 +205,43 @@ async def scenario() -> Any: _assert_fallback_record(record, expected) +def test_one_budget_bounds_the_embedding_waits_of_a_preparation_across_scopes(tmp_path: Path) -> None: + """The budget is the preparation's, not each search's. + + A preparation searches its Scope and every context reference. If each search waited a full budget, a slow + provider would hold a three-Scope preparation for three budgets, past the one-second Hook request of some hosts, + and the Host would drop the context entirely instead of keeping the lexical matches. + """ + + budget = 0.2 + embedding = _TimedQueryEmbedding(delay=5.0) + + async def scenario() -> Any: + async with _runtime(tmp_path / "shared.db", budget, embedding) as runtime: + assert runtime.scopes is not None + assert runtime.atomic_memory is not None + references = tuple([await _scope_with_memory(runtime, f"reference-{index}") for index in range(2)]) + current = await runtime.scopes.create( + ScopeDraft( + title="Current", + summary="Two context references", + idempotency_key="current", + context_references=references, + ) + ) + await runtime.atomic_memory.for_scope(current.scope_id).create(( + AtomicMemoryContent(kind="fact", text=_MEMORY), + )) + return await runtime.context.for_scope(current.scope_id).prepare(_request()) + + prepared = asyncio.run(scenario()) + + assert prepared.status == "ready" + assert _MEMORY in (prepared.content or "") + assert prepared.reason == "inference_timeout" + assert sum(embedding.waits) < budget * 1.5 + + def test_a_sufficient_budget_waits_for_the_query_embedding(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: embedding = _SlowQueryEmbedding(delay=0.3) diff --git a/tests/e2e/test_recall_sufficiency_gate.py b/tests/e2e/test_recall_sufficiency_gate.py index 1922897499..aecf9dd1ab 100644 --- a/tests/e2e/test_recall_sufficiency_gate.py +++ b/tests/e2e/test_recall_sufficiency_gate.py @@ -44,6 +44,7 @@ from powercontext.builtin.artifacts.topic_memory import ( TopicMemoryContent, TopicMemoryDraft, + chunk_topic_memory_detail, prepare_topic_memory_projection, ) from powercontext.builtin.inference import EmbeddingResult, InferenceUnavailableError @@ -118,6 +119,7 @@ async def counted( admission: Any, reuse: Any, topic_reuse: Any, + embedding_budget: Any, excluded_memory: Any, ) -> Any: log.calls.append({"families": set(families), "admission": admission}) @@ -132,6 +134,7 @@ async def counted( admission=admission, reuse=reuse, topic_reuse=topic_reuse, + embedding_budget=embedding_budget, excluded_memory=excluded_memory, ) log.calls[-1]["memory_recoverable"] = result.memory_recoverable @@ -184,21 +187,33 @@ async def _seed(runtime: BuiltinRuntime, scope_id: str, texts: list[str]) -> Non await runtime.atomic_memory.for_scope(scope_id).create(tuple(_entry(text) for text in map(_embedded_body, texts))) -async def _seed_topic_memories(runtime: BuiltinRuntime, scope_id: str, count: int) -> None: +async def _seed_topic_memories( + runtime: BuiltinRuntime, scope_id: str, count: int, *, embedding_model: Any = None +) -> None: provider = cast(Any, runtime._provider) - async with provider.database.transaction() as connection: - for position in range(count): - content = TopicMemoryContent( - title=f"alpha beta topic {position:02d}", - summary=f"Useful evidence {position}", - detail="unrelated detail", + for position in range(count): + content = TopicMemoryContent( + title=f"alpha beta topic {position:02d}", + summary=f"Useful evidence {position}", + detail="unrelated detail", + ) + projection = prepare_topic_memory_projection(content) + if embedding_model is not None: + chunks = chunk_topic_memory_detail(content.detail) + embedded = await embedding_model.embed((f"{content.title}\n{content.summary}", *(c.text for c in chunks))) + projection = prepare_topic_memory_projection( + content, + topic_embedding=embedded.vectors[0], + chunk_embeddings=embedded.vectors[1:], + embedding_profile=embedding_model.profile, ) + async with provider.database.transaction() as connection: await provider.repositories.topic_memories.publish_create( connection, scope_id, f"topic-over-cap-{position:02d}", TopicMemoryDraft(content=content), - prepare_topic_memory_projection(content), + projection, ) @@ -796,6 +811,61 @@ async def scenario() -> None: asyncio.run(scenario()) +def test_work_between_embedding_waits_does_not_spend_the_preparation_budget(tmp_path) -> None: + """Only waiting for a query embedding is charged to ``recall_embedding_timeout_seconds``. + + Atomic Memory and Topic Memory each embed the query here, and a slow Experience recall runs between them. Both + providers answer well inside the budget; charging the Experience recall too would leave Topic Memory nothing and + drop its semantic channel. + """ + + class SlowQueryEmbedding: + profile = EmbeddingProfile(profile_id="query", model="query", dimension=2) + + def __init__(self) -> None: + self.answered = False + + async def embed(self, texts: tuple[str, ...], /) -> EmbeddingResult: + if _QUERY in texts: + await asyncio.sleep(0.15) + self.answered = True + return EmbeddingResult(vectors=tuple((1.0, 0.0) for _ in texts)) + + async def slow_experience(_scope_id: str, _query: str, _limit: int, **_kwargs: Any) -> ExperienceSearchOutcome: + await asyncio.sleep(0.3) + return ExperienceSearchOutcome() + + memory_embedding = SlowQueryEmbedding() + topic_embedding = SlowQueryEmbedding() + + async def scenario() -> Any: + async with _runtime( + tmp_path / "embedding-budget.db", + RuntimeConfig(recall_embedding_timeout_seconds=0.4), + embedding_model=memory_embedding, + ) as runtime: + scope_id = await _create_scope(runtime, "embedding-budget") + await _seed(runtime, scope_id, ["alpha beta gamma evidence"]) + await _seed_topic_memories(runtime, scope_id, 1, embedding_model=memory_embedding) + assert runtime._topic_memory_search is not None + runtime._topic_memory_searcher = TopicMemorySearcher( + search=runtime._topic_memory_search, + embedding_model=topic_embedding, + observer=runtime._topic_memory_search_observer, + ) + # A separate Topic Memory model makes it embed the query itself instead of reusing the Atomic vector. + runtime._topic_memory_embedding_model = topic_embedding + runtime._experience_recall = slow_experience + request = PrepareContextRequest.model_validate({"query": _QUERY, "max_bytes": 8000, "explain": True}) + return await runtime.context.for_scope(scope_id).prepare(request) + + prepared = asyncio.run(scenario()) + + assert memory_embedding.answered + assert topic_embedding.answered + assert prepared.reason is None + + def test_empty_scope_is_no_content_and_does_not_expand(tmp_path, monkeypatch) -> None: log = _RecallRoundLog() log.install(monkeypatch) @@ -900,6 +970,7 @@ async def changed_head( admission: Any, reuse: Any, topic_reuse: Any, + embedding_budget: Any, excluded_memory: Any, ) -> Any: nonlocal calls @@ -914,6 +985,7 @@ async def changed_head( admission=admission, reuse=reuse, topic_reuse=topic_reuse, + embedding_budget=embedding_budget, excluded_memory=excluded_memory, ) raise AtomicMemoryConflictError("Memory changed while expanding retrieval") # noqa: TRY003 @@ -1000,6 +1072,7 @@ async def fails_on_second_expansion( admission: Any, reuse: Any, topic_reuse: Any, + embedding_budget: Any, excluded_memory: Any, ) -> Any: nonlocal calls @@ -1015,6 +1088,7 @@ async def fails_on_second_expansion( admission=admission, reuse=reuse, topic_reuse=topic_reuse, + embedding_budget=embedding_budget, excluded_memory=excluded_memory, ) if calls == 2: From 4afbcd6006ae7a4313bbf60ab353eec8215b8d4f Mon Sep 17 00:00:00 2001 From: Fengzdadi <453788063@qq.com> Date: Sun, 25 Oct 2026 03:10:39 -0400 Subject: [PATCH 2/2] test(context): pin one provider call and one fallback event per spent budget Once the first search of a preparation spends the embedding budget, the later searches must not call the provider or log another fallback event, or memory.search.embedding_fallback would count waits that never reached the provider as provider timeouts. Summing the waits could not tell a skipped embedding from a very short one. --- tests/e2e/test_recall_embedding_timeout.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/e2e/test_recall_embedding_timeout.py b/tests/e2e/test_recall_embedding_timeout.py index a1ae8862f4..3a533bd699 100644 --- a/tests/e2e/test_recall_embedding_timeout.py +++ b/tests/e2e/test_recall_embedding_timeout.py @@ -205,12 +205,15 @@ async def scenario() -> Any: _assert_fallback_record(record, expected) -def test_one_budget_bounds_the_embedding_waits_of_a_preparation_across_scopes(tmp_path: Path) -> None: +def test_one_budget_bounds_the_embedding_waits_of_a_preparation_across_scopes( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: """The budget is the preparation's, not each search's. A preparation searches its Scope and every context reference. If each search waited a full budget, a slow provider would hold a three-Scope preparation for three budgets, past the one-second Hook request of some hosts, - and the Host would drop the context entirely instead of keeping the lexical matches. + and the Host would drop the context entirely instead of keeping the lexical matches. Once the first search has + spent the budget, the later ones do not call the provider, so the fallback event counts only the real timeout. """ budget = 0.2 @@ -239,7 +242,11 @@ async def scenario() -> Any: assert prepared.status == "ready" assert _MEMORY in (prepared.content or "") assert prepared.reason == "inference_timeout" - assert sum(embedding.waits) < budget * 1.5 + # This is the external inference-attempt budget for one request. + assert len(embedding.waits) == 1 + assert embedding.waits[0] < budget * 1.5 + (record,) = _fallbacks(caplog) + _assert_fallback_record(record, "inference_timeout") def test_a_sufficient_budget_waits_for_the_query_embedding(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: