Compare commits
20
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e3fe1719f4 | ||
|
|
e9cce69650 | ||
|
|
2fed7adf6b | ||
|
|
7308ed9df2 | ||
|
|
1c0cbf6d05 | ||
|
|
71854a3510 | ||
|
|
fcb1bcd866 | ||
|
|
6305e4cea8 | ||
|
|
a4b0969248 | ||
|
|
e0920311a1 | ||
|
|
e6426f6fd0 | ||
|
|
99e3c8d9d7 | ||
|
|
896a9c6c83 | ||
|
|
193fabeeb4 | ||
|
|
0aeb3e8760 | ||
|
|
70aaa7be85 | ||
|
|
18046152ab | ||
|
|
7b946a28ef | ||
|
|
e5ce0cd964 | ||
|
|
d8837b4fb2 |
No files matched your search
@@ -4,6 +4,28 @@ Hy Memory and MemHub are versioned and released independently.
|
||||
|
||||
## Unreleased
|
||||
|
||||
- The SDK memory page now reads its list, count, source and time labels and
|
||||
graph from one page response for one source filter, so rows, totals and the
|
||||
map no longer disagree. Memories with no recorded source are shown and
|
||||
filterable as "unidentified source"; library search keeps only memories
|
||||
that contain the words within the SDK's relevance window.
|
||||
- Metadata repair v2 follows `governed_memory_evidence` across every
|
||||
pre-cutover bundle, which restores the source of memories the cutover
|
||||
backfilled without one, and moves impossible future conversation times back
|
||||
to the write time.
|
||||
- Reused the existing three-column memory map for SDK-authority list/detail
|
||||
views and removed the retired "original evidence" surface. SDK-native status,
|
||||
source, conversation time and creation time now render without fabricated
|
||||
revisions; source and text filters are applied by the paged Core endpoint.
|
||||
- SDK memory search now calls the SDK search contract instead of listing the
|
||||
entire library in Core. Corrupt empty records are skipped and diagnosed
|
||||
without mutating data or failing a page, and post-cutover metadata repair
|
||||
restores source and original times from the verified migration bundle.
|
||||
- Kept Cursor's WAL change token separate from the file mtime so a hash can no
|
||||
longer appear as a conversation date centuries in the future.
|
||||
- A trusted packaged Runtime now repairs a dirty same-build installation by
|
||||
staging a clean replacement, and `runtime_service_not_ready` is rendered as
|
||||
user-facing startup guidance instead of an internal code.
|
||||
- Fixed legacy-to-SDK startup migration to adopt native Coding Experience and
|
||||
Working Fact IDs instead of importing them as Chat layers. Raw migration now
|
||||
rebuilds bounded user/assistant turns from the latest three days, excludes
|
||||
|
||||
@@ -15,6 +15,8 @@
|
||||
3. MemHub **不存 memory 实体**;**可以存 `memory_id`** 做组织/回执。原文进入 **SDK raw**。
|
||||
4. SDK raw 只保留滚动三天;一条 raw 对应一个完整 turn,只含真人 user 与
|
||||
assistant,tool 不持久化。提取后的高层 memory 不受该窗口影响。
|
||||
5. `source_raw_ids` 是 SDK 内部溯源关系,不是产品 Evidence 合同;记忆详情不提供
|
||||
“原始依据”或按 raw ID 展开原文的入口。
|
||||
|
||||
生产入口通过共享 Runtime 租约取得公开 `HyMemoryClient`;MemHub 只在调用前执行
|
||||
模型配置、密钥和 harness/model 门禁,不维护第二套 memory 正文权威。
|
||||
|
||||
@@ -3180,7 +3180,8 @@ async function runCommand(
|
||||
migration.code ?? "sdk_migration_incomplete",
|
||||
);
|
||||
}
|
||||
migrated = migration.state === "cutover";
|
||||
migrated =
|
||||
migration.state === "cutover" || migration.state === "complete";
|
||||
await applyPreparedServiceInstall(preparedService);
|
||||
serviceApplied = true;
|
||||
// The receipt exists so a later command can find the data dir without
|
||||
@@ -3665,7 +3666,8 @@ async function runCommand(
|
||||
}
|
||||
// An older Runtime refuses a migrated home, so from here a failure
|
||||
// keeps the new version instead of rolling back to one that cannot start.
|
||||
migrated = migration.state === "cutover";
|
||||
migrated =
|
||||
migration.state === "cutover" || migration.state === "complete";
|
||||
await restartUserService(layout, serviceOptions);
|
||||
// A restart that returns without error only means launchd/systemd
|
||||
// accepted the command. O-22: the first `upgrade` implementation
|
||||
@@ -4131,6 +4133,15 @@ function notReadyAfterMigration(command: "install" | "upgrade") {
|
||||
};
|
||||
}
|
||||
|
||||
const SDK_MIGRATION_FAILURES: Record<string, string> = {
|
||||
sdk_migration_disk_space_insufficient:
|
||||
"磁盘空间不足,记忆数据迁移没有完成,已回滚到上一版本。请腾出与记忆数据同等大小的空间后重试;原有记忆数据不会丢失。",
|
||||
sdk_migration_memory_unreadable:
|
||||
"核对发现新版本还读不到你的全部记忆,没有切换数据源,已回滚到上一版本。请把这个错误码反馈给我们;原有记忆数据不会丢失。",
|
||||
sdk_migration_model_profile_required:
|
||||
"迁移需要先有可用的模型配置,用来核对新版本能否读回记忆。已回滚到上一版本;请先用 memhub models setup 完成配置后再升级。",
|
||||
};
|
||||
|
||||
function sdkMigrationFailure(command: "install" | "upgrade", code: string) {
|
||||
return {
|
||||
schema_version: 1,
|
||||
@@ -4139,9 +4150,8 @@ function sdkMigrationFailure(command: "install" | "upgrade", code: string) {
|
||||
status: "error",
|
||||
code,
|
||||
message:
|
||||
code === "sdk_migration_disk_space_insufficient"
|
||||
? "磁盘空间不足,记忆数据迁移没有完成,已回滚到上一版本。请腾出与记忆数据同等大小的空间后重试;原有记忆数据不会丢失。"
|
||||
: "记忆数据迁移没有完成,已回滚到上一版本。请重试;如果仍然失败,请把这个错误码反馈给我们。原有记忆数据不会丢失。",
|
||||
SDK_MIGRATION_FAILURES[code] ??
|
||||
"记忆数据迁移没有完成,已回滚到上一版本。请重试;如果仍然失败,请把这个错误码反馈给我们。原有记忆数据不会丢失。",
|
||||
};
|
||||
}
|
||||
|
||||
@@ -4861,10 +4871,10 @@ async function runMaintenanceRestore(options: {
|
||||
const instanceId = randomUUID();
|
||||
const child = spawn(options.coreCommand, ["serve", "--maintenance"], {
|
||||
env: {
|
||||
// Do not pass the caller's Runtime tuning variables (for example
|
||||
// retrieval-index flags) into a maintenance process: a restore must
|
||||
// start the product's minimal maintenance configuration, not inherit
|
||||
// an unrelated experiment from the shell.
|
||||
// Do not pass the caller's Runtime tuning variables into a
|
||||
// maintenance process: a restore must start the product's minimal
|
||||
// maintenance configuration, not inherit an unrelated experiment from
|
||||
// the shell.
|
||||
PATH: process.env.PATH,
|
||||
HOME: process.env.HOME,
|
||||
LANG: process.env.LANG,
|
||||
|
||||
@@ -13,7 +13,8 @@ dependencies = [
|
||||
"filelock>=3.32.4",
|
||||
# Public SDK version validated against MemHub's permission and recall contracts.
|
||||
"hy-memory==2.0.5",
|
||||
# Retrieval's derived vector index directly consumes these locked versions.
|
||||
# chromadb arrives through the HY SDK; locked here so an unrelated upgrade
|
||||
# cannot move the SDK's store. Built-in embedding consumes numpy directly.
|
||||
"chromadb==1.5.9",
|
||||
"numpy==2.5.2",
|
||||
"httpx>=0.28.1",
|
||||
|
||||
@@ -103,7 +103,9 @@ from one_memory.hy_sdk_types import (
|
||||
HySdkMemoryRecord,
|
||||
HySdkRawPage,
|
||||
HySdkRawRecord,
|
||||
HySdkReadError,
|
||||
HySdkUserProfile,
|
||||
library_keyword_hits,
|
||||
parse_sdk_search,
|
||||
)
|
||||
from one_memory.imports.access import (
|
||||
@@ -841,8 +843,6 @@ def create_app(
|
||||
secret_store: SecretStore | None = None,
|
||||
) -> FastAPI:
|
||||
runtime_settings = settings or load_settings()
|
||||
if runtime_settings.retrieval_index_enabled and not manage_runtime:
|
||||
raise ValueError("retrieval_runtime_ownership_required")
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
|
||||
@@ -881,6 +881,28 @@ def create_app(
|
||||
app.state.sdk_data_plane_migrations = SdkDataPlaneMigrationStore(database_path)
|
||||
app.state.memory_authority = app.state.sdk_data_plane_migrations.authority()
|
||||
if app.state.memory_authority is MemoryAuthority.SDK:
|
||||
from one_memory.hy_sdk_migration import (
|
||||
repair_sdk_metadata_from_bundle,
|
||||
)
|
||||
|
||||
try:
|
||||
repair = await asyncio.to_thread(
|
||||
repair_sdk_metadata_from_bundle,
|
||||
data_dir=runtime_settings.data_dir,
|
||||
database_path=database_path,
|
||||
)
|
||||
if repair.vector_rows or repair.coding_rows:
|
||||
logger.info(
|
||||
"sdk_metadata_repair_completed mapped=%d vector=%d coding=%d",
|
||||
repair.mapped,
|
||||
repair.vector_rows,
|
||||
repair.coding_rows,
|
||||
)
|
||||
except Exception as error:
|
||||
logger.warning(
|
||||
"sdk_metadata_repair_failed error_type=%s",
|
||||
type(error).__name__,
|
||||
)
|
||||
app.state.memory_governance_store = RetiredGovernanceStore(database_path)
|
||||
app.state.evidence_store = RetiredEvidenceStore(database_path)
|
||||
else:
|
||||
@@ -1456,22 +1478,7 @@ def create_app(
|
||||
app.state.resource_diagnostics = watched_task(
|
||||
asyncio.get_running_loop(), _emit_resource_snapshots
|
||||
)
|
||||
app.state.retrieval_runtime = None
|
||||
try:
|
||||
if runtime_settings.retrieval_index_enabled:
|
||||
from one_memory.retrieval.index_provisioning import provision_retrieval_index
|
||||
from one_memory.retrieval.runtime import ExistingRetrievalRuntime
|
||||
|
||||
# The runtime only resumes an index someone initialized. Turning
|
||||
# the feature on is that someone; otherwise enabling it would
|
||||
# leave every search on the lexical lane with nothing to say why.
|
||||
await asyncio.to_thread(provision_retrieval_index, runtime_settings.data_dir)
|
||||
app.state.retrieval_runtime = ExistingRetrievalRuntime(
|
||||
runtime_settings.data_dir,
|
||||
minimum_free_bytes=runtime_settings.retrieval_minimum_free_bytes,
|
||||
service_levels=app.state.service_levels,
|
||||
)
|
||||
await app.state.retrieval_runtime.start()
|
||||
# Startup writes are done. Only now let the derived-content
|
||||
# rebuild take the single write checkout: it is a long
|
||||
# background job, and on a large library it would otherwise
|
||||
@@ -1498,15 +1505,6 @@ def create_app(
|
||||
outcome="started",
|
||||
instance_id=runtime_settings.instance_id,
|
||||
)
|
||||
retrieval_stop_error: BaseException | None = None
|
||||
if app.state.retrieval_runtime is not None:
|
||||
try:
|
||||
await app.state.retrieval_runtime.stop()
|
||||
except (Exception, asyncio.CancelledError) as error:
|
||||
# Preserve failure while cleaning other owned services.
|
||||
# Experimental: an exceptional live-thread failure still
|
||||
# needs an outer ownership fence before production enablement.
|
||||
retrieval_stop_error = error
|
||||
app.state.main_database_report.cancel()
|
||||
app.state.loop_heartbeat.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
@@ -1564,8 +1562,6 @@ def create_app(
|
||||
# Waited for off the loop: this is the task whose whole point
|
||||
# is that the loop never waits on a disk.
|
||||
await asyncio.to_thread(flush_diagnostics)
|
||||
if retrieval_stop_error is not None:
|
||||
raise retrieval_stop_error
|
||||
|
||||
app = FastAPI(
|
||||
title="MemHub Core",
|
||||
@@ -3086,7 +3082,8 @@ def create_app(
|
||||
offset=0,
|
||||
order="desc",
|
||||
)
|
||||
page = HySdkMemoryPage.from_sdk(raw, limit=1, offset=0)
|
||||
page, skipped = HySdkMemoryPage.from_sdk_skipping_invalid(raw, limit=1, offset=0)
|
||||
_log_skipped_sdk_memories(skipped)
|
||||
except Exception as error:
|
||||
raise _sdk_http_error(error, "hy_memory_list_failed") from None
|
||||
return SdkHomeSummary(
|
||||
@@ -3130,10 +3127,10 @@ def create_app(
|
||||
store = getattr(app.state, "sdk_data_plane_migrations", None)
|
||||
if store is None:
|
||||
store = SdkDataPlaneMigrationStore(_evidence_store(app).database_path)
|
||||
marker = store.cutover()
|
||||
authority, migration_id = store.authority_state()
|
||||
return MemoryAuthorityResponse(
|
||||
authority="sdk" if marker is not None else "legacy",
|
||||
migration_id=marker.migration_id if marker is not None else None,
|
||||
authority="sdk" if authority is MemoryAuthority.SDK else "legacy",
|
||||
migration_id=migration_id,
|
||||
)
|
||||
|
||||
@app.get(
|
||||
@@ -3143,22 +3140,221 @@ def create_app(
|
||||
)
|
||||
async def sdk_memory_library(
|
||||
statuses: list[str] = Query(default=["active"]), # noqa: B008
|
||||
source: str | None = Query(default=None, min_length=1),
|
||||
query: str | None = Query(default=None, max_length=512),
|
||||
offset: int = Query(default=0, ge=0),
|
||||
limit: int = Query(default=100, ge=1, le=500),
|
||||
) -> HySdkMemoryPage:
|
||||
try:
|
||||
client = await _sdk_client_handle(app).client()
|
||||
search_query = (query or "").strip()
|
||||
if search_query:
|
||||
if set(statuses) != {"active"}:
|
||||
raise HTTPException(
|
||||
status_code=422,
|
||||
detail="sdk_search_supports_active_status_only",
|
||||
)
|
||||
# SDK search has no count-only contract and, on the local store,
|
||||
# no keyword channel. Ask for one bounded relevance window with
|
||||
# no score floor, then keep the rows that contain the words, so
|
||||
# a word in the library is found and a word in none of it finds
|
||||
# nothing. The whole library is never listed in Core.
|
||||
# A larger window made one search take long enough to freeze
|
||||
# the page; 20 is the bounded window the library searches.
|
||||
fetch = 20
|
||||
search_raw = await asyncio.to_thread(
|
||||
client.search,
|
||||
search_query,
|
||||
user_ids=["local"],
|
||||
agent_ids=[],
|
||||
include_sources=[source] if source else None,
|
||||
limit=fetch,
|
||||
min_score=0.0,
|
||||
schema_min_score=0.0,
|
||||
schema_limit=fetch,
|
||||
# The library lists stored memories; proactive intentions
|
||||
# are a recall-time lane the list never shows.
|
||||
intention_limit=0,
|
||||
scene="normal",
|
||||
reader=HY_MEMORY_READER,
|
||||
)
|
||||
hit_ids, window = library_keyword_hits(search_raw, search_query)
|
||||
rows = await _get_sdk_rows(client, list(hit_ids[offset : offset + limit]))
|
||||
memories: list[HySdkMemoryRecord] = []
|
||||
skipped = 0
|
||||
for row in rows:
|
||||
if row is None:
|
||||
continue
|
||||
try:
|
||||
record = HySdkMemoryRecord.from_sdk(row)
|
||||
except HySdkReadError:
|
||||
skipped += 1
|
||||
continue
|
||||
if record.status == "active":
|
||||
memories.append(record)
|
||||
_log_skipped_sdk_memories(skipped)
|
||||
return HySdkMemoryPage(
|
||||
memories=tuple(memories),
|
||||
total=len(hit_ids),
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
search_window=fetch if window >= fetch else None,
|
||||
)
|
||||
raw = await asyncio.to_thread(
|
||||
client.list_memories,
|
||||
user_id="local",
|
||||
statuses=list(statuses),
|
||||
include_sources=[source] if source else None,
|
||||
offset=offset,
|
||||
limit=limit,
|
||||
)
|
||||
return HySdkMemoryPage.from_sdk(raw, limit=limit, offset=offset)
|
||||
page, skipped = HySdkMemoryPage.from_sdk_skipping_invalid(
|
||||
raw, limit=limit, offset=offset
|
||||
)
|
||||
_log_skipped_sdk_memories(skipped)
|
||||
return page
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as error:
|
||||
raise _sdk_http_error(error, "hy_memory_list_failed") from None
|
||||
|
||||
@app.get(
|
||||
"/api/v2/memories/graph",
|
||||
response_model=MemoryGraphResponse,
|
||||
dependencies=protected,
|
||||
)
|
||||
async def sdk_memory_graph(
|
||||
limit: int = Query(default=160, ge=1, le=300),
|
||||
offset: int = Query(default=0, ge=0),
|
||||
source: str | None = Query(default=None, min_length=1),
|
||||
) -> MemoryGraphResponse:
|
||||
"""The memory map on SDK authority: SDK memories plus the Kuzu overlay.
|
||||
|
||||
Memories come from the public client; only the overlay reads Kuzu
|
||||
(`hy_graph_projection`, the visualization exception). Shape matches
|
||||
the legacy map so the same view draws both.
|
||||
"""
|
||||
try:
|
||||
client = await _sdk_client_handle(app).client()
|
||||
# The map draws the page the list shows, so it reads the same
|
||||
# source filter and window rather than the newest memories overall.
|
||||
raw = await asyncio.to_thread(
|
||||
client.list_memories,
|
||||
user_id="local",
|
||||
include_sources=[source] if source else None,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
order="desc",
|
||||
)
|
||||
page, skipped = HySdkMemoryPage.from_sdk_skipping_invalid(
|
||||
raw, limit=limit, offset=offset
|
||||
)
|
||||
_log_skipped_sdk_memories(skipped)
|
||||
except Exception as error:
|
||||
raise _sdk_http_error(error, "hy_memory_list_failed") from None
|
||||
vdb = raw.get("vdb") if isinstance(raw, dict) else None
|
||||
total = vdb.get("total") if isinstance(vdb, dict) else None
|
||||
memories = [memory for memory in page.memories if memory.layer != "l1_raw"]
|
||||
catalog = await asyncio.to_thread(_source_catalog, app)
|
||||
nodes: dict[str, MemoryGraphNode] = {}
|
||||
edges: list[MemoryGraphEdge] = []
|
||||
sessions: dict[str, list[str]] = {}
|
||||
for memory in memories:
|
||||
memory_node_id = f"memory:{memory.memory_id}"
|
||||
nodes[memory_node_id] = MemoryGraphNode(
|
||||
node_id=memory_node_id,
|
||||
node_type="memory",
|
||||
label=memory.content,
|
||||
memory_id=memory.memory_id,
|
||||
kind=memory.layer,
|
||||
status=memory.status,
|
||||
)
|
||||
# The SDK fills in "default_session" when a writer gave none (the
|
||||
# cutover backfill did so for every legacy memory); it names no
|
||||
# conversation, so it must not chain those memories together.
|
||||
if memory.session_id and memory.session_id != "default_session":
|
||||
sessions.setdefault(memory.session_id, []).append(memory_node_id)
|
||||
if memory.source and memory.source != "default":
|
||||
source_node_id = f"source:{memory.source}"
|
||||
registration = catalog.get(memory.source)
|
||||
nodes[source_node_id] = MemoryGraphNode(
|
||||
node_id=source_node_id,
|
||||
node_type="source",
|
||||
label=(registration.descriptor.display_name if registration else memory.source),
|
||||
source_id=memory.source,
|
||||
)
|
||||
edges.append(
|
||||
MemoryGraphEdge(
|
||||
edge_id=f"{source_node_id}->{memory_node_id}",
|
||||
source_node_id=source_node_id,
|
||||
target_node_id=memory_node_id,
|
||||
relation="evidence_source",
|
||||
origin="provenance",
|
||||
)
|
||||
)
|
||||
for _session_id, node_ids in sorted(sessions.items()):
|
||||
for left, right in zip(node_ids, node_ids[1:], strict=False):
|
||||
edges.append(
|
||||
MemoryGraphEdge(
|
||||
edge_id=f"{left}->{right}:same_conversation",
|
||||
source_node_id=left,
|
||||
target_node_id=right,
|
||||
relation="same_conversation",
|
||||
origin="same_conversation",
|
||||
)
|
||||
)
|
||||
hy_graph = None
|
||||
# Aggregate nodes carry no per-source provenance yet, so a revoked
|
||||
# source withholds the overlay, as on the legacy map.
|
||||
if not await asyncio.to_thread(_read_access(app).has_revoked_sources):
|
||||
try:
|
||||
hy_graph = await _sdk_memory_runtime(app).graph_snapshot(limit=limit)
|
||||
except HySdkError:
|
||||
hy_graph = None
|
||||
if hy_graph is not None:
|
||||
for node in hy_graph.nodes:
|
||||
node_id = f"memory:{node.node_id}"
|
||||
if node_id in nodes:
|
||||
continue
|
||||
node_id = f"hy:{node.node_id}"
|
||||
nodes.setdefault(
|
||||
node_id,
|
||||
MemoryGraphNode(
|
||||
node_id=node_id,
|
||||
node_type="topic" if node.layer == "TOPIC" else "schema",
|
||||
label=node.content,
|
||||
confidence=node.confidence,
|
||||
),
|
||||
)
|
||||
for edge in hy_graph.edges:
|
||||
source_node_id, target_node_id = (
|
||||
f"memory:{edge.source_node_id}"
|
||||
if f"memory:{edge.source_node_id}" in nodes
|
||||
else f"hy:{edge.source_node_id}",
|
||||
f"memory:{edge.target_node_id}"
|
||||
if f"memory:{edge.target_node_id}" in nodes
|
||||
else f"hy:{edge.target_node_id}",
|
||||
)
|
||||
if source_node_id not in nodes or target_node_id not in nodes:
|
||||
continue
|
||||
edges.append(
|
||||
MemoryGraphEdge(
|
||||
edge_id=f"{source_node_id}->{target_node_id}:{edge.relation}",
|
||||
source_node_id=source_node_id,
|
||||
target_node_id=target_node_id,
|
||||
relation=edge.relation,
|
||||
weight=edge.weight,
|
||||
origin="hy_relation",
|
||||
)
|
||||
)
|
||||
return MemoryGraphResponse(
|
||||
total_memories=total if isinstance(total, int) else len(memories),
|
||||
projected_memories=len(memories),
|
||||
nodes=tuple(nodes.values()),
|
||||
edges=tuple(edges),
|
||||
hy_graph_available=hy_graph is not None and hy_graph.consulted,
|
||||
)
|
||||
|
||||
@app.get(
|
||||
"/api/v2/user-profile",
|
||||
response_model=SdkUserProfileResponse,
|
||||
@@ -4484,9 +4680,8 @@ class _LazyASGIApp:
|
||||
`create_app()` must not run at import time. The managed `serve()` entry
|
||||
imports `create_app` from this module before it can construct its own app
|
||||
with `manage_runtime=True`; an eager module-level app was therefore built
|
||||
first with the default `manage_runtime=False`. With retrieval index enabled
|
||||
that default app hit `retrieval_runtime_ownership_required` and killed the
|
||||
process before the managed app could bind a port.
|
||||
first with the default `manage_runtime=False`, doing startup work for an
|
||||
app that never serves.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
@@ -5463,6 +5658,26 @@ def _sdk_client_handle(app: FastAPI) -> SdkClientHandle:
|
||||
return handle
|
||||
|
||||
|
||||
def _log_skipped_sdk_memories(count: int) -> None:
|
||||
if count:
|
||||
logger.warning("sdk_memory_invalid_rows_skipped count=%d", count)
|
||||
|
||||
|
||||
# `client.get` funnels into the SDK's single loop thread, so more workers only
|
||||
# queue there while holding the default executor every other request needs.
|
||||
_SDK_GET_CONCURRENCY = 8
|
||||
|
||||
|
||||
async def _get_sdk_rows(client: Any, memory_ids: list[str]) -> list[object]:
|
||||
gate = asyncio.Semaphore(_SDK_GET_CONCURRENCY)
|
||||
|
||||
async def one(memory_id: str) -> object:
|
||||
async with gate:
|
||||
return await asyncio.to_thread(client.get, memory_id)
|
||||
|
||||
return list(await asyncio.gather(*(one(memory_id) for memory_id in memory_ids)))
|
||||
|
||||
|
||||
async def _get_sdk_memory(
|
||||
handle: SdkClientHandle,
|
||||
memory_id: str,
|
||||
|
||||
@@ -335,6 +335,8 @@ def sdk_migrate(options: Sequence[str], stream: TextIO) -> int:
|
||||
"bundle": "正在备份记忆数据……",
|
||||
"prepare": "正在整理待处理的写入任务……",
|
||||
"cutover": "正在切换数据源……",
|
||||
"vectors": "正在接回旧版本的向量数据……",
|
||||
"verify": "正在用新版本核对记忆能否读回……",
|
||||
"cleanup": "正在删除旧表……",
|
||||
}
|
||||
|
||||
@@ -356,13 +358,17 @@ def sdk_migrate(options: Sequence[str], stream: TextIO) -> int:
|
||||
progress=progress,
|
||||
)
|
||||
except Exception as error:
|
||||
from one_memory.sdk_migration_verify import MemoryReadbackFailed
|
||||
|
||||
code = migration_failure_code(error)
|
||||
document = {"schema_version": 1, "state": "failed", "code": code}
|
||||
_write_restore(stream, as_json, document, f"迁移未完成:{code}")
|
||||
failed: dict[str, object] = {"schema_version": 1, "state": "failed", "code": code}
|
||||
if isinstance(error, MemoryReadbackFailed):
|
||||
failed["readback"] = error.report.as_document()
|
||||
_write_restore(stream, as_json, failed, f"迁移未完成:{code}")
|
||||
return 1
|
||||
messages = {
|
||||
"not_required": "不需要迁移。",
|
||||
"complete": "迁移已完成,旧表已清理。",
|
||||
"complete": "迁移已完成;遗留的旧表或向量已处理。",
|
||||
"cutover": "迁移完成:记忆数据已以 SDK 为准。",
|
||||
}
|
||||
_write_restore(stream, as_json, result.as_document(), messages[result.state])
|
||||
@@ -390,16 +396,23 @@ def sdk_migration_status(options: Sequence[str], stream: TextIO) -> int:
|
||||
SdkDataPlaneMigrationStore,
|
||||
)
|
||||
from one_memory.sdk_migration_runner import abandoned_staging_present
|
||||
from one_memory.sdk_migration_verify import read_active_profile
|
||||
from one_memory.sdk_vector_reattach import legacy_vector_collections
|
||||
from one_memory.target_schema import legacy_tables_present
|
||||
|
||||
authority = SdkDataPlaneMigrationStore(database_path).authority()
|
||||
if authority is MemoryAuthority.SDK:
|
||||
profile = read_active_profile(database_path)
|
||||
leftovers = bool(
|
||||
legacy_tables_present(database_path)
|
||||
or abandoned_staging_present(
|
||||
data_dir=settings.data_dir,
|
||||
namespace=HY_MEMORY_NAMESPACE,
|
||||
)
|
||||
or (
|
||||
profile is not None
|
||||
and legacy_vector_collections(settings.data_dir, profile)
|
||||
)
|
||||
)
|
||||
state = "required" if leftovers else "not_required"
|
||||
else:
|
||||
|
||||
@@ -23,10 +23,6 @@ class Settings(BaseSettings):
|
||||
live_sync_enabled: bool = True
|
||||
automatic_memory_processing: bool = True
|
||||
imports_enabled: bool = True
|
||||
# Experimental and unreleased: ownership/error shutdown still requires acceptance.
|
||||
# Resume only explicitly initialized generations; never migrate at startup.
|
||||
retrieval_index_enabled: bool = False
|
||||
retrieval_minimum_free_bytes: int = Field(default=256 * 1024 * 1024, ge=0)
|
||||
# Bounded metadata only: query, original text and answers are never stored.
|
||||
call_event_retention_days: int | None = Field(default=30, ge=1, le=365)
|
||||
call_event_max_events: int | None = Field(default=10_000, ge=1, le=100_000)
|
||||
|
||||
@@ -47,6 +47,19 @@ class ContextSearchRequest(BaseModel):
|
||||
return self
|
||||
|
||||
|
||||
class SdkContextSearchRequest(ContextSearchRequest):
|
||||
"""`POST /api/v2/context/search`: the three SDK recall lanes are set apart.
|
||||
|
||||
`limit` is the normal lane's top-k and the cap on returned items;
|
||||
`schema_limit` and `intention_limit` are the SDK's schema (L6) and
|
||||
proactive intention (L7) lanes. A caller's values are passed through.
|
||||
"""
|
||||
|
||||
limit: int = Field(default=10, ge=1, le=20, strict=True)
|
||||
schema_limit: int = Field(default=3, ge=0, le=20, strict=True)
|
||||
intention_limit: int = Field(default=1, ge=0, le=20, strict=True)
|
||||
|
||||
|
||||
class ContextMessage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
|
||||
@@ -34,11 +34,12 @@ from one_memory.context_models import (
|
||||
MemoryContextItem,
|
||||
RawContextItem,
|
||||
SdkContextItem,
|
||||
SdkContextSearchRequest,
|
||||
SdkContextSearchResponse,
|
||||
SdkMemoryContextItem,
|
||||
SdkRawContextItem,
|
||||
)
|
||||
from one_memory.context_raw import revalidate_raw_context
|
||||
from one_memory.context_raw import read_raw_context, revalidate_raw_context
|
||||
from one_memory.context_search import search_context
|
||||
from one_memory.diagnostic_log import diagnostic_event, diagnostic_exception
|
||||
from one_memory.external_memory_access import external_access
|
||||
@@ -50,7 +51,6 @@ from one_memory.hy_sdk_types import (
|
||||
parse_sdk_search,
|
||||
)
|
||||
from one_memory.imports.storage import EvidenceStore
|
||||
from one_memory.retrieval.raw_lane import read_raw_indexed_context
|
||||
from one_memory.semantic_deadline import hy_wait_budget, search_within
|
||||
|
||||
CALL_EVENT_RESPONSE_BUDGET_MS = 250
|
||||
@@ -99,12 +99,11 @@ def context_router(dependencies: Sequence[Depends]) -> APIRouter:
|
||||
read = asyncio.create_task(
|
||||
in_read_pool(
|
||||
raw_read_pool(),
|
||||
read_raw_indexed_context,
|
||||
read_raw_context,
|
||||
database,
|
||||
value,
|
||||
source_ids=authority.raw_sources,
|
||||
deadline=lane_end,
|
||||
index=getattr(request.app.state, "retrieval_runtime", None),
|
||||
)
|
||||
)
|
||||
done, _ = await asyncio.wait({read}, timeout=max(0, lane_end - time.monotonic()))
|
||||
@@ -441,7 +440,7 @@ def sdk_context_router(dependencies: Sequence[Depends]) -> APIRouter:
|
||||
|
||||
@router.post("/search", response_model=SdkContextSearchResponse)
|
||||
async def search(
|
||||
value: ContextSearchRequest,
|
||||
value: SdkContextSearchRequest,
|
||||
request: Request,
|
||||
response: Response,
|
||||
) -> SdkContextSearchResponse:
|
||||
@@ -516,8 +515,8 @@ def sdk_context_router(dependencies: Sequence[Depends]) -> APIRouter:
|
||||
limit=value.limit,
|
||||
min_score=sdk_relevance_floor(),
|
||||
schema_min_score=sdk_relevance_floor(),
|
||||
schema_limit=value.limit,
|
||||
intention_limit=value.limit,
|
||||
schema_limit=value.schema_limit,
|
||||
intention_limit=value.intention_limit,
|
||||
scene="normal",
|
||||
reader=HY_MEMORY_READER,
|
||||
)
|
||||
|
||||
@@ -22,7 +22,7 @@ from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, Protocol, cast
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
from one_memory.background_work import interactive_read_pool
|
||||
from one_memory.context_expand import ExpansionAuthority, ExpansionRefSigner
|
||||
@@ -36,23 +36,8 @@ from one_memory.retrieval.models import (
|
||||
RawEvidence,
|
||||
RetrievalPath,
|
||||
RetrievalResponse,
|
||||
TimeRange,
|
||||
)
|
||||
from one_memory.retrieval.query_plan import plan_local_query
|
||||
from one_memory.retrieval.raw_candidates import RawCandidateBatch
|
||||
from one_memory.retrieval.raw_evidence import materialize_raw_candidates
|
||||
|
||||
|
||||
class RawIndexLookup(Protocol):
|
||||
def lookup_raw(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
source_ids: frozenset[str],
|
||||
deadline: float,
|
||||
time_range: TimeRange | None = None,
|
||||
) -> tuple[RawCandidateBatch, bool]: ...
|
||||
|
||||
|
||||
ExpandRefs = tuple[ExpansionRefSigner, ExpansionAuthority]
|
||||
|
||||
@@ -106,11 +91,6 @@ class LaneResult:
|
||||
coverage_limited: bool = False
|
||||
|
||||
|
||||
def _semantic_only_raw(item: RawEvidence) -> bool:
|
||||
"""Semantic overlap alone is a candidate, not a confirmed recall item."""
|
||||
return set(item.matched_by) == {"semantic"}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StageTimings:
|
||||
raw_ms: float = 0.0
|
||||
@@ -191,12 +171,6 @@ def _managed(
|
||||
return selected, truncated
|
||||
|
||||
|
||||
def _group_missing_requirements(item: RawEvidence | GovernedEvidence) -> list[str]:
|
||||
if isinstance(item, RawEvidence) and _semantic_only_raw(item):
|
||||
return ["relevance_confirmation"]
|
||||
return []
|
||||
|
||||
|
||||
def _path(
|
||||
status: str,
|
||||
reason: str | None,
|
||||
@@ -296,8 +270,7 @@ def read_lexical_raw_lane(
|
||||
("lexical", "object") if plan.object_keys else ("lexical",),
|
||||
truncated,
|
||||
# Starvation means: this lane matched, then dropped the item whole. An
|
||||
# ordinary no-match must stay empty, and an unconfigured semantic lane
|
||||
# is not starvation.
|
||||
# ordinary no-match must stay empty.
|
||||
starved and not items,
|
||||
probe_saturated,
|
||||
)
|
||||
@@ -564,97 +537,6 @@ def warm_raw_read_path(
|
||||
return False
|
||||
|
||||
|
||||
def read_semantic_raw_lane(
|
||||
index: RawIndexLookup | None,
|
||||
database_path: Path,
|
||||
query: CoordinatorQuery,
|
||||
*,
|
||||
source_ids: frozenset[str],
|
||||
deadline: float,
|
||||
) -> LaneResult:
|
||||
if not source_ids:
|
||||
return LaneResult((), "unavailable", "permission_required", (), False)
|
||||
# No derived index is the no-model baseline, not a failure: the lexical lane
|
||||
# must still answer. Semantic is genuinely unsearched in that state, so it
|
||||
# must not claim a channel or a coverage result it never produced.
|
||||
if index is None:
|
||||
return LaneResult((), "not_searched", None, (), False)
|
||||
try:
|
||||
batch, unpublished = index.lookup_raw(
|
||||
query.query,
|
||||
source_ids=source_ids,
|
||||
deadline=deadline,
|
||||
time_range=plan_local_query(query.query).time_range,
|
||||
)
|
||||
except Exception:
|
||||
return LaneResult((), "unavailable", "index_unavailable", ("semantic",), False)
|
||||
if not isinstance(batch, RawCandidateBatch):
|
||||
return LaneResult((), "unavailable", "index_unavailable", ("semantic",), False)
|
||||
if batch.reason == "not_authorized":
|
||||
return LaneResult((), "unavailable", "permission_required", ("semantic",), False)
|
||||
if batch.reason == "deadline":
|
||||
return LaneResult((), "unavailable", "timeout_budget", ("semantic",), True)
|
||||
if batch.reason in {"vector_unavailable", "index_unavailable"}:
|
||||
return LaneResult((), "unavailable", batch.reason, ("semantic",), False)
|
||||
store, generation = _index_state(index, database_path)
|
||||
from one_memory.retrieval.index_store import RetrievalIndexStore
|
||||
|
||||
store = cast("RetrievalIndexStore | None", store)
|
||||
if store is None or generation is None:
|
||||
return LaneResult((), "unavailable", "index_unavailable", ("semantic",), False)
|
||||
try:
|
||||
materialized = materialize_raw_candidates(
|
||||
store,
|
||||
batch,
|
||||
generation_id=generation,
|
||||
source_ids=source_ids,
|
||||
deadline=deadline,
|
||||
character_budget=query.character_budget,
|
||||
)
|
||||
except Exception:
|
||||
return LaneResult((), "unavailable", "index_unavailable", ("semantic",), False)
|
||||
if materialized.reason == "deadline":
|
||||
return LaneResult((), "unavailable", "timeout_budget", ("semantic",), True)
|
||||
if materialized.reason in {"index_unavailable", "not_authorized"}:
|
||||
return LaneResult((), "unavailable", materialized.reason, ("semantic",), False)
|
||||
complete = materialized.complete and not unpublished
|
||||
if materialized.reason == "evidence_changed":
|
||||
status: Literal["ok", "empty", "unavailable", "not_searched"] = "unavailable"
|
||||
reason = "evidence_changed"
|
||||
elif not materialized.items and not complete:
|
||||
status, reason = "unavailable", "index_incomplete"
|
||||
elif materialized.reason == "character_budget":
|
||||
status, reason = "ok", None
|
||||
else:
|
||||
status, reason = ("ok" if materialized.items else "empty"), None
|
||||
return LaneResult(
|
||||
tuple(materialized.items),
|
||||
status,
|
||||
reason,
|
||||
("semantic",),
|
||||
materialized.reason == "character_budget",
|
||||
)
|
||||
|
||||
|
||||
def _index_state(index: RawIndexLookup, database_path: Path) -> tuple[object | None, str | None]:
|
||||
store = getattr(index, "_store", None)
|
||||
generation = None
|
||||
snapshot = getattr(index, "snapshot", None)
|
||||
if snapshot is not None:
|
||||
try:
|
||||
generation = snapshot().generation_id
|
||||
except Exception:
|
||||
generation = None
|
||||
if store is None:
|
||||
try:
|
||||
from one_memory.retrieval.index_store import RetrievalIndexStore
|
||||
|
||||
store = RetrievalIndexStore(database_path)
|
||||
except Exception:
|
||||
store = None
|
||||
return store, generation
|
||||
|
||||
|
||||
def read_memory_lane(
|
||||
database_path: Path,
|
||||
query: CoordinatorQuery,
|
||||
@@ -857,7 +739,6 @@ def coordinate_context_v2(
|
||||
query: CoordinatorQuery,
|
||||
*,
|
||||
database_path: Path,
|
||||
index: RawIndexLookup | None = None,
|
||||
deadline: float,
|
||||
expand_refs: ExpandRefs | None = None,
|
||||
) -> RetrievalResponse:
|
||||
@@ -871,62 +752,29 @@ def coordinate_context_v2(
|
||||
lexical = read_lexical_raw_lane(
|
||||
database_path, query, source_ids=authority.raw_sources, deadline=deadline
|
||||
)
|
||||
semantic = read_semantic_raw_lane(
|
||||
index, database_path, query, source_ids=authority.raw_sources, deadline=deadline
|
||||
)
|
||||
timings = StageTimings(raw_ms=(time.monotonic() - raw_started) * 1000)
|
||||
combined: dict[tuple[str, str], RawEvidence] = {}
|
||||
for item in (*lexical.items, *semantic.items):
|
||||
if not isinstance(item, RawEvidence):
|
||||
continue
|
||||
identity = _identity(item)
|
||||
previous = combined.get(identity)
|
||||
if previous is None or _semantic_only_raw(previous):
|
||||
combined[identity] = item
|
||||
# Semantic similarity is a candidate signal, not a verified recall
|
||||
# result. Keep semantic-only candidates visible in coverage, but do not
|
||||
# release their bodies for host injection unless the item is
|
||||
# independently confirmed by the lexical/object lane. A global distance
|
||||
# threshold is deliberately not used: the measured positive and
|
||||
# negative distributions overlap.
|
||||
semantic_only = {_identity(item) for item in combined.values() if _semantic_only_raw(item)}
|
||||
confirmed = {_identity(item) for item in lexical.items if isinstance(item, RawEvidence)}
|
||||
deferred_raw = semantic_only - confirmed
|
||||
if deferred_raw:
|
||||
combined = {
|
||||
identity: item
|
||||
for identity, item in combined.items()
|
||||
if identity not in deferred_raw
|
||||
}
|
||||
channels = tuple(dict.fromkeys((*lexical.channels, *semantic.channels)))
|
||||
for item in lexical.items:
|
||||
if isinstance(item, RawEvidence):
|
||||
combined.setdefault(_identity(item), item)
|
||||
if lexical.status == "unavailable":
|
||||
status: Literal["ok", "empty", "unavailable", "not_searched", "degraded"] = (
|
||||
"unavailable"
|
||||
)
|
||||
reason = lexical.reason
|
||||
elif semantic.status == "unavailable":
|
||||
status = "unavailable"
|
||||
reason = semantic.reason
|
||||
elif combined:
|
||||
status, reason = "ok", semantic.reason if semantic.status != "ok" else None
|
||||
elif deferred_raw:
|
||||
status, reason = "degraded", "relevance_confirmation"
|
||||
elif lexical.status == "empty":
|
||||
status, reason = ("empty" if semantic.status != "not_searched" else "ok"), None
|
||||
elif combined or lexical.status == "empty":
|
||||
status, reason = "ok", None
|
||||
else:
|
||||
status, reason = "empty", None
|
||||
truncated = lexical.truncated or semantic.truncated
|
||||
raw = LaneResult(
|
||||
tuple(combined.values()),
|
||||
status,
|
||||
reason,
|
||||
channels,
|
||||
truncated,
|
||||
lexical.channels,
|
||||
lexical.truncated,
|
||||
lexical.starved,
|
||||
lexical.coverage_limited,
|
||||
)
|
||||
if deferred_raw:
|
||||
object.__setattr__(raw, "coverage_limited", True)
|
||||
elif "raw" in query.kinds and not authority.raw_allowed:
|
||||
raw = LaneResult((), "unavailable", "permission_required", (), False)
|
||||
if "memory" in query.kinds:
|
||||
@@ -995,9 +843,7 @@ def coordinate_context_v2(
|
||||
return _path("unavailable", raw.reason, raw.channels, "not_searched", False, 0)
|
||||
if dropped and not final_raw:
|
||||
return _path("unavailable", "permission_required", raw.channels, "unavailable", True, 0)
|
||||
# Lexical-only is complete when the lexical lane actually answered. An
|
||||
# unconfigured semantic lane is a capability absence, not proof that
|
||||
# less than the authorized range was searched.
|
||||
# The raw lane is complete when the lexical lane actually answered.
|
||||
complete = (
|
||||
raw.status in {"ok", "empty"}
|
||||
and not dropped
|
||||
@@ -1022,9 +868,7 @@ def coordinate_context_v2(
|
||||
"count_budget"
|
||||
if truncated
|
||||
else "coverage_budget"
|
||||
if raw.coverage_limited and not deferred_raw
|
||||
else "relevance_confirmation"
|
||||
if deferred_raw
|
||||
if raw.coverage_limited
|
||||
else "character_budget"
|
||||
)
|
||||
return _path(
|
||||
@@ -1112,7 +956,7 @@ def coordinate_context_v2(
|
||||
"group_id": f"group-{index}",
|
||||
"answer_target": query.intent,
|
||||
"decision_status": "missing_confirmation",
|
||||
"missing_requirements": _group_missing_requirements(item),
|
||||
"missing_requirements": [],
|
||||
"object_refs": _object_refs(item, plan.object_keys),
|
||||
"items": [group_item.model_dump(mode="json") for group_item in (item,)],
|
||||
}
|
||||
@@ -1185,7 +1029,6 @@ async def search_context_v2(
|
||||
query: CoordinatorQuery,
|
||||
*,
|
||||
database_path: Path,
|
||||
index: RawIndexLookup | None = None,
|
||||
deadline: float,
|
||||
expand_refs: ExpandRefs | None = None,
|
||||
) -> RetrievalResponse:
|
||||
@@ -1198,7 +1041,6 @@ async def search_context_v2(
|
||||
authority,
|
||||
query,
|
||||
database_path=database_path,
|
||||
index=index,
|
||||
deadline=deadline,
|
||||
expand_refs=expand_refs,
|
||||
),
|
||||
|
||||
@@ -1,13 +1,20 @@
|
||||
"""Offline snapshot and restore of the HY SDK data plane for the App migration.
|
||||
"""Offline maintenance of the HY SDK data plane for App migrations.
|
||||
|
||||
Runtime product code must use public ``HyMemoryClient`` methods. The startup
|
||||
migration bundle is the isolated exception: all writers are stopped, so it may
|
||||
copy and restore the SDK-owned local stores as files.
|
||||
migration boundary is the isolated exception: all writers are stopped, so it
|
||||
may copy, restore and repair SDK-owned local stores.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
from collections.abc import Iterator
|
||||
from contextlib import closing
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from hy_memory import ( # type: ignore[import-untyped]
|
||||
DataPlaneRestore,
|
||||
@@ -16,6 +23,28 @@ from hy_memory import ( # type: ignore[import-untyped]
|
||||
snapshot_local_data_plane,
|
||||
)
|
||||
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.hy_runtime_config import HY_MEMORY_NAMESPACE
|
||||
|
||||
# v2 reads the backfill's governed_memory_evidence links; v1 missed them.
|
||||
_METADATA_REPAIR_VERSION = 2
|
||||
_BATCH = 500
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RepairMetadata:
|
||||
source: str | None
|
||||
memory_at: int | None
|
||||
created_at: int | None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MetadataRepairResult:
|
||||
migration_id: str | None
|
||||
mapped: int
|
||||
vector_rows: int
|
||||
coding_rows: int
|
||||
|
||||
|
||||
def snapshot_data_plane(
|
||||
*,
|
||||
@@ -45,3 +74,400 @@ def restore_data_plane(
|
||||
writers_stopped=writers_stopped,
|
||||
protective_root=protective_root,
|
||||
)
|
||||
|
||||
|
||||
def repair_sdk_metadata_from_bundle(
|
||||
*,
|
||||
data_dir: Path,
|
||||
database_path: Path,
|
||||
now: datetime | None = None,
|
||||
) -> MetadataRepairResult:
|
||||
"""Restore metadata lost by historical cutovers from their verified bundles."""
|
||||
migration_id = _sdk_migration_id(database_path)
|
||||
if migration_id is None:
|
||||
return MetadataRepairResult(None, 0, 0, 0)
|
||||
receipt = data_dir / f".sdk-metadata-repair-v{_METADATA_REPAIR_VERSION}-{migration_id}.json"
|
||||
if receipt.is_file():
|
||||
document = json.loads(receipt.read_text(encoding="utf-8"))
|
||||
return MetadataRepairResult(
|
||||
migration_id,
|
||||
int(document.get("mapped", 0)),
|
||||
int(document.get("vector_rows", 0)),
|
||||
int(document.get("coding_rows", 0)),
|
||||
)
|
||||
# A home can carry several cutovers (v1 and then v2). A later bundle may
|
||||
# have been taken after an earlier cutover had already retired the legacy
|
||||
# tables, so every bundle is read and their links are pooled.
|
||||
legacy = sorted(
|
||||
(data_dir / "backups").glob("*-pre-cutover.sdk-bundle/management/one-memory.sqlite3")
|
||||
)
|
||||
if not legacy:
|
||||
# No bundle yet is not a finished repair: try again on the next start.
|
||||
return MetadataRepairResult(migration_id, 0, 0, 0)
|
||||
current = now or datetime.now(UTC)
|
||||
pools: dict[str, _Pool] = {}
|
||||
for path in legacy:
|
||||
_merge_pools(pools, _legacy_pools(path, now=current))
|
||||
mapping = _decide(pools)
|
||||
vector_rows = _repair_vector(data_dir, mapping, now=current)
|
||||
coding_rows = _repair_coding(data_dir, mapping)
|
||||
result = MetadataRepairResult(
|
||||
migration_id,
|
||||
len(mapping),
|
||||
vector_rows,
|
||||
coding_rows,
|
||||
)
|
||||
temporary = receipt.with_suffix(".tmp")
|
||||
temporary.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"schema_version": 1,
|
||||
"migration_id": migration_id,
|
||||
"mapped": result.mapped,
|
||||
"vector_rows": vector_rows,
|
||||
"coding_rows": coding_rows,
|
||||
},
|
||||
sort_keys=True,
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
temporary.replace(receipt)
|
||||
return result
|
||||
|
||||
|
||||
def _sdk_migration_id(database_path: Path) -> str | None:
|
||||
with main_database(database_path).read(row_factory=sqlite3.Row) as database:
|
||||
row = database.execute(
|
||||
"SELECT authority,migration_id FROM memory_authority_state WHERE singleton=1"
|
||||
).fetchone()
|
||||
if row is None or row["authority"] != "sdk":
|
||||
return None
|
||||
value = row["migration_id"]
|
||||
return str(value) if isinstance(value, str) and value else None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Pool:
|
||||
"""Everything the legacy library knew about one SDK memory."""
|
||||
|
||||
sources: set[str] = field(default_factory=set)
|
||||
times: list[datetime] = field(default_factory=list)
|
||||
created: list[datetime] = field(default_factory=list)
|
||||
|
||||
|
||||
def _legacy_metadata(path: Path, *, now: datetime) -> dict[str, RepairMetadata]:
|
||||
return _decide(_legacy_pools(path, now=now))
|
||||
|
||||
|
||||
def _legacy_pools(path: Path, *, now: datetime) -> dict[str, _Pool]:
|
||||
"""Link each SDK memory id to the legacy evidence it was made from.
|
||||
|
||||
Three legacy paths wrote memories the SDK now holds:
|
||||
|
||||
- the cutover backfill (`import_exact`) copied every governed memory into the
|
||||
SDK under its legacy id, or under the id `memory_task_outputs` recorded,
|
||||
without a source or a conversation time; `governed_memory_evidence` is the
|
||||
link from that memory to its evidence;
|
||||
- extraction tasks recorded their evidence as `support_refs`;
|
||||
- the write queue recorded the SDK ids each job produced.
|
||||
"""
|
||||
with closing(sqlite3.connect(f"file:{path}?mode=ro", uri=True)) as database:
|
||||
database.row_factory = sqlite3.Row
|
||||
tables = {
|
||||
str(row[0])
|
||||
for row in database.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
||||
}
|
||||
if "evidence" not in tables:
|
||||
return {}
|
||||
links: dict[str, set[str]] = {}
|
||||
created: dict[str, list[datetime]] = {}
|
||||
canonical: dict[str, str] = {}
|
||||
if "memory_task_outputs" in tables:
|
||||
for row in database.execute("SELECT body FROM memory_task_outputs"):
|
||||
try:
|
||||
body = json.loads(str(row["body"]))
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
legacy_id = body.get("memory_id")
|
||||
canonical_id = body.get("hy_memory_id")
|
||||
refs = body.get("support_refs")
|
||||
if not isinstance(legacy_id, str) or not isinstance(canonical_id, str):
|
||||
continue
|
||||
canonical[legacy_id] = canonical_id
|
||||
links.setdefault(canonical_id, set()).update(
|
||||
str(ref["evidence_id"])
|
||||
for ref in (refs if isinstance(refs, list) else ())
|
||||
if isinstance(ref, dict) and isinstance(ref.get("evidence_id"), str)
|
||||
)
|
||||
if "governed_memories" in tables:
|
||||
for row in database.execute(
|
||||
"SELECT memory_id,created_at,updated_at FROM governed_memories"
|
||||
):
|
||||
legacy_id = str(row["memory_id"])
|
||||
canonical_id = canonical.get(legacy_id, legacy_id)
|
||||
links.setdefault(canonical_id, set())
|
||||
legacy_created = _plausible_time(
|
||||
row["created_at"], fallback=row["updated_at"], now=now
|
||||
)
|
||||
if legacy_created is not None:
|
||||
created.setdefault(canonical_id, []).append(legacy_created)
|
||||
if "governed_memory_evidence" in tables:
|
||||
for row in database.execute(
|
||||
"SELECT memory_id,evidence_id FROM governed_memory_evidence"
|
||||
):
|
||||
legacy_id = str(row["memory_id"])
|
||||
links.setdefault(canonical.get(legacy_id, legacy_id), set()).add(
|
||||
str(row["evidence_id"])
|
||||
)
|
||||
if "memory_write_jobs" in tables and "sdk_memory_ids_json" in _columns(
|
||||
database, "memory_write_jobs"
|
||||
):
|
||||
job_evidence: dict[str, set[str]] = {}
|
||||
if "memory_write_job_evidence" in tables:
|
||||
for row in database.execute(
|
||||
"SELECT job_id,evidence_id FROM memory_write_job_evidence"
|
||||
):
|
||||
job_evidence.setdefault(str(row["job_id"]), set()).add(str(row["evidence_id"]))
|
||||
for row in database.execute(
|
||||
"SELECT job_id,sdk_memory_ids_json FROM memory_write_jobs "
|
||||
"WHERE sdk_memory_ids_json IS NOT NULL"
|
||||
):
|
||||
try:
|
||||
memory_ids = json.loads(str(row["sdk_memory_ids_json"]))
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
for memory_id in memory_ids if isinstance(memory_ids, list) else ():
|
||||
if isinstance(memory_id, str) and memory_id:
|
||||
links.setdefault(memory_id, set()).update(
|
||||
job_evidence.get(str(row["job_id"]), ())
|
||||
)
|
||||
evidence: dict[str, tuple[str, datetime | None]] = {}
|
||||
wanted = sorted({identifier for refs in links.values() for identifier in refs})
|
||||
for batch in _batches(wanted):
|
||||
placeholders = ",".join("?" for _ in batch)
|
||||
for row in database.execute(
|
||||
"SELECT evidence_id,source_id,created_at,imported_at "
|
||||
f"FROM evidence WHERE evidence_id IN ({placeholders})",
|
||||
batch,
|
||||
):
|
||||
evidence[str(row["evidence_id"])] = (
|
||||
str(row["source_id"]),
|
||||
_plausible_time(row["created_at"], fallback=row["imported_at"], now=now),
|
||||
)
|
||||
pools: dict[str, _Pool] = {}
|
||||
for canonical_id, refs in links.items():
|
||||
pool = pools.setdefault(canonical_id, _Pool())
|
||||
for identifier in refs:
|
||||
if identifier not in evidence:
|
||||
continue
|
||||
source, timestamp = evidence[identifier]
|
||||
if source and source != "default":
|
||||
pool.sources.add(source)
|
||||
if timestamp is not None:
|
||||
pool.times.append(timestamp)
|
||||
pool.created.extend(created.get(canonical_id, ()))
|
||||
return pools
|
||||
|
||||
|
||||
def _merge_pools(into: dict[str, _Pool], other: dict[str, _Pool]) -> None:
|
||||
for canonical_id, pool in other.items():
|
||||
merged = into.setdefault(canonical_id, _Pool())
|
||||
merged.sources |= pool.sources
|
||||
merged.times.extend(pool.times)
|
||||
merged.created.extend(pool.created)
|
||||
|
||||
|
||||
def _decide(pools: dict[str, _Pool]) -> dict[str, RepairMetadata]:
|
||||
# Several legacy memories can fold into one SDK memory. Their evidence is
|
||||
# pooled before deciding, so a disagreeing source stays unknown instead of
|
||||
# whichever link happened to come last.
|
||||
return {
|
||||
canonical_id: RepairMetadata(
|
||||
source=next(iter(pool.sources)) if len(pool.sources) == 1 else None,
|
||||
memory_at=int(max(pool.times).timestamp()) if pool.times else None,
|
||||
created_at=int(min(pool.created).timestamp()) if pool.created else None,
|
||||
)
|
||||
for canonical_id, pool in pools.items()
|
||||
}
|
||||
|
||||
|
||||
def _columns(database: sqlite3.Connection, table: str) -> set[str]:
|
||||
return {str(row[1]) for row in database.execute(f"PRAGMA table_info({table})")}
|
||||
|
||||
|
||||
def _plausible_time(
|
||||
value: object,
|
||||
*,
|
||||
fallback: object,
|
||||
now: datetime,
|
||||
) -> datetime | None:
|
||||
for candidate in (value, fallback):
|
||||
if not isinstance(candidate, str) or not candidate:
|
||||
continue
|
||||
try:
|
||||
parsed = datetime.fromisoformat(candidate)
|
||||
except ValueError:
|
||||
continue
|
||||
if parsed.tzinfo is None:
|
||||
parsed = parsed.replace(tzinfo=UTC)
|
||||
parsed = parsed.astimezone(UTC)
|
||||
if parsed <= now + timedelta(days=1):
|
||||
return parsed
|
||||
return None
|
||||
|
||||
|
||||
def _repair_vector(
|
||||
data_dir: Path,
|
||||
mapping: dict[str, RepairMetadata],
|
||||
*,
|
||||
now: datetime,
|
||||
) -> int:
|
||||
directory = data_dir / "hy-memory" / HY_MEMORY_NAMESPACE / "vector"
|
||||
if not (directory / "chroma.sqlite3").is_file():
|
||||
return 0
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
|
||||
client = chromadb.PersistentClient(
|
||||
path=str(directory),
|
||||
settings=Settings(anonymized_telemetry=False),
|
||||
)
|
||||
# The SDK client is not open yet (startup has not created its handle), so
|
||||
# this is the only writer; release the store before the SDK takes it.
|
||||
try:
|
||||
repaired = 0
|
||||
ids = tuple(mapping)
|
||||
for collection in client.list_collections():
|
||||
for batch in _batches(ids):
|
||||
page = collection.get(ids=list(batch), include=["metadatas"])
|
||||
updated_ids: list[str] = []
|
||||
updated_metadata: list[dict[str, object]] = []
|
||||
for memory_id, metadata in zip(
|
||||
page.get("ids") or (),
|
||||
page.get("metadatas") or (),
|
||||
strict=True,
|
||||
):
|
||||
if not isinstance(metadata, dict):
|
||||
continue
|
||||
desired = mapping.get(str(memory_id))
|
||||
if desired is None:
|
||||
continue
|
||||
next_metadata = dict(metadata)
|
||||
changed = False
|
||||
if desired.source and next_metadata.get("source") in {
|
||||
None,
|
||||
"",
|
||||
"default",
|
||||
}:
|
||||
next_metadata["source"] = desired.source
|
||||
changed = True
|
||||
if (
|
||||
desired.memory_at is not None
|
||||
and next_metadata.get("memory_at") != desired.memory_at
|
||||
):
|
||||
next_metadata["memory_at"] = desired.memory_at
|
||||
changed = True
|
||||
if (
|
||||
desired.created_at is not None
|
||||
and next_metadata.get("gmt_created") != desired.created_at
|
||||
):
|
||||
next_metadata["gmt_created"] = desired.created_at
|
||||
changed = True
|
||||
if changed:
|
||||
updated_ids.append(str(memory_id))
|
||||
updated_metadata.append(next_metadata)
|
||||
if updated_ids:
|
||||
collection.update(ids=updated_ids, metadatas=updated_metadata)
|
||||
repaired += len(updated_ids)
|
||||
repaired += _settle_future_times(collection, now=now)
|
||||
finally:
|
||||
client.close()
|
||||
return repaired
|
||||
|
||||
|
||||
def _settle_future_times(collection: Any, *, now: datetime) -> int:
|
||||
"""Move conversation times that are still in the future back to the write time.
|
||||
|
||||
A historical Cursor bug stored a file hash as the message time (years like
|
||||
2254). The SDK lists newest first by `memory_at`, so those rows hold the
|
||||
first page. Rows the bundles could date were fixed above; for the rest the
|
||||
write time is the closest true moment that is known.
|
||||
"""
|
||||
cutoff = int((now + timedelta(days=1)).timestamp())
|
||||
page = collection.get(where={"memory_at": {"$gt": cutoff}}, include=["metadatas"])
|
||||
updated_ids: list[str] = []
|
||||
updated_metadata: list[dict[str, object]] = []
|
||||
for memory_id, metadata in zip(page.get("ids") or (), page.get("metadatas") or (), strict=True):
|
||||
if not isinstance(metadata, dict):
|
||||
continue
|
||||
created = metadata.get("gmt_created")
|
||||
settled = (
|
||||
created
|
||||
if isinstance(created, int) and not isinstance(created, bool) and created <= cutoff
|
||||
else int(now.timestamp())
|
||||
)
|
||||
updated_ids.append(str(memory_id))
|
||||
updated_metadata.append({**metadata, "memory_at": settled})
|
||||
if updated_ids:
|
||||
collection.update(ids=updated_ids, metadatas=updated_metadata)
|
||||
return len(updated_ids)
|
||||
|
||||
|
||||
def _repair_coding(
|
||||
data_dir: Path,
|
||||
mapping: dict[str, RepairMetadata],
|
||||
) -> int:
|
||||
root = data_dir / "hy-memory" / HY_MEMORY_NAMESPACE
|
||||
repaired = 0
|
||||
for name in ("coding.sqlite3", "coding_cv.sqlite3"):
|
||||
path = root / name
|
||||
if not path.is_file():
|
||||
continue
|
||||
with closing(sqlite3.connect(path)) as database:
|
||||
tables = {
|
||||
str(row[0])
|
||||
for row in database.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
||||
}
|
||||
for table, identifier in (
|
||||
("coding_memory_meta", "experience_id"),
|
||||
("working_fact", "fact_id"),
|
||||
):
|
||||
if table not in tables:
|
||||
continue
|
||||
rows = database.execute(
|
||||
f"SELECT {identifier},source,created_at FROM {table}"
|
||||
).fetchall()
|
||||
for memory_id, source, created_at in rows:
|
||||
desired = mapping.get(str(memory_id))
|
||||
if desired is None:
|
||||
continue
|
||||
updates: list[str] = []
|
||||
values: list[object] = []
|
||||
if desired.source and source in {None, "", "default"}:
|
||||
updates.append("source=?")
|
||||
values.append(desired.source)
|
||||
if desired.created_at is not None:
|
||||
desired_created_at = (
|
||||
datetime.fromtimestamp(desired.created_at, tz=UTC)
|
||||
.replace(tzinfo=None)
|
||||
.isoformat()
|
||||
)
|
||||
if created_at != desired_created_at:
|
||||
updates.append("created_at=?")
|
||||
values.append(desired_created_at)
|
||||
if not updates:
|
||||
continue
|
||||
values.append(memory_id)
|
||||
database.execute(
|
||||
f"UPDATE {table} SET {','.join(updates)} WHERE {identifier}=?",
|
||||
values,
|
||||
)
|
||||
repaired += 1
|
||||
database.commit()
|
||||
return repaired
|
||||
|
||||
|
||||
def _batches[T](values: list[T] | tuple[T, ...]) -> Iterator[tuple[T, ...]]:
|
||||
for start in range(0, len(values), _BATCH):
|
||||
yield tuple(values[start : start + _BATCH])
|
||||
@@ -237,15 +237,9 @@ def build_sdk_config(
|
||||
|
||||
root = data_dir.resolve() / "hy-memory" / "governed-v2"
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
collection_identity = "\0".join(
|
||||
(
|
||||
embedding_provider.value,
|
||||
embedding_endpoint or "local",
|
||||
embedding_model,
|
||||
str(dimensions or "auto"),
|
||||
)
|
||||
collection = _collection_base(
|
||||
embedding_provider, embedding_endpoint, embedding_model, dimensions
|
||||
)
|
||||
collection = "memhub_" + hashlib.sha256(collection_identity.encode()).hexdigest()[:16]
|
||||
|
||||
return MemoryConfig.from_dict(
|
||||
{
|
||||
@@ -348,6 +342,49 @@ def _generation_access(
|
||||
return endpoint, api_key
|
||||
|
||||
|
||||
def _collection_base(
|
||||
provider: ProviderKind,
|
||||
endpoint: str | None,
|
||||
model: str,
|
||||
dimensions: int | None,
|
||||
) -> str:
|
||||
identity = "\0".join((provider.value, endpoint or "local", model, str(dimensions or "auto")))
|
||||
return "memhub_" + hashlib.sha256(identity.encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
def vector_collection_bases(profile: ModelProfileSnapshot) -> tuple[str, tuple[str, ...]]:
|
||||
"""The SDK vector collection base for this embedding space, and the older ones.
|
||||
|
||||
The SDK derives every vector collection (chat, Coding keys, Working facts,
|
||||
each variant) by suffixing this base. Builds before #145 named the same
|
||||
embedding space `one_memory_<12 hex>`, with or without a declared width;
|
||||
the offline migration renames those so their vectors stay readable.
|
||||
"""
|
||||
provider = profile.embedding_provider_kind or profile.provider_kind
|
||||
local = provider is ProviderKind.LOCAL_ONNX
|
||||
model = profile.embedding_model or (BUILTIN_EMBEDDING_MODEL if local else None)
|
||||
if not model:
|
||||
raise HySdkRuntimeError("hy_sdk_embedding_model_required")
|
||||
endpoint: str | None
|
||||
dimensions: int | None
|
||||
if local:
|
||||
endpoint, dimensions = None, profile.embedding_dimensions or BUILTIN_EMBEDDING_DIMENSIONS
|
||||
legacy_endpoint = "local-onnx"
|
||||
else:
|
||||
endpoint = _openai_endpoint(provider, profile.embedding_base_url or profile.base_url)
|
||||
dimensions = profile.embedding_dimensions
|
||||
legacy_endpoint = endpoint
|
||||
legacy_model = profile.embedding_model or "bge-small-zh-v1.5"
|
||||
legacy = tuple(
|
||||
"one_memory_"
|
||||
+ hashlib.sha256(
|
||||
":".join((provider.value, legacy_endpoint, legacy_model, *declared)).encode()
|
||||
).hexdigest()[:12]
|
||||
for declared in ((), ("auto",), (str(profile.embedding_dimensions or "auto"),))
|
||||
)
|
||||
return _collection_base(provider, endpoint, model, dimensions), tuple(dict.fromkeys(legacy))
|
||||
|
||||
|
||||
def _openai_endpoint(provider: ProviderKind, endpoint: str | None) -> str:
|
||||
value = (endpoint or "").rstrip("/")
|
||||
if not value:
|
||||
|
||||
@@ -92,9 +92,35 @@ class HySdkMemoryPage(BaseModel):
|
||||
total: int = Field(ge=0)
|
||||
limit: int = Field(ge=1)
|
||||
offset: int = Field(ge=0)
|
||||
# Set on a search whose relevance window was full: matches were looked for
|
||||
# only among this many most relevant memories, so `total` may undercount.
|
||||
search_window: int | None = None
|
||||
|
||||
@classmethod
|
||||
def from_sdk(cls, raw: object, *, limit: int, offset: int) -> HySdkMemoryPage:
|
||||
page, _skipped = cls._parse(raw, limit=limit, offset=offset, skip_invalid=False)
|
||||
return page
|
||||
|
||||
@classmethod
|
||||
def from_sdk_skipping_invalid(
|
||||
cls, raw: object, *, limit: int, offset: int
|
||||
) -> tuple[HySdkMemoryPage, int]:
|
||||
"""Parse a page, leaving out rows outside the public record contract.
|
||||
|
||||
Returns the page and how many rows were left out, so the caller can
|
||||
report them. The rows themselves are never touched.
|
||||
"""
|
||||
return cls._parse(raw, limit=limit, offset=offset, skip_invalid=True)
|
||||
|
||||
@classmethod
|
||||
def _parse(
|
||||
cls,
|
||||
raw: object,
|
||||
*,
|
||||
limit: int,
|
||||
offset: int,
|
||||
skip_invalid: bool,
|
||||
) -> tuple[HySdkMemoryPage, int]:
|
||||
if not isinstance(raw, Mapping):
|
||||
raise HySdkReadError("hy_memory_list_shape_unknown")
|
||||
rows: list[object] = []
|
||||
@@ -115,15 +141,27 @@ class HySdkMemoryPage(BaseModel):
|
||||
total += bucket_total
|
||||
rows.extend(bucket[rows_name])
|
||||
deduplicated: dict[str, HySdkMemoryRecord] = {}
|
||||
skipped = 0
|
||||
for row in rows:
|
||||
record = HySdkMemoryRecord.from_sdk(row)
|
||||
try:
|
||||
record = HySdkMemoryRecord.from_sdk(row)
|
||||
except HySdkReadError:
|
||||
if not skip_invalid:
|
||||
raise
|
||||
skipped += 1
|
||||
continue
|
||||
deduplicated.setdefault(record.memory_id, record)
|
||||
return cls(
|
||||
page = cls(
|
||||
memories=tuple(deduplicated.values()),
|
||||
# `total` is the SDK's stable result-set count. Subtracting only
|
||||
# invalid rows observed on this page would make it change as the
|
||||
# caller paged. Until the SDK exposes a quarantined count, keep the
|
||||
# authoritative total and skip malformed rows only in this page.
|
||||
total=total,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
return page, skipped
|
||||
|
||||
|
||||
class HySdkRawRecord(BaseModel):
|
||||
@@ -289,6 +327,38 @@ def parse_sdk_search(raw: object, *, error_code: str) -> tuple[SemanticMemoryHit
|
||||
return tuple(found)
|
||||
|
||||
|
||||
def library_keyword_hits(raw: object, query: str) -> tuple[tuple[str, ...], int]:
|
||||
"""Ids in an SDK search window whose text contains every query term.
|
||||
|
||||
The local SDK store has no keyword channel (Chroma's keyword search is
|
||||
empty), so search is pure vector similarity: a word in the library can
|
||||
score below the floor, and a word in none of it still scores above. The
|
||||
library search asks the SDK for its most relevant window and keeps only the
|
||||
rows that actually contain the words, most relevant first. Returns those ids
|
||||
and how many distinct rows the window held.
|
||||
"""
|
||||
rows = _reply_rows(raw, "hy_memory_search_shape_unknown")
|
||||
terms = query.casefold().split()
|
||||
seen: set[str] = set()
|
||||
matched: list[str] = []
|
||||
for item in sorted(rows, key=_score, reverse=True):
|
||||
memory_id = item.get("memory_id") or item.get("id")
|
||||
if not isinstance(memory_id, str) or not memory_id or memory_id in seen:
|
||||
continue
|
||||
seen.add(memory_id)
|
||||
content = item.get("content")
|
||||
tags = item.get("tags")
|
||||
text = " ".join(
|
||||
[
|
||||
content if isinstance(content, str) else "",
|
||||
*(tag for tag in (tags if isinstance(tags, list) else ()) if isinstance(tag, str)),
|
||||
]
|
||||
).casefold()
|
||||
if all(term in text for term in terms):
|
||||
matched.append(memory_id)
|
||||
return tuple(matched), len(seen)
|
||||
|
||||
|
||||
def parse_sdk_raw_search(raw: object) -> tuple[HySdkRawHit, ...]:
|
||||
return tuple(
|
||||
HySdkRawHit.from_sdk_search(row)
|
||||
|
||||
@@ -50,6 +50,10 @@ class AuthorizedSourceFile:
|
||||
size: int
|
||||
modified_ns: int
|
||||
changed_ns: int
|
||||
# Source-specific change token (for example Cursor's WAL identity). It is
|
||||
# cursor metadata only; modified_ns remains the real filesystem mtime used
|
||||
# as a fallback conversation timestamp.
|
||||
cursor_version: str | None = None
|
||||
root_identity: str | None = None
|
||||
requested_root: Path | None = None
|
||||
access_check: Callable[[], None] | None = field(default=None, repr=False, compare=False)
|
||||
|
||||
@@ -644,7 +644,7 @@ def is_unchanged(source_file: AuthorizedSourceFile, cursors: dict[str, SourceCur
|
||||
return (
|
||||
previous is not None
|
||||
and previous.read_to_end
|
||||
and previous.file_identity == source_file.file_identity
|
||||
and previous.file_identity == _cursor_file_identity(source_file)
|
||||
and previous.observed_size == source_file.size
|
||||
and previous.observed_modified_ns == source_file.modified_ns
|
||||
# Chmod alone leaves size, mtime and inode untouched, so a file that
|
||||
@@ -691,7 +691,7 @@ def touched_cursor(
|
||||
previous = options.cursors.get(source_file.safe_id)
|
||||
keep = (
|
||||
previous is not None
|
||||
and previous.file_identity == source_file.file_identity
|
||||
and _same_file(previous, source_file)
|
||||
and previous.observed_size <= source_file.size
|
||||
)
|
||||
# A kept cursor means the next scan resumes exactly where this one ended:
|
||||
@@ -871,7 +871,7 @@ def resume_checkpoint(
|
||||
previous = options.cursors.get(source_file.safe_id)
|
||||
if previous is None:
|
||||
return None
|
||||
if previous.file_identity != source_file.file_identity:
|
||||
if not _same_file(previous, source_file):
|
||||
return None
|
||||
if source_file.size < previous.observed_size:
|
||||
return None
|
||||
@@ -1046,7 +1046,7 @@ def source_cursor(
|
||||
) -> SourceCursor:
|
||||
return SourceCursor(
|
||||
file_key=source_file.safe_id,
|
||||
file_identity=source_file.file_identity,
|
||||
file_identity=_cursor_file_identity(source_file),
|
||||
checkpoint=checkpoint,
|
||||
checkpoint_anchor=checkpoint_anchor,
|
||||
observed_size=source_file.size,
|
||||
@@ -1054,6 +1054,27 @@ def source_cursor(
|
||||
)
|
||||
|
||||
|
||||
def _cursor_file_identity(source_file: AuthorizedSourceFile) -> str:
|
||||
return (
|
||||
f"{source_file.file_identity}{_VERSION_SEPARATOR}{source_file.cursor_version}"
|
||||
if source_file.cursor_version
|
||||
else source_file.file_identity
|
||||
)
|
||||
|
||||
|
||||
_VERSION_SEPARATOR = ":v:"
|
||||
|
||||
|
||||
def _same_file(previous: SourceCursor, source_file: AuthorizedSourceFile) -> bool:
|
||||
"""Whether a stored cursor still points into this file.
|
||||
|
||||
The stored identity may carry a change token (Cursor's WAL version). The
|
||||
token says the file changed, which `is_unchanged` handles; it does not mean
|
||||
the file was replaced, so resuming compares the bare `dev:ino` only.
|
||||
"""
|
||||
return previous.file_identity.partition(_VERSION_SEPARATOR)[0] == source_file.file_identity
|
||||
|
||||
|
||||
_POSITION = re.compile(r"(?<=:)\d+(?=:|$)")
|
||||
|
||||
|
||||
|
||||
@@ -533,8 +533,8 @@ def _with_write_ahead_log(database: AuthorizedSourceFile) -> AuthorizedSourceFil
|
||||
version = hashlib.blake2b(
|
||||
f"{database.modified_ns}:{journal.st_mtime_ns}:{journal.st_size}:{journal.st_ino}".encode(),
|
||||
digest_size=8,
|
||||
).digest()
|
||||
return replace(database, modified_ns=int.from_bytes(version, "big") & ((1 << 63) - 1))
|
||||
).hexdigest()
|
||||
return replace(database, cursor_version=version)
|
||||
|
||||
|
||||
def _live_baseline(listing: _CursorListing) -> dict[str, SourceCursor]:
|
||||
|
||||
@@ -10,16 +10,12 @@ weights alone.
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from importlib import import_module
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from one_memory.embedding_assets import (
|
||||
BUILTIN_EMBEDDING_DIMENSIONS,
|
||||
BUILTIN_EMBEDDING_MAX_TOKENS,
|
||||
EMBEDDING_ASSETS,
|
||||
AssetVerificationError,
|
||||
MissingEmbeddingAssetsError,
|
||||
assets_present,
|
||||
resolve_asset_dir,
|
||||
@@ -34,28 +30,6 @@ _FETCH_HINT = (
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TokenOffset:
|
||||
"""Half-open Python Unicode code point offsets; synthetic specials use [0,0)."""
|
||||
|
||||
start_cp: int
|
||||
end_cp: int
|
||||
special: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TokenizedText:
|
||||
"""Untruncated, unpadded encoding, including tokenizer postprocessor specials."""
|
||||
|
||||
token_ids: tuple[int, ...]
|
||||
offsets: tuple[TokenOffset, ...]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Defensive copies also keep runtime callers from retaining mutable lists.
|
||||
object.__setattr__(self, "token_ids", tuple(self.token_ids))
|
||||
object.__setattr__(self, "offsets", tuple(self.offsets))
|
||||
|
||||
|
||||
class LocalEmbedder:
|
||||
"""Encodes text with the built-in model, loading it on first use."""
|
||||
|
||||
@@ -68,43 +42,6 @@ class LocalEmbedder:
|
||||
self._session: Any | None = None
|
||||
self._tokenizer: Any | None = None
|
||||
self._input_names: frozenset[str] = frozenset()
|
||||
self._offset_lock = threading.Lock()
|
||||
self._offset_tokenizer: Any | None = None
|
||||
|
||||
def tokenize_with_offsets(self, text: str) -> TokenizedText:
|
||||
"""Tokenize original text without encode()'s truncation or batch padding.
|
||||
|
||||
The Python tokenizers binding returns code point offsets into the input,
|
||||
even when its model normalizer changes spelling. This separate lazy
|
||||
tokenizer uses the exact locked asset; it never loads an ONNX session or
|
||||
mutates the tokenizer shared by encode(). Offsets are not UTF-16 indices.
|
||||
"""
|
||||
if not isinstance(text, str):
|
||||
raise TypeError("embedding_offsets_require_text")
|
||||
with self._offset_lock:
|
||||
if self._offset_tokenizer is None:
|
||||
directory = resolve_asset_dir()
|
||||
path = directory / "tokenizer.json"
|
||||
if not path.is_file():
|
||||
raise MissingEmbeddingAssetsError(f"{_FETCH_HINT} Looked in {directory}.")
|
||||
data = path.read_bytes()
|
||||
expected = next(a.sha256 for a in EMBEDDING_ASSETS if a.filename == path.name)
|
||||
if sha256(data).hexdigest() != expected:
|
||||
raise AssetVerificationError("tokenizer.json: locked digest mismatch")
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
tokenizer = Tokenizer.from_str(data.decode("utf-8"))
|
||||
tokenizer.no_truncation()
|
||||
tokenizer.no_padding()
|
||||
self._offset_tokenizer = tokenizer
|
||||
encoding = self._offset_tokenizer.encode(text, add_special_tokens=True)
|
||||
offsets = tuple(
|
||||
TokenOffset(start, end, bool(special))
|
||||
for (start, end), special in zip(
|
||||
encoding.offsets, encoding.special_tokens_mask, strict=True
|
||||
)
|
||||
)
|
||||
return TokenizedText(tuple(encoding.ids), offsets)
|
||||
|
||||
def _load(self) -> tuple[Any, Any]:
|
||||
# Double-checked under a lock: memory processing encodes from worker
|
||||
|
||||
@@ -2217,34 +2217,6 @@ CREATE TABLE IF NOT EXISTS retrieval_chunks (
|
||||
"""
|
||||
|
||||
|
||||
def retrieval_documents_v2_ddl(dialect: str = "mysql") -> str:
|
||||
"""The v2 documents table under its temporary `_v2` name.
|
||||
|
||||
`index_store.migrate_v2` creates this, copies rows from the v1 table, drops
|
||||
the v1 table and renames. The rename is the only reason a `_v2` name exists
|
||||
at all, and it is derived from the single definition so the copy-across
|
||||
cannot drift from what fresh databases get.
|
||||
"""
|
||||
documents = next(
|
||||
statement
|
||||
for statement in _split_statements(retrieval_index_ddl(dialect))
|
||||
if "CREATE TABLE IF NOT EXISTS retrieval_documents (" in statement
|
||||
)
|
||||
return documents.replace(
|
||||
"CREATE TABLE IF NOT EXISTS retrieval_documents (",
|
||||
"CREATE TABLE retrieval_documents_v2 (",
|
||||
)
|
||||
|
||||
|
||||
def retrieval_build_maps_ddl(dialect: str = "mysql") -> str:
|
||||
"""The build-maps table added by v2."""
|
||||
return next(
|
||||
statement
|
||||
for statement in _split_statements(retrieval_index_ddl(dialect))
|
||||
if "retrieval_build_maps" in statement
|
||||
)
|
||||
|
||||
|
||||
async def _033_retrieval_index(cursor: Any) -> None:
|
||||
for statement in _split_statements(retrieval_index_ddl(_dialect(cursor))):
|
||||
await cursor.execute(statement)
|
||||
|
||||
@@ -22,6 +22,13 @@ class RetiredEvidenceStore:
|
||||
del limit
|
||||
return 0
|
||||
|
||||
def release_sdk_staging(self, evidence_ids: tuple[str, ...]) -> int:
|
||||
# The write drain releases staged Evidence after every SDK write. On
|
||||
# SDK authority nothing is staged; without this, `__getattr__` raised
|
||||
# past `getattr(..., None)` and every completed write logged an error.
|
||||
del evidence_ids
|
||||
return 0
|
||||
|
||||
def get(self, _identifier: str) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@@ -1,167 +0,0 @@
|
||||
"""Single-Evidence identity chunking; no reads, persistence, normalization or grants.
|
||||
|
||||
Callers must first establish eligibility/sanitization/permission. This module
|
||||
does not strip wrappers or secrets and must not be used as a content safety gate.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol
|
||||
|
||||
from one_memory.embedding_assets import BUILTIN_EMBEDDING_MAX_TOKENS
|
||||
from one_memory.local_embedding import TokenizedText, TokenOffset
|
||||
|
||||
|
||||
class OffsetTokenizer(Protocol):
|
||||
def tokenize_with_offsets(self, text: str) -> TokenizedText: ...
|
||||
|
||||
|
||||
class ChunkingError(ValueError):
|
||||
"""Invalid offsets, uncovered text or insufficient encoding budget; never truncate."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TextChunk:
|
||||
ordinal: int
|
||||
start_cp: int
|
||||
end_cp: int
|
||||
token_count: int
|
||||
text: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Span:
|
||||
start: int
|
||||
end: int
|
||||
tokens: int
|
||||
|
||||
|
||||
def _spans(text: str, encoded: TokenizedText) -> tuple[_Span, ...]:
|
||||
if not isinstance(encoded, TokenizedText) or len(encoded.token_ids) != len(encoded.offsets):
|
||||
raise ChunkingError("chunk_token_offset_count_mismatch")
|
||||
spans: list[_Span] = []
|
||||
previous_start = previous_end = 0
|
||||
for token_id, offset in zip(encoded.token_ids, encoded.offsets, strict=True):
|
||||
if type(token_id) is not int or token_id < 0 or not isinstance(offset, TokenOffset):
|
||||
raise ChunkingError("chunk_invalid_token_offset")
|
||||
if (
|
||||
type(offset.start_cp) is not int
|
||||
or type(offset.end_cp) is not int
|
||||
or type(offset.special) is not bool
|
||||
):
|
||||
raise ChunkingError("chunk_invalid_token_offset")
|
||||
start, end = offset.start_cp, offset.end_cp
|
||||
if offset.special:
|
||||
if start != 0 or end != 0:
|
||||
raise ChunkingError("chunk_invalid_special_offset")
|
||||
continue
|
||||
if not 0 <= start < end <= len(text) or start < previous_start or end < previous_end:
|
||||
raise ChunkingError("chunk_invalid_offset_order_or_bounds")
|
||||
if start > previous_end and not text[previous_end:start].isspace():
|
||||
raise ChunkingError("chunk_uncovered_nonwhitespace")
|
||||
# Overlapping subpieces (including several pieces for one code point)
|
||||
# form an indivisible interval; slicing inside it loses token identity.
|
||||
if spans and start < spans[-1].end:
|
||||
prior = spans[-1]
|
||||
spans[-1] = _Span(prior.start, end, prior.tokens + 1)
|
||||
else:
|
||||
spans.append(_Span(start, end, 1))
|
||||
previous_start, previous_end = start, end
|
||||
if text[previous_end:] and not text[previous_end:].isspace():
|
||||
raise ChunkingError("chunk_uncovered_nonwhitespace")
|
||||
return tuple(spans)
|
||||
|
||||
|
||||
def _encoding(text: str, tokenizer: OffsetTokenizer) -> tuple[TokenizedText, tuple[_Span, ...]]:
|
||||
try:
|
||||
result = tokenizer.tokenize_with_offsets(text)
|
||||
except Exception as error:
|
||||
raise ChunkingError("chunk_tokenization_failed") from error
|
||||
return result, _spans(text, result)
|
||||
|
||||
|
||||
def chunk_evidence(
|
||||
text: str,
|
||||
tokenizer: OffsetTokenizer,
|
||||
*,
|
||||
prefix: str = "",
|
||||
target_tokens: int = 384,
|
||||
overlap_tokens: int = 64,
|
||||
max_tokens: int = BUILTIN_EMBEDDING_MAX_TOKENS,
|
||||
) -> tuple[TextChunk, ...]:
|
||||
"""Chunk one original with code point [start,end), preserving exact slices.
|
||||
|
||||
target/overlap count original non-special tokens. Overlap is up to the
|
||||
requested token count at indivisible offset boundaries; forward progress is
|
||||
mandatory. token_count is re-tokenized prefix + chunk.text INCLUDING specials,
|
||||
never an additive estimate. Prefix is caller-provided, unchanged, and omitted
|
||||
from chunk.text/offsets. Empty/whitespace-only inputs have no index chunks.
|
||||
"""
|
||||
if not isinstance(text, str) or not isinstance(prefix, str):
|
||||
raise ChunkingError("chunk_requires_text")
|
||||
if (
|
||||
any(type(value) is not int for value in (target_tokens, overlap_tokens, max_tokens))
|
||||
or not 0 <= overlap_tokens < target_tokens
|
||||
or not 1 <= max_tokens <= BUILTIN_EMBEDDING_MAX_TOKENS
|
||||
or target_tokens > max_tokens
|
||||
):
|
||||
raise ChunkingError("chunk_invalid_token_budget")
|
||||
if not text or text.isspace():
|
||||
return ()
|
||||
_, spans = _encoding(text, tokenizer)
|
||||
if not spans:
|
||||
raise ChunkingError("chunk_uncovered_nonwhitespace")
|
||||
chunks: list[TextChunk] = []
|
||||
start_index = 0
|
||||
previous_end_cp = 0
|
||||
while start_index < len(spans):
|
||||
end_index = start_index
|
||||
body_tokens = 0
|
||||
while end_index < len(spans):
|
||||
candidate = body_tokens + spans[end_index].tokens
|
||||
if end_index > start_index and candidate > target_tokens:
|
||||
break
|
||||
body_tokens = candidate
|
||||
end_index += 1
|
||||
# Prefer the last paragraph boundary in the latter half of the target.
|
||||
# Never trade away an entire overlap window or split a token interval.
|
||||
if end_index < len(spans):
|
||||
count = 0
|
||||
preferred = None
|
||||
for boundary in range(start_index + 1, end_index + 1):
|
||||
count += spans[boundary - 1].tokens
|
||||
if count >= max(target_tokens // 2, overlap_tokens + 1):
|
||||
left, right = spans[boundary - 1].end, spans[boundary].start
|
||||
if "\n" in text[max(spans[boundary - 1].start, left - 1) : right]:
|
||||
preferred = boundary
|
||||
if preferred is not None:
|
||||
end_index = preferred
|
||||
start_cp = 0 if start_index == 0 else spans[start_index].start
|
||||
while end_index > start_index:
|
||||
# Include gaps verbatim, including whitespace and combining marks;
|
||||
# source coverage is never reconstructed by joining decoded tokens.
|
||||
end_cp = spans[end_index].start if end_index < len(spans) else len(text)
|
||||
encoded, _ = _encoding(prefix + text[start_cp:end_cp], tokenizer)
|
||||
token_count = len(encoded.token_ids)
|
||||
if token_count <= max_tokens:
|
||||
break
|
||||
end_index -= 1
|
||||
else:
|
||||
raise ChunkingError("chunk_indivisible_span_exceeds_token_budget")
|
||||
if end_cp <= previous_end_cp or start_cp > previous_end_cp:
|
||||
raise ChunkingError("chunk_coverage_cannot_progress")
|
||||
chunks.append(TextChunk(len(chunks), start_cp, end_cp, token_count, text[start_cp:end_cp]))
|
||||
previous_end_cp = end_cp
|
||||
if end_index == len(spans):
|
||||
break
|
||||
next_start = end_index
|
||||
overlap = 0
|
||||
while next_start > start_index + 1:
|
||||
candidate = overlap + spans[next_start - 1].tokens
|
||||
if candidate > overlap_tokens:
|
||||
break
|
||||
next_start -= 1
|
||||
overlap = candidate
|
||||
start_index = next_start
|
||||
if previous_end_cp != len(text):
|
||||
raise ChunkingError("chunk_incomplete_coverage")
|
||||
return tuple(chunks)
|
||||
@@ -1,92 +0,0 @@
|
||||
"""Shared qualification for identity-mapped raw indexing and candidate reads."""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from one_memory.conversation_processing import _is_hy_input_turn
|
||||
from one_memory.imports.models import EvidenceContentKind, EvidenceRecord
|
||||
from one_memory.imports.sources.parsing import is_host_generated_user_turn, strip_host_wrappers
|
||||
from one_memory.input_sanitization import sanitize_model_input
|
||||
from one_memory.retrieval.chunking import TextChunk
|
||||
from one_memory.retrieval.mapped_content import PreparedMappedContent, prepare_mapped_content
|
||||
|
||||
|
||||
def identity_content_ready(record: EvidenceRecord) -> bool:
|
||||
# Reuse pure existing hygiene checks, without invoking HY. Nonidentity
|
||||
# cleaning requires a source map and remains explicitly uncovered for now.
|
||||
return (
|
||||
record.content_kind is EvidenceContentKind.RAW_CONVERSATION
|
||||
and record.role in {"user", "assistant"}
|
||||
and _is_hy_input_turn(record)
|
||||
and not is_host_generated_user_turn(record.verbatim)
|
||||
and strip_host_wrappers(record.verbatim) == record.verbatim.strip()
|
||||
and sanitize_model_input(record.verbatim) == record.verbatim
|
||||
)
|
||||
|
||||
|
||||
def index_content(
|
||||
record: EvidenceRecord, *, normalizer_version: int, chunker_version: int
|
||||
) -> PreparedMappedContent | None:
|
||||
"""Generation-owned content semantics, never silently upgrade identity generations."""
|
||||
if (normalizer_version, chunker_version) == (1, 1):
|
||||
if not identity_content_ready(record):
|
||||
raise ValueError("content_mapping_required")
|
||||
return None
|
||||
if (normalizer_version, chunker_version) == (2, 2):
|
||||
return prepare_mapped_content(record)
|
||||
raise ValueError("unsupported_content_generation")
|
||||
|
||||
|
||||
def validate_chunk_coverage(
|
||||
record: EvidenceRecord, chunks: Sequence[TextChunk], prepared: PreparedMappedContent | None
|
||||
) -> None:
|
||||
"""Verify exact slices and complete citation-safe coverage without a tokenizer call."""
|
||||
if len(chunks) > 8192 or (prepared is None and not chunks):
|
||||
raise ValueError("invalid_chunk_manifest")
|
||||
if prepared is None:
|
||||
ranges = [(0, len(record.verbatim))]
|
||||
else:
|
||||
ranges = []
|
||||
for span in prepared.metadata.segments:
|
||||
if span.kind == "copy":
|
||||
assert span.source_start_cp is not None and span.source_end_cp is not None
|
||||
ranges.append((span.source_start_cp, span.source_end_cp))
|
||||
range_index = 0
|
||||
covered_end = ranges[0][0] if ranges else 0
|
||||
previous_end = -1
|
||||
for index, chunk in enumerate(chunks):
|
||||
if not isinstance(chunk, TextChunk) or any(
|
||||
type(value) is not int
|
||||
for value in (chunk.ordinal, chunk.start_cp, chunk.end_cp, chunk.token_count)
|
||||
):
|
||||
raise ValueError("invalid_chunk_manifest")
|
||||
while range_index < len(ranges) and chunk.start_cp >= ranges[range_index][1]:
|
||||
if record.verbatim[covered_end : ranges[range_index][1]].strip():
|
||||
raise ValueError("incomplete_chunk_manifest")
|
||||
range_index += 1
|
||||
covered_end = ranges[range_index][0] if range_index < len(ranges) else 0
|
||||
if range_index >= len(ranges):
|
||||
raise ValueError("invalid_chunk_manifest")
|
||||
left, right = ranges[range_index]
|
||||
if (
|
||||
chunk.ordinal != index
|
||||
or not left <= chunk.start_cp < chunk.end_cp <= right
|
||||
or chunk.end_cp <= previous_end
|
||||
or not 1 <= chunk.token_count <= 512
|
||||
or record.verbatim[chunk.start_cp : chunk.end_cp] != chunk.text
|
||||
or (
|
||||
chunk.start_cp > covered_end
|
||||
and (prepared is None or record.verbatim[covered_end : chunk.start_cp].strip())
|
||||
)
|
||||
):
|
||||
raise ValueError("invalid_chunk_manifest")
|
||||
covered_end = previous_end = chunk.end_cp
|
||||
while range_index < len(ranges):
|
||||
if prepared is None:
|
||||
if covered_end != ranges[range_index][1]:
|
||||
raise ValueError("incomplete_chunk_manifest")
|
||||
elif record.verbatim[covered_end : ranges[range_index][1]].strip():
|
||||
raise ValueError("incomplete_chunk_manifest")
|
||||
range_index += 1
|
||||
covered_end = ranges[range_index][0] if range_index < len(ranges) else 0
|
||||
if prepared is not None and (not chunks) != (prepared.state == "no_indexable_content"):
|
||||
raise ValueError("invalid_chunk_manifest")
|
||||
@@ -1,96 +0,0 @@
|
||||
"""Publish one active generation for the built-in encoder, when asked to.
|
||||
|
||||
The runtime deliberately only resumes an index somebody already initialized, so
|
||||
without this the feature could be switched on and still never index anything.
|
||||
Provisioning stays separate from `start()` for the same reason it exists: it is
|
||||
the one place allowed to create, and it creates only for a first enable.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.embedding_assets import assets_present
|
||||
from one_memory.retrieval.index_store import RetrievalIndexStore
|
||||
from one_memory.retrieval.index_worker import LocalRetrievalEncoder
|
||||
from one_memory.retrieval.vector_store import RawVectorStore
|
||||
|
||||
DATABASE_NAME = "one-memory.sqlite3"
|
||||
VECTOR_DIR_NAME = "retrieval-vectors"
|
||||
|
||||
|
||||
def provision_retrieval_index(data_dir: Path) -> str | None:
|
||||
"""The active generation id for the built-in encoder, creating it once.
|
||||
|
||||
Returns None whenever the right answer is to leave the data alone: no
|
||||
built-in encoder, no database yet, an index from an older schema, or one
|
||||
already published for a different encoder. None is not an error; the
|
||||
lexical lane works without an index and says so through `degraded`.
|
||||
"""
|
||||
database = data_dir / DATABASE_NAME
|
||||
if not assets_present() or not database.is_file():
|
||||
return None
|
||||
identity = LocalRetrievalEncoder().identity
|
||||
store = RetrievalIndexStore(database)
|
||||
try:
|
||||
with main_database(database).write() as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
if (existing := _published(db)) is not None:
|
||||
return existing if _matches(store, db, existing, identity) else None
|
||||
version = _schema_version(db)
|
||||
if version == 1:
|
||||
# Someone's v1 index. Migrating it is maintenance, not startup.
|
||||
return None
|
||||
if version is None:
|
||||
store.initialize_v2(db)
|
||||
generation = store.create_generation(db, identity)
|
||||
store.activate_generation(generation)
|
||||
_ensure_vectors(data_dir / VECTOR_DIR_NAME, generation, identity.dimensions)
|
||||
except Exception:
|
||||
logging.getLogger(__name__).warning("retrieval_index_provision_failed", exc_info=True)
|
||||
return None
|
||||
return generation
|
||||
|
||||
|
||||
def _schema_version(db: sqlite3.Connection) -> int | None:
|
||||
if not db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='retrieval_index_state'"
|
||||
).fetchone():
|
||||
return None
|
||||
row = db.execute(
|
||||
"SELECT schema_version FROM retrieval_index_state WHERE singleton=1"
|
||||
).fetchone()
|
||||
return None if row is None else int(row[0])
|
||||
|
||||
|
||||
def _published(db: sqlite3.Connection) -> str | None:
|
||||
if not db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='retrieval_index_generations'"
|
||||
).fetchone():
|
||||
return None
|
||||
rows = db.execute(
|
||||
"SELECT generation_id FROM retrieval_index_generations WHERE state='active' LIMIT 2"
|
||||
).fetchall()
|
||||
return str(rows[0][0]) if len(rows) == 1 else None
|
||||
|
||||
|
||||
def _matches(
|
||||
store: RetrievalIndexStore,
|
||||
db: sqlite3.Connection,
|
||||
generation: str,
|
||||
identity: object,
|
||||
) -> bool:
|
||||
try:
|
||||
return store._identity(db, generation) == identity
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _ensure_vectors(directory: Path, generation: str, dimensions: int) -> None:
|
||||
if directory.exists():
|
||||
return
|
||||
vectors = RawVectorStore.create(directory, generation_id=generation, dimensions=dimensions)
|
||||
vectors.close()
|
||||
@@ -1,856 +0,0 @@
|
||||
"""Explicit derived-index transactions. No constructor migration or model invocation."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import sqlite3
|
||||
import time
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import asdict, dataclass, replace
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
from one_memory import ingest_journal as journal
|
||||
from one_memory.context_raw import _check, _connection
|
||||
from one_memory.database import take_the_write_lock_first
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.imports.models import EvidenceRecord
|
||||
from one_memory.imports.storage import EvidenceStore
|
||||
from one_memory.memory_server.migrations import (
|
||||
_split_statements,
|
||||
retrieval_build_maps_ddl,
|
||||
retrieval_documents_v2_ddl,
|
||||
retrieval_index_ddl,
|
||||
)
|
||||
from one_memory.retrieval.chunking import TextChunk
|
||||
from one_memory.retrieval.content import index_content, validate_chunk_coverage
|
||||
from one_memory.retrieval.mapped_content import ContentMap, PreparedMappedContent
|
||||
from one_memory.retrieval.models import TimeRange
|
||||
|
||||
|
||||
def digest(value: object) -> str:
|
||||
return hashlib.sha256(
|
||||
json.dumps(
|
||||
value, sort_keys=True, ensure_ascii=False, separators=(",", ":"), allow_nan=False
|
||||
).encode()
|
||||
).hexdigest()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EncoderIdentity:
|
||||
encoder_id: str
|
||||
encoder_digest: str
|
||||
tokenizer_digest: str
|
||||
dimensions: int
|
||||
chunker_version: int = 1
|
||||
normalizer_version: int = 1
|
||||
query_template_version: int = 1
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not self.encoder_id or len(self.encoder_id) > 128:
|
||||
raise ValueError("invalid_encoder")
|
||||
for value in (self.encoder_digest, self.tokenizer_digest):
|
||||
if len(value) != 64 or any(c not in "0123456789abcdef" for c in value):
|
||||
raise ValueError("invalid_encoder_digest")
|
||||
for number in (
|
||||
self.dimensions,
|
||||
self.chunker_version,
|
||||
self.normalizer_version,
|
||||
self.query_template_version,
|
||||
):
|
||||
if type(number) is not int or not 1 <= number <= 1000000:
|
||||
raise ValueError("invalid_encoder_number")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class IndexLease:
|
||||
generation_id: str
|
||||
evidence_id: str
|
||||
build_id: str
|
||||
event_seq: int
|
||||
owner: str
|
||||
token: str
|
||||
expires_at: float
|
||||
grant_fingerprint: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PublishedChunk:
|
||||
chunk_id: str
|
||||
build_id: str
|
||||
evidence_id: str
|
||||
revision: int
|
||||
start_cp: int
|
||||
end_cp: int
|
||||
ordinal: int
|
||||
source_id: str
|
||||
event_time: datetime | None
|
||||
|
||||
|
||||
def chunk_identity(lease: IndexLease, chunk: TextChunk, map_digest: str | None = None) -> str:
|
||||
values = [lease.build_id, chunk.ordinal, chunk.start_cp, chunk.end_cp, digest(chunk.text)]
|
||||
return digest(values if map_digest is None else [*values, map_digest])
|
||||
|
||||
|
||||
@contextmanager
|
||||
def transaction(db: sqlite3.Connection) -> Iterator[None]:
|
||||
if not db.in_transaction:
|
||||
raise RuntimeError("retrieval_requires_caller_transaction")
|
||||
name = "retrieval_" + uuid4().hex
|
||||
previous = db.row_factory
|
||||
db.row_factory = sqlite3.Row
|
||||
db.execute(f"SAVEPOINT {name}")
|
||||
try:
|
||||
yield
|
||||
except BaseException:
|
||||
db.execute(f"ROLLBACK TO {name}")
|
||||
db.execute(f"RELEASE {name}")
|
||||
raise
|
||||
else:
|
||||
db.execute(f"RELEASE {name}")
|
||||
finally:
|
||||
db.row_factory = previous
|
||||
|
||||
|
||||
class RetrievalIndexStore:
|
||||
def __init__(self, database_path: Path):
|
||||
self.database_path = database_path
|
||||
|
||||
@contextmanager
|
||||
def _db(
|
||||
self, *, write: bool = False, deadline: float | None = None
|
||||
) -> Iterator[sqlite3.Connection]:
|
||||
if deadline is not None:
|
||||
if write:
|
||||
raise ValueError("deadline_connection_is_read_only")
|
||||
with _connection(self.database_path, deadline) as db:
|
||||
db.execute("PRAGMA query_only=ON")
|
||||
yield db
|
||||
return
|
||||
database = main_database(self.database_path)
|
||||
with (
|
||||
database.write(create=False, foreign_keys=True, row_factory=sqlite3.Row)
|
||||
if write
|
||||
else database.read(foreign_keys=True, row_factory=sqlite3.Row)
|
||||
) as db:
|
||||
try:
|
||||
if write:
|
||||
take_the_write_lock_first(db)
|
||||
else:
|
||||
db.execute("PRAGMA query_only=ON")
|
||||
db.execute("BEGIN")
|
||||
yield db
|
||||
if write:
|
||||
db.commit()
|
||||
except BaseException:
|
||||
db.rollback()
|
||||
raise
|
||||
|
||||
def initialize(self, db: sqlite3.Connection) -> None:
|
||||
"""Caller owns maintenance/exclusive writer admission; rollback includes baseline."""
|
||||
self._initialize(db, version=1)
|
||||
|
||||
def initialize_v2(self, db: sqlite3.Connection) -> None:
|
||||
"""Initialize a fresh store directly at the folded v2 schema.
|
||||
|
||||
Startup provisioning uses this path so a brand-new database does not
|
||||
create a v1 schema and immediately migrate it. Explicit maintenance
|
||||
callers keep `initialize()`'s v1 baseline contract instead.
|
||||
"""
|
||||
self._initialize(db, version=2)
|
||||
|
||||
def _initialize(self, db: sqlite3.Connection, *, version: int) -> None:
|
||||
with transaction(db):
|
||||
exists = db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE name='retrieval_index_state'"
|
||||
).fetchone()
|
||||
if exists:
|
||||
self._store_id(db)
|
||||
elif db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE name IN ('retrieval_documents',"
|
||||
"'retrieval_chunks','retrieval_index_generations','retrieval_consumers')"
|
||||
).fetchone():
|
||||
raise ValueError("retrieval_partial_initialization")
|
||||
journal.initialize(db)
|
||||
# Migration 33 owns the schema. Explicit initialization keeps the
|
||||
# maintenance-facing v1 baseline so `migrate_v2` remains observable;
|
||||
# startup provisioning may instead ask for the folded v2 directly.
|
||||
for sql in _split_statements(retrieval_index_ddl("sqlite", version=version)):
|
||||
db.execute(sql)
|
||||
row = db.execute(
|
||||
"SELECT schema_version FROM retrieval_index_state WHERE singleton=1"
|
||||
).fetchone()
|
||||
if row is None:
|
||||
db.execute(
|
||||
"INSERT INTO retrieval_index_state VALUES (1,?,?)",
|
||||
(version, uuid4().hex),
|
||||
)
|
||||
elif row[0] not in (1, 2):
|
||||
raise ValueError("retrieval_schema_mismatch")
|
||||
|
||||
def migrate_v2(self, db: sqlite3.Connection) -> None:
|
||||
"""Explicit maintenance transaction only; readers/worker never invoke this."""
|
||||
with transaction(db):
|
||||
self._store_id(db)
|
||||
version = db.execute("SELECT schema_version FROM retrieval_index_state").fetchone()[0]
|
||||
if version == 2:
|
||||
db.execute(
|
||||
"SELECT build_id,source_map_json,map_digest,coverage_state "
|
||||
"FROM retrieval_build_maps LIMIT 0"
|
||||
)
|
||||
return
|
||||
if (
|
||||
db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='retrieval_build_maps'"
|
||||
).fetchone()
|
||||
is not None
|
||||
):
|
||||
# A v1 index must not already own a v2-only table. Treat this
|
||||
# as a foreign/partial schema instead of letting `IF NOT
|
||||
# EXISTS` silently accept a table with the wrong columns.
|
||||
raise sqlite3.OperationalError("retrieval_build_maps already exists")
|
||||
# Both tables come from migration 33, named/derived there so this
|
||||
# module holds no DDL text of its own.
|
||||
db.execute(retrieval_documents_v2_ddl("sqlite"))
|
||||
db.execute("INSERT INTO retrieval_documents_v2 SELECT * FROM retrieval_documents")
|
||||
db.execute("DROP TABLE retrieval_documents")
|
||||
db.execute("ALTER TABLE retrieval_documents_v2 RENAME TO retrieval_documents")
|
||||
db.execute(
|
||||
"CREATE INDEX retrieval_jobs ON retrieval_documents(generation_id,state,event_seq)"
|
||||
)
|
||||
db.execute(retrieval_build_maps_ddl("sqlite"))
|
||||
db.execute("UPDATE retrieval_index_state SET schema_version=2 WHERE singleton=1")
|
||||
|
||||
@staticmethod
|
||||
def _store_id(db: sqlite3.Connection) -> str:
|
||||
row = db.execute(
|
||||
"SELECT schema_version,data_store_id FROM retrieval_index_state WHERE singleton=1"
|
||||
).fetchone()
|
||||
if row is None or row[0] not in (1, 2) or not isinstance(row[1], str) or len(row[1]) != 32:
|
||||
raise ValueError("retrieval_not_initialized")
|
||||
return str(row[1])
|
||||
|
||||
def create_generation(self, db: sqlite3.Connection, identity: EncoderIdentity) -> str:
|
||||
with transaction(db):
|
||||
self._check_content_version(db, identity)
|
||||
key = digest([self._store_id(db), asdict(identity)])
|
||||
encoded = json.dumps(asdict(identity), sort_keys=True)
|
||||
row = db.execute(
|
||||
"SELECT identity_json FROM retrieval_index_generations WHERE generation_id=?",
|
||||
(key,),
|
||||
).fetchone()
|
||||
if row is not None and row[0] != encoded:
|
||||
raise ValueError("generation_identity_changed")
|
||||
if row is None:
|
||||
db.execute(
|
||||
"INSERT INTO retrieval_index_generations VALUES (?,?,'building',?)",
|
||||
(key, encoded, time.time()),
|
||||
)
|
||||
db.execute("INSERT INTO retrieval_consumers VALUES (?,0,0)", (key,))
|
||||
return key
|
||||
|
||||
def consume(self, db: sqlite3.Connection, generation_id: str, *, limit: int = 256) -> int:
|
||||
if type(limit) is not int or not 1 <= limit <= 256:
|
||||
raise ValueError("invalid_consume_limit")
|
||||
with transaction(db):
|
||||
store_id = self._store_id(db)
|
||||
generation = db.execute(
|
||||
"SELECT state FROM retrieval_index_generations WHERE generation_id=?",
|
||||
(generation_id,),
|
||||
).fetchone()
|
||||
if generation is None or generation[0] == "retired":
|
||||
raise ValueError("generation_not_writable")
|
||||
cursor = db.execute(
|
||||
"SELECT after_seq FROM retrieval_consumers WHERE generation_id=?", (generation_id,)
|
||||
).fetchone()
|
||||
if cursor is None:
|
||||
raise ValueError("consumer_missing")
|
||||
rows = db.execute(
|
||||
"SELECT seq FROM evidence_ingest_events WHERE seq>? ORDER BY seq LIMIT ?",
|
||||
(cursor[0], limit),
|
||||
).fetchall()
|
||||
for row in rows:
|
||||
event = journal.event(db, row[0])
|
||||
if event is None:
|
||||
raise ValueError("journal_event_missing")
|
||||
if event.user_id != "local":
|
||||
continue
|
||||
build = digest(
|
||||
[
|
||||
store_id,
|
||||
generation_id,
|
||||
event.evidence_id,
|
||||
event.incarnation,
|
||||
event.evidence_revision,
|
||||
event.fingerprint,
|
||||
]
|
||||
)
|
||||
state = "deleted" if event.kind == "delete" else "pending"
|
||||
db.execute(
|
||||
"""INSERT INTO retrieval_documents
|
||||
(generation_id,evidence_id,source_id,event_seq,build_id,state)
|
||||
VALUES (?,?,?,?,?,?) ON CONFLICT(generation_id,evidence_id) DO UPDATE SET
|
||||
source_id=excluded.source_id,event_seq=excluded.event_seq,build_id=excluded.build_id,
|
||||
state=excluded.state,active_build_id=NULL,lease_owner=NULL,lease_token=NULL,
|
||||
lease_until=NULL,grant_fingerprint=NULL
|
||||
WHERE excluded.event_seq>retrieval_documents.event_seq""",
|
||||
(generation_id, event.evidence_id, event.source_id, event.seq, build, state),
|
||||
)
|
||||
if rows:
|
||||
db.execute(
|
||||
"UPDATE retrieval_consumers SET after_seq=? WHERE generation_id=?",
|
||||
(rows[-1][0], generation_id),
|
||||
)
|
||||
return len(rows)
|
||||
|
||||
def consume_pending(self, generation_id: str, *, limit: int = 256) -> int:
|
||||
with self._db(write=True) as db:
|
||||
return self.consume(db, generation_id, limit=limit)
|
||||
|
||||
def encoder_identity(self, generation_id: str) -> EncoderIdentity:
|
||||
with self._db() as db:
|
||||
return self._identity(db, generation_id)
|
||||
|
||||
@staticmethod
|
||||
def _check_content_version(db: sqlite3.Connection, identity: EncoderIdentity) -> None:
|
||||
versions = (identity.normalizer_version, identity.chunker_version)
|
||||
if versions not in ((1, 1), (2, 2)):
|
||||
raise ValueError("unsupported_content_generation")
|
||||
if versions == (2, 2):
|
||||
row = db.execute(
|
||||
"SELECT schema_version FROM retrieval_index_state WHERE singleton=1"
|
||||
).fetchone()
|
||||
if row is None or row[0] != 2:
|
||||
raise ValueError("retrieval_schema2_required")
|
||||
|
||||
def _identity(self, db: sqlite3.Connection, generation_id: str) -> EncoderIdentity:
|
||||
row = db.execute(
|
||||
"SELECT identity_json FROM retrieval_index_generations WHERE generation_id=?",
|
||||
(generation_id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise ValueError("generation_missing")
|
||||
identity = EncoderIdentity(**json.loads(row[0]))
|
||||
if digest([self._store_id(db), asdict(identity)]) != generation_id:
|
||||
raise ValueError("generation_identity_changed")
|
||||
self._check_content_version(db, identity)
|
||||
return identity
|
||||
|
||||
def _prepare(
|
||||
self, db: sqlite3.Connection, generation_id: str, record: EvidenceRecord
|
||||
) -> PreparedMappedContent | None:
|
||||
identity = self._identity(db, generation_id)
|
||||
return index_content(
|
||||
record,
|
||||
normalizer_version=identity.normalizer_version,
|
||||
chunker_version=identity.chunker_version,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _grant(db: sqlite3.Connection, source_id: str) -> str | None:
|
||||
row = db.execute(
|
||||
"""SELECT root_path,candidate_kind,granted_at,root_identity,revision
|
||||
FROM source_grants
|
||||
WHERE user_id='local' AND source_id=? AND revoked_at IS NULL""",
|
||||
(source_id,),
|
||||
).fetchone()
|
||||
return digest(tuple(row)) if row and row[3] and row[4] > 0 else None
|
||||
|
||||
def claim(
|
||||
self, generation_id: str, owner: str, *, now: float | None = None, lease_seconds: int = 60
|
||||
) -> IndexLease | None:
|
||||
now = time.time() if now is None else now
|
||||
if not owner or len(owner) > 128 or not math.isfinite(now) or not 1 <= lease_seconds <= 300:
|
||||
raise ValueError("invalid_lease")
|
||||
with self._db(write=True) as db:
|
||||
generation = db.execute(
|
||||
"SELECT state FROM retrieval_index_generations WHERE generation_id=?",
|
||||
(generation_id,),
|
||||
).fetchone()
|
||||
if generation is None or generation[0] == "retired":
|
||||
raise ValueError("generation_not_writable")
|
||||
rows = db.execute(
|
||||
"""SELECT * FROM retrieval_documents WHERE generation_id=? AND
|
||||
(state='pending' OR (state='leased' AND lease_until<=?))
|
||||
ORDER BY event_seq LIMIT 256""",
|
||||
(generation_id, now),
|
||||
).fetchall()
|
||||
for row in rows:
|
||||
event = journal.event(db, row["event_seq"])
|
||||
grant = self._grant(db, row["source_id"])
|
||||
if event is None or grant is None or not journal.matches_current(db, event):
|
||||
db.execute(
|
||||
"UPDATE retrieval_documents SET state='blocked',active_build_id=NULL "
|
||||
"WHERE generation_id=? AND evidence_id=?",
|
||||
(generation_id, row["evidence_id"]),
|
||||
)
|
||||
continue
|
||||
token = uuid4().hex
|
||||
until = now + lease_seconds
|
||||
db.execute(
|
||||
"""UPDATE retrieval_documents SET state='leased',lease_owner=?,lease_token=?,
|
||||
lease_until=?,grant_fingerprint=?,attempts=attempts+1
|
||||
WHERE generation_id=? AND evidence_id=?""",
|
||||
(owner, token, until, grant, generation_id, row["evidence_id"]),
|
||||
)
|
||||
return IndexLease(
|
||||
generation_id,
|
||||
row["evidence_id"],
|
||||
row["build_id"],
|
||||
row["event_seq"],
|
||||
owner,
|
||||
token,
|
||||
until,
|
||||
grant,
|
||||
)
|
||||
return None
|
||||
|
||||
def _validate_lease(
|
||||
self, db: sqlite3.Connection, lease: IndexLease, now: float
|
||||
) -> journal.IngestEvent:
|
||||
if not math.isfinite(now):
|
||||
raise ValueError("invalid_lease_time")
|
||||
generation = db.execute(
|
||||
"SELECT state FROM retrieval_index_generations WHERE generation_id=?",
|
||||
(lease.generation_id,),
|
||||
).fetchone()
|
||||
if generation is None or generation[0] == "retired":
|
||||
raise ValueError("lease_generation_retired")
|
||||
row = db.execute(
|
||||
"SELECT * FROM retrieval_documents WHERE generation_id=? AND evidence_id=?",
|
||||
(lease.generation_id, lease.evidence_id),
|
||||
).fetchone()
|
||||
if (
|
||||
row is None
|
||||
or row["state"] != "leased"
|
||||
or (
|
||||
row["build_id"],
|
||||
row["event_seq"],
|
||||
row["lease_owner"],
|
||||
row["lease_token"],
|
||||
row["lease_until"],
|
||||
row["grant_fingerprint"],
|
||||
)
|
||||
!= (
|
||||
lease.build_id,
|
||||
lease.event_seq,
|
||||
lease.owner,
|
||||
lease.token,
|
||||
lease.expires_at,
|
||||
lease.grant_fingerprint,
|
||||
)
|
||||
or lease.expires_at <= now
|
||||
):
|
||||
raise ValueError("lease_changed_or_expired")
|
||||
event = journal.event(db, lease.event_seq)
|
||||
if (
|
||||
event is None
|
||||
or event.user_id != "local"
|
||||
or event.evidence_id != lease.evidence_id
|
||||
or not journal.matches_current(db, event)
|
||||
):
|
||||
raise ValueError("lease_evidence_changed")
|
||||
if self._grant(db, event.source_id) != lease.grant_fingerprint:
|
||||
raise ValueError("lease_source_authority_changed")
|
||||
return event
|
||||
|
||||
def read_lease(self, lease: IndexLease, *, now: float | None = None) -> EvidenceRecord:
|
||||
with self._db() as db:
|
||||
self._validate_lease(db, lease, time.time() if now is None else now)
|
||||
record = EvidenceStore._get_with_connection(db, lease.evidence_id)
|
||||
if record is None or len(record.verbatim) > 2_000_000:
|
||||
raise ValueError("lease_input_unavailable_or_too_large")
|
||||
return record
|
||||
|
||||
def renew(
|
||||
self, lease: IndexLease, *, now: float | None = None, lease_seconds: int = 60
|
||||
) -> IndexLease:
|
||||
now = time.time() if now is None else now
|
||||
if type(lease_seconds) is not int or not 1 <= lease_seconds <= 300:
|
||||
raise ValueError("invalid_lease")
|
||||
with self._db(write=True) as db:
|
||||
self._validate_lease(db, lease, now)
|
||||
expires = max(lease.expires_at, now + lease_seconds)
|
||||
db.execute(
|
||||
"UPDATE retrieval_documents SET lease_until=? "
|
||||
"WHERE generation_id=? AND evidence_id=?",
|
||||
(expires, lease.generation_id, lease.evidence_id),
|
||||
)
|
||||
return replace(lease, expires_at=expires)
|
||||
|
||||
def block(self, lease: IndexLease) -> None:
|
||||
"""Release only this claim; preserve uncertain derived vectors for later cleanup."""
|
||||
with self._db(write=True) as db:
|
||||
db.execute(
|
||||
"""UPDATE retrieval_documents SET state='blocked',active_build_id=NULL,
|
||||
lease_owner=NULL,lease_token=NULL,lease_until=NULL
|
||||
WHERE generation_id=? AND evidence_id=? AND state='leased'
|
||||
AND build_id=? AND lease_token=? AND lease_owner=?""",
|
||||
(lease.generation_id, lease.evidence_id, lease.build_id, lease.token, lease.owner),
|
||||
)
|
||||
|
||||
def retry_blocked(self, generation_id: str) -> int:
|
||||
"""Bounded, durable round-robin retry/regrant pass; never creates grants."""
|
||||
changed = 0
|
||||
with self._db(write=True) as db:
|
||||
generation = db.execute(
|
||||
"SELECT g.state,c.retry_after_seq FROM retrieval_index_generations g "
|
||||
"JOIN retrieval_consumers c USING(generation_id) WHERE generation_id=?",
|
||||
(generation_id,),
|
||||
).fetchone()
|
||||
if generation is None or generation[0] == "retired":
|
||||
raise ValueError("generation_not_writable")
|
||||
rows = db.execute(
|
||||
"SELECT * FROM retrieval_documents WHERE generation_id=? "
|
||||
"AND state IN ('blocked','ready','excluded') AND event_seq>? "
|
||||
"ORDER BY event_seq LIMIT 256",
|
||||
(generation_id, generation[1]),
|
||||
).fetchall()
|
||||
for row in rows:
|
||||
event = journal.event(db, row["event_seq"])
|
||||
grant = self._grant(db, row["source_id"])
|
||||
if row["state"] in ("ready", "excluded") and grant == row["grant_fingerprint"]:
|
||||
continue
|
||||
if (
|
||||
event
|
||||
and event.source_id == row["source_id"]
|
||||
and grant
|
||||
and journal.matches_current(db, event)
|
||||
):
|
||||
db.execute(
|
||||
"UPDATE retrieval_documents SET state='pending',active_build_id=NULL,"
|
||||
"lease_owner=NULL,lease_token=NULL,lease_until=NULL,grant_fingerprint=NULL "
|
||||
"WHERE generation_id=? AND evidence_id=?",
|
||||
(generation_id, row["evidence_id"]),
|
||||
)
|
||||
changed += 1
|
||||
db.execute(
|
||||
"UPDATE retrieval_consumers SET retry_after_seq=? WHERE generation_id=?",
|
||||
(rows[-1]["event_seq"] if len(rows) == 256 else 0, generation_id),
|
||||
)
|
||||
return changed
|
||||
|
||||
def publish(
|
||||
self,
|
||||
lease: IndexLease,
|
||||
chunks: Sequence[TextChunk],
|
||||
*,
|
||||
vectors_verified: bool,
|
||||
now: float | None = None,
|
||||
source_map: ContentMap | None = None,
|
||||
checkpoint: Callable[[], None] | None = None,
|
||||
) -> tuple[str, ...]:
|
||||
"""Caller verifies actual immutable vectors before entering this short transaction."""
|
||||
if vectors_verified is not True:
|
||||
raise ValueError("vectors_not_verified")
|
||||
now = time.time() if now is None else now
|
||||
if not math.isfinite(now):
|
||||
raise ValueError("invalid_lease_time")
|
||||
if len(chunks) > 8192:
|
||||
raise ValueError("invalid_chunk_manifest")
|
||||
with self._db(write=True) as db:
|
||||
if checkpoint is not None:
|
||||
checkpoint()
|
||||
event = self._validate_lease(db, lease, now)
|
||||
record = EvidenceStore._get_with_connection(db, lease.evidence_id)
|
||||
if record is None or len(record.verbatim) > 2_000_000:
|
||||
raise ValueError("lease_evidence_missing")
|
||||
prepared = self._prepare(db, lease.generation_id, record)
|
||||
if source_map != (prepared.metadata if prepared is not None else None):
|
||||
raise ValueError("source_map_mismatch")
|
||||
validate_chunk_coverage(record, chunks, prepared)
|
||||
map_digest = source_map.digest if source_map is not None else None
|
||||
keys = []
|
||||
manifest = []
|
||||
for index, chunk in enumerate(chunks):
|
||||
text_hash = digest(chunk.text)
|
||||
key = chunk_identity(lease, chunk, map_digest)
|
||||
keys.append(key)
|
||||
manifest.append(
|
||||
(
|
||||
key,
|
||||
lease.generation_id,
|
||||
lease.evidence_id,
|
||||
lease.build_id,
|
||||
index,
|
||||
chunk.start_cp,
|
||||
chunk.end_cp,
|
||||
chunk.token_count,
|
||||
text_hash,
|
||||
)
|
||||
)
|
||||
existing = db.execute(
|
||||
"SELECT * FROM retrieval_chunks WHERE build_id=? ORDER BY ordinal",
|
||||
(lease.build_id,),
|
||||
).fetchall()
|
||||
if existing:
|
||||
if [tuple(row) for row in existing] != manifest:
|
||||
raise ValueError("immutable_chunk_manifest_conflict")
|
||||
else:
|
||||
db.executemany("INSERT INTO retrieval_chunks VALUES (?,?,?,?,?,?,?,?,?)", manifest)
|
||||
state = "ready" if chunks else "excluded"
|
||||
if source_map is not None:
|
||||
values = (
|
||||
lease.build_id,
|
||||
source_map.canonical_json(),
|
||||
map_digest,
|
||||
"indexable" if chunks else "excluded",
|
||||
)
|
||||
prior = db.execute(
|
||||
"SELECT * FROM retrieval_build_maps WHERE build_id=?", (lease.build_id,)
|
||||
).fetchone()
|
||||
if prior is not None and tuple(prior) != values:
|
||||
raise ValueError("immutable_source_map_conflict")
|
||||
if prior is None:
|
||||
db.execute("INSERT INTO retrieval_build_maps VALUES (?,?,?,?)", values)
|
||||
if not journal.matches_current(db, event):
|
||||
raise ValueError("lease_evidence_changed")
|
||||
db.execute(
|
||||
"UPDATE retrieval_documents SET state=?,active_build_id=build_id,"
|
||||
"lease_owner=NULL,lease_token=NULL,lease_until=NULL "
|
||||
"WHERE generation_id=? AND evidence_id=?",
|
||||
(state, lease.generation_id, lease.evidence_id),
|
||||
)
|
||||
if checkpoint is not None:
|
||||
checkpoint()
|
||||
return tuple(keys)
|
||||
|
||||
def activate_generation(self, generation_id: str) -> None:
|
||||
with self._db(write=True) as db:
|
||||
row = db.execute(
|
||||
"SELECT state FROM retrieval_index_generations WHERE generation_id=?",
|
||||
(generation_id,),
|
||||
).fetchone()
|
||||
cursor = db.execute(
|
||||
"SELECT after_seq FROM retrieval_consumers WHERE generation_id=?", (generation_id,)
|
||||
).fetchone()
|
||||
if (
|
||||
row is None
|
||||
or row[0] == "retired"
|
||||
or cursor is None
|
||||
or cursor[0] != journal.watermark(db)
|
||||
):
|
||||
raise ValueError("generation_not_caught_up")
|
||||
for document in db.execute(
|
||||
"SELECT * FROM retrieval_documents WHERE generation_id=? "
|
||||
"AND state NOT IN ('ready','excluded','deleted')",
|
||||
(generation_id,),
|
||||
):
|
||||
if self._grant(db, document["source_id"]) is not None:
|
||||
raise ValueError("generation_incomplete")
|
||||
for document in db.execute(
|
||||
"SELECT * FROM retrieval_documents WHERE generation_id=? "
|
||||
"AND state IN ('ready','excluded')",
|
||||
(generation_id,),
|
||||
):
|
||||
if self._grant(db, document["source_id"]) is not None:
|
||||
event = journal.event(db, document["event_seq"])
|
||||
if (
|
||||
not event
|
||||
or not journal.matches_current(db, event)
|
||||
or self._grant(db, document["source_id"]) != document["grant_fingerprint"]
|
||||
):
|
||||
raise ValueError("generation_stale")
|
||||
record = EvidenceStore._get_with_connection(db, document["evidence_id"])
|
||||
if record is None:
|
||||
raise ValueError("generation_stale")
|
||||
self.validate_document(db, document, record)
|
||||
db.execute(
|
||||
"UPDATE retrieval_index_generations SET state='retired' "
|
||||
"WHERE state='active' AND generation_id!=?",
|
||||
(generation_id,),
|
||||
)
|
||||
db.execute(
|
||||
"UPDATE retrieval_index_generations SET state='active' WHERE generation_id=?",
|
||||
(generation_id,),
|
||||
)
|
||||
|
||||
def active_chunks(
|
||||
self,
|
||||
generation_id: str,
|
||||
*,
|
||||
source_ids: frozenset[str] | None = None,
|
||||
time_range: TimeRange | None = None,
|
||||
deadline: float | None = None,
|
||||
require_active: bool = False,
|
||||
) -> tuple[PublishedChunk, ...]:
|
||||
"""Return current chunks after SQL-level source/time eligibility filtering.
|
||||
|
||||
Grant fingerprints, journal freshness, and manifests are still checked
|
||||
below in the same snapshot. The SQL predicate keeps unauthorized rows
|
||||
from filling an in-memory candidate set before those checks run.
|
||||
"""
|
||||
if source_ids is not None and not source_ids:
|
||||
return ()
|
||||
with self._db(deadline=deadline) as db:
|
||||
generation = db.execute(
|
||||
"SELECT state FROM retrieval_index_generations WHERE generation_id=?",
|
||||
(generation_id,),
|
||||
).fetchone()
|
||||
if generation is None or generation[0] != "active":
|
||||
if require_active:
|
||||
raise ValueError("active_generation_changed")
|
||||
return ()
|
||||
clauses = ["d.generation_id=?", "d.state IN ('ready','excluded')"]
|
||||
parameters: list[object] = [generation_id]
|
||||
join = ""
|
||||
if source_ids is not None:
|
||||
clauses.append("d.source_id IN (SELECT value FROM json_each(?))")
|
||||
parameters.append(json.dumps(sorted(source_ids)))
|
||||
if time_range is not None:
|
||||
join = "JOIN evidence e ON e.evidence_id=d.evidence_id "
|
||||
clauses.extend(
|
||||
(
|
||||
"e.created_at IS NOT NULL",
|
||||
"julianday(e.created_at) >= julianday(?)",
|
||||
"julianday(e.created_at) < julianday(?)",
|
||||
)
|
||||
)
|
||||
parameters.extend((time_range.start.isoformat(), time_range.end.isoformat()))
|
||||
output = []
|
||||
for row in db.execute(
|
||||
"SELECT d.* FROM retrieval_documents d "
|
||||
+ join
|
||||
+ f"WHERE {' AND '.join(clauses)} ORDER BY d.event_seq,d.evidence_id",
|
||||
parameters,
|
||||
):
|
||||
if deadline is not None:
|
||||
_check(deadline)
|
||||
grant = self._grant(db, row["source_id"])
|
||||
if grant is None or grant != row["grant_fingerprint"]:
|
||||
continue
|
||||
event = journal.event(db, row["event_seq"])
|
||||
if (
|
||||
event is None
|
||||
or event.source_id != row["source_id"]
|
||||
or not journal.matches_current(db, event)
|
||||
):
|
||||
continue
|
||||
record = EvidenceStore._get_with_connection(db, row["evidence_id"])
|
||||
if (
|
||||
record is None
|
||||
or record.user_id != "local"
|
||||
or record.source_id != row["source_id"]
|
||||
):
|
||||
continue
|
||||
event_time = record.created_at
|
||||
if event_time is not None:
|
||||
event_time = event_time.astimezone(UTC) if event_time.tzinfo else None
|
||||
if time_range is not None and (
|
||||
event_time is None or not time_range.start <= event_time < time_range.end
|
||||
):
|
||||
continue
|
||||
chunks = self.validate_document(db, row, record, deadline=deadline)
|
||||
for chunk in chunks:
|
||||
if deadline is not None:
|
||||
_check(deadline)
|
||||
output.append(
|
||||
PublishedChunk(
|
||||
chunk["chunk_id"],
|
||||
row["build_id"],
|
||||
row["evidence_id"],
|
||||
event.evidence_revision,
|
||||
chunk["start_cp"],
|
||||
chunk["end_cp"],
|
||||
chunk["ordinal"],
|
||||
record.source_id,
|
||||
event_time,
|
||||
)
|
||||
)
|
||||
return tuple(output)
|
||||
|
||||
def validate_document(
|
||||
self,
|
||||
db: sqlite3.Connection,
|
||||
document: sqlite3.Row,
|
||||
record: EvidenceRecord,
|
||||
*,
|
||||
deadline: float | None = None,
|
||||
) -> tuple[sqlite3.Row, ...]:
|
||||
"""Same-snapshot complete manifest/map check shared by candidate and final read."""
|
||||
if deadline is not None:
|
||||
_check(deadline)
|
||||
if document["active_build_id"] != document["build_id"]:
|
||||
raise ValueError("published_build_binding_invalid")
|
||||
prepared = self._prepare(db, document["generation_id"], record)
|
||||
map_digest = None
|
||||
if prepared is not None:
|
||||
source_map = db.execute(
|
||||
"SELECT source_map_json,map_digest,coverage_state FROM retrieval_build_maps "
|
||||
"WHERE build_id=?",
|
||||
(document["build_id"],),
|
||||
).fetchone()
|
||||
expected_state = "excluded" if prepared.state == "no_indexable_content" else "ready"
|
||||
if (
|
||||
source_map is None
|
||||
or tuple(source_map)
|
||||
!= (
|
||||
prepared.metadata.canonical_json(),
|
||||
prepared.metadata.digest,
|
||||
"excluded" if expected_state == "excluded" else "indexable",
|
||||
)
|
||||
or document["state"] != expected_state
|
||||
):
|
||||
raise ValueError("published_source_map_invalid")
|
||||
map_digest = prepared.metadata.digest
|
||||
elif document["state"] != "ready":
|
||||
raise ValueError("published_manifest_invalid")
|
||||
rows = tuple(
|
||||
db.execute(
|
||||
"SELECT * FROM retrieval_chunks WHERE generation_id=? AND evidence_id=? "
|
||||
"AND build_id=? ORDER BY ordinal LIMIT 8193",
|
||||
(document["generation_id"], record.evidence_id, document["build_id"]),
|
||||
).fetchall()
|
||||
)
|
||||
chunks = []
|
||||
for row in rows:
|
||||
if deadline is not None:
|
||||
_check(deadline)
|
||||
start, end = row["start_cp"], row["end_cp"]
|
||||
if (
|
||||
type(start) is not int
|
||||
or type(end) is not int
|
||||
or not 0 <= start < end <= len(record.verbatim)
|
||||
):
|
||||
raise ValueError("published_manifest_invalid")
|
||||
text = record.verbatim[start:end]
|
||||
text_hash = digest(text)
|
||||
identity = [document["build_id"], row["ordinal"], start, end, text_hash]
|
||||
if map_digest is not None:
|
||||
identity.append(map_digest)
|
||||
if row["text_hash"] != text_hash or row["chunk_id"] != digest(identity):
|
||||
raise ValueError("published_manifest_invalid")
|
||||
chunks.append(TextChunk(row["ordinal"], start, end, row["token_count"], text))
|
||||
try:
|
||||
validate_chunk_coverage(record, chunks, prepared)
|
||||
except ValueError as error:
|
||||
raise ValueError("published_manifest_incomplete_or_invalid") from error
|
||||
if deadline is not None:
|
||||
_check(deadline)
|
||||
return rows
|
||||
|
||||
def status(self, generation_id: str) -> dict[str, int]:
|
||||
with self._db() as db:
|
||||
row = db.execute(
|
||||
"SELECT after_seq FROM retrieval_consumers WHERE generation_id=?", (generation_id,)
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise ValueError("generation_missing")
|
||||
result = {
|
||||
"after_seq": int(row[0]),
|
||||
"pending": 0,
|
||||
"leased": 0,
|
||||
"ready": 0,
|
||||
"blocked": 0,
|
||||
"deleted": 0,
|
||||
"excluded": 0,
|
||||
}
|
||||
for state, count in db.execute(
|
||||
"SELECT state,count(*) FROM retrieval_documents "
|
||||
"WHERE generation_id=? GROUP BY state",
|
||||
(generation_id,),
|
||||
):
|
||||
result[state] = count
|
||||
return result
|
||||
@@ -1,200 +0,0 @@
|
||||
"""Explicit ownership of one admitted index worker; no migration or resource close."""
|
||||
|
||||
import asyncio
|
||||
import math
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal, Protocol
|
||||
|
||||
from one_memory.retrieval.index_worker import IndexWorkResult
|
||||
from one_memory.service_levels import ServiceLevel, service_level_context
|
||||
|
||||
|
||||
class IndexWorker(Protocol):
|
||||
def run_once(self, *, cancel: threading.Event | None = None) -> IndexWorkResult: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class IndexSupervisorSnapshot:
|
||||
state: Literal["new", "running", "stopping", "stopped"]
|
||||
last_result: str | None
|
||||
error_code: str | None
|
||||
slices: int
|
||||
failures: int
|
||||
thread_alive: bool
|
||||
|
||||
|
||||
class RetrievalIndexSupervisor:
|
||||
"""Single-use supervisor. Caller owns admission and closes vectors only after stop.
|
||||
|
||||
Cancellation is cooperative: stop/wait cannot finish until a blocking native
|
||||
run_once returns. The lifecycle owner must drain before shutting its event loop.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
worker: IndexWorker,
|
||||
*,
|
||||
idle_wait_seconds: float = 5,
|
||||
failure_wait_seconds: float = 5,
|
||||
yield_wait_seconds: float = 0.05,
|
||||
max_consecutive_slices: int = 8,
|
||||
background_should_yield: Callable[[], bool] | None = None,
|
||||
):
|
||||
for interval in (idle_wait_seconds, failure_wait_seconds, yield_wait_seconds):
|
||||
if (
|
||||
isinstance(interval, bool)
|
||||
or not isinstance(interval, (int, float))
|
||||
or not math.isfinite(interval)
|
||||
or not 0.001 <= interval <= 300
|
||||
):
|
||||
raise ValueError("index_supervisor_invalid_config")
|
||||
if type(max_consecutive_slices) is not int or not 1 <= max_consecutive_slices <= 256:
|
||||
raise ValueError("index_supervisor_invalid_config")
|
||||
self._worker = worker
|
||||
self._idle_wait = idle_wait_seconds
|
||||
self._failure_wait = failure_wait_seconds
|
||||
self._yield_wait = yield_wait_seconds
|
||||
self._max_slices = max_consecutive_slices
|
||||
self._background_should_yield = background_should_yield or (lambda: False)
|
||||
self._condition = threading.Condition()
|
||||
self._cancel = threading.Event()
|
||||
self._wake = False
|
||||
self._thread: threading.Thread | None = None
|
||||
self._state: Literal["new", "running", "stopping", "stopped"] = "new"
|
||||
self._last_result: str | None = None
|
||||
self._error_code: str | None = None
|
||||
self._slices = self._failures = 0
|
||||
|
||||
def start(self) -> None:
|
||||
with self._condition:
|
||||
if self._state != "new":
|
||||
raise RuntimeError("index_supervisor_already_started")
|
||||
thread = threading.Thread(target=self._run, name="retrieval-index", daemon=False)
|
||||
self._thread = thread
|
||||
self._state = "running"
|
||||
try:
|
||||
thread.start()
|
||||
except Exception:
|
||||
self._thread = None
|
||||
self._state = "stopped"
|
||||
self._error_code = "index_supervisor_start_failed"
|
||||
raise RuntimeError("index_supervisor_start_failed") from None
|
||||
|
||||
def trigger(self) -> bool:
|
||||
with self._condition:
|
||||
if self._state != "running":
|
||||
return False
|
||||
self._wake = True
|
||||
self._condition.notify_all()
|
||||
return True
|
||||
|
||||
def snapshot(self) -> IndexSupervisorSnapshot:
|
||||
with self._condition:
|
||||
return IndexSupervisorSnapshot(
|
||||
self._state,
|
||||
self._last_result,
|
||||
self._error_code,
|
||||
self._slices,
|
||||
self._failures,
|
||||
self._thread is not None and self._thread.is_alive(),
|
||||
)
|
||||
|
||||
def _request_stop(self) -> None:
|
||||
with self._condition:
|
||||
self._cancel.set()
|
||||
self._wake = False
|
||||
if self._state == "new":
|
||||
self._state = "stopped"
|
||||
elif self._state == "running":
|
||||
self._state = "stopping"
|
||||
self._condition.notify_all()
|
||||
|
||||
async def stop(self) -> None:
|
||||
self._request_stop()
|
||||
await self.wait()
|
||||
|
||||
async def wait(self) -> None:
|
||||
"""Cancellation requests shutdown, drains the real thread, then propagates."""
|
||||
cancelled = False
|
||||
while True:
|
||||
with self._condition:
|
||||
thread = self._thread
|
||||
if thread is None or not thread.is_alive():
|
||||
if thread is not None:
|
||||
thread.join()
|
||||
break
|
||||
try:
|
||||
# No executor waiter or to_thread future to orphan on cancellation.
|
||||
await asyncio.sleep(0.01)
|
||||
except asyncio.CancelledError:
|
||||
cancelled = True
|
||||
self._request_stop()
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
def _pause(self, seconds: float, *, wakeable: bool) -> None:
|
||||
deadline = time.monotonic() + seconds
|
||||
with self._condition:
|
||||
while not self._cancel.is_set():
|
||||
if wakeable and self._wake:
|
||||
break
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
break
|
||||
self._condition.wait(remaining)
|
||||
self._wake = False
|
||||
|
||||
def _run(self) -> None:
|
||||
consecutive = 0
|
||||
try:
|
||||
while not self._cancel.is_set():
|
||||
if self._background_should_yield():
|
||||
self._pause(self._yield_wait, wakeable=False)
|
||||
continue
|
||||
result: str | None = None
|
||||
error: str | None = None
|
||||
try:
|
||||
with service_level_context(ServiceLevel.BACKGROUND):
|
||||
value = self._worker.run_once(cancel=self._cancel)
|
||||
if not isinstance(value, IndexWorkResult) or value.state not in {
|
||||
"idle",
|
||||
"ready",
|
||||
"excluded",
|
||||
"blocked",
|
||||
"cancelled",
|
||||
"deferred",
|
||||
}:
|
||||
error = "index_worker_invalid_result"
|
||||
else:
|
||||
result = value.state
|
||||
if result == "blocked":
|
||||
error = "index_worker_blocked"
|
||||
except BaseException:
|
||||
# Worker/provider exceptions must never escape via thread stderr.
|
||||
error = "index_worker_failed"
|
||||
with self._condition:
|
||||
self._last_result, self._error_code = result, error
|
||||
self._slices += 1
|
||||
if error is not None:
|
||||
self._failures += 1
|
||||
if self._cancel.is_set():
|
||||
break
|
||||
if error is not None:
|
||||
consecutive = 0
|
||||
self._pause(self._failure_wait, wakeable=False)
|
||||
elif result in {"ready", "excluded"}:
|
||||
consecutive += 1
|
||||
if consecutive >= self._max_slices:
|
||||
consecutive = 0
|
||||
self._pause(self._yield_wait, wakeable=False)
|
||||
else:
|
||||
consecutive = 0
|
||||
self._pause(self._idle_wait, wakeable=True)
|
||||
finally:
|
||||
with self._condition:
|
||||
self._state = "stopped"
|
||||
self._wake = False
|
||||
self._condition.notify_all()
|
||||
@@ -1,262 +0,0 @@
|
||||
"""Bounded raw-index work. No migration, scheduler, HY or remote model invocation.
|
||||
|
||||
Runtime must admit a single worker, own its thread/process, and drain that owner
|
||||
before closing the vector client. Cancellation here is cooperative between
|
||||
native calls, not a claim that ONNX/Chroma can be preempted in a Python thread.
|
||||
"""
|
||||
|
||||
import math
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal, Protocol
|
||||
|
||||
from one_memory.embedding_assets import (
|
||||
BUILTIN_EMBEDDING_MODEL,
|
||||
EMBEDDING_ASSETS,
|
||||
verify_assets,
|
||||
)
|
||||
from one_memory.local_embedding import LocalEmbedder, TokenizedText
|
||||
from one_memory.retrieval.chunking import ChunkingError, chunk_evidence
|
||||
from one_memory.retrieval.content import index_content
|
||||
from one_memory.retrieval.index_store import (
|
||||
EncoderIdentity,
|
||||
IndexLease,
|
||||
RetrievalIndexStore,
|
||||
chunk_identity,
|
||||
)
|
||||
from one_memory.retrieval.mapped_content import MappedContentError, chunk_mapped_content
|
||||
from one_memory.retrieval.vector_store import RawVectorRecord, RawVectorStore, VectorStoreError
|
||||
|
||||
|
||||
class RetrievalEncoder(Protocol):
|
||||
@property
|
||||
def identity(self) -> EncoderIdentity: ...
|
||||
|
||||
def tokenize_with_offsets(self, text: str) -> TokenizedText: ...
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
|
||||
class LocalRetrievalEncoder:
|
||||
"""Pinned local assets; mapped content v2 by default, explicit v1 for old generations."""
|
||||
|
||||
def __init__(self, embedder: LocalEmbedder | None = None, *, content_version: int = 2):
|
||||
if type(content_version) is not int or content_version not in (1, 2):
|
||||
raise ValueError("unsupported_content_generation")
|
||||
self._embedder = embedder if embedder is not None else LocalEmbedder()
|
||||
digests = {asset.filename: asset.sha256 for asset in EMBEDDING_ASSETS}
|
||||
self.identity = EncoderIdentity(
|
||||
BUILTIN_EMBEDDING_MODEL,
|
||||
digests["model_int8.onnx"],
|
||||
digests["tokenizer.json"],
|
||||
self._embedder.dimensions,
|
||||
chunker_version=content_version,
|
||||
normalizer_version=content_version,
|
||||
)
|
||||
self._verified = False
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def _verify(self) -> None:
|
||||
with self._lock:
|
||||
if not self._verified:
|
||||
verify_assets()
|
||||
self._verified = True
|
||||
|
||||
def tokenize_with_offsets(self, text: str) -> TokenizedText:
|
||||
self._verify()
|
||||
return self._embedder.tokenize_with_offsets(text)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
self._verify()
|
||||
if len(texts) > 8 or any(
|
||||
not 1 <= len(self.tokenize_with_offsets(text).token_ids) <= 512 for text in texts
|
||||
):
|
||||
raise ValueError("encoder_batch_or_token_budget")
|
||||
with self._lock:
|
||||
return self._embedder.encode(texts)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class IndexWorkResult:
|
||||
state: Literal["idle", "ready", "excluded", "blocked", "cancelled", "deferred"]
|
||||
reason: str | None = None
|
||||
chunk_count: int = 0
|
||||
|
||||
|
||||
class _Interrupted(Exception):
|
||||
def __init__(self, state: Literal["cancelled", "deferred"]):
|
||||
self.state = state
|
||||
|
||||
|
||||
class RetrievalIndexWorker:
|
||||
def __init__(
|
||||
self,
|
||||
store: RetrievalIndexStore,
|
||||
vectors: RawVectorStore,
|
||||
encoder: RetrievalEncoder,
|
||||
*,
|
||||
generation_id: str,
|
||||
owner: str,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
wall_clock: Callable[[], float] = time.time,
|
||||
max_job_seconds: float = 120,
|
||||
):
|
||||
if (
|
||||
not owner
|
||||
or len(owner) > 128
|
||||
or isinstance(max_job_seconds, bool)
|
||||
or not math.isfinite(max_job_seconds)
|
||||
or not 0 < max_job_seconds <= 300
|
||||
):
|
||||
raise ValueError("invalid_worker_budget")
|
||||
if store.encoder_identity(generation_id) != encoder.identity:
|
||||
raise ValueError("worker_encoder_identity_mismatch")
|
||||
if (
|
||||
vectors.generation_id != generation_id
|
||||
or vectors.dimensions != encoder.identity.dimensions
|
||||
):
|
||||
raise ValueError("worker_vector_identity_mismatch")
|
||||
self.store, self.vectors, self.encoder = store, vectors, encoder
|
||||
self.generation_id, self.owner = generation_id, owner
|
||||
self.clock, self.wall_clock = clock, wall_clock
|
||||
self.max_job_seconds = max_job_seconds
|
||||
self._run_lock = threading.Lock()
|
||||
|
||||
def run_once(self, *, cancel: threading.Event | None = None) -> IndexWorkResult:
|
||||
"""Consume <=256 events and handle <=1 document; never activate a generation.
|
||||
|
||||
Interrupted jobs retain their finite lease for restart recovery. Failed
|
||||
jobs are blocked for explicit retry; arbitrary provider exceptions and
|
||||
original text never enter the public result or a log.
|
||||
"""
|
||||
if not self._run_lock.acquire(blocking=False):
|
||||
raise RuntimeError("index_worker_already_running")
|
||||
lease: IndexLease | None = None
|
||||
deadline = self.clock() + self.max_job_seconds
|
||||
|
||||
def check() -> None:
|
||||
if cancel is not None and cancel.is_set():
|
||||
raise _Interrupted("cancelled")
|
||||
now = self.clock()
|
||||
if not math.isfinite(now) or now >= deadline:
|
||||
raise _Interrupted("deferred")
|
||||
|
||||
def boundary() -> None:
|
||||
nonlocal lease
|
||||
check()
|
||||
assert lease is not None
|
||||
# This revalidates the original's current fingerprint and authority,
|
||||
# not just a heartbeat on an otherwise stale claim.
|
||||
lease = self.store.renew(lease, now=self.wall_clock())
|
||||
check()
|
||||
|
||||
def publish_checkpoint() -> None:
|
||||
check()
|
||||
assert lease is not None
|
||||
now = self.wall_clock()
|
||||
if not math.isfinite(now) or now >= lease.expires_at:
|
||||
raise ValueError("lease_changed_or_expired")
|
||||
|
||||
try:
|
||||
check()
|
||||
self.store.consume_pending(self.generation_id)
|
||||
check()
|
||||
lease = self.store.claim(self.generation_id, self.owner, now=self.wall_clock())
|
||||
if lease is None:
|
||||
return IndexWorkResult("idle")
|
||||
record = self.store.read_lease(lease, now=self.wall_clock())
|
||||
check()
|
||||
identity = self.encoder.identity
|
||||
try:
|
||||
prepared = index_content(
|
||||
record,
|
||||
normalizer_version=identity.normalizer_version,
|
||||
chunker_version=identity.chunker_version,
|
||||
)
|
||||
except ValueError:
|
||||
self.store.block(lease)
|
||||
return IndexWorkResult("blocked", "content_mapping_required")
|
||||
|
||||
worker_encoder = self.encoder
|
||||
|
||||
class CheckedTokenizer:
|
||||
def tokenize_with_offsets(self, text: str) -> TokenizedText:
|
||||
check()
|
||||
result = worker_encoder.tokenize_with_offsets(text)
|
||||
check()
|
||||
return result
|
||||
|
||||
chunks = (
|
||||
chunk_mapped_content(prepared, CheckedTokenizer())
|
||||
if prepared is not None
|
||||
else chunk_evidence(record.verbatim, CheckedTokenizer())
|
||||
)
|
||||
if (not chunks and prepared is None) or len(chunks) > 8192:
|
||||
raise ChunkingError("invalid_chunk_count")
|
||||
source_map = prepared.metadata if prepared is not None else None
|
||||
map_digest = source_map.digest if source_map is not None else None
|
||||
boundary()
|
||||
expected: list[RawVectorRecord] = []
|
||||
for start in range(0, len(chunks), 8):
|
||||
batch = chunks[start : start + 8]
|
||||
boundary()
|
||||
encoded = self.encoder.encode_documents([chunk.text for chunk in batch])
|
||||
boundary()
|
||||
if len(encoded) != len(batch):
|
||||
raise ValueError("encoder_count_mismatch")
|
||||
records = [
|
||||
RawVectorRecord(
|
||||
chunk_identity(lease, chunk, map_digest),
|
||||
lease.build_id,
|
||||
lease.generation_id,
|
||||
tuple(vector),
|
||||
)
|
||||
for chunk, vector in zip(batch, encoded, strict=True)
|
||||
]
|
||||
self.vectors.upsert(records)
|
||||
boundary()
|
||||
expected.extend(records)
|
||||
self._verify(expected, boundary)
|
||||
assert lease is not None
|
||||
check()
|
||||
self.store.publish(
|
||||
lease,
|
||||
chunks,
|
||||
vectors_verified=True,
|
||||
now=self.wall_clock(),
|
||||
source_map=source_map,
|
||||
checkpoint=publish_checkpoint,
|
||||
)
|
||||
return IndexWorkResult("ready" if chunks else "excluded", chunk_count=len(chunks))
|
||||
except _Interrupted as error:
|
||||
return IndexWorkResult(
|
||||
error.state, "cancelled" if error.state == "cancelled" else "deadline"
|
||||
)
|
||||
except Exception as error:
|
||||
cause: BaseException | None = error
|
||||
while cause is not None:
|
||||
if isinstance(cause, _Interrupted):
|
||||
return IndexWorkResult(
|
||||
cause.state, "cancelled" if cause.state == "cancelled" else "deadline"
|
||||
)
|
||||
cause = cause.__cause__
|
||||
if lease is None:
|
||||
raise
|
||||
self.store.block(lease)
|
||||
reason = (
|
||||
"chunking_unavailable"
|
||||
if isinstance(error, (ChunkingError, MappedContentError))
|
||||
else "vector_unavailable"
|
||||
if isinstance(error, VectorStoreError)
|
||||
else "index_build_unavailable"
|
||||
)
|
||||
return IndexWorkResult("blocked", reason)
|
||||
finally:
|
||||
self._run_lock.release()
|
||||
|
||||
def _verify(self, records: Sequence[RawVectorRecord], boundary: Callable[[], None]) -> None:
|
||||
for start in range(0, len(records), 8):
|
||||
self.vectors.verify(records[start : start + 8])
|
||||
boundary()
|
||||
@@ -1,349 +0,0 @@
|
||||
"""Pure raw-content admission and exact source maps; no grants, writes or HY calls."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from typing import Literal
|
||||
|
||||
from one_memory.conversation_processing import _is_hy_input_turn
|
||||
from one_memory.embedding_assets import BUILTIN_EMBEDDING_MAX_TOKENS
|
||||
from one_memory.imports.models import (
|
||||
MAX_VERBATIM_CHARACTERS,
|
||||
EvidenceContentKind,
|
||||
EvidenceRecord,
|
||||
)
|
||||
from one_memory.imports.sources.parsing import (
|
||||
is_host_generated_user_turn,
|
||||
unwrap_host_envelope_with_positions,
|
||||
)
|
||||
from one_memory.input_sanitization import sanitize_model_input_with_spans
|
||||
from one_memory.local_embedding import TokenizedText
|
||||
from one_memory.retrieval.chunking import ChunkingError, OffsetTokenizer, TextChunk, chunk_evidence
|
||||
|
||||
MAP_VERSION = "raw-content-map-v1"
|
||||
PARSER_VERSION = "host-envelope-positions-v1"
|
||||
SANITIZER_VERSION = "model-input-spans-v1"
|
||||
MAX_MAPPED_CHUNKS = 8192
|
||||
|
||||
|
||||
class MappedContentError(ValueError):
|
||||
"""Fixed reason only; never include original or sanitized content."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MappedSegment:
|
||||
kind: Literal["copy", "synthetic"]
|
||||
start_cp: int
|
||||
end_cp: int
|
||||
source_start_cp: int | None
|
||||
source_end_cp: int | None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExcludedRange:
|
||||
start_cp: int
|
||||
end_cp: int
|
||||
reason: Literal["envelope_or_trim", "sanitizer"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ContentMap:
|
||||
map_version: str
|
||||
parser_version: str
|
||||
sanitizer_version: str
|
||||
source_length_cp: int
|
||||
body_start_cp: int | None
|
||||
body_end_cp: int | None
|
||||
sanitized_length_cp: int
|
||||
sanitized_sha256: str
|
||||
segments: tuple[MappedSegment, ...]
|
||||
excluded: tuple[ExcludedRange, ...]
|
||||
|
||||
def canonical_json(self) -> str:
|
||||
return json.dumps(asdict(self), sort_keys=True, separators=(",", ":"), allow_nan=False)
|
||||
|
||||
@property
|
||||
def digest(self) -> str:
|
||||
return hashlib.sha256(self.canonical_json().encode()).hexdigest()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PreparedMappedContent:
|
||||
actual_text: str = field(repr=False)
|
||||
metadata: ContentMap
|
||||
state: Literal["ready", "no_indexable_content"]
|
||||
|
||||
|
||||
def _validate(content: PreparedMappedContent, original: str | None = None) -> None:
|
||||
meta, text = content.metadata, content.actual_text
|
||||
invalid = "mapped_invalid_source_map"
|
||||
if (
|
||||
(meta.map_version, meta.parser_version, meta.sanitizer_version)
|
||||
!= (MAP_VERSION, PARSER_VERSION, SANITIZER_VERSION)
|
||||
or type(meta.source_length_cp) is not int
|
||||
or not 0 <= meta.source_length_cp <= MAX_VERBATIM_CHARACTERS
|
||||
or type(meta.sanitized_length_cp) is not int
|
||||
or meta.sanitized_length_cp != len(text)
|
||||
or meta.sanitized_sha256 != hashlib.sha256(text.encode()).hexdigest()
|
||||
or type(meta.segments) is not tuple
|
||||
or type(meta.excluded) is not tuple
|
||||
):
|
||||
raise MappedContentError(invalid)
|
||||
body_start, body_end = meta.body_start_cp, meta.body_end_cp
|
||||
if body_start is None or body_end is None:
|
||||
if body_start is not None or body_end is not None or text or meta.segments:
|
||||
raise MappedContentError(invalid)
|
||||
elif (
|
||||
type(body_start) is not int
|
||||
or type(body_end) is not int
|
||||
or not 0 <= body_start < body_end <= meta.source_length_cp
|
||||
):
|
||||
raise MappedContentError(invalid)
|
||||
cursor = source_cursor = 0
|
||||
has_content = False
|
||||
partition: list[tuple[int, int]] = []
|
||||
for span in meta.segments:
|
||||
if (
|
||||
not isinstance(span, MappedSegment)
|
||||
or type(span.start_cp) is not int
|
||||
or type(span.end_cp) is not int
|
||||
or not cursor == span.start_cp < span.end_cp <= len(text)
|
||||
):
|
||||
raise MappedContentError(invalid)
|
||||
cursor = span.end_cp
|
||||
if span.kind == "synthetic":
|
||||
if span.source_start_cp is not None or span.source_end_cp is not None:
|
||||
raise MappedContentError(invalid)
|
||||
continue
|
||||
left, right = span.source_start_cp, span.source_end_cp
|
||||
if (
|
||||
span.kind != "copy"
|
||||
or type(left) is not int
|
||||
or type(right) is not int
|
||||
or body_start is None
|
||||
or body_end is None
|
||||
or not body_start <= left < right <= body_end
|
||||
or left < source_cursor
|
||||
or right - left != span.end_cp - span.start_cp
|
||||
):
|
||||
raise MappedContentError(invalid)
|
||||
if original is not None and text[span.start_cp : span.end_cp] != original[left:right]:
|
||||
raise MappedContentError(invalid)
|
||||
has_content |= bool(text[span.start_cp : span.end_cp].strip())
|
||||
source_cursor = right
|
||||
partition.append((left, right))
|
||||
if cursor != len(text):
|
||||
raise MappedContentError(invalid)
|
||||
excluded_end = 0
|
||||
for excluded in meta.excluded:
|
||||
if (
|
||||
not isinstance(excluded, ExcludedRange)
|
||||
or type(excluded.start_cp) is not int
|
||||
or type(excluded.end_cp) is not int
|
||||
or not excluded_end <= excluded.start_cp < excluded.end_cp <= meta.source_length_cp
|
||||
or excluded.reason not in {"envelope_or_trim", "sanitizer"}
|
||||
):
|
||||
raise MappedContentError(invalid)
|
||||
excluded_end = excluded.end_cp
|
||||
partition.append((excluded.start_cp, excluded.end_cp))
|
||||
source_cursor = 0
|
||||
for left, right in sorted(partition):
|
||||
if left != source_cursor:
|
||||
raise MappedContentError(invalid)
|
||||
source_cursor = right
|
||||
if source_cursor != meta.source_length_cp or content.state != (
|
||||
"ready" if has_content else "no_indexable_content"
|
||||
):
|
||||
raise MappedContentError(invalid)
|
||||
|
||||
|
||||
def prepare_mapped_content(record: EvidenceRecord) -> PreparedMappedContent:
|
||||
"""Map a qualified body to Evidence.verbatim code points, without modifying it.
|
||||
|
||||
The returned text is transient; only metadata is suitable for an outer manifest.
|
||||
That manifest must independently bind Evidence identity/revision/hash and grants.
|
||||
"""
|
||||
if record.content_kind is not EvidenceContentKind.RAW_CONVERSATION or record.role not in {
|
||||
"user",
|
||||
"assistant",
|
||||
}:
|
||||
raise MappedContentError("mapped_content_ineligible")
|
||||
original = record.verbatim
|
||||
if len(original) > MAX_VERBATIM_CHARACTERS:
|
||||
raise MappedContentError("mapped_source_limit_exceeded")
|
||||
envelope = unwrap_host_envelope_with_positions(original)
|
||||
if envelope.status == "ambiguous":
|
||||
raise MappedContentError("mapped_ambiguous_envelope")
|
||||
body, start, end = envelope.text, envelope.start_cp, envelope.end_cp
|
||||
segments: list[MappedSegment] = []
|
||||
excluded: list[ExcludedRange] = []
|
||||
text = ""
|
||||
if body is None:
|
||||
if start is not None or end is not None:
|
||||
raise MappedContentError("mapped_invalid_envelope")
|
||||
if original:
|
||||
excluded.append(ExcludedRange(0, len(original), "envelope_or_trim"))
|
||||
else:
|
||||
if (
|
||||
type(start) is not int
|
||||
or type(end) is not int
|
||||
or not 0 <= start < end <= len(original)
|
||||
or original[start:end] != body
|
||||
):
|
||||
raise MappedContentError("mapped_invalid_envelope")
|
||||
if is_host_generated_user_turn(original) or is_host_generated_user_turn(body):
|
||||
raise MappedContentError("mapped_generated_content")
|
||||
if not _is_hy_input_turn(record.model_copy(update={"verbatim": body})):
|
||||
raise MappedContentError("mapped_body_hygiene_rejected")
|
||||
sanitized = sanitize_model_input_with_spans(body)
|
||||
text = sanitized.text
|
||||
cursor = 0
|
||||
for span in sanitized.spans:
|
||||
if span.start_cp > cursor:
|
||||
segments.append(MappedSegment("synthetic", cursor, span.start_cp, None, None))
|
||||
copied = MappedSegment(
|
||||
"copy",
|
||||
span.start_cp,
|
||||
span.end_cp,
|
||||
start + span.source_start_cp,
|
||||
start + span.source_end_cp,
|
||||
)
|
||||
# Only source AND output adjacency can make one citation-safe interval.
|
||||
if (
|
||||
segments
|
||||
and segments[-1].kind == "copy"
|
||||
and (
|
||||
segments[-1].end_cp == copied.start_cp
|
||||
and segments[-1].source_end_cp == copied.source_start_cp
|
||||
)
|
||||
):
|
||||
previous = segments.pop()
|
||||
copied = MappedSegment(
|
||||
"copy",
|
||||
previous.start_cp,
|
||||
copied.end_cp,
|
||||
previous.source_start_cp,
|
||||
copied.source_end_cp,
|
||||
)
|
||||
segments.append(copied)
|
||||
cursor = span.end_cp
|
||||
if cursor < len(text):
|
||||
segments.append(MappedSegment("synthetic", cursor, len(text), None, None))
|
||||
if start:
|
||||
excluded.append(ExcludedRange(0, start, "envelope_or_trim"))
|
||||
excluded.extend(
|
||||
ExcludedRange(start + left, start + right, "sanitizer")
|
||||
for left, right in sanitized.excluded
|
||||
)
|
||||
if end < len(original):
|
||||
excluded.append(ExcludedRange(end, len(original), "envelope_or_trim"))
|
||||
metadata = ContentMap(
|
||||
MAP_VERSION,
|
||||
PARSER_VERSION,
|
||||
SANITIZER_VERSION,
|
||||
len(original),
|
||||
start,
|
||||
end,
|
||||
len(text),
|
||||
hashlib.sha256(text.encode()).hexdigest(),
|
||||
tuple(segments),
|
||||
tuple(excluded),
|
||||
)
|
||||
prepared = PreparedMappedContent(
|
||||
text,
|
||||
metadata,
|
||||
"ready"
|
||||
if any(
|
||||
span.kind == "copy" and text[span.start_cp : span.end_cp].strip() for span in segments
|
||||
)
|
||||
else "no_indexable_content",
|
||||
)
|
||||
_validate(prepared, original)
|
||||
return prepared
|
||||
|
||||
|
||||
class _BoundedTokenizer:
|
||||
"""Count successful chunk re-encodings without copying the chunking algorithm."""
|
||||
|
||||
def __init__(self, tokenizer: OffsetTokenizer, remaining: int, max_tokens: int):
|
||||
self.tokenizer = tokenizer
|
||||
self.remaining = remaining
|
||||
self.max_tokens = max_tokens
|
||||
self.initial = True
|
||||
|
||||
def tokenize_with_offsets(self, text: str) -> TokenizedText:
|
||||
if not self.initial and self.remaining == 0:
|
||||
raise MappedContentError("mapped_chunk_limit_exceeded")
|
||||
result = self.tokenizer.tokenize_with_offsets(text)
|
||||
# chunk_evidence first encodes the whole copy span, then prefix+candidate;
|
||||
# an in-budget candidate emits exactly one chunk. Shrink retries do not.
|
||||
if (
|
||||
not self.initial
|
||||
and isinstance(result, TokenizedText)
|
||||
and len(result.token_ids) <= self.max_tokens
|
||||
):
|
||||
self.remaining -= 1
|
||||
self.initial = False
|
||||
return result
|
||||
|
||||
|
||||
def chunk_mapped_content(
|
||||
content: PreparedMappedContent,
|
||||
tokenizer: OffsetTokenizer,
|
||||
*,
|
||||
prefix: str = "",
|
||||
target_tokens: int = 384,
|
||||
overlap_tokens: int = 64,
|
||||
max_tokens: int = BUILTIN_EMBEDDING_MAX_TOKENS,
|
||||
max_chunks: int = MAX_MAPPED_CHUNKS,
|
||||
) -> tuple[TextChunk, ...]:
|
||||
"""Chunk each continuous copy independently; never encode synthetic placeholders."""
|
||||
if type(max_chunks) is not int or not 1 <= max_chunks <= MAX_MAPPED_CHUNKS:
|
||||
raise MappedContentError("mapped_invalid_chunk_limit")
|
||||
_validate(content)
|
||||
# Reuse budget validation even for an explicitly nonindexable record.
|
||||
chunk_evidence(
|
||||
"",
|
||||
tokenizer,
|
||||
prefix=prefix,
|
||||
target_tokens=target_tokens,
|
||||
overlap_tokens=overlap_tokens,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
chunks: list[TextChunk] = []
|
||||
for span in content.metadata.segments:
|
||||
if span.kind != "copy":
|
||||
continue
|
||||
text = content.actual_text[span.start_cp : span.end_cp]
|
||||
if not text.strip():
|
||||
continue
|
||||
if len(chunks) == max_chunks:
|
||||
raise MappedContentError("mapped_chunk_limit_exceeded")
|
||||
try:
|
||||
pieces = chunk_evidence(
|
||||
text,
|
||||
_BoundedTokenizer(tokenizer, max_chunks - len(chunks), max_tokens),
|
||||
prefix=prefix,
|
||||
target_tokens=target_tokens,
|
||||
overlap_tokens=overlap_tokens,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
except ChunkingError as error:
|
||||
if isinstance(error.__cause__, MappedContentError):
|
||||
raise error.__cause__ from None
|
||||
raise
|
||||
if len(chunks) + len(pieces) > max_chunks:
|
||||
raise MappedContentError("mapped_chunk_limit_exceeded")
|
||||
assert span.source_start_cp is not None # Validated copy span.
|
||||
ordinal = len(chunks)
|
||||
chunks.extend(
|
||||
TextChunk(
|
||||
ordinal + i,
|
||||
span.source_start_cp + piece.start_cp,
|
||||
span.source_start_cp + piece.end_cp,
|
||||
piece.token_count,
|
||||
piece.text,
|
||||
)
|
||||
for i, piece in enumerate(pieces)
|
||||
)
|
||||
return tuple(chunks)
|
||||
@@ -1,147 +0,0 @@
|
||||
"""Internal raw semantic candidate lookup. No route, raw body release or full coverage claim."""
|
||||
|
||||
import math
|
||||
import sqlite3
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
from one_memory.context_raw import _check
|
||||
from one_memory.retrieval.index_store import PublishedChunk, RetrievalIndexStore
|
||||
from one_memory.retrieval.models import RetrievalPath, TimeRange
|
||||
from one_memory.retrieval.vector_store import RawVectorStore, VectorStoreError
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RawSemanticCandidate:
|
||||
chunk: PublishedChunk
|
||||
distance: float
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RawCandidateBatch:
|
||||
hits: tuple[RawSemanticCandidate, ...] = ()
|
||||
complete: bool = False
|
||||
reason: str | None = None
|
||||
coverage: Literal["indexed_only"] = "indexed_only"
|
||||
time_range: TimeRange | None = None
|
||||
|
||||
|
||||
def read_raw_semantic_candidates(
|
||||
store: RetrievalIndexStore,
|
||||
vectors: RawVectorStore,
|
||||
query_vector: Sequence[float],
|
||||
*,
|
||||
generation_id: str,
|
||||
source_ids: frozenset[str],
|
||||
deadline: float,
|
||||
limit: int = 64,
|
||||
time_range: TimeRange | None = None,
|
||||
) -> RawCandidateBatch:
|
||||
"""Caller resolves host authority; explicit source IDs can only narrow it.
|
||||
|
||||
complete means the eligible published chunks were checked and searched, not
|
||||
that every authorized original is indexed. The future coordinator must still
|
||||
revalidate host identity and all raw/memory data together before body release.
|
||||
"""
|
||||
if not isinstance(source_ids, frozenset) or any(
|
||||
not isinstance(source, str) or not 1 <= len(source) <= 128 for source in source_ids
|
||||
):
|
||||
raise ValueError("explicit_resolved_sources_required")
|
||||
if type(limit) is not int or not 1 <= limit <= 64:
|
||||
raise ValueError("invalid_candidate_limit")
|
||||
if isinstance(deadline, bool) or not math.isfinite(deadline):
|
||||
raise ValueError("invalid_candidate_deadline")
|
||||
if not source_ids:
|
||||
return RawCandidateBatch(reason="not_authorized")
|
||||
if vectors.generation_id != generation_id:
|
||||
raise ValueError("candidate_generation_mismatch")
|
||||
try:
|
||||
_check(deadline)
|
||||
eligible = store.active_chunks(
|
||||
generation_id, source_ids=source_ids, time_range=time_range, deadline=deadline
|
||||
)
|
||||
if not eligible:
|
||||
return RawCandidateBatch(
|
||||
complete=True, reason="no_eligible_chunks", time_range=time_range
|
||||
)
|
||||
before = {chunk.chunk_id: chunk for chunk in eligible}
|
||||
expected_vectors = [(chunk.chunk_id, chunk.build_id) for chunk in eligible]
|
||||
# When one query can return the entire allowlist, that response itself
|
||||
# is the integrity check. Avoid a duplicate Chroma get-before-query on
|
||||
# the hot small/medium-authority path. Larger allowlists still need the
|
||||
# explicit complete verification because search returns only top-K.
|
||||
verify_before_search = len(eligible) > limit
|
||||
if verify_before_search and not vectors.verify_candidates(
|
||||
expected_vectors, deadline=deadline
|
||||
):
|
||||
return RawCandidateBatch(reason="deadline")
|
||||
result = vectors.search(query_vector, tuple(before), limit, deadline)
|
||||
if not verify_before_search and {hit.chunk_id for hit in result.hits} != set(before):
|
||||
raise VectorStoreError("vector_record_missing")
|
||||
_check(deadline)
|
||||
current = {
|
||||
chunk.chunk_id: chunk
|
||||
for chunk in store.active_chunks(
|
||||
generation_id, source_ids=source_ids, time_range=time_range, deadline=deadline
|
||||
)
|
||||
}
|
||||
hits = []
|
||||
for hit in result.hits:
|
||||
if (
|
||||
hit.chunk_id not in before
|
||||
or hit.generation_id != generation_id
|
||||
or hit.build_id != before[hit.chunk_id].build_id
|
||||
):
|
||||
raise VectorStoreError("candidate_identity_mismatch")
|
||||
if current.get(hit.chunk_id) == before[hit.chunk_id]:
|
||||
hits.append(RawSemanticCandidate(before[hit.chunk_id], hit.distance))
|
||||
_check(deadline)
|
||||
unchanged = before == current
|
||||
return RawCandidateBatch(
|
||||
tuple(hits),
|
||||
result.complete and unchanged,
|
||||
"evidence_changed" if not unchanged else result.reason,
|
||||
time_range=time_range,
|
||||
)
|
||||
except TimeoutError:
|
||||
return RawCandidateBatch(reason="deadline")
|
||||
except VectorStoreError:
|
||||
return RawCandidateBatch(reason="vector_unavailable")
|
||||
except (sqlite3.Error, OSError, ValueError):
|
||||
return RawCandidateBatch(reason="index_unavailable")
|
||||
|
||||
|
||||
def compose_raw_retrieval_path(
|
||||
batch: RawCandidateBatch,
|
||||
*,
|
||||
unpublished_authorized: bool,
|
||||
lexical_hit_count: int,
|
||||
) -> RetrievalPath:
|
||||
"""Map channel results to a path. Indexed-only search is never complete empty."""
|
||||
if type(lexical_hit_count) is not int or lexical_hit_count < 0:
|
||||
raise ValueError("invalid_lexical_hit_count")
|
||||
hits = len(batch.hits) + lexical_hit_count
|
||||
incomplete = unpublished_authorized or not batch.complete
|
||||
channels: tuple[Literal["lexical", "semantic"], ...] = ("lexical", "semantic")
|
||||
if hits:
|
||||
return RetrievalPath(
|
||||
status="ok",
|
||||
returned_count=min(hits, 32),
|
||||
searched_channels=channels,
|
||||
coverage="partial" if incomplete else "complete",
|
||||
)
|
||||
if incomplete:
|
||||
return RetrievalPath(
|
||||
status="degraded",
|
||||
reason="indexed_only",
|
||||
returned_count=0,
|
||||
searched_channels=channels,
|
||||
coverage="partial",
|
||||
)
|
||||
return RetrievalPath(
|
||||
status="empty",
|
||||
returned_count=0,
|
||||
searched_channels=channels,
|
||||
coverage="complete",
|
||||
)
|
||||
@@ -1,189 +0,0 @@
|
||||
"""Materialize internal semantic candidates from current authorized Evidence.
|
||||
|
||||
Not an arbitrary-ID endpoint or signed expansion capability. The combined v2
|
||||
coordinator must revalidate this lane with memory and host identity at release.
|
||||
"""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC
|
||||
|
||||
from one_memory import ingest_journal as journal
|
||||
from one_memory.context_raw import _check, _connection
|
||||
from one_memory.imports.models import EvidenceRecord
|
||||
from one_memory.imports.storage import EvidenceStore
|
||||
from one_memory.retrieval.index_store import RetrievalIndexStore
|
||||
from one_memory.retrieval.models import EvidenceSpan, RawEvidence
|
||||
from one_memory.retrieval.raw_candidates import RawCandidateBatch
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RawEvidenceBatch:
|
||||
items: tuple[RawEvidence, ...] = ()
|
||||
complete: bool = False
|
||||
reason: str | None = None
|
||||
used_characters: int = 0
|
||||
|
||||
|
||||
def _merge(intervals: list[tuple[int, int]]) -> list[tuple[int, int]]:
|
||||
output: list[tuple[int, int]] = []
|
||||
for start, end in sorted(intervals):
|
||||
if output and start <= output[-1][1]:
|
||||
output[-1] = (output[-1][0], max(end, output[-1][1]))
|
||||
else:
|
||||
output.append((start, end))
|
||||
return output
|
||||
|
||||
|
||||
def materialize_raw_candidates(
|
||||
store: RetrievalIndexStore,
|
||||
candidates: RawCandidateBatch,
|
||||
*,
|
||||
generation_id: str,
|
||||
source_ids: frozenset[str],
|
||||
deadline: float,
|
||||
character_budget: int = 6000,
|
||||
) -> RawEvidenceBatch:
|
||||
if not isinstance(source_ids, frozenset) or any(
|
||||
not isinstance(source, str) or not 1 <= len(source) <= 128 for source in source_ids
|
||||
):
|
||||
raise ValueError("explicit_resolved_sources_required")
|
||||
if type(character_budget) is not int or not 1 <= character_budget <= 20000:
|
||||
raise ValueError("invalid_evidence_budget")
|
||||
if len(candidates.hits) > 64:
|
||||
raise ValueError("candidate_budget_exceeded")
|
||||
if not source_ids:
|
||||
return RawEvidenceBatch(reason="not_authorized")
|
||||
try:
|
||||
_check(deadline)
|
||||
expected = {hit.chunk.chunk_id: hit.chunk for hit in candidates.hits}
|
||||
if len(expected) != len(candidates.hits):
|
||||
raise ValueError("duplicate_candidates")
|
||||
if not expected:
|
||||
# Empty/excluded builds still carry an integrity claim. Validate
|
||||
# current authorized completed manifests, never only nonempty hits.
|
||||
current_chunks = store.active_chunks(
|
||||
generation_id,
|
||||
source_ids=source_ids,
|
||||
time_range=candidates.time_range,
|
||||
deadline=deadline,
|
||||
require_active=True,
|
||||
)
|
||||
_check(deadline)
|
||||
if current_chunks and candidates.reason == "no_eligible_chunks":
|
||||
return RawEvidenceBatch(reason="evidence_changed")
|
||||
return RawEvidenceBatch(complete=candidates.complete, reason=candidates.reason)
|
||||
changed = False
|
||||
grouped: dict[str, tuple[EvidenceRecord, list[tuple[int, int]]]] = {}
|
||||
with _connection(store.database_path, deadline) as db:
|
||||
db.execute("PRAGMA query_only=ON")
|
||||
rows = {
|
||||
row["chunk_id"]: row
|
||||
for row in db.execute(
|
||||
"""SELECT c.*, d.event_seq,d.grant_fingerprint,d.source_id
|
||||
FROM retrieval_chunks c JOIN retrieval_documents d
|
||||
ON d.generation_id=c.generation_id AND d.evidence_id=c.evidence_id
|
||||
AND d.active_build_id=c.build_id AND d.build_id=c.build_id
|
||||
JOIN retrieval_index_generations g ON g.generation_id=c.generation_id
|
||||
WHERE c.generation_id=? AND g.state='active' AND d.state='ready'
|
||||
AND c.chunk_id IN (SELECT value FROM json_each(?))""",
|
||||
(generation_id, json.dumps(tuple(expected))),
|
||||
)
|
||||
}
|
||||
for key, chunk in expected.items():
|
||||
_check(deadline)
|
||||
row = rows.get(key)
|
||||
if row is None or row["source_id"] not in source_ids:
|
||||
changed = True
|
||||
continue
|
||||
event = journal.event(db, row["event_seq"])
|
||||
grant = store._grant(db, row["source_id"])
|
||||
if (
|
||||
event is None
|
||||
or event.user_id != "local"
|
||||
or event.source_id != row["source_id"]
|
||||
or event.evidence_id != chunk.evidence_id
|
||||
or event.evidence_revision != chunk.revision
|
||||
or grant is None
|
||||
or grant != row["grant_fingerprint"]
|
||||
or not journal.matches_current(db, event)
|
||||
):
|
||||
changed = True
|
||||
continue
|
||||
record = EvidenceStore._get_with_connection(db, chunk.evidence_id)
|
||||
if (
|
||||
record is None
|
||||
or record.user_id != "local"
|
||||
or record.source_id != row["source_id"]
|
||||
):
|
||||
changed = True
|
||||
continue
|
||||
event_time = record.created_at
|
||||
if event_time is not None:
|
||||
event_time = event_time.astimezone(UTC) if event_time.tzinfo else None
|
||||
if (
|
||||
(row["build_id"], row["ordinal"], row["start_cp"], row["end_cp"])
|
||||
!= (chunk.build_id, chunk.ordinal, chunk.start_cp, chunk.end_cp)
|
||||
or (record.source_id, event_time) != (chunk.source_id, chunk.event_time)
|
||||
or not 0 <= chunk.start_cp < chunk.end_cp <= len(record.verbatim)
|
||||
):
|
||||
changed = True
|
||||
continue
|
||||
if record.evidence_id not in grouped:
|
||||
document = db.execute(
|
||||
"SELECT * FROM retrieval_documents WHERE generation_id=? AND evidence_id=?",
|
||||
(generation_id, record.evidence_id),
|
||||
).fetchone()
|
||||
if document is None:
|
||||
raise ValueError("published_document_missing")
|
||||
store.validate_document(db, document, record, deadline=deadline)
|
||||
grouped[record.evidence_id] = (record, [])
|
||||
grouped[record.evidence_id][1].append((chunk.start_cp, chunk.end_cp))
|
||||
items = []
|
||||
used = 0
|
||||
truncated = False
|
||||
for record, intervals in grouped.values():
|
||||
_check(deadline)
|
||||
spans: list[EvidenceSpan] = []
|
||||
for start, end in _merge(intervals):
|
||||
if len(spans) >= 4 or end - start > min(20000, character_budget - used):
|
||||
truncated = True
|
||||
continue
|
||||
spans.append(
|
||||
EvidenceSpan(start_cp=start, end_cp=end, text=record.verbatim[start:end])
|
||||
)
|
||||
used += end - start
|
||||
if not spans:
|
||||
continue
|
||||
event_time = record.created_at
|
||||
if event_time is not None:
|
||||
event_time = event_time.astimezone(UTC) if event_time.tzinfo else None
|
||||
items.append(
|
||||
RawEvidence(
|
||||
evidence_id=record.evidence_id,
|
||||
revision=record.revision,
|
||||
source_id=record.source_id,
|
||||
conversation_id=record.conversation_id,
|
||||
role=record.role,
|
||||
written_by="person" if record.role == "user" else "agent",
|
||||
event_time=event_time,
|
||||
spans=tuple(spans),
|
||||
matched_by=("semantic",),
|
||||
)
|
||||
)
|
||||
_check(deadline)
|
||||
return RawEvidenceBatch(
|
||||
tuple(items),
|
||||
candidates.complete and not changed and not truncated,
|
||||
"evidence_changed"
|
||||
if changed
|
||||
else "character_budget"
|
||||
if truncated
|
||||
else candidates.reason,
|
||||
used,
|
||||
)
|
||||
except TimeoutError:
|
||||
return RawEvidenceBatch(reason="deadline")
|
||||
except (sqlite3.Error, OSError, ValueError):
|
||||
return RawEvidenceBatch(reason="index_unavailable")
|
||||
@@ -1,181 +0,0 @@
|
||||
"""Fuse lexical raw context with an optional derived index. No new HTTP schema."""
|
||||
|
||||
import sqlite3
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Protocol
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from one_memory.context_models import (
|
||||
ContextLaneResult,
|
||||
ContextSearchRequest,
|
||||
ContextSemanticStatus,
|
||||
RawContextItem,
|
||||
)
|
||||
from one_memory.context_raw import (
|
||||
_METADATA,
|
||||
_check,
|
||||
_connection,
|
||||
_execute,
|
||||
_message,
|
||||
_readable,
|
||||
read_raw_context,
|
||||
)
|
||||
from one_memory.retrieval.models import TimeRange
|
||||
from one_memory.retrieval.raw_candidates import RawCandidateBatch, compose_raw_retrieval_path
|
||||
|
||||
|
||||
class RawIndexLookup(Protocol):
|
||||
def lookup_raw(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
source_ids: frozenset[str] | None,
|
||||
deadline: float,
|
||||
time_range: TimeRange | None = None,
|
||||
) -> tuple[RawCandidateBatch, bool]: ...
|
||||
|
||||
|
||||
def fuse_raw_lane(
|
||||
lexical: ContextLaneResult,
|
||||
batch: RawCandidateBatch,
|
||||
*,
|
||||
unpublished_authorized: bool,
|
||||
extra_items: tuple[RawContextItem, ...] = (),
|
||||
) -> ContextLaneResult:
|
||||
"""v1 has no coverage field; degraded + lexical_only stands in for partial index."""
|
||||
semantic_status: ContextSemanticStatus = (
|
||||
"timed_out"
|
||||
if batch.reason == "deadline"
|
||||
else "failed"
|
||||
if batch.reason in {"index_unavailable", "vector_unavailable"}
|
||||
else "completed"
|
||||
)
|
||||
index_complete = batch.complete and not unpublished_authorized
|
||||
timed_out = semantic_status == "timed_out"
|
||||
if lexical.status not in {"ok", "empty", "degraded"}:
|
||||
return lexical.model_copy(
|
||||
update={
|
||||
"searched_channels": ("lexical", "semantic"),
|
||||
"semantic_status": semantic_status,
|
||||
"index_complete": index_complete,
|
||||
"timed_out": timed_out,
|
||||
}
|
||||
)
|
||||
seen = {item.anchor_evidence_id for item in lexical.items if isinstance(item, RawContextItem)}
|
||||
extras = tuple(item for item in extra_items if isinstance(item, RawContextItem))
|
||||
# Semantic similarity is a candidate signal, not a verified recall result.
|
||||
# Release only semantic items that the lexical lane independently confirmed
|
||||
# for the same evidence identity. Unconfirmed semantic hits stay in the
|
||||
# coverage receipt as partial/degraded and must not reach host injection.
|
||||
unconfirmed_semantic = any(item.anchor_evidence_id not in seen for item in extras)
|
||||
items = lexical.items
|
||||
path = compose_raw_retrieval_path(
|
||||
batch,
|
||||
unpublished_authorized=unpublished_authorized,
|
||||
lexical_hit_count=len(lexical.items),
|
||||
)
|
||||
status = path.status
|
||||
if unconfirmed_semantic:
|
||||
status = "degraded"
|
||||
if lexical.status == "degraded" and status == "ok":
|
||||
status = "degraded"
|
||||
return ContextLaneResult(
|
||||
status=status,
|
||||
reason=(
|
||||
"lexical_only" if unconfirmed_semantic and lexical.reason is None else lexical.reason
|
||||
),
|
||||
items=items[:21],
|
||||
truncated_by=lexical.truncated_by,
|
||||
searched_channels=("lexical", "semantic"),
|
||||
semantic_status=semantic_status,
|
||||
index_complete=index_complete,
|
||||
timed_out=timed_out,
|
||||
)
|
||||
|
||||
|
||||
def load_raw_items_from_candidates(
|
||||
database_path: Path,
|
||||
batch: RawCandidateBatch,
|
||||
*,
|
||||
source_ids: frozenset[str] | None,
|
||||
deadline: float,
|
||||
character_budget: int = 6000,
|
||||
) -> tuple[RawContextItem, ...]:
|
||||
"""Release authorized Evidence bodies for semantic IDs. Not a vector dump."""
|
||||
ids = tuple(dict.fromkeys(hit.chunk.evidence_id for hit in batch.hits))
|
||||
if not ids or source_ids == frozenset():
|
||||
return ()
|
||||
readable, authority = _readable(source_ids)
|
||||
body_limit = min(20000, max(1, character_budget))
|
||||
items: list[RawContextItem] = []
|
||||
try:
|
||||
with _connection(database_path, deadline) as db:
|
||||
for evidence_id in ids[:21]:
|
||||
_check(deadline)
|
||||
row = _execute(
|
||||
db,
|
||||
deadline,
|
||||
f"""
|
||||
SELECT {_METADATA}, substr(e.verbatim,1,?) AS text,
|
||||
length(e.verbatim)>? AS body_truncated
|
||||
FROM evidence e
|
||||
WHERE {readable} AND e.evidence_id=?
|
||||
""",
|
||||
[body_limit, body_limit, *authority, evidence_id],
|
||||
).fetchone()
|
||||
if row is None:
|
||||
continue
|
||||
try:
|
||||
message = _message(row)
|
||||
items.append(
|
||||
RawContextItem(
|
||||
source_id=row["source_id"],
|
||||
conversation_id=row["conversation_id"],
|
||||
anchor_evidence_id=message.evidence_id,
|
||||
messages=(message,),
|
||||
truncated=message.truncated,
|
||||
)
|
||||
)
|
||||
except (ValidationError, ValueError, TypeError):
|
||||
continue
|
||||
except TimeoutError:
|
||||
raise
|
||||
except (sqlite3.Error, OSError):
|
||||
return ()
|
||||
return tuple(items)
|
||||
|
||||
|
||||
def read_raw_indexed_context(
|
||||
database_path: Path,
|
||||
request: ContextSearchRequest,
|
||||
*,
|
||||
source_ids: frozenset[str] | None,
|
||||
deadline: float,
|
||||
index: RawIndexLookup | None = None,
|
||||
load_semantic_items: Callable[..., tuple[RawContextItem, ...]] | None = None,
|
||||
) -> ContextLaneResult:
|
||||
lexical = read_raw_context(database_path, request, source_ids=source_ids, deadline=deadline)
|
||||
if index is None:
|
||||
return lexical
|
||||
try:
|
||||
batch, unpublished = index.lookup_raw(
|
||||
request.query, source_ids=source_ids, deadline=deadline
|
||||
)
|
||||
except Exception:
|
||||
batch, unpublished = RawCandidateBatch(reason="index_unavailable"), True
|
||||
extra: tuple[RawContextItem, ...] = ()
|
||||
if batch.hits:
|
||||
extra = (
|
||||
load_semantic_items(database_path, batch, source_ids=source_ids, deadline=deadline)
|
||||
if load_semantic_items is not None
|
||||
else load_raw_items_from_candidates(
|
||||
database_path,
|
||||
batch,
|
||||
source_ids=source_ids,
|
||||
deadline=deadline,
|
||||
character_budget=request.character_budget,
|
||||
)
|
||||
)
|
||||
return fuse_raw_lane(lexical, batch, unpublished_authorized=unpublished, extra_items=extra)
|
||||
@@ -1,379 +0,0 @@
|
||||
"""Resume an explicitly configured local index; never create or migrate one."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
import stat
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.retrieval.index_store import EncoderIdentity, RetrievalIndexStore
|
||||
from one_memory.retrieval.index_supervisor import IndexSupervisorSnapshot, RetrievalIndexSupervisor
|
||||
from one_memory.retrieval.index_worker import LocalRetrievalEncoder, RetrievalIndexWorker
|
||||
from one_memory.retrieval.models import TimeRange
|
||||
from one_memory.retrieval.raw_candidates import RawCandidateBatch, read_raw_semantic_candidates
|
||||
from one_memory.retrieval.vector_store import RawVectorStore, VectorStoreError
|
||||
from one_memory.service_levels import ServiceLevel, ServiceLevelScheduler, ServiceResource
|
||||
|
||||
type RuntimeState = Literal[
|
||||
"new", "starting", "running", "not_initialized", "failed", "stopping", "stopped"
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExistingRetrievalSnapshot:
|
||||
state: RuntimeState
|
||||
error_code: str | None
|
||||
generation_id: str | None
|
||||
supervisor: IndexSupervisorSnapshot | None
|
||||
|
||||
|
||||
class _AdmissionError(Exception):
|
||||
def __init__(self, code: str, *, missing: bool = False):
|
||||
self.code, self.missing = code, missing
|
||||
|
||||
|
||||
async def _native[T](operation: Callable[[], T]) -> T:
|
||||
"""Owned, joined native call. Only retained internal lifecycle tasks invoke this."""
|
||||
result: list[T] = []
|
||||
errors: list[BaseException] = []
|
||||
|
||||
def run() -> None:
|
||||
try:
|
||||
result.append(operation())
|
||||
except BaseException as error:
|
||||
errors.append(error)
|
||||
|
||||
thread = threading.Thread(target=run, name="retrieval-resource", daemon=False)
|
||||
thread.start()
|
||||
while thread.is_alive():
|
||||
await asyncio.sleep(0.01)
|
||||
thread.join()
|
||||
if errors:
|
||||
raise errors[0]
|
||||
return result[0]
|
||||
|
||||
|
||||
class ExistingRetrievalRuntime:
|
||||
"""Caller guarantees managed Core single ownership and explicit configuration.
|
||||
|
||||
Construct without I/O; start and stop on the owning event loop. A retained
|
||||
lifecycle task shields native work from caller cancellation until cleanup.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_dir: Path,
|
||||
*,
|
||||
minimum_free_bytes: int = 256 * 1024 * 1024,
|
||||
service_levels: ServiceLevelScheduler | None = None,
|
||||
):
|
||||
if type(minimum_free_bytes) is not int or minimum_free_bytes < 0:
|
||||
raise ValueError("retrieval_runtime_invalid_disk_limit")
|
||||
self._directory = Path(data_dir)
|
||||
self._minimum_free = minimum_free_bytes
|
||||
self._service_levels = service_levels or ServiceLevelScheduler()
|
||||
self._state: ExistingRetrievalSnapshot = ExistingRetrievalSnapshot("new", None, None, None)
|
||||
self._vectors: RawVectorStore | None = None
|
||||
self._store: RetrievalIndexStore | None = None
|
||||
self._encoder: LocalRetrievalEncoder | None = None
|
||||
self._supervisor: RetrievalIndexSupervisor | None = None
|
||||
self._last_supervisor_snapshot: IndexSupervisorSnapshot | None = None
|
||||
self._start_task: asyncio.Task[bool] | None = None
|
||||
self._stop_task: asyncio.Task[None] | None = None
|
||||
self._stop_requested = False
|
||||
|
||||
def snapshot(self) -> ExistingRetrievalSnapshot:
|
||||
return ExistingRetrievalSnapshot(
|
||||
self._state.state,
|
||||
self._state.error_code,
|
||||
self._state.generation_id,
|
||||
self._supervisor.snapshot()
|
||||
if self._supervisor is not None
|
||||
else self._last_supervisor_snapshot,
|
||||
)
|
||||
|
||||
def _set(
|
||||
self, state: RuntimeState, code: str | None = None, generation: str | None = None
|
||||
) -> None:
|
||||
self._state = ExistingRetrievalSnapshot(
|
||||
state,
|
||||
code,
|
||||
generation if generation is not None else self._state.generation_id,
|
||||
None,
|
||||
)
|
||||
|
||||
def trigger(self) -> bool:
|
||||
return bool(
|
||||
self._state.state == "running" and self._supervisor and self._supervisor.trigger()
|
||||
)
|
||||
|
||||
def lookup_raw(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
source_ids: frozenset[str] | None,
|
||||
deadline: float,
|
||||
time_range: TimeRange | None = None,
|
||||
) -> tuple[RawCandidateBatch, bool]:
|
||||
"""Read-only candidate lookup. Does not release original bodies."""
|
||||
store, encoder, vectors, generation = (
|
||||
self._store,
|
||||
self._encoder,
|
||||
self._vectors,
|
||||
self._state.generation_id,
|
||||
)
|
||||
if (
|
||||
self._state.state != "running"
|
||||
or store is None
|
||||
or encoder is None
|
||||
or vectors is None
|
||||
or generation is None
|
||||
):
|
||||
return RawCandidateBatch(reason="index_unavailable"), True
|
||||
counts = store.status(generation)
|
||||
unpublished = counts["pending"] + counts["leased"] + counts["blocked"] > 0
|
||||
resolved = source_ids
|
||||
if resolved is None:
|
||||
resolved = frozenset(
|
||||
chunk.source_id for chunk in store.active_chunks(generation, deadline=deadline)
|
||||
)
|
||||
if not resolved:
|
||||
return (
|
||||
RawCandidateBatch(complete=not unpublished, reason="no_eligible_chunks"),
|
||||
unpublished,
|
||||
)
|
||||
try:
|
||||
vector = encoder.encode_documents([query.strip()])[0]
|
||||
except Exception:
|
||||
return RawCandidateBatch(reason="index_unavailable"), True
|
||||
batch = read_raw_semantic_candidates(
|
||||
store,
|
||||
vectors,
|
||||
vector,
|
||||
generation_id=generation,
|
||||
source_ids=resolved,
|
||||
deadline=deadline,
|
||||
time_range=time_range,
|
||||
)
|
||||
return batch, unpublished
|
||||
|
||||
async def start(self) -> bool:
|
||||
if self._state.state != "new" or self._stop_requested:
|
||||
raise RuntimeError("retrieval_runtime_already_started")
|
||||
self._set("starting")
|
||||
self._start_task = asyncio.create_task(self._start())
|
||||
try:
|
||||
return await asyncio.shield(self._start_task)
|
||||
except asyncio.CancelledError:
|
||||
# Retain and drain opening resources, then release them, before cancellation escapes.
|
||||
await self.stop()
|
||||
raise
|
||||
|
||||
async def stop(self) -> None:
|
||||
self._stop_requested = True
|
||||
if self._stop_task is None:
|
||||
self._stop_task = asyncio.create_task(self._stop())
|
||||
cancelled = False
|
||||
while not self._stop_task.done():
|
||||
try:
|
||||
await asyncio.shield(self._stop_task)
|
||||
except asyncio.CancelledError:
|
||||
cancelled = True
|
||||
self._stop_task.result()
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
async def _start(self) -> bool:
|
||||
try:
|
||||
store, encoder, vectors, generation = await _native(self._open)
|
||||
self._store = store
|
||||
self._encoder = encoder
|
||||
self._vectors = vectors
|
||||
self._set("starting", generation=generation)
|
||||
if self._stop_requested:
|
||||
return False
|
||||
worker = RetrievalIndexWorker(
|
||||
store, vectors, encoder, generation_id=generation, owner="managed-core-retrieval"
|
||||
)
|
||||
self._supervisor = RetrievalIndexSupervisor(
|
||||
worker,
|
||||
background_should_yield=lambda: self._service_levels.should_yield(
|
||||
ServiceResource.INDEX,
|
||||
ServiceLevel.BACKGROUND,
|
||||
),
|
||||
)
|
||||
self._supervisor.start()
|
||||
self._set("running")
|
||||
return True
|
||||
except _AdmissionError as error:
|
||||
self._set("not_initialized" if error.missing else "failed", error.code)
|
||||
except Exception:
|
||||
self._set("failed", "retrieval_runtime_start_failed")
|
||||
await self._release()
|
||||
return False
|
||||
|
||||
async def _release(self) -> None:
|
||||
stop_failed = False
|
||||
if self._supervisor is not None:
|
||||
try:
|
||||
await self._supervisor.stop()
|
||||
except Exception:
|
||||
stop_failed = True
|
||||
supervisor_snapshot = self._supervisor.snapshot()
|
||||
if supervisor_snapshot.thread_alive:
|
||||
self._set("failed", "retrieval_runtime_stop_failed")
|
||||
raise RuntimeError("retrieval_runtime_stop_failed") from None
|
||||
# The supervisor owns the worker, which owns the native vector
|
||||
# client. Retain only the immutable receipt once its thread drains.
|
||||
self._last_supervisor_snapshot = supervisor_snapshot
|
||||
self._supervisor = None
|
||||
if self._vectors is not None:
|
||||
try:
|
||||
await _native(self._vectors.close)
|
||||
except Exception:
|
||||
self._set("failed", "retrieval_runtime_close_failed")
|
||||
raise RuntimeError("retrieval_runtime_close_failed") from None
|
||||
self._vectors = None
|
||||
self._store = None
|
||||
self._encoder = None
|
||||
if stop_failed:
|
||||
self._set("failed", "retrieval_runtime_stop_failed")
|
||||
raise RuntimeError("retrieval_runtime_stop_failed") from None
|
||||
|
||||
async def _stop(self) -> None:
|
||||
if self._start_task is not None:
|
||||
try:
|
||||
await asyncio.shield(self._start_task)
|
||||
except Exception:
|
||||
# Startup can itself fail during cleanup. Do not lose resource ownership.
|
||||
self._set("failed", "retrieval_runtime_start_failed")
|
||||
self._set("stopping", self._state.error_code)
|
||||
await self._release()
|
||||
if self._state.error_code == "retrieval_runtime_close_failed":
|
||||
return
|
||||
self._set("stopped", self._state.error_code)
|
||||
|
||||
@staticmethod
|
||||
def _path(path: Path, *, directory: bool) -> tuple[int, int]:
|
||||
if any(part.is_symlink() for part in (path, *path.parents)):
|
||||
raise _AdmissionError("retrieval_runtime_unsafe_path")
|
||||
try:
|
||||
info = path.stat()
|
||||
except FileNotFoundError:
|
||||
raise _AdmissionError("retrieval_runtime_not_initialized", missing=True) from None
|
||||
if not (stat.S_ISDIR(info.st_mode) if directory else stat.S_ISREG(info.st_mode)):
|
||||
raise _AdmissionError("retrieval_runtime_unsafe_path")
|
||||
if not directory and info.st_nlink != 1:
|
||||
raise _AdmissionError("retrieval_runtime_unsafe_path")
|
||||
if info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) & 0o022:
|
||||
raise _AdmissionError("retrieval_runtime_unsafe_path")
|
||||
return info.st_dev, info.st_ino
|
||||
|
||||
def _sidecars(self, database: Path) -> None:
|
||||
for suffix in ("-wal", "-shm", "-journal"):
|
||||
path = database.with_name(database.name + suffix)
|
||||
if path.is_symlink() or path.exists():
|
||||
self._path(path, directory=False)
|
||||
|
||||
def _inspect(self, database: Path) -> tuple[RetrievalIndexStore, EncoderIdentity, str]:
|
||||
with main_database(database.absolute()).read(busy_timeout_ms=1_000) as db:
|
||||
db.row_factory = sqlite3.Row
|
||||
db.execute("PRAGMA query_only=ON")
|
||||
db.execute("BEGIN")
|
||||
if not db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE name='retrieval_index_state'"
|
||||
).fetchone():
|
||||
raise _AdmissionError("retrieval_runtime_not_initialized", missing=True)
|
||||
state = db.execute(
|
||||
"SELECT schema_version FROM retrieval_index_state WHERE singleton=1"
|
||||
).fetchone()
|
||||
if state is not None and state[0] == 1:
|
||||
raise _AdmissionError("retrieval_runtime_not_initialized", missing=True)
|
||||
if state is None or state[0] != 2:
|
||||
raise _AdmissionError("retrieval_runtime_invalid_index")
|
||||
rows = db.execute(
|
||||
"SELECT generation_id,identity_json FROM retrieval_index_generations "
|
||||
"WHERE state='active' LIMIT 2"
|
||||
).fetchall()
|
||||
if not rows:
|
||||
raise _AdmissionError("retrieval_runtime_not_initialized", missing=True)
|
||||
if len(rows) != 1:
|
||||
raise _AdmissionError("retrieval_runtime_invalid_index")
|
||||
store = RetrievalIndexStore(database)
|
||||
generation = rows[0]["generation_id"]
|
||||
identity = store._identity(db, generation)
|
||||
raw = json.loads(rows[0]["identity_json"])
|
||||
if raw != asdict(identity) or set(raw) != set(asdict(identity)):
|
||||
raise _AdmissionError("retrieval_runtime_invalid_index")
|
||||
# Check required schema shape without repairing it or reading original bodies.
|
||||
db.execute(
|
||||
"SELECT build_id,source_map_json,map_digest,coverage_state "
|
||||
"FROM retrieval_build_maps LIMIT 0"
|
||||
)
|
||||
db.execute("SELECT after_seq,retry_after_seq FROM retrieval_consumers LIMIT 0")
|
||||
db.execute("SELECT lease_token,active_build_id FROM retrieval_documents LIMIT 0")
|
||||
db.execute("SELECT chunk_id,text_hash FROM retrieval_chunks LIMIT 0")
|
||||
return store, identity, generation
|
||||
|
||||
def _open(self) -> tuple[RetrievalIndexStore, LocalRetrievalEncoder, RawVectorStore, str]:
|
||||
if sys.platform == "win32" or not hasattr(os, "getuid"):
|
||||
raise _AdmissionError("retrieval_runtime_platform_unsupported")
|
||||
directory = self._directory.absolute()
|
||||
database, vector_dir = directory / "one-memory.sqlite3", directory / "retrieval-vectors"
|
||||
root_pin = self._path(directory, directory=True)
|
||||
database_pin = self._path(database, directory=False)
|
||||
self._sidecars(database)
|
||||
try:
|
||||
store, identity, generation = self._inspect(database)
|
||||
except (ValueError, TypeError, sqlite3.Error):
|
||||
raise _AdmissionError("retrieval_runtime_invalid_index") from None
|
||||
vector_pin = self._path(vector_dir, directory=True)
|
||||
vector_database = vector_dir / "chroma.sqlite3"
|
||||
vector_db_pin = self._path(vector_database, directory=False)
|
||||
self._sidecars(vector_database)
|
||||
|
||||
def recheck() -> None:
|
||||
if (
|
||||
self._path(directory, directory=True) != root_pin
|
||||
or self._path(database, directory=False) != database_pin
|
||||
or self._path(vector_dir, directory=True) != vector_pin
|
||||
or self._path(vector_database, directory=False) != vector_db_pin
|
||||
):
|
||||
raise _AdmissionError("retrieval_runtime_index_changed")
|
||||
self._sidecars(database)
|
||||
self._sidecars(vector_database)
|
||||
|
||||
if shutil.disk_usage(directory).free < self._minimum_free:
|
||||
raise _AdmissionError("retrieval_runtime_low_disk")
|
||||
encoder = LocalRetrievalEncoder(content_version=identity.normalizer_version)
|
||||
if encoder.identity != identity:
|
||||
raise _AdmissionError("retrieval_runtime_encoder_mismatch")
|
||||
recheck()
|
||||
try:
|
||||
vectors = RawVectorStore.open_writer_existing(
|
||||
vector_dir, generation_id=generation, dimensions=identity.dimensions
|
||||
)
|
||||
# Retain ownership immediately so all later failures use the same cleanup path.
|
||||
self._vectors = vectors
|
||||
# A generation change during the native open is not authorization to resume another.
|
||||
recheck()
|
||||
if self._inspect(database)[1:] != (identity, generation):
|
||||
raise _AdmissionError("retrieval_runtime_index_changed")
|
||||
recheck()
|
||||
return store, encoder, vectors, generation
|
||||
except BaseException as error:
|
||||
if isinstance(error, VectorStoreError):
|
||||
if str(error) in {"vector_store_not_found", "vector_collection_not_found"}:
|
||||
raise _AdmissionError(
|
||||
"retrieval_runtime_not_initialized", missing=True
|
||||
) from None
|
||||
raise _AdmissionError("retrieval_runtime_vector_unavailable") from None
|
||||
raise
|
||||
@@ -1,534 +0,0 @@
|
||||
"""Independent raw-vector adapter. Eligibility is supplied by Core's SQLite snapshot.
|
||||
|
||||
No HY collection, document text, embedding function, migration/repair or authority
|
||||
store lives here. Chroma PersistentClient is NOT physically read-only: even the
|
||||
existing-collection handle can maintain backend metadata/index files. A caller
|
||||
requiring filesystem read-only operation must not use this backend.
|
||||
|
||||
Scoring is an exact pass over the eligible vectors, not a global ANN top-K that
|
||||
is filtered afterwards. Chroma's filtered query does not score every candidate
|
||||
in its `where $in` allowlist: at 10k stored chunks an explicit 64-ID allowlist
|
||||
returned 26 of 64 while reporting the query as complete, and 0 of 64 after the
|
||||
writer handle was reopened. A post-filtered global search would be neither
|
||||
permission-correct nor recall-correct, so the eligible embeddings are read
|
||||
back and ranked here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
import sqlite3
|
||||
import struct
|
||||
import threading
|
||||
from collections.abc import Callable, Sequence
|
||||
from contextlib import closing
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from numbers import Real
|
||||
from pathlib import Path
|
||||
from time import monotonic
|
||||
from typing import Any, Literal
|
||||
|
||||
ALLOWLIST_BATCH_SIZE = 2048
|
||||
|
||||
|
||||
class VectorStoreError(RuntimeError):
|
||||
"""Bounded failure code; callers must not treat corruption as complete empty search."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RawVectorRecord:
|
||||
chunk_id: str
|
||||
build_id: str
|
||||
generation_id: str
|
||||
vector: tuple[float, ...]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "vector", tuple(self.vector))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VectorHit:
|
||||
chunk_id: str
|
||||
build_id: str
|
||||
generation_id: str
|
||||
distance: float
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VectorSearchResult:
|
||||
hits: tuple[VectorHit, ...]
|
||||
complete: bool
|
||||
eligible_count: int
|
||||
searched_count: int
|
||||
reason: Literal["deadline"] | None
|
||||
|
||||
|
||||
class RawVectorStore:
|
||||
"""Unopened constructor has no I/O. Use create/open_existing explicitly.
|
||||
|
||||
Directory must be a dedicated, trusted Runtime retrieval-vectors directory,
|
||||
never the HY data directory. Only create/open_writer_existing enable upsert.
|
||||
Core must serialize writers across handles/processes; this instance's lock
|
||||
is not a persistent writer lease. close() releases only this Chroma client,
|
||||
retaining other live clients' references.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
directory: Path,
|
||||
*,
|
||||
generation_id: str,
|
||||
dimensions: int,
|
||||
clock: Callable[[], float] = monotonic,
|
||||
) -> None:
|
||||
_identifier(generation_id)
|
||||
if type(dimensions) is not int or not 1 <= dimensions <= 65536:
|
||||
raise VectorStoreError("vector_invalid_dimensions")
|
||||
self.directory = Path(directory)
|
||||
self.generation_id = generation_id
|
||||
self.dimensions = dimensions
|
||||
self.collection_name = "raw_" + sha256(generation_id.encode()).hexdigest()
|
||||
self._clock = clock
|
||||
self._client: Any = None
|
||||
self._collection: Any = None
|
||||
self._writable = False
|
||||
self._closed = False
|
||||
self._lock = threading.RLock()
|
||||
|
||||
@classmethod
|
||||
def create(cls, directory: Path, *, generation_id: str, dimensions: int) -> RawVectorStore:
|
||||
"""Explicit writer initialization. Refuses to reuse an existing collection."""
|
||||
value = cls(directory, generation_id=generation_id, dimensions=dimensions)
|
||||
value._open(create=True)
|
||||
return value
|
||||
|
||||
@classmethod
|
||||
def open_existing(
|
||||
cls,
|
||||
directory: Path,
|
||||
*,
|
||||
generation_id: str,
|
||||
dimensions: int,
|
||||
) -> RawVectorStore:
|
||||
"""No logical index writes; persistent Chroma may still write maintenance files."""
|
||||
value = cls(directory, generation_id=generation_id, dimensions=dimensions)
|
||||
value._open(create=False)
|
||||
return value
|
||||
|
||||
@classmethod
|
||||
def open_writer_existing(
|
||||
cls,
|
||||
directory: Path,
|
||||
*,
|
||||
generation_id: str,
|
||||
dimensions: int,
|
||||
) -> RawVectorStore:
|
||||
"""Explicit worker recovery; verifies existing collection, never creates/repairs."""
|
||||
value = cls(directory, generation_id=generation_id, dimensions=dimensions)
|
||||
value._open(create=False, writable=True)
|
||||
return value
|
||||
|
||||
def _metadata(self) -> dict[str, str | int]:
|
||||
return {
|
||||
"adapter": "one-memory-raw-v1",
|
||||
"generation_id": self.generation_id,
|
||||
"dimensions": self.dimensions,
|
||||
"distance": "cosine",
|
||||
"normalization": "l2-float32-v1",
|
||||
}
|
||||
|
||||
def _preflight(self, create: bool) -> None:
|
||||
if any(part.is_symlink() for part in (self.directory, *self.directory.parents)):
|
||||
raise VectorStoreError("vector_unsafe_directory")
|
||||
if not self.directory.exists():
|
||||
if not create:
|
||||
raise VectorStoreError("vector_store_not_found")
|
||||
self.directory.mkdir(parents=True)
|
||||
if not self.directory.is_dir():
|
||||
raise VectorStoreError("vector_unsafe_directory")
|
||||
database = self.directory / "chroma.sqlite3"
|
||||
if database.is_symlink():
|
||||
raise VectorStoreError("vector_unsafe_database")
|
||||
if not database.is_file():
|
||||
if not create:
|
||||
raise VectorStoreError("vector_store_not_found")
|
||||
if any(self.directory.iterdir()):
|
||||
raise VectorStoreError("vector_unknown_directory")
|
||||
return
|
||||
if database.stat().st_nlink != 1:
|
||||
raise VectorStoreError("vector_unsafe_database")
|
||||
# Refuse missing/corrupt collections BEFORE PersistentClient can perform
|
||||
# startup maintenance. This checks the installed Chroma 1.5.9 schema.
|
||||
try:
|
||||
with closing(sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True)) as db:
|
||||
found = db.execute(
|
||||
"SELECT 1 FROM collections WHERE name=?", (self.collection_name,)
|
||||
).fetchone()
|
||||
if not create and found is None:
|
||||
raise VectorStoreError("vector_collection_not_found")
|
||||
if create and found is not None:
|
||||
raise VectorStoreError("vector_collection_exists")
|
||||
except sqlite3.Error:
|
||||
raise VectorStoreError("vector_store_unavailable") from None
|
||||
|
||||
def _open(self, *, create: bool, writable: bool = False) -> None:
|
||||
self._preflight(create)
|
||||
try:
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
|
||||
self._client = chromadb.PersistentClient(
|
||||
path=str(self.directory), settings=Settings(anonymized_telemetry=False)
|
||||
)
|
||||
if create:
|
||||
self._collection = self._client.create_collection(
|
||||
self.collection_name,
|
||||
embedding_function=None,
|
||||
metadata=self._metadata(),
|
||||
# Unit vectors + inner product give cosine distance without
|
||||
# cosine HNSW's lossy normalization/de-normalization on get.
|
||||
configuration={"hnsw": {"space": "ip"}},
|
||||
)
|
||||
else:
|
||||
self._collection = self._client.get_collection(
|
||||
self.collection_name, embedding_function=None
|
||||
)
|
||||
if (
|
||||
self._collection.metadata != self._metadata()
|
||||
or type(self._collection.metadata.get("dimensions")) is not int
|
||||
):
|
||||
raise VectorStoreError("vector_collection_identity_mismatch")
|
||||
config = self._collection.configuration
|
||||
if (
|
||||
config.get("hnsw", {}).get("space") != "ip"
|
||||
or config.get("embedding_function") is not None
|
||||
):
|
||||
raise VectorStoreError("vector_collection_configuration_mismatch")
|
||||
self._writable = create or writable
|
||||
except Exception as error:
|
||||
self.close()
|
||||
if isinstance(error, VectorStoreError):
|
||||
raise
|
||||
raise VectorStoreError("vector_store_unavailable") from None
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
if self._client is not None:
|
||||
self._client.close()
|
||||
self._client = self._collection = None
|
||||
|
||||
def _require_open(self) -> Any:
|
||||
if self._closed:
|
||||
raise VectorStoreError("vector_store_closed")
|
||||
if self._collection is None:
|
||||
raise VectorStoreError("vector_store_not_open")
|
||||
return self._collection
|
||||
|
||||
def _records(self, records: Sequence[RawVectorRecord]) -> tuple[RawVectorRecord, ...]:
|
||||
result = []
|
||||
seen: set[str] = set()
|
||||
for record in records:
|
||||
if not isinstance(record, RawVectorRecord):
|
||||
raise VectorStoreError("vector_invalid_record")
|
||||
for identifier in (record.chunk_id, record.build_id, record.generation_id):
|
||||
_identifier(identifier)
|
||||
if record.generation_id != self.generation_id or record.chunk_id in seen:
|
||||
raise VectorStoreError("vector_invalid_record_identity")
|
||||
seen.add(record.chunk_id)
|
||||
result.append(
|
||||
RawVectorRecord(
|
||||
record.chunk_id,
|
||||
record.build_id,
|
||||
record.generation_id,
|
||||
_vector(record.vector, self.dimensions),
|
||||
)
|
||||
)
|
||||
return tuple(result)
|
||||
|
||||
def _read_records(self, ids: Sequence[str]) -> dict[str, tuple[RawVectorRecord, str]]:
|
||||
if not ids:
|
||||
return {}
|
||||
collection = self._require_open()
|
||||
result: dict[str, tuple[RawVectorRecord, str]] = {}
|
||||
for start in range(0, len(ids), ALLOWLIST_BATCH_SIZE):
|
||||
batch = ids[start : start + ALLOWLIST_BATCH_SIZE]
|
||||
try:
|
||||
data = collection.get(ids=list(batch), include=["embeddings", "metadatas"])
|
||||
returned, vectors, metadatas = data["ids"], data["embeddings"], data["metadatas"]
|
||||
if (
|
||||
vectors is None
|
||||
or metadatas is None
|
||||
or not len(returned) == len(vectors) == len(metadatas)
|
||||
):
|
||||
raise VectorStoreError("vector_invalid_record_response")
|
||||
for key, vector, metadata in zip(returned, vectors, metadatas, strict=True):
|
||||
if key not in batch or key in result:
|
||||
raise VectorStoreError("vector_ineligible_record")
|
||||
result[key] = self._decode(key, vector, metadata)
|
||||
except VectorStoreError:
|
||||
raise
|
||||
except Exception:
|
||||
raise VectorStoreError("vector_read_unavailable") from None
|
||||
return result
|
||||
|
||||
def _decode(
|
||||
self, key: str, vector: Sequence[float], metadata: Any
|
||||
) -> tuple[RawVectorRecord, str]:
|
||||
if not isinstance(metadata, dict) or set(metadata) != {
|
||||
"chunk_id",
|
||||
"build_id",
|
||||
"generation_id",
|
||||
"vector_sha256",
|
||||
"input_sha256",
|
||||
}:
|
||||
raise VectorStoreError("vector_invalid_metadata")
|
||||
if metadata["chunk_id"] != key or metadata["generation_id"] != self.generation_id:
|
||||
raise VectorStoreError("vector_metadata_identity_mismatch")
|
||||
_identifier(key)
|
||||
_identifier(metadata["build_id"])
|
||||
if not isinstance(metadata["input_sha256"], str) or not re.fullmatch(
|
||||
r"[0-9a-f]{64}", metadata["input_sha256"]
|
||||
):
|
||||
raise VectorStoreError("vector_invalid_input_digest")
|
||||
normalized = _vector(vector, self.dimensions)
|
||||
if metadata["vector_sha256"] != _digest(normalized):
|
||||
raise VectorStoreError("vector_content_mismatch")
|
||||
return RawVectorRecord(key, metadata["build_id"], self.generation_id, normalized), metadata[
|
||||
"input_sha256"
|
||||
]
|
||||
|
||||
def upsert(self, records: Sequence[RawVectorRecord]) -> None:
|
||||
"""Idempotent immutable insertion; add-only at Chroma, never overwrite IDs.
|
||||
|
||||
Equality is canonical float32 input identity plus explicit normalized
|
||||
storage identity. Conflict is checked before
|
||||
any add. Backend failure may leave a partial new batch: caller must verify
|
||||
the entire build before publication. Existing IDs are never modified.
|
||||
"""
|
||||
with self._lock:
|
||||
collection = self._require_open()
|
||||
if not self._writable:
|
||||
raise VectorStoreError("vector_store_read_only")
|
||||
incoming = self._records(records)
|
||||
existing = self._read_records([record.chunk_id for record in incoming])
|
||||
for record in incoming:
|
||||
if record.chunk_id in existing and not _matches(existing[record.chunk_id], record):
|
||||
raise VectorStoreError("vector_immutable_conflict")
|
||||
new = [record for record in incoming if record.chunk_id not in existing]
|
||||
for start in range(0, len(new), ALLOWLIST_BATCH_SIZE):
|
||||
batch = new[start : start + ALLOWLIST_BATCH_SIZE]
|
||||
try:
|
||||
collection.add(
|
||||
ids=[record.chunk_id for record in batch],
|
||||
embeddings=[list(_unit_vector(record.vector)) for record in batch],
|
||||
metadatas=[
|
||||
{
|
||||
"chunk_id": record.chunk_id,
|
||||
"build_id": record.build_id,
|
||||
"generation_id": record.generation_id,
|
||||
"vector_sha256": _digest(_unit_vector(record.vector)),
|
||||
"input_sha256": _digest(record.vector),
|
||||
}
|
||||
for record in batch
|
||||
],
|
||||
)
|
||||
except Exception:
|
||||
raise VectorStoreError("vector_write_unavailable") from None
|
||||
self.verify(incoming)
|
||||
|
||||
def verify(self, records: Sequence[RawVectorRecord]) -> None:
|
||||
"""Verify every expected ID, metadata, dimensions and exact float32 vector."""
|
||||
with self._lock:
|
||||
self._require_open()
|
||||
expected = self._records(records)
|
||||
actual = self._read_records([record.chunk_id for record in expected])
|
||||
for record in expected:
|
||||
if record.chunk_id not in actual:
|
||||
raise VectorStoreError("vector_record_missing")
|
||||
if not _matches(actual[record.chunk_id], record):
|
||||
raise VectorStoreError("vector_immutable_conflict")
|
||||
|
||||
def verify_candidates(
|
||||
self,
|
||||
candidates: Sequence[tuple[str, str]],
|
||||
*,
|
||||
deadline: float,
|
||||
) -> bool:
|
||||
"""Read-only presence/integrity check for (chunk_id, expected_build_id).
|
||||
|
||||
Costs O(candidates * dimensions), not a bounded top-K lookup. False means
|
||||
deadline exhaustion, never proof of absence. Each read covers at most
|
||||
2048 IDs; synchronous backend calls cannot be interrupted. This is not a
|
||||
snapshot across Chroma and SQLite: Core still owns final eligibility
|
||||
checks and writer serialization. Chroma may perform maintenance writes.
|
||||
"""
|
||||
with self._lock:
|
||||
self._require_open()
|
||||
if (
|
||||
isinstance(deadline, bool)
|
||||
or not isinstance(deadline, Real)
|
||||
or not math.isfinite(deadline)
|
||||
):
|
||||
raise VectorStoreError("vector_invalid_deadline")
|
||||
if not isinstance(candidates, Sequence) or isinstance(candidates, (str, bytes)):
|
||||
raise VectorStoreError("vector_invalid_candidates")
|
||||
expected: list[tuple[str, str]] = []
|
||||
seen: set[str] = set()
|
||||
for pair in candidates:
|
||||
if (
|
||||
not isinstance(pair, Sequence)
|
||||
or isinstance(pair, (str, bytes))
|
||||
or len(pair) != 2
|
||||
):
|
||||
raise VectorStoreError("vector_invalid_candidates")
|
||||
key, build = pair
|
||||
_identifier(key)
|
||||
_identifier(build)
|
||||
if key in seen:
|
||||
raise VectorStoreError("vector_duplicate_candidates")
|
||||
seen.add(key)
|
||||
expected.append((key, build))
|
||||
for start in range(0, len(expected), ALLOWLIST_BATCH_SIZE):
|
||||
if self._clock() >= deadline:
|
||||
return False
|
||||
batch = expected[start : start + ALLOWLIST_BATCH_SIZE]
|
||||
# Bound this call to one batch so it cannot read past the next
|
||||
# deadline checkpoint via _read_records' internal loop.
|
||||
actual = self._read_records([key for key, _build in batch])
|
||||
for key, build in batch:
|
||||
if key not in actual:
|
||||
raise VectorStoreError("vector_record_missing")
|
||||
if actual[key][0].build_id != build:
|
||||
raise VectorStoreError("vector_metadata_identity_mismatch")
|
||||
if self._clock() >= deadline:
|
||||
return False
|
||||
return True
|
||||
|
||||
def search(
|
||||
self,
|
||||
query_vector: Sequence[float],
|
||||
eligible_ids: Sequence[str],
|
||||
limit: int,
|
||||
deadline: float,
|
||||
) -> VectorSearchResult:
|
||||
"""Rank exactly the eligible candidates; synchronous backend cannot be preempted.
|
||||
|
||||
deadline is absolute monotonic seconds. A late batch is discarded; prior
|
||||
on-time batches may return partial. complete means every eligible
|
||||
vector was scored, NOT complete index coverage: only chunks the caller
|
||||
already resolved as eligible and published are searched.
|
||||
|
||||
The eligible embeddings are read back in bounded batches and ranked
|
||||
here, because Chroma's own filtered query does not score every ID in
|
||||
its allowlist. A global top-K filtered afterwards would rank
|
||||
unauthorized vectors first and let them crowd the eligible ones out.
|
||||
"""
|
||||
with self._lock:
|
||||
self._require_open()
|
||||
vector = _unit_vector(_vector(query_vector, self.dimensions))
|
||||
if type(limit) is not int or not 1 <= limit <= ALLOWLIST_BATCH_SIZE:
|
||||
raise VectorStoreError("vector_invalid_limit")
|
||||
if (
|
||||
isinstance(deadline, bool)
|
||||
or not isinstance(deadline, Real)
|
||||
or not math.isfinite(deadline)
|
||||
):
|
||||
raise VectorStoreError("vector_invalid_deadline")
|
||||
if isinstance(eligible_ids, (str, bytes)):
|
||||
raise VectorStoreError("vector_invalid_allowlist")
|
||||
ids = tuple(eligible_ids)
|
||||
for key in ids:
|
||||
_identifier(key)
|
||||
if len(set(ids)) != len(ids):
|
||||
raise VectorStoreError("vector_duplicate_allowlist")
|
||||
hits: list[VectorHit] = []
|
||||
searched = 0
|
||||
for start in range(0, len(ids), ALLOWLIST_BATCH_SIZE):
|
||||
if self._clock() >= deadline:
|
||||
break
|
||||
batch = ids[start : start + ALLOWLIST_BATCH_SIZE]
|
||||
try:
|
||||
read = self._read_records(batch)
|
||||
except VectorStoreError as error:
|
||||
# A missing or corrupt eligible vector is not "no hits".
|
||||
raise VectorStoreError("vector_query_unavailable") from error
|
||||
if self._clock() >= deadline:
|
||||
break
|
||||
try:
|
||||
missing = [key for key in batch if key not in read]
|
||||
if missing:
|
||||
# The caller resolved these as eligible and published,
|
||||
# so an absent vector is a broken build, not a miss.
|
||||
raise VectorStoreError("vector_query_unavailable")
|
||||
validated = [
|
||||
VectorHit(
|
||||
key,
|
||||
read[key][0].build_id,
|
||||
read[key][0].generation_id,
|
||||
# Stored vectors are unit length, so the inner
|
||||
# product is the cosine similarity.
|
||||
1.0 - _dot(read[key][0].vector, vector),
|
||||
)
|
||||
for key in batch
|
||||
]
|
||||
except VectorStoreError:
|
||||
raise
|
||||
except Exception:
|
||||
raise VectorStoreError("vector_invalid_query_response") from None
|
||||
if self._clock() >= deadline:
|
||||
break
|
||||
searched += len(batch)
|
||||
hits.extend(validated)
|
||||
hits.sort(key=lambda hit: (hit.distance, hit.chunk_id))
|
||||
del hits[limit:]
|
||||
complete = searched == len(ids)
|
||||
return VectorSearchResult(
|
||||
tuple(hits), complete, len(ids), searched, None if complete else "deadline"
|
||||
)
|
||||
|
||||
|
||||
def _dot(left: Sequence[float], right: Sequence[float]) -> float:
|
||||
return math.fsum(a * b for a, b in zip(left, right, strict=True))
|
||||
|
||||
|
||||
def _identifier(value: Any) -> None:
|
||||
if not isinstance(value, str) or not re.fullmatch(r"[a-zA-Z0-9][a-zA-Z0-9_.:-]{0,255}", value):
|
||||
raise VectorStoreError("vector_invalid_identifier")
|
||||
|
||||
|
||||
def _vector(value: Sequence[float], dimensions: int) -> tuple[float, ...]:
|
||||
try:
|
||||
if len(value) != dimensions or any(
|
||||
isinstance(v, bool) or not isinstance(v, Real) or not math.isfinite(v) for v in value
|
||||
):
|
||||
raise VectorStoreError("vector_invalid_values_or_dimensions")
|
||||
packed = struct.pack(f"<{dimensions}f", *value)
|
||||
normalized = struct.unpack(f"<{dimensions}f", packed)
|
||||
if not all(math.isfinite(v) for v in normalized) or not any(normalized):
|
||||
raise VectorStoreError("vector_invalid_values_or_dimensions")
|
||||
return normalized
|
||||
except VectorStoreError:
|
||||
raise
|
||||
except (ValueError, TypeError, OverflowError, struct.error):
|
||||
raise VectorStoreError("vector_invalid_values_or_dimensions") from None
|
||||
|
||||
|
||||
def _digest(vector: tuple[float, ...]) -> str:
|
||||
return sha256(struct.pack(f"<{len(vector)}f", *vector)).hexdigest()
|
||||
|
||||
|
||||
def _unit_vector(vector: tuple[float, ...]) -> tuple[float, ...]:
|
||||
length = math.sqrt(math.fsum(component * component for component in vector))
|
||||
return _vector(tuple(component / length for component in vector), len(vector))
|
||||
|
||||
|
||||
def _matches(actual: tuple[RawVectorRecord, str], expected: RawVectorRecord) -> bool:
|
||||
stored, input_digest = actual
|
||||
return input_digest == _digest(expected.vector) and stored == RawVectorRecord(
|
||||
expected.chunk_id, expected.build_id, expected.generation_id, _unit_vector(expected.vector)
|
||||
)
|
||||
@@ -50,18 +50,29 @@ class SdkDataPlaneMigrationStore:
|
||||
return datetime.now(UTC).isoformat()
|
||||
|
||||
def authority(self) -> MemoryAuthority:
|
||||
return self.authority_state()[0]
|
||||
|
||||
def authority_state(self) -> tuple[MemoryAuthority, str | None]:
|
||||
"""The surviving singleton: authority and the migration that set it.
|
||||
|
||||
This is the only source of authority. The detailed cutover receipt is
|
||||
dropped after cleanup and never exists on a fresh home.
|
||||
"""
|
||||
try:
|
||||
with main_database(self.database_path).read() as connection:
|
||||
row = connection.execute(
|
||||
"SELECT authority FROM memory_authority_state WHERE singleton=1"
|
||||
"SELECT authority,migration_id FROM memory_authority_state WHERE singleton=1"
|
||||
).fetchone()
|
||||
except sqlite3.OperationalError as error:
|
||||
if "no such table" not in str(error).lower():
|
||||
raise
|
||||
else:
|
||||
if row is not None:
|
||||
return MemoryAuthority(str(row[0]))
|
||||
return MemoryAuthority.SDK if self.cutover() is not None else MemoryAuthority.LEGACY
|
||||
return MemoryAuthority(str(row[0])), (str(row[1]) if row[1] else None)
|
||||
marker = self.cutover()
|
||||
if marker is not None:
|
||||
return MemoryAuthority.SDK, marker.migration_id
|
||||
return MemoryAuthority.LEGACY, None
|
||||
|
||||
def cutover(self) -> CutoverMarker | None:
|
||||
try:
|
||||
|
||||
@@ -15,9 +15,10 @@ import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
from collections.abc import Callable, Iterator
|
||||
from contextlib import contextmanager
|
||||
from contextlib import AbstractContextManager, contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.imports.models import EvidenceRecord
|
||||
@@ -35,6 +36,7 @@ from one_memory.memory_server.durable_queue import (
|
||||
DurableWriteQueue,
|
||||
)
|
||||
from one_memory.memory_server.unified_entry import UnifiedWriteIntake
|
||||
from one_memory.model_profiles import ModelProfileSnapshot
|
||||
from one_memory.sdk_data_plane_migration import (
|
||||
CutoverMarker,
|
||||
DataPlaneMigrationError,
|
||||
@@ -42,6 +44,20 @@ from one_memory.sdk_data_plane_migration import (
|
||||
SdkDataPlaneMigrationStore,
|
||||
)
|
||||
from one_memory.sdk_migration_bundle import inspect_sdk_migration_bundle
|
||||
from one_memory.sdk_migration_verify import (
|
||||
MemoryReadbackFailed,
|
||||
ReadbackReport,
|
||||
legacy_memory_ids_from,
|
||||
legacy_memory_ids_from_bundle,
|
||||
read_active_profile,
|
||||
runtime_sdk_client,
|
||||
verify_readback,
|
||||
)
|
||||
from one_memory.sdk_vector_reattach import (
|
||||
ReattachedCollection,
|
||||
legacy_vector_collections,
|
||||
reattach_legacy_vectors,
|
||||
)
|
||||
from one_memory.target_schema import drop_legacy_tables, legacy_tables_present
|
||||
|
||||
PREPARED_PHASE = "prepared"
|
||||
@@ -341,6 +357,8 @@ class StartupMigrationResult:
|
||||
bundle_id: str | None = None
|
||||
preparation: CutoverPreparation | None = None
|
||||
dropped_tables: tuple[str, ...] = ()
|
||||
reattached: tuple[ReattachedCollection, ...] = ()
|
||||
readback: ReadbackReport | None = None
|
||||
|
||||
def as_document(self) -> dict[str, object]:
|
||||
document: dict[str, object] = {"schema_version": 1, "state": self.state}
|
||||
@@ -349,6 +367,18 @@ class StartupMigrationResult:
|
||||
document["bundle_id"] = self.bundle_id
|
||||
if self.preparation is not None:
|
||||
document.update(self.preparation.as_document())
|
||||
if self.reattached:
|
||||
document["reattached_vectors"] = [
|
||||
{
|
||||
"collection": item.current,
|
||||
"legacy_count": item.legacy_count,
|
||||
"merged_count": item.merged_count,
|
||||
"final_count": item.final_count,
|
||||
}
|
||||
for item in self.reattached
|
||||
]
|
||||
if self.readback is not None:
|
||||
document["readback"] = self.readback.as_document()
|
||||
if self.dropped_tables:
|
||||
document["dropped_tables"] = list(self.dropped_tables)
|
||||
return document
|
||||
@@ -415,6 +445,15 @@ def _ensure_bundle(*, data_dir: Path, bundle_id: str, runtime_build: str) -> dic
|
||||
return inspect_sdk_migration_bundle(data_dir=data_dir, bundle_id=bundle_id)
|
||||
|
||||
|
||||
def _vacuum(database_path: Path) -> None:
|
||||
"""Return the dropped tables' pages to the disk; a full disk only skips it."""
|
||||
try:
|
||||
with main_database(database_path).write(autocommit=True) as database:
|
||||
database.execute("VACUUM")
|
||||
except sqlite3.OperationalError:
|
||||
return
|
||||
|
||||
|
||||
def run_startup_migration(
|
||||
*,
|
||||
data_dir: Path,
|
||||
@@ -422,24 +461,59 @@ def run_startup_migration(
|
||||
sdk_version: str,
|
||||
namespace: str,
|
||||
progress: Callable[[str], None] = lambda _step: None,
|
||||
open_client: Callable[[ModelProfileSnapshot], AbstractContextManager[Any]] | None = None,
|
||||
) -> StartupMigrationResult:
|
||||
"""Bring one stopped home to SDK authority. Safe to repeat after any failure.
|
||||
|
||||
The caller must stop Runtime and keep its service from restarting.
|
||||
`open_client` yields a public SDK client configured as the Runtime will be;
|
||||
tests replace it.
|
||||
"""
|
||||
database_path = data_dir / "one-memory.sqlite3"
|
||||
if not database_path.is_file():
|
||||
return StartupMigrationResult(state="not_required")
|
||||
|
||||
def client_for(profile: ModelProfileSnapshot) -> AbstractContextManager[Any]:
|
||||
if open_client is not None:
|
||||
return open_client(profile)
|
||||
return runtime_sdk_client(database_path, data_dir, profile)
|
||||
|
||||
store = SdkDataPlaneMigrationStore(database_path)
|
||||
if store.authority() is MemoryAuthority.SDK:
|
||||
if not legacy_tables_present(database_path) and not abandoned_staging_present(
|
||||
authority, authority_migration = store.authority_state()
|
||||
profile = read_active_profile(database_path)
|
||||
if authority is MemoryAuthority.SDK:
|
||||
vectors_waiting = profile is not None and bool(legacy_vector_collections(data_dir, profile))
|
||||
leftovers = bool(legacy_tables_present(database_path)) or abandoned_staging_present(
|
||||
data_dir=data_dir, namespace=namespace
|
||||
):
|
||||
)
|
||||
if not vectors_waiting and not leftovers:
|
||||
return StartupMigrationResult(state="not_required")
|
||||
progress("cleanup")
|
||||
reattached: tuple[ReattachedCollection, ...] = ()
|
||||
readback: ReadbackReport | None = None
|
||||
if vectors_waiting:
|
||||
assert profile is not None
|
||||
progress("vectors")
|
||||
reattached = reattach_legacy_vectors(data_dir, profile)
|
||||
legacy = legacy_memory_ids_from_bundle(data_dir, authority_migration)
|
||||
if legacy is not None and not legacy.empty:
|
||||
# Authority already moved, so nothing is left to protect by
|
||||
# refusing; record what the repaired home reads back.
|
||||
progress("verify")
|
||||
with client_for(profile) as client:
|
||||
readback = verify_readback(client, legacy)
|
||||
dropped: tuple[str, ...] = ()
|
||||
if leftovers:
|
||||
progress("cleanup")
|
||||
dropped = resume_legacy_cleanup(data_dir=data_dir, database_path=database_path)
|
||||
if dropped:
|
||||
_vacuum(database_path)
|
||||
return StartupMigrationResult(
|
||||
state="complete",
|
||||
dropped_tables=resume_legacy_cleanup(data_dir=data_dir, database_path=database_path),
|
||||
migration_id=authority_migration,
|
||||
bundle_id=f"{authority_migration}-pre-cutover" if authority_migration else None,
|
||||
dropped_tables=dropped,
|
||||
reattached=reattached,
|
||||
readback=readback,
|
||||
)
|
||||
|
||||
from one_memory.database_access import close_main_databases
|
||||
@@ -447,6 +521,11 @@ def run_startup_migration(
|
||||
from one_memory.memory_server.migrations import MIGRATIONS
|
||||
|
||||
migration_id, bundle_id = _resumable_identity(database_path)
|
||||
legacy = legacy_memory_ids_from(database_path)
|
||||
if profile is None and not legacy.empty:
|
||||
# Without the embedding profile the Runtime's collections cannot be
|
||||
# named, so neither reattached nor read back. Legacy stays intact.
|
||||
raise DataPlaneMigrationError("sdk_migration_model_profile_required")
|
||||
with capture_fence(data_dir):
|
||||
progress("bundle")
|
||||
manifest = _ensure_bundle(
|
||||
@@ -472,6 +551,19 @@ def run_startup_migration(
|
||||
backup_bundle_id=bundle_id,
|
||||
)
|
||||
preparation = prepare_legacy_work_for_cutover(database_path)
|
||||
reattached = ()
|
||||
readback = None
|
||||
if profile is not None:
|
||||
progress("vectors")
|
||||
reattached = reattach_legacy_vectors(data_dir, profile)
|
||||
if not legacy.empty:
|
||||
assert profile is not None
|
||||
progress("verify")
|
||||
with client_for(profile) as client:
|
||||
readback = verify_readback(client, legacy)
|
||||
if not readback.passed:
|
||||
store.set_phase(migration_id, phase="readback_failed", state="blocked")
|
||||
raise MemoryReadbackFailed(readback)
|
||||
store.set_phase(migration_id, phase=PREPARED_PHASE, state="ready")
|
||||
progress("cutover")
|
||||
commit_sdk_cutover(
|
||||
@@ -483,10 +575,14 @@ def run_startup_migration(
|
||||
)
|
||||
progress("cleanup")
|
||||
dropped = resume_legacy_cleanup(data_dir=data_dir, database_path=database_path)
|
||||
if dropped:
|
||||
_vacuum(database_path)
|
||||
return StartupMigrationResult(
|
||||
state="cutover",
|
||||
migration_id=migration_id,
|
||||
bundle_id=bundle_id,
|
||||
preparation=preparation,
|
||||
dropped_tables=dropped,
|
||||
reattached=reattached,
|
||||
readback=readback,
|
||||
)
|
||||
@@ -0,0 +1,232 @@
|
||||
"""Read legacy memory IDs back through the SDK before (and after) the cutover.
|
||||
|
||||
The cutover is only safe if the SDK, configured exactly as the new Runtime
|
||||
will configure it, returns the memories the legacy mirror lists. Checking that
|
||||
the backup bundle is complete says nothing about it: build 20 dropped the
|
||||
legacy tables while its Runtime read an empty vector collection.
|
||||
|
||||
Reads go through the public `HyMemoryClient` (`list_memories`, `get`).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
from collections.abc import Iterator
|
||||
from contextlib import closing, contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.model_profiles import ModelProfileSnapshot
|
||||
from one_memory.sdk_data_plane_migration import DataPlaneMigrationError
|
||||
|
||||
_PRODUCTIVITY_KINDS = frozenset({"coding", "working"})
|
||||
_PRODUCTIVITY_LAYERS = frozenset(
|
||||
{"coding_experience", "working_experience", "coding_fact", "working_fact"}
|
||||
)
|
||||
_SDK_STATUSES = ("active", "superseded", "shadow", "negated", "conflicted", "archived")
|
||||
_PAGE = 500
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LegacyMemoryIds:
|
||||
chat: tuple[str, ...]
|
||||
productivity: tuple[str, ...]
|
||||
users: tuple[str, ...]
|
||||
|
||||
@property
|
||||
def empty(self) -> bool:
|
||||
return not self.chat and not self.productivity
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ReadbackReport:
|
||||
chat_expected: int
|
||||
chat_readable: int
|
||||
productivity_expected: int
|
||||
productivity_readable: int
|
||||
|
||||
@property
|
||||
def passed(self) -> bool:
|
||||
# Every chat memory is an SDK node by construction, so all must read
|
||||
# back. The SDK's Coding reconciler deletes experiences it retires, so
|
||||
# a partial productivity count is normal; none at all is not.
|
||||
return self.chat_readable == self.chat_expected and (
|
||||
self.productivity_expected == 0 or self.productivity_readable > 0
|
||||
)
|
||||
|
||||
def as_document(self) -> dict[str, object]:
|
||||
return {
|
||||
"chat_expected": self.chat_expected,
|
||||
"chat_readable": self.chat_readable,
|
||||
"productivity_expected": self.productivity_expected,
|
||||
"productivity_readable": self.productivity_readable,
|
||||
"passed": self.passed,
|
||||
}
|
||||
|
||||
|
||||
class MemoryReadbackFailed(DataPlaneMigrationError):
|
||||
def __init__(self, report: ReadbackReport) -> None:
|
||||
super().__init__("sdk_migration_memory_unreadable")
|
||||
self.report = report
|
||||
|
||||
|
||||
def _tables(connection: sqlite3.Connection) -> set[str]:
|
||||
return {
|
||||
str(row[0])
|
||||
for row in connection.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
||||
}
|
||||
|
||||
|
||||
def read_legacy_memory_ids(connection: sqlite3.Connection) -> LegacyMemoryIds:
|
||||
"""Confirmed legacy memories, by the SDK id the old bridge recorded."""
|
||||
tables = _tables(connection)
|
||||
if "governed_memories" not in tables:
|
||||
return LegacyMemoryIds(chat=(), productivity=(), users=())
|
||||
mapped: dict[str, tuple[str, str | None]] = {}
|
||||
if "memory_task_outputs" in tables:
|
||||
for (body,) in connection.execute("SELECT body FROM memory_task_outputs"):
|
||||
try:
|
||||
document = json.loads(body)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
legacy_id, sdk_id = document.get("memory_id"), document.get("hy_memory_id")
|
||||
if isinstance(legacy_id, str) and isinstance(sdk_id, str):
|
||||
layer = document.get("layer")
|
||||
mapped[legacy_id] = (sdk_id, layer if isinstance(layer, str) else None)
|
||||
chat: set[str] = set()
|
||||
productivity: set[str] = set()
|
||||
users: set[str] = set()
|
||||
for memory_id, user_id, kind in connection.execute(
|
||||
"SELECT memory_id,user_id,kind FROM governed_memories WHERE status='confirmed'"
|
||||
):
|
||||
sdk_id, layer = mapped.get(str(memory_id), (str(memory_id), None))
|
||||
users.add(str(user_id))
|
||||
if str(kind) in _PRODUCTIVITY_KINDS or layer in _PRODUCTIVITY_LAYERS:
|
||||
productivity.add(sdk_id)
|
||||
else:
|
||||
chat.add(sdk_id)
|
||||
return LegacyMemoryIds(
|
||||
chat=tuple(sorted(chat)),
|
||||
productivity=tuple(sorted(productivity)),
|
||||
users=tuple(sorted(users)),
|
||||
)
|
||||
|
||||
|
||||
def legacy_memory_ids_from(database_path: Path) -> LegacyMemoryIds:
|
||||
with main_database(database_path).read() as connection:
|
||||
return read_legacy_memory_ids(connection)
|
||||
|
||||
|
||||
def legacy_memory_ids_from_bundle(
|
||||
data_dir: Path, migration_id: str | None
|
||||
) -> LegacyMemoryIds | None:
|
||||
"""The same list from the pre-cutover bundle, for a home already migrated."""
|
||||
if not migration_id:
|
||||
return None
|
||||
database = (
|
||||
data_dir
|
||||
/ "backups"
|
||||
/ f"{migration_id}-pre-cutover.sdk-bundle"
|
||||
/ "management"
|
||||
/ "one-memory.sqlite3"
|
||||
)
|
||||
if not database.is_file():
|
||||
return None
|
||||
with closing(sqlite3.connect(f"file:{database}?mode=ro", uri=True)) as connection:
|
||||
return read_legacy_memory_ids(connection)
|
||||
|
||||
|
||||
def read_active_profile(database_path: Path) -> ModelProfileSnapshot | None:
|
||||
"""The active model profile, read only; the status probe runs beside Core."""
|
||||
with main_database(database_path).read(row_factory=sqlite3.Row) as connection:
|
||||
if not {"model_profiles", "model_profile_versions"} <= _tables(connection):
|
||||
return None
|
||||
row = connection.execute(
|
||||
"SELECT version.* FROM model_profiles profile "
|
||||
"JOIN model_profile_versions version "
|
||||
"ON version.profile_id=profile.profile_id AND version.revision=profile.revision "
|
||||
"WHERE profile.is_active=1"
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return ModelProfileSnapshot(
|
||||
profile_id=row["profile_id"],
|
||||
revision=row["revision"],
|
||||
provider_kind=row["provider_kind"],
|
||||
base_url=row["base_url"],
|
||||
model=row["model"],
|
||||
embedding_model=row["embedding_model"],
|
||||
embedding_provider_kind=row["embedding_provider_kind"],
|
||||
embedding_base_url=row["embedding_base_url"],
|
||||
embedding_dimensions=row["embedding_dimensions"],
|
||||
timeout_seconds=row["timeout_seconds"],
|
||||
max_concurrency=row["max_concurrency"],
|
||||
context_window=row["context_window"],
|
||||
has_api_key=row["secret_ref"] is not None,
|
||||
disable_thinking=bool(row["disable_thinking"]),
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def runtime_sdk_client(
|
||||
database_path: Path, data_dir: Path, profile: ModelProfileSnapshot
|
||||
) -> Iterator[Any]:
|
||||
"""A public SDK client configured exactly as the Runtime configures it."""
|
||||
from one_memory.hy_sdk_runtime import acquire_client
|
||||
from one_memory.model_profiles import ModelProfileStore
|
||||
from one_memory.secrets import KeyringSecretStore, SecretStoreUnavailableError
|
||||
|
||||
try:
|
||||
secret = ModelProfileStore(database_path).resolve_snapshot_secret(
|
||||
profile.profile_id, profile.revision, KeyringSecretStore()
|
||||
)
|
||||
except SecretStoreUnavailableError:
|
||||
raise DataPlaneMigrationError("sdk_migration_secure_storage_unavailable") from None
|
||||
if profile.has_api_key and secret is None:
|
||||
raise DataPlaneMigrationError("sdk_migration_model_secret_unavailable")
|
||||
lease = acquire_client(profile, api_key=secret, data_dir=data_dir)
|
||||
try:
|
||||
yield lease.client
|
||||
finally:
|
||||
lease.close()
|
||||
|
||||
|
||||
def verify_readback(client: Any, legacy: LegacyMemoryIds) -> ReadbackReport:
|
||||
listed: set[str] = set()
|
||||
for user_id in legacy.users or ("local",):
|
||||
offset = 0
|
||||
while True:
|
||||
raw = client.list_memories(
|
||||
user_id=user_id,
|
||||
statuses=list(_SDK_STATUSES),
|
||||
limit=_PAGE,
|
||||
offset=offset,
|
||||
order="asc",
|
||||
)
|
||||
bucket = raw.get("vdb") if isinstance(raw, dict) else None
|
||||
rows = bucket.get("memories") if isinstance(bucket, dict) else None
|
||||
total = bucket.get("total") if isinstance(bucket, dict) else None
|
||||
if not isinstance(rows, list) or not isinstance(total, int):
|
||||
raise DataPlaneMigrationError("sdk_migration_memory_list_shape_unknown")
|
||||
listed.update(
|
||||
str(row["memory_id"])
|
||||
for row in rows
|
||||
if isinstance(row, dict) and row.get("memory_id")
|
||||
)
|
||||
offset += len(rows)
|
||||
if not rows or offset >= total:
|
||||
break
|
||||
|
||||
def readable(memory_id: str) -> bool:
|
||||
return memory_id in listed or client.get(memory_id) is not None
|
||||
|
||||
return ReadbackReport(
|
||||
chat_expected=len(legacy.chat),
|
||||
chat_readable=sum(1 for memory_id in legacy.chat if readable(memory_id)),
|
||||
productivity_expected=len(legacy.productivity),
|
||||
productivity_readable=sum(1 for memory_id in legacy.productivity if readable(memory_id)),
|
||||
)
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
"""Offline reattachment of SDK vector collections named by builds before #145.
|
||||
|
||||
Those builds named a profile's embedding space `one_memory_<12 hex>`; #145
|
||||
renamed it `memhub_<16 hex>` and dropped the step that carried the old
|
||||
collections over, so every later Runtime opened an empty space beside the old
|
||||
one. Runtime code only reaches vectors through the SDK. This module runs from
|
||||
`sdk-migrate` while the Runtime is stopped, through Chroma's public client.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
from collections.abc import Iterator
|
||||
from contextlib import closing
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from one_memory.hy_sdk_runtime import vector_collection_bases
|
||||
from one_memory.model_profiles import ModelProfileSnapshot
|
||||
from one_memory.sdk_data_plane_migration import DataPlaneMigrationError
|
||||
|
||||
_BATCH = 500
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ReattachedCollection:
|
||||
legacy: str
|
||||
current: str
|
||||
legacy_count: int
|
||||
merged_count: int
|
||||
final_count: int
|
||||
|
||||
|
||||
def _vector_directory(data_dir: Path) -> Path:
|
||||
return data_dir / "hy-memory" / "governed-v2" / "vector"
|
||||
|
||||
|
||||
def _stored_names(data_dir: Path) -> tuple[str, ...]:
|
||||
database = _vector_directory(data_dir) / "chroma.sqlite3"
|
||||
if not database.is_file():
|
||||
return ()
|
||||
with closing(sqlite3.connect(f"file:{database}?mode=ro", uri=True)) as connection:
|
||||
try:
|
||||
rows = connection.execute("SELECT name FROM collections").fetchall()
|
||||
except sqlite3.Error:
|
||||
return ()
|
||||
return tuple(sorted(str(row[0]) for row in rows))
|
||||
|
||||
|
||||
def legacy_vector_collections(
|
||||
data_dir: Path, profile: ModelProfileSnapshot
|
||||
) -> tuple[tuple[str, str], ...]:
|
||||
"""(legacy name, current name) pairs still waiting to be reattached.
|
||||
|
||||
Read-only, so `sdk-migration-status` can ask it.
|
||||
"""
|
||||
current, legacy_bases = vector_collection_bases(profile)
|
||||
pairs: list[tuple[str, str]] = []
|
||||
for name in _stored_names(data_dir):
|
||||
for base in legacy_bases:
|
||||
if name == base or name.startswith(f"{base}_"):
|
||||
pairs.append((name, current + name[len(base) :]))
|
||||
break
|
||||
return tuple(pairs)
|
||||
|
||||
|
||||
def _records(collection: Any) -> Iterator[dict[str, Any]]:
|
||||
offset = 0
|
||||
while True:
|
||||
page = collection.get(
|
||||
include=["embeddings", "documents", "metadatas"],
|
||||
limit=_BATCH,
|
||||
offset=offset,
|
||||
)
|
||||
ids = list(page["ids"])
|
||||
if not ids:
|
||||
return
|
||||
yield {
|
||||
"ids": ids,
|
||||
"embeddings": page["embeddings"],
|
||||
"documents": page["documents"],
|
||||
"metadatas": page["metadatas"],
|
||||
}
|
||||
offset += len(ids)
|
||||
|
||||
|
||||
def reattach_legacy_vectors(
|
||||
data_dir: Path, profile: ModelProfileSnapshot
|
||||
) -> tuple[ReattachedCollection, ...]:
|
||||
"""Rename legacy collections to the current names, keeping newer writes.
|
||||
|
||||
A current collection that already exists (written after the rename broke)
|
||||
is merged into the legacy one first; on an id present in both, the newer
|
||||
current record wins. Callers must have stopped every SDK writer.
|
||||
"""
|
||||
pairs = legacy_vector_collections(data_dir, profile)
|
||||
if not pairs:
|
||||
return ()
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
|
||||
client = chromadb.PersistentClient(
|
||||
path=str(_vector_directory(data_dir)),
|
||||
settings=Settings(anonymized_telemetry=False),
|
||||
)
|
||||
existing = set(_stored_names(data_dir))
|
||||
done: list[ReattachedCollection] = []
|
||||
for legacy_name, current_name in pairs:
|
||||
legacy = client.get_collection(legacy_name)
|
||||
legacy_count = legacy.count()
|
||||
merged = 0
|
||||
if current_name in existing:
|
||||
current = client.get_collection(current_name)
|
||||
legacy_ids = set(_ids(legacy))
|
||||
current_ids: set[str] = set()
|
||||
for batch in _records(current):
|
||||
legacy.upsert(**batch)
|
||||
current_ids.update(batch["ids"])
|
||||
merged = len(current_ids)
|
||||
if legacy.count() != len(legacy_ids | current_ids):
|
||||
raise DataPlaneMigrationError("sdk_migration_vector_merge_incomplete")
|
||||
client.delete_collection(current_name)
|
||||
legacy.modify(name=current_name)
|
||||
final = client.get_collection(current_name).count()
|
||||
if final < legacy_count:
|
||||
raise DataPlaneMigrationError("sdk_migration_vector_reattach_incomplete")
|
||||
existing.discard(legacy_name)
|
||||
existing.add(current_name)
|
||||
done.append(
|
||||
ReattachedCollection(
|
||||
legacy=legacy_name,
|
||||
current=current_name,
|
||||
legacy_count=legacy_count,
|
||||
merged_count=merged,
|
||||
final_count=final,
|
||||
)
|
||||
)
|
||||
return tuple(done)
|
||||
|
||||
|
||||
def _ids(collection: Any) -> Iterator[str]:
|
||||
offset = 0
|
||||
while True:
|
||||
ids = list(collection.get(include=[], limit=_BATCH, offset=offset)["ids"])
|
||||
if not ids:
|
||||
return
|
||||
yield from ids
|
||||
offset += len(ids)
|
||||
@@ -139,7 +139,6 @@
|
||||
| `test_live_watch.py` | `contract` | 覆盖 a new turn lands well inside the ten second target、start prime catches a change that predates the watcher、a run that predates the change is followed by a fresh one 等 15 个场景。 |
|
||||
| `test_live_watch_grants.py` | `integration` | 覆盖 explicit sync grant starts watch without restart and revocation stops it 场景。 |
|
||||
| `test_local_embedding.py` | `contract` | Guards for the built-in embedding model. |
|
||||
| `test_local_embedding_offsets.py` | `contract` | 覆盖 real locked tokenizer offsets are unicode codepoints、known emoji and combining offsets are not utf8 or utf16 units、offsets are not truncated or padded and do not create onnx session 等 9 个场景。 |
|
||||
| `test_local_onnx_provider.py` | `contract` | Expect local ONNX embedding to be the default and to feed Hy Memory. |
|
||||
| `test_log_file.py` | `contract` | 覆盖 what the watchdog says survives a dock launch、the log lives where the data does、a machine left running for months does not fill up 等 4 个场景。 |
|
||||
| `test_log_privacy.py` | `integration` | The file log is shipped to us in bug reports; it must not carry originals. |
|
||||
@@ -232,21 +231,8 @@
|
||||
| `test_recall_scope_control.py` | `integration` | 覆盖 external scope cannot override shared or host workspace boundary、scope only mutation is versioned without changing status or reindexing 场景。 |
|
||||
| `test_recall_scope_off_loop.py` | `contract` | The workspace lookup behind an external recall must not run on the loop. |
|
||||
| `test_recent_first_import.py` | `contract` | Recent-file priority survives restart through a selection-bound frozen copy. |
|
||||
| `test_retrieval_chunking.py` | `contract` | 覆盖 default 384 body 64 overlap and exact source slices、all prefix and special tokens count toward strict 512 budget、real tokenizer covers long chinese middle and tail object locations 等 20 个场景。 |
|
||||
| `test_retrieval_eligibility_benchmark.py` | `contract` | 覆盖 scale probe proves composite authority index changes the plan 场景。 |
|
||||
| `test_retrieval_index_store.py` | `contract` | 覆盖 constructor and reads do not create database or schema、explicit initialization rolls back and never grants model policy、consumer and jobs commit together 等 19 个场景。 |
|
||||
| `test_retrieval_index_supervisor.py` | `contract` | 覆盖 stop waits for native exit and does not close resources、async cancellation repeatedly waits for native thread、idle no trigger no spin and duplicate start rejected 等 14 个场景。 |
|
||||
| `test_retrieval_index_worker.py` | `contract` | 覆盖 tick indexes complete original and reopen can search、change during native encoding never publishes、cancel after encode does not publish and is reclaimable 等 10 个场景。 |
|
||||
| `test_retrieval_lifespan.py` | `integration` | 覆盖 retrieval defaults off and does not create index、enabling retrieval requires runtime ownership、index enabled serve imports without building an unmanaged app 等 6 个场景。 |
|
||||
| `test_retrieval_mapped_content.py` | `contract` | 覆盖 raw authors keep exact body and do not mutate evidence、raw and author qualification rejects without tokenizing、ambiguous generated and body hygiene are not admitted 等 22 个场景。 |
|
||||
| `test_retrieval_mapped_integration.py` | `contract` | 覆盖 empty materialization rechecks excluded map、completed excluded missing active binding is not empty success、old empty completion is invalidated when a new indexed record arrives 等 24 个场景。 |
|
||||
| `test_retrieval_provisioning.py` | `contract` | The index has to come up through the product, not through a test helper. |
|
||||
| `test_retrieval_query_plan.py` | `contract` | Local QueryPlan: extract URL/path objects. |
|
||||
| `test_retrieval_raw_candidates.py` | `contract` | 覆盖 positive only current authorized published chunks are searched、empty authority never reads database or vector、unapproved source never enters vector query 等 15 个场景。 |
|
||||
| `test_retrieval_raw_evidence.py` | `contract` | 覆盖 materialized text is exact current original、stale candidates do not authorize original read、budget does not cut a chunk into fake partial object 等 6 个场景。 |
|
||||
| `test_retrieval_raw_lane.py` | `contract` | v1 raw lane must not claim complete empty when the derived index is partial. |
|
||||
| `test_retrieval_runtime.py` | `contract` | 覆盖 constructor and missing start create nothing、stop releases supervisor worker but preserves immutable snapshot、not initialized or invalid configuration never creates index 等 21 个场景。 |
|
||||
| `test_retrieval_vector_store.py` | `contract` | 覆盖 constructor does no io and unopened calls do not create、open existing missing store or collection never creates、real chroma collection has no embedding function or documents 等 32 个场景。 |
|
||||
| `test_runtime.py` | `integration` | 覆盖 runtime writes and removes ownership receipt、runtime rejects second owner for same data directory、nonpublishing maintenance owner still excludes a second core 等 4 个场景。 |
|
||||
| `test_runtime_client.py` | `contract` | 覆盖 runtime inspection authenticates the discovered instance、runtime inspection rejects untrusted receipts without probing、runtime inspection distinguishes stopped and identity mismatch 场景。 |
|
||||
| `test_sanitization_source_map.py` | `contract` | 覆盖 mapping matches existing sanitizer and exact original slices、synthetic replacements never claim original locations、repeated identical suffix is mapped by positions not string search 等 6 个场景。 |
|
||||
@@ -328,4 +314,4 @@
|
||||
| `test_workspace_recall.py` | `contract` | Workspace memories were unreachable from every host; these pin the way in. |
|
||||
| `test_write_queue_lease_renewal.py` | `contract` | Rule 8: a lease must outlive the work it covers, and be renewed while it runs. |
|
||||
|
||||
共 323 个测试文件。
|
||||
共 309 个测试文件。
|
||||
@@ -14,8 +14,8 @@ from pathlib import Path
|
||||
import pytest
|
||||
|
||||
PRIMARY_MARKERS = ("unit", "contract", "integration", "performance", "live")
|
||||
EXPECTED_TEST_FILE_COUNT = 323
|
||||
EXPECTED_TEST_FILES_SHA256 = "10a2327bd047a281511980f5b558c1ea64d857f0fff550e4bd96cb53800205ba"
|
||||
EXPECTED_TEST_FILE_COUNT = 309
|
||||
EXPECTED_TEST_FILES_SHA256 = "ef6165c249971f5a2ab71a0c49f6d859c8c0f3612e3f749fba0acce0eab513ea"
|
||||
|
||||
# Pure transformations, validation models, parity checks and state machines.
|
||||
# This list is intentionally small: uncertain files default to contract.
|
||||
|
||||
@@ -1,132 +0,0 @@
|
||||
import socket
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import FrozenInstanceError
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory import local_embedding
|
||||
from one_memory.embedding_assets import AssetVerificationError, resolve_asset_dir
|
||||
from one_memory.local_embedding import LocalEmbedder, TokenizedText
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_network(monkeypatch):
|
||||
def refused(*args, **kwargs):
|
||||
pytest.fail("local tokenizer/encoder attempted network")
|
||||
|
||||
monkeypatch.setattr(socket.socket, "connect", refused)
|
||||
monkeypatch.setattr(socket, "create_connection", refused)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text", ["中文😀e\u0301末尾", "👩💻链接", "e\u0301 cafe\u0301", "URL路径 /Case/末尾"]
|
||||
)
|
||||
def test_real_locked_tokenizer_offsets_are_unicode_codepoints(text):
|
||||
embedder = LocalEmbedder()
|
||||
tokenized = embedder.tokenize_with_offsets(text)
|
||||
assert isinstance(tokenized, TokenizedText)
|
||||
assert isinstance(tokenized.token_ids, tuple) and isinstance(tokenized.offsets, tuple)
|
||||
ordinary = [offset for offset in tokenized.offsets if not offset.special]
|
||||
assert ordinary
|
||||
for offset in ordinary:
|
||||
assert 0 <= offset.start_cp < offset.end_cp <= len(text)
|
||||
assert text[offset.start_cp : offset.end_cp]
|
||||
assert all(
|
||||
any(o.start_cp <= i < o.end_cp for o in ordinary)
|
||||
for i, character in enumerate(text)
|
||||
if not character.isspace()
|
||||
)
|
||||
assert embedder._session is None and embedder._tokenizer is None
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
ordinary[0].start_cp = 10
|
||||
|
||||
|
||||
def test_known_emoji_and_combining_offsets_are_not_utf8_or_utf16_units():
|
||||
value = LocalEmbedder().tokenize_with_offsets("中文😀e\u0301末尾")
|
||||
offsets = [(o.start_cp, o.end_cp) for o in value.offsets if not o.special]
|
||||
assert offsets == [(0, 1), (1, 2), (2, 5), (5, 6), (6, 7)]
|
||||
|
||||
|
||||
def test_offsets_are_not_truncated_or_padded_and_do_not_create_onnx_session(monkeypatch):
|
||||
monkeypatch.setattr(local_embedding, "import_module", lambda *args: pytest.fail("ONNX load"))
|
||||
embedder = LocalEmbedder()
|
||||
long = embedder.tokenize_with_offsets("中" * 1600)
|
||||
short = embedder.tokenize_with_offsets("中")
|
||||
empty = embedder.tokenize_with_offsets("")
|
||||
assert len(long.token_ids) == 1602
|
||||
assert len(short.token_ids) == 3
|
||||
assert len(empty.token_ids) == 2
|
||||
assert all(o.special for o in empty.offsets)
|
||||
assert long.offsets[-2].end_cp == 1600
|
||||
assert embedder._session is None
|
||||
|
||||
|
||||
def test_parallel_offset_calls_share_one_lazy_tokenizer_without_shared_encoder_mutation():
|
||||
embedder = LocalEmbedder()
|
||||
texts = ["中" * 700, "😀中e\u0301", "short"] * 8
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
actual = list(pool.map(embedder.tokenize_with_offsets, texts))
|
||||
assert actual == [embedder.tokenize_with_offsets(text) for text in texts]
|
||||
assert embedder._session is None and embedder._tokenizer is None
|
||||
|
||||
|
||||
def test_offset_tokenizer_does_not_modify_encode_state_or_vectors():
|
||||
embedder = LocalEmbedder()
|
||||
baseline = embedder.encode(["原有编码必须保持", "短句"])
|
||||
session, tokenizer = embedder._session, embedder._tokenizer
|
||||
truncation, padding = tokenizer.truncation.copy(), tokenizer.padding.copy()
|
||||
with ThreadPoolExecutor(max_workers=4) as pool:
|
||||
tokenized = pool.submit(embedder.tokenize_with_offsets, "中" * 1100)
|
||||
encoded = pool.submit(embedder.encode, ["原有编码必须保持", "短句"])
|
||||
assert len(tokenized.result().token_ids) == 1102
|
||||
assert encoded.result() == baseline
|
||||
assert embedder._session is session and embedder._tokenizer is tokenizer
|
||||
assert tokenizer.truncation == truncation and tokenizer.padding == padding
|
||||
assert len(tokenizer.encode("中" * 1100).ids) == 512
|
||||
|
||||
|
||||
def test_offset_assets_reject_unlocked_tokenizer_without_loading_weights(tmp_path, monkeypatch):
|
||||
(tmp_path / "tokenizer.json").write_bytes(b'{"invalid":"not locked"}')
|
||||
monkeypatch.setattr(local_embedding, "resolve_asset_dir", lambda: tmp_path)
|
||||
with pytest.raises(AssetVerificationError):
|
||||
LocalEmbedder().tokenize_with_offsets("中")
|
||||
|
||||
|
||||
def test_offset_tokenization_uses_only_tokenizer_asset(tmp_path, monkeypatch):
|
||||
# Tokenization must not require another ONNX session or even the weights file.
|
||||
(tmp_path / "tokenizer.json").write_bytes((resolve_asset_dir() / "tokenizer.json").read_bytes())
|
||||
monkeypatch.setattr(local_embedding, "resolve_asset_dir", lambda: tmp_path)
|
||||
assert len(LocalEmbedder().tokenize_with_offsets("中文").token_ids) == 4
|
||||
|
||||
|
||||
def test_real_adjacent_wordpieces_and_long_ascii_unknown_keep_original_offsets():
|
||||
embedder = LocalEmbedder()
|
||||
word = "unaffordable"
|
||||
parts = [o for o in embedder.tokenize_with_offsets(word).offsets if not o.special]
|
||||
assert len(parts) > 1
|
||||
assert parts[0].start_cp == 0 and parts[-1].end_cp == len(word)
|
||||
assert all(a.end_cp == b.start_cp for a, b in zip(parts, parts[1:], strict=False))
|
||||
long_word = "x" * 6000
|
||||
unknown = embedder.tokenize_with_offsets(long_word)
|
||||
offsets = [o for o in unknown.offsets if not o.special]
|
||||
assert len(offsets) == 1 and unknown.token_ids[1] == 100 # locked BERT [UNK]
|
||||
assert (offsets[0].start_cp, offsets[0].end_cp) == (0, 6000)
|
||||
|
||||
|
||||
def test_parallel_initial_load_constructs_one_offsets_tokenizer(monkeypatch):
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
calls = []
|
||||
original = Tokenizer.from_str
|
||||
|
||||
def counted(value):
|
||||
calls.append(1)
|
||||
return original(value)
|
||||
|
||||
monkeypatch.setattr(Tokenizer, "from_str", counted)
|
||||
embedder = LocalEmbedder()
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
results = list(pool.map(embedder.tokenize_with_offsets, ["中😀e\u0301"] * 16))
|
||||
assert len(calls) == 1
|
||||
assert all(value == results[0] for value in results)
|
||||
assert embedder._session is None
|
||||
@@ -87,9 +87,24 @@ ALLOWED = {
|
||||
"SQLite's own defaults, read once from a connection that opens no file",
|
||||
opened_with='":memory:"',
|
||||
),
|
||||
("retrieval/vector_store.py", "RawVectorStore._preflight"): Allowed(
|
||||
"Chroma's own chroma.sqlite3, read-only, checked before Chroma opens it",
|
||||
opened_with='"?mode=ro"',
|
||||
("sdk_migration_verify.py", "legacy_memory_ids_from_bundle"): Allowed(
|
||||
"the pre-cutover bundle's copy of the legacy database, read-only, to list "
|
||||
"the memory ids a migrated home must still read back",
|
||||
opened_with="?mode=ro",
|
||||
),
|
||||
("hy_sdk_migration.py", "_legacy_pools"): Allowed(
|
||||
"each pre-cutover bundle's copy of the legacy database, read-only, to "
|
||||
"recover the source and times historical cutovers dropped",
|
||||
opened_with="?mode=ro",
|
||||
),
|
||||
("hy_sdk_migration.py", "_repair_coding"): Allowed(
|
||||
"the SDK's own Coding/Working stores, not the main database: the one-time "
|
||||
"metadata repair writes them at startup, before the SDK client opens"
|
||||
),
|
||||
("sdk_vector_reattach.py", "_stored_names"): Allowed(
|
||||
"Chroma's own chroma.sqlite3, read-only, to list collection names before "
|
||||
"the offline migration opens Chroma",
|
||||
opened_with="?mode=ro",
|
||||
),
|
||||
("imports/sqlite_access.py", "open_sqlite_readonly"): Allowed(
|
||||
"a host app's own SQLite (Cursor state.vscdb / store.db): one direct "
|
||||
|
||||
@@ -278,3 +278,216 @@ def test_sdk_user_profile_http_returns_native_assembled_profile(tmp_path):
|
||||
}
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.fresh_sdk
|
||||
def test_sdk_memory_library_uses_sdk_search_and_skips_empty_rows_without_mutation(
|
||||
tmp_path,
|
||||
caplog,
|
||||
):
|
||||
async def scenario():
|
||||
app = application(tmp_path)
|
||||
list_calls = []
|
||||
search_calls = []
|
||||
deleted = []
|
||||
records = {
|
||||
"release": {
|
||||
"memory_id": "release",
|
||||
"content": "Release plan",
|
||||
"layer": "l2_fact",
|
||||
"status": "active",
|
||||
"user_id": "local",
|
||||
"agent_id": "all-apps",
|
||||
"session_id": "s1",
|
||||
"source": "cursor",
|
||||
"tags": ["shipping"],
|
||||
},
|
||||
"retired": {
|
||||
"memory_id": "retired",
|
||||
"content": "Release plan, first draft",
|
||||
"layer": "l2_fact",
|
||||
"status": "superseded",
|
||||
"user_id": "local",
|
||||
"agent_id": "all-apps",
|
||||
"session_id": "s1",
|
||||
"source": "cursor",
|
||||
"tags": [],
|
||||
},
|
||||
"default": {
|
||||
"memory_id": "default",
|
||||
"content": "Unknown source",
|
||||
"layer": "l2_fact",
|
||||
"status": "active",
|
||||
"user_id": "local",
|
||||
"agent_id": "all-apps",
|
||||
"session_id": "s3",
|
||||
"source": "default",
|
||||
"tags": [],
|
||||
},
|
||||
}
|
||||
|
||||
class Client:
|
||||
def list_memories(self, **kwargs):
|
||||
list_calls.append(kwargs)
|
||||
rows = [
|
||||
records["release"],
|
||||
records["default"],
|
||||
{
|
||||
"memory_id": "empty",
|
||||
"content": "",
|
||||
"layer": "l2_fact",
|
||||
"status": "active",
|
||||
"user_id": "local",
|
||||
"agent_id": "all-apps",
|
||||
"session_id": "s2",
|
||||
"source": "workbuddy",
|
||||
"tags": [],
|
||||
},
|
||||
]
|
||||
include_sources = kwargs.get("include_sources") or []
|
||||
if include_sources:
|
||||
rows = [row for row in rows if row["source"] in include_sources]
|
||||
return {
|
||||
"vdb": {
|
||||
"memories": rows,
|
||||
"total": len(rows),
|
||||
"limit": kwargs["limit"],
|
||||
"offset": kwargs["offset"],
|
||||
}
|
||||
}
|
||||
|
||||
def search(self, query, **kwargs):
|
||||
search_calls.append((query, kwargs))
|
||||
return {
|
||||
"memories": {
|
||||
"schema": [],
|
||||
"proactive": [],
|
||||
"normal": [
|
||||
{
|
||||
"memory_id": "release",
|
||||
"content": "Release plan",
|
||||
"score": 0.9,
|
||||
},
|
||||
{
|
||||
"memory_id": "retired",
|
||||
"content": "Release plan, first draft",
|
||||
"score": 0.8,
|
||||
},
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
def get(self, memory_id):
|
||||
return records.get(memory_id)
|
||||
|
||||
def delete(self, memory_id):
|
||||
deleted.append(memory_id)
|
||||
return {"success": True, "deleted_count": 1}
|
||||
|
||||
sdk = Client()
|
||||
|
||||
class Handle:
|
||||
async def client(self):
|
||||
return sdk
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
app.state.sdk_client_handle = Handle()
|
||||
async with client_for(app) as client:
|
||||
response = await client.get(
|
||||
"/api/v2/memory-library",
|
||||
params={"source": "cursor", "query": "release"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
# `total` counts the window's keyword matches; the superseded
|
||||
# hit matches but is not listed as an active memory.
|
||||
assert body["total"] == 2
|
||||
assert body["search_window"] is None
|
||||
assert [item["memory_id"] for item in body["memories"]] == ["release"]
|
||||
assert search_calls[0][0] == "release"
|
||||
assert search_calls[0][1]["include_sources"] == ["cursor"]
|
||||
assert search_calls[0][1]["limit"] == 20
|
||||
assert search_calls[0][1]["intention_limit"] == 0
|
||||
assert search_calls[0][1]["min_score"] == 0.0
|
||||
assert list_calls == []
|
||||
|
||||
# The SDK still returns its nearest rows for a word that is in
|
||||
# none of them; the library reports no match instead.
|
||||
unmatched = await client.get(
|
||||
"/api/v2/memory-library", params={"query": "zzz nowhere"}
|
||||
)
|
||||
assert unmatched.json()["total"] == 0
|
||||
assert unmatched.json()["memories"] == []
|
||||
|
||||
listed = await client.get(
|
||||
"/api/v2/memory-library",
|
||||
params={"source": "workbuddy", "limit": 100},
|
||||
)
|
||||
assert listed.status_code == 200
|
||||
assert listed.json()["memories"] == []
|
||||
assert listed.json()["total"] == 1
|
||||
assert deleted == []
|
||||
assert list_calls[0]["limit"] == 100
|
||||
assert "sdk_memory_invalid_rows_skipped count=1" in caplog.messages
|
||||
|
||||
graph = await client.get("/api/v2/memories/graph")
|
||||
assert graph.status_code == 200
|
||||
node_ids = {node["node_id"] for node in graph.json()["nodes"]}
|
||||
assert "source:cursor" in node_ids
|
||||
assert "source:default" not in node_ids
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.fresh_sdk
|
||||
def test_sdk_context_search_sets_each_recall_lane_apart(tmp_path):
|
||||
"""Normal, schema and intention lanes default to 10 / 3 / 1 and pass through."""
|
||||
|
||||
async def scenario():
|
||||
app = application(tmp_path)
|
||||
calls = []
|
||||
|
||||
class Client:
|
||||
def search(self, query, **kwargs):
|
||||
calls.append(kwargs)
|
||||
return {"memories": {"schema": [], "proactive": [], "normal": []}}
|
||||
|
||||
def get(self, memory_id):
|
||||
return None
|
||||
|
||||
sdk = Client()
|
||||
|
||||
class Handle:
|
||||
async def client(self):
|
||||
return sdk
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
app.state.sdk_client_handle = Handle()
|
||||
async with client_for(app) as client:
|
||||
defaulted = await client.post(
|
||||
"/api/v2/context/search",
|
||||
json={"query": "release", "kinds": ["memory"]},
|
||||
)
|
||||
chosen = await client.post(
|
||||
"/api/v2/context/search",
|
||||
json={
|
||||
"query": "release",
|
||||
"kinds": ["memory"],
|
||||
"limit": 5,
|
||||
"schema_limit": 0,
|
||||
"intention_limit": 2,
|
||||
},
|
||||
)
|
||||
assert defaulted.status_code == 200
|
||||
assert chosen.status_code == 200
|
||||
assert [
|
||||
(call["limit"], call["schema_limit"], call["intention_limit"]) for call in calls
|
||||
] == [(10, 3, 1), (5, 0, 2)]
|
||||
|
||||
asyncio.run(scenario())
|
||||
@@ -85,7 +85,6 @@ async def application(tmp_path: Path):
|
||||
live_sync_enabled=False,
|
||||
live_sync_watch_enabled=False,
|
||||
automatic_memory_processing=False,
|
||||
retrieval_index_enabled=False,
|
||||
)
|
||||
app = create_app(settings, source_catalog={}, secret_store=secrets)
|
||||
async with (
|
||||
|
||||
@@ -592,7 +592,6 @@ def test_probe_api_optional_body_lookup_and_authentication(tmp_path, monkeypatch
|
||||
automatic_memory_processing=False,
|
||||
live_sync_enabled=False,
|
||||
live_sync_watch_enabled=False,
|
||||
retrieval_index_enabled=False,
|
||||
),
|
||||
source_catalog={},
|
||||
secret_store=secrets,
|
||||
|
||||
@@ -136,6 +136,64 @@ def test_cursor_store_timestamp_envelope_keeps_only_the_user_query_text(tmp_path
|
||||
assert meta["turn_event_at"] == "2026-09-22T23:34:00+00:00"
|
||||
|
||||
|
||||
def test_cursor_wal_change_token_does_not_become_a_future_conversation_time(
|
||||
tmp_path,
|
||||
) -> None:
|
||||
root = tmp_path / ".cursor"
|
||||
database = root / "chats" / "workspace" / "wal-time" / "store.db"
|
||||
database.parent.mkdir(parents=True)
|
||||
connection = sqlite3.connect(database)
|
||||
connection.execute("PRAGMA journal_mode=WAL")
|
||||
connection.execute("PRAGMA wal_autocheckpoint=0")
|
||||
connection.execute("CREATE TABLE blobs (id TEXT, data TEXT)")
|
||||
connection.execute(
|
||||
"INSERT INTO blobs VALUES (?, ?)",
|
||||
("message", json.dumps({"role": "assistant", "content": "answer"})),
|
||||
)
|
||||
connection.commit()
|
||||
stamp = 1_788_220_800
|
||||
os.utime(database, (stamp, stamp))
|
||||
try:
|
||||
records = evidence(
|
||||
asyncio.run(scan(CursorUnifiedSourceAdapter(grant=grant("cursor", root))))
|
||||
)
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
assert len(records) == 1
|
||||
assert int(records[0].evidence.created_at.timestamp()) == stamp
|
||||
|
||||
|
||||
def test_cursor_wal_store_resumes_from_its_rowid_after_new_frames(tmp_path) -> None:
|
||||
"""The WAL change token marks a store changed; it must not reset its checkpoint."""
|
||||
root = tmp_path / ".cursor"
|
||||
database = root / "chats" / "workspace" / "wal-resume" / "store.db"
|
||||
database.parent.mkdir(parents=True)
|
||||
connection = sqlite3.connect(database)
|
||||
connection.execute("PRAGMA journal_mode=WAL")
|
||||
connection.execute("PRAGMA wal_autocheckpoint=0")
|
||||
connection.execute("CREATE TABLE blobs (id TEXT, data TEXT)")
|
||||
connection.execute(
|
||||
"INSERT INTO blobs VALUES (?, ?)",
|
||||
("old", json.dumps({"id": "old", "role": "user", "content": "old row"})),
|
||||
)
|
||||
connection.commit()
|
||||
adapter = CursorUnifiedSourceAdapter(grant=grant("cursor", root))
|
||||
try:
|
||||
baseline = asyncio.run(adapter.live_baseline())
|
||||
assert evidence(asyncio.run(scan(adapter, cursors=baseline))) == []
|
||||
connection.execute(
|
||||
"INSERT INTO blobs VALUES (?, ?)",
|
||||
("new", json.dumps({"id": "new", "role": "user", "content": "new row"})),
|
||||
)
|
||||
connection.commit()
|
||||
appended = asyncio.run(scan(adapter, cursors=baseline))
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
assert [item.evidence.verbatim for item in evidence(appended)] == ["new row"]
|
||||
|
||||
|
||||
def test_cursor_transcript_timestamp_envelope_keeps_only_the_user_query_text(tmp_path) -> None:
|
||||
root = tmp_path / ".cursor"
|
||||
transcript = (
|
||||
|
||||
@@ -1,15 +1,13 @@
|
||||
"""Search imported conversations with no generative model configured at all.
|
||||
|
||||
Every other retrieval test either builds the index by hand or stubs the
|
||||
encoder, so none of them answers the question a user would ask: I imported my
|
||||
history and I never set up a model, can One Memory find things I said in
|
||||
different words? This goes through the product's own entry points only:
|
||||
authorize the source, run an import job, search `/api/v1/context/search`.
|
||||
This answers the question a user would ask: I imported my history and I
|
||||
never set up a model, can One Memory find things I said? This goes through the
|
||||
product's own entry points only: authorize the source, run an import job,
|
||||
search `/api/v1/context/search`.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import pytest
|
||||
@@ -18,7 +16,6 @@ from pydantic import SecretStr
|
||||
|
||||
from one_memory.api import create_app
|
||||
from one_memory.config import Settings
|
||||
from one_memory.embedding_assets import assets_present
|
||||
from one_memory.imports.sources.registry import build_source_catalog
|
||||
from one_memory.secrets import InMemorySecretStore
|
||||
from tests.source_consent import authorized_request, thaw_source
|
||||
@@ -29,14 +26,7 @@ LICENCE = "我们最后决定 One Memory 用 MIT 协议发布,不再考虑 GPL
|
||||
BACKUP = "备份先写到本地磁盘,再同步到对象存储。"
|
||||
SYNONYM = "开源软件的授权条款最后怎么挑的"
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.integration,
|
||||
pytest.mark.serial,
|
||||
pytest.mark.requires_embedding_assets,
|
||||
pytest.mark.skipif(
|
||||
not assets_present(), reason="built-in embedding weights required for the semantic lane"
|
||||
),
|
||||
]
|
||||
pytestmark = [pytest.mark.integration, pytest.mark.serial]
|
||||
|
||||
|
||||
def transcript(root, name, *texts):
|
||||
@@ -57,7 +47,7 @@ def transcript(root, name, *texts):
|
||||
(root / f"rollout-{name}.jsonl").write_text("\n".join(lines) + "\n")
|
||||
|
||||
|
||||
def application(tmp_path, *, index: bool):
|
||||
def application(tmp_path):
|
||||
home = tmp_path / "home"
|
||||
root = home / ".codex/sessions"
|
||||
root.mkdir(parents=True)
|
||||
@@ -73,10 +63,8 @@ def application(tmp_path, *, index: bool):
|
||||
api_token=SecretStr(TOKEN),
|
||||
automatic_memory_processing=False,
|
||||
live_sync_enabled=False,
|
||||
retrieval_index_enabled=index,
|
||||
),
|
||||
source_catalog={"codex": build_source_catalog(home=home, environment={})["codex"]},
|
||||
manage_runtime=index,
|
||||
secret_store=InMemorySecretStore(),
|
||||
)
|
||||
|
||||
@@ -106,20 +94,6 @@ async def search(client, query, *, attempts=1):
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
|
||||
def lexical_only(app, query):
|
||||
"""What the word-matching reader alone returns, for the same request."""
|
||||
from one_memory.context_models import ContextSearchRequest
|
||||
from one_memory.context_raw import read_raw_context
|
||||
|
||||
result = read_raw_context(
|
||||
app.state.evidence_store.database_path,
|
||||
ContextSearchRequest(query=query, kinds=("raw",)),
|
||||
source_ids=frozenset({"codex"}),
|
||||
deadline=time.monotonic() + 5,
|
||||
)
|
||||
return tuple(result.items)
|
||||
|
||||
|
||||
def texts(body):
|
||||
return [
|
||||
message["text"]
|
||||
@@ -131,7 +105,7 @@ def texts(body):
|
||||
|
||||
def test_shared_words_are_found_without_any_model_or_index(tmp_path):
|
||||
async def scenario():
|
||||
app = application(tmp_path, index=False)
|
||||
app = application(tmp_path)
|
||||
async with (
|
||||
app.router.lifespan_context(app),
|
||||
AsyncClient(
|
||||
@@ -144,38 +118,8 @@ def test_shared_words_are_found_without_any_model_or_index(tmp_path):
|
||||
assert any(LICENCE in text for text in texts(body))
|
||||
|
||||
missed = await search(client, SYNONYM)
|
||||
assert not texts(missed), (
|
||||
"the lexical lane already answers this phrasing, so the semantic "
|
||||
"test below would prove nothing"
|
||||
)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_a_semantic_only_candidate_is_not_released_without_confirmation(tmp_path):
|
||||
async def scenario():
|
||||
app = application(tmp_path, index=True)
|
||||
async with (
|
||||
app.router.lifespan_context(app),
|
||||
AsyncClient(
|
||||
transport=ASGITransport(app=app), base_url="http://test", headers=AUTH
|
||||
) as client,
|
||||
):
|
||||
assert app.state.retrieval_runtime.snapshot().state == "running"
|
||||
await imported(app, client)
|
||||
deadline = time.monotonic() + 30
|
||||
while time.monotonic() < deadline:
|
||||
snapshot = app.state.retrieval_runtime.snapshot().supervisor
|
||||
if snapshot is not None and snapshot.slices >= 4:
|
||||
break
|
||||
await asyncio.sleep(1)
|
||||
assert snapshot is not None and snapshot.slices >= 4 and snapshot.failures == 0
|
||||
body = await search(client, SYNONYM)
|
||||
assert body["paths"]["raw"]["reason"] == "lexical_only", body
|
||||
assert not texts(body), "a semantic-only candidate must not release original text"
|
||||
assert not lexical_only(app, SYNONYM), "lexical alone answers this; nothing is proven"
|
||||
assert any(LICENCE in text for text in texts(await search(client, "MIT 协议"))), (
|
||||
"the index must not cost the lexical lane its own hits"
|
||||
)
|
||||
# Different words are not a lexical match, and there is no private
|
||||
# vector index to pretend otherwise.
|
||||
assert not texts(missed)
|
||||
|
||||
asyncio.run(scenario())
|
||||
@@ -1,273 +0,0 @@
|
||||
from dataclasses import FrozenInstanceError
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.local_embedding import LocalEmbedder, TokenizedText, TokenOffset
|
||||
from one_memory.retrieval.chunking import ChunkingError, chunk_evidence
|
||||
|
||||
|
||||
class Characters:
|
||||
def tokenize_with_offsets(self, text):
|
||||
offsets = (
|
||||
TokenOffset(0, 0, True),
|
||||
*(TokenOffset(i, i + 1) for i in range(len(text))),
|
||||
TokenOffset(0, 0, True),
|
||||
)
|
||||
return TokenizedText(tuple(range(len(offsets))), offsets)
|
||||
|
||||
|
||||
def assert_coverage(text, chunks, tokenizer, prefix=""):
|
||||
assert chunks
|
||||
assert [chunk.ordinal for chunk in chunks] == list(range(len(chunks)))
|
||||
covered = bytearray(len(text))
|
||||
previous_start = -1
|
||||
for chunk in chunks:
|
||||
assert 0 <= chunk.start_cp < chunk.end_cp <= len(text)
|
||||
assert chunk.start_cp > previous_start
|
||||
previous_start = chunk.start_cp
|
||||
assert chunk.text == text[chunk.start_cp : chunk.end_cp]
|
||||
assert chunk.token_count == len(
|
||||
tokenizer.tokenize_with_offsets(prefix + chunk.text).token_ids
|
||||
)
|
||||
assert chunk.token_count <= 512
|
||||
covered[chunk.start_cp : chunk.end_cp] = b"\x01" * (chunk.end_cp - chunk.start_cp)
|
||||
assert all(covered[i] for i, c in enumerate(text) if not c.isspace())
|
||||
assert chunks[-1].end_cp == len(text)
|
||||
|
||||
|
||||
def test_default_384_body_64_overlap_and_exact_source_slices():
|
||||
text = "中" * 1000
|
||||
chunks = chunk_evidence(text, Characters())
|
||||
assert [(c.start_cp, c.end_cp) for c in chunks] == [(0, 384), (320, 704), (640, 1000)]
|
||||
assert chunks[0].token_count == 386
|
||||
assert_coverage(text, chunks, Characters())
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
chunks[0].text = "replacement"
|
||||
|
||||
|
||||
def test_all_prefix_and_special_tokens_count_toward_strict_512_budget():
|
||||
prefix = "P" * 200
|
||||
text = "中" * 900
|
||||
chunks = chunk_evidence(text, Characters(), prefix=prefix)
|
||||
assert chunks[0].end_cp == 310
|
||||
assert chunks[0].token_count == 512
|
||||
assert_coverage(text, chunks, Characters(), prefix)
|
||||
with pytest.raises(ChunkingError, match="budget"):
|
||||
chunk_evidence("中", Characters(), prefix="P" * 510)
|
||||
|
||||
|
||||
def test_real_tokenizer_covers_long_chinese_middle_and_tail_object_locations():
|
||||
middle = "https://example.test/Middle?Key=Value#Anchor"
|
||||
tail = "/CaseSensitive/路径/final.txt"
|
||||
text = "中文😀e\u0301\n" + "中" * 740 + "\n" + middle + "\n" + "尾" * 800 + "\n" + tail
|
||||
tokenizer = LocalEmbedder()
|
||||
chunks = chunk_evidence(text, tokenizer)
|
||||
assert len(chunks) > 3
|
||||
assert_coverage(text, chunks, tokenizer)
|
||||
for value in (middle, tail):
|
||||
location = text.index(value)
|
||||
assert any(c.start_cp <= location and c.end_cp >= location + len(value) for c in chunks)
|
||||
assert chunks[0].text.startswith("中文😀e\u0301\n")
|
||||
|
||||
|
||||
def test_prefers_paragraph_boundary_near_target_without_losing_text():
|
||||
text = "中" * 300 + "\n\n" + "尾" * 500
|
||||
chunks = chunk_evidence(text, Characters())
|
||||
assert chunks[0].end_cp == 302
|
||||
assert_coverage(text, chunks, Characters())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text", ["", " \n\t\r "])
|
||||
def test_empty_or_whitespace_has_no_index_chunks(text):
|
||||
assert chunk_evidence(text, Characters()) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text", ["😀" * 300, "e\u0301" * 400, "甲\n乙\t丙", "原样 <user_query>原文</user_query>"]
|
||||
)
|
||||
def test_unicode_and_wrapper_text_identity_is_preserved(text):
|
||||
chunks = chunk_evidence(text, Characters())
|
||||
assert_coverage(text, chunks, Characters())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"offsets,ids",
|
||||
[
|
||||
((TokenOffset(0, 4),), (1,)),
|
||||
((TokenOffset(-1, 1),), (1,)),
|
||||
((TokenOffset(1, 2), TokenOffset(0, 1)), (1, 2)),
|
||||
((TokenOffset(0, 1), TokenOffset(2, 3)), (1, 2)),
|
||||
((TokenOffset(0, 0),), (1,)),
|
||||
((TokenOffset(0, 1, True),), (1,)),
|
||||
((TokenOffset(0, 3),), (1, 2)),
|
||||
((), ()),
|
||||
],
|
||||
)
|
||||
def test_bad_or_incomplete_offsets_fail_explicitly(offsets, ids):
|
||||
class Broken:
|
||||
def tokenize_with_offsets(self, text):
|
||||
return TokenizedText(ids, offsets)
|
||||
|
||||
with pytest.raises(ChunkingError):
|
||||
chunk_evidence("abc", Broken())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs",
|
||||
[
|
||||
{"target_tokens": 64},
|
||||
{"overlap_tokens": -1},
|
||||
{"max_tokens": 513},
|
||||
{"max_tokens": True},
|
||||
{"target_tokens": 0},
|
||||
{"overlap_tokens": 384},
|
||||
],
|
||||
)
|
||||
def test_invalid_budgets_fail_before_tokenizing(kwargs):
|
||||
with pytest.raises(ChunkingError):
|
||||
chunk_evidence("content", Characters(), **kwargs)
|
||||
|
||||
|
||||
def test_retokenizes_actual_prefix_plus_slice_instead_of_assuming_additive_counts():
|
||||
class ContextSensitive(Characters):
|
||||
def tokenize_with_offsets(self, text):
|
||||
result = super().tokenize_with_offsets(text)
|
||||
if text.startswith("P"):
|
||||
# Same valid offsets may carry multiple pieces. Budget must include them all.
|
||||
return TokenizedText(
|
||||
result.token_ids * 2,
|
||||
tuple(o for o in result.offsets if o.special)
|
||||
+ tuple(o for o in result.offsets if o.special)
|
||||
+ tuple(o for o in result.offsets if not o.special for _ in range(2)),
|
||||
)
|
||||
return result
|
||||
|
||||
tokenizer = ContextSensitive()
|
||||
chunks = chunk_evidence("中" * 600, tokenizer, prefix="P")
|
||||
assert_coverage("中" * 600, chunks, tokenizer, "P")
|
||||
|
||||
|
||||
def test_overlapping_token_offsets_are_not_split_inside_a_character():
|
||||
class Repeated(Characters):
|
||||
def tokenize_with_offsets(self, text):
|
||||
return TokenizedText(
|
||||
tuple(range(2 * len(text))),
|
||||
tuple(TokenOffset(i, i + 1) for i in range(len(text)) for _ in range(2)),
|
||||
)
|
||||
|
||||
tokenizer = Repeated()
|
||||
text = "😀" * 500
|
||||
chunks = chunk_evidence(text, tokenizer)
|
||||
assert_coverage(text, chunks, tokenizer)
|
||||
|
||||
|
||||
def test_single_indivisible_offset_over_budget_is_not_truncated():
|
||||
class Indivisible:
|
||||
def tokenize_with_offsets(self, text):
|
||||
return TokenizedText(
|
||||
tuple(range(600)), tuple(TokenOffset(0, len(text)) for _ in range(600))
|
||||
)
|
||||
|
||||
with pytest.raises(ChunkingError, match="budget"):
|
||||
chunk_evidence("😀", Indivisible())
|
||||
|
||||
|
||||
def test_real_long_continuous_ascii_unknown_preserves_all_6000_characters():
|
||||
tokenizer = LocalEmbedder()
|
||||
text = " \t" + "x" * 6000 + "\n末尾链接 https://example.test/Case?A=B#C\t"
|
||||
chunks = chunk_evidence(text, tokenizer)
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].text == text
|
||||
assert_coverage(text, chunks, tokenizer)
|
||||
|
||||
|
||||
def test_real_subword_and_whitespace_boundaries_have_no_gaps():
|
||||
tokenizer = LocalEmbedder()
|
||||
text = " \t\n" + ("unaffordable embeddings\t中\u0301文 a\u0301b\n\n" * 180) + "\t\n "
|
||||
prefix = "本段用于检索:"
|
||||
chunks = chunk_evidence(text, tokenizer, prefix=prefix)
|
||||
assert len(chunks) > 3
|
||||
assert chunks[0].start_cp == 0 and chunks[-1].end_cp == len(text)
|
||||
assert_coverage(text, chunks, tokenizer, prefix)
|
||||
for first, second in zip(chunks, chunks[1:], strict=False):
|
||||
assert second.start_cp <= first.end_cp
|
||||
|
||||
|
||||
def test_partially_overlapping_offsets_for_combining_cluster_are_indivisible():
|
||||
class Combining(Characters):
|
||||
def tokenize_with_offsets(self, text):
|
||||
offsets = []
|
||||
i = 0
|
||||
while i < len(text):
|
||||
if text.startswith("e\u0301", i):
|
||||
offsets.extend((TokenOffset(i, i + 2), TokenOffset(i + 1, i + 2)))
|
||||
i += 2
|
||||
else:
|
||||
offsets.append(TokenOffset(i, i + 1))
|
||||
i += 1
|
||||
return TokenizedText(tuple(range(len(offsets))), tuple(offsets))
|
||||
|
||||
text = "中" * 383 + "e\u0301" + "尾" * 500
|
||||
tokenizer = Combining()
|
||||
chunks = chunk_evidence(text, tokenizer)
|
||||
assert_coverage(text, chunks, tokenizer)
|
||||
assert all(c.start_cp != 384 and c.end_cp != 384 for c in chunks)
|
||||
|
||||
|
||||
def test_sparse_whitespace_gaps_are_retained_not_rebuilt_from_tokens():
|
||||
tokenizer = LocalEmbedder()
|
||||
text = " \n" + "中" * 380 + "\n \t\n" + "尾" * 380 + " \t\r\n"
|
||||
chunks = chunk_evidence(text, tokenizer)
|
||||
assert_coverage(text, chunks, tokenizer)
|
||||
assert chunks[0].text.startswith(" \n")
|
||||
assert chunks[-1].text.endswith(" \t\r\n")
|
||||
|
||||
|
||||
def test_real_prefix_retokenization_never_exceeds_encoder_window():
|
||||
tokenizer = LocalEmbedder()
|
||||
prefix = "检索前缀" * 50
|
||||
text = ("中文tokenization😀e\u0301\n" * 120) + "https://example.test/Tail"
|
||||
chunks = chunk_evidence(text, tokenizer, prefix=prefix)
|
||||
assert len(chunks) > 1
|
||||
assert_coverage(text, chunks, tokenizer, prefix)
|
||||
|
||||
|
||||
def test_tokenizer_failure_is_explicit_not_partial_output():
|
||||
class Broken:
|
||||
def tokenize_with_offsets(self, text):
|
||||
raise RuntimeError("tokenizer unavailable")
|
||||
|
||||
with pytest.raises(ChunkingError, match="tokenization_failed"):
|
||||
chunk_evidence("中", Broken())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text", ["e\u0301", "\u0301", "中\u0301文", "\n\r\t中\n尾\n", "👩💻", "👩👩👧👦", "✈️"]
|
||||
)
|
||||
def test_locked_normalizer_supported_unicode_examples_cover_original(text):
|
||||
tokenizer = LocalEmbedder()
|
||||
assert_coverage(text, chunk_evidence(text, tokenizer), tokenizer)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text", ["中\u200d文", "\u200d", "中\u200b文", "中\x00文", "\x00"])
|
||||
def test_locked_normalizer_dropped_nonwhitespace_is_explicitly_unsupported(text):
|
||||
tokenizer = LocalEmbedder()
|
||||
offsets = tokenizer.tokenize_with_offsets(text).offsets
|
||||
assert any(
|
||||
not character.isspace()
|
||||
and not any(
|
||||
not offset.special and offset.start_cp <= i < offset.end_cp for offset in offsets
|
||||
)
|
||||
for i, character in enumerate(text)
|
||||
)
|
||||
with pytest.raises(ChunkingError, match="uncovered_nonwhitespace"):
|
||||
chunk_evidence(text, tokenizer)
|
||||
|
||||
|
||||
def test_another_normalizer_dropping_trailing_combining_mark_is_not_reported_complete():
|
||||
class NormalizesAccent:
|
||||
def tokenize_with_offsets(self, text):
|
||||
return TokenizedText((1,), (TokenOffset(0, 1),))
|
||||
|
||||
with pytest.raises(ChunkingError, match="uncovered_nonwhitespace"):
|
||||
chunk_evidence("e\u0301", NormalizesAccent())
|
||||
@@ -1,38 +0,0 @@
|
||||
import importlib.util
|
||||
import sqlite3
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.performance, pytest.mark.serial]
|
||||
|
||||
SCRIPT = Path(__file__).resolve().parents[3] / "scripts" / "bench_retrieval_eligibility_scale.py"
|
||||
if not SCRIPT.exists():
|
||||
pytest.skip(
|
||||
f"benchmark script is not part of this checkout: {SCRIPT.name}",
|
||||
allow_module_level=True,
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location("bench_retrieval_eligibility_scale", SCRIPT)
|
||||
benchmark = importlib.util.module_from_spec(spec)
|
||||
sys.modules["bench_retrieval_eligibility_scale"] = benchmark
|
||||
assert spec.loader is not None
|
||||
spec.loader.exec_module(benchmark)
|
||||
|
||||
|
||||
def test_scale_probe_proves_composite_authority_index_changes_the_plan(tmp_path):
|
||||
db = sqlite3.connect(tmp_path / "probe.sqlite3", cached_statements=0)
|
||||
try:
|
||||
benchmark.create_schema(db)
|
||||
benchmark.fill_to(db, 0, 10_000)
|
||||
before = benchmark.measure(db, benchmark.FILTERED_SQL, ("g", "allowed"), repeat=3)
|
||||
db.execute(
|
||||
"CREATE INDEX retrieval_authority_jobs "
|
||||
"ON retrieval_documents(generation_id,state,source_id,event_seq)"
|
||||
)
|
||||
after = benchmark.measure(db, benchmark.FILTERED_SQL, ("g", "allowed"), repeat=3)
|
||||
finally:
|
||||
db.close()
|
||||
assert before["rows"] == after["rows"] == 100
|
||||
assert not any("source_id" in step for step in before["plan"])
|
||||
assert any("retrieval_authority_jobs" in step for step in after["plan"])
|
||||
@@ -1,452 +0,0 @@
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory import ingest_journal
|
||||
from one_memory.imports.grants import SourceGrantStore
|
||||
from one_memory.imports.models import EvidenceInput, EvidenceOrigin, OriginMetadata, WrittenBy
|
||||
from one_memory.imports.sources.registry import RootCandidateKind
|
||||
from one_memory.imports.storage import EvidenceStore
|
||||
from one_memory.memory_server.config_state import ConfigStore
|
||||
from one_memory.retrieval.chunking import TextChunk
|
||||
from one_memory.retrieval.index_store import EncoderIdentity, RetrievalIndexStore
|
||||
from one_memory.retrieval.models import TimeRange
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fixture(tmp_path):
|
||||
path = tmp_path / "data.sqlite3"
|
||||
evidence = EvidenceStore(path)
|
||||
grants = SourceGrantStore(path, home=tmp_path)
|
||||
source = tmp_path / "source"
|
||||
source.mkdir()
|
||||
grants.grant(
|
||||
user_id="local", source_id="codex", root=source, candidate_kind=RootCandidateKind.DEFAULT
|
||||
)
|
||||
store = RetrievalIndexStore(path)
|
||||
identity = EncoderIdentity("test", "a" * 64, "b" * 64, 3)
|
||||
return path, evidence, grants, store, identity
|
||||
|
||||
|
||||
def put(evidence, text="original text", message="one"):
|
||||
return evidence.put(
|
||||
EvidenceInput(
|
||||
source_id="codex",
|
||||
conversation_id="thread",
|
||||
message_id=message,
|
||||
role="user",
|
||||
written_by=WrittenBy.PERSON,
|
||||
verbatim=text,
|
||||
origin_meta={EvidenceOrigin.FILE: OriginMetadata()},
|
||||
)
|
||||
).record
|
||||
|
||||
|
||||
def initialize(f):
|
||||
path, _, _, store, identity = f
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
store.initialize(db)
|
||||
return store.create_generation(db, identity)
|
||||
|
||||
|
||||
def test_constructor_and_reads_do_not_create_database_or_schema(tmp_path):
|
||||
store = RetrievalIndexStore(tmp_path / "missing")
|
||||
assert not store.database_path.exists()
|
||||
with pytest.raises(sqlite3.Error):
|
||||
store.status("missing")
|
||||
assert not store.database_path.exists()
|
||||
|
||||
|
||||
def test_explicit_initialization_rolls_back_and_never_grants_model_policy(fixture):
|
||||
path, evidence, _, store, _ = fixture
|
||||
row = put(evidence)
|
||||
with closing(sqlite3.connect(path)) as db:
|
||||
with pytest.raises(RuntimeError):
|
||||
store.initialize(db)
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
store.initialize(db)
|
||||
event = ingest_journal.current(db, row.evidence_id)
|
||||
assert not event.generation_eligible and event.kind == "migration_baseline"
|
||||
db.rollback()
|
||||
assert not db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE name='retrieval_index_state'"
|
||||
).fetchone()
|
||||
assert not db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE name='evidence_ingest_events'"
|
||||
).fetchone()
|
||||
initialize(fixture)
|
||||
with closing(sqlite3.connect(path)) as db:
|
||||
assert not db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE name='continuous_policy_history'"
|
||||
).fetchone()
|
||||
|
||||
|
||||
def test_consumer_and_jobs_commit_together(fixture):
|
||||
path, evidence, _, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
row = put(evidence)
|
||||
with closing(sqlite3.connect(path)) as db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
assert store.consume(db, generation) == 1
|
||||
assert db.execute("SELECT COUNT(*) FROM retrieval_documents").fetchone()[0] == 1
|
||||
db.rollback()
|
||||
assert store.status(generation)["after_seq"] == 0
|
||||
assert store.consume_pending(generation) == 1
|
||||
assert store.consume_pending(generation) == 0
|
||||
lease = store.claim(generation, "worker", now=100)
|
||||
assert lease.evidence_id == row.evidence_id
|
||||
|
||||
|
||||
def test_lease_expiry_old_owner_and_publish_manifest(fixture):
|
||||
_, evidence, _, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
row = put(evidence)
|
||||
store.consume_pending(generation)
|
||||
lease = store.claim(generation, "one", now=100, lease_seconds=10)
|
||||
assert store.claim(generation, "two", now=105) is None
|
||||
replacement = store.claim(generation, "two", now=111)
|
||||
chunks = (TextChunk(0, 0, len(row.verbatim), 4, row.verbatim),)
|
||||
with pytest.raises(ValueError, match="lease"):
|
||||
store.publish(lease, chunks, vectors_verified=True, now=112)
|
||||
with pytest.raises(ValueError, match="vectors"):
|
||||
store.publish(replacement, chunks, vectors_verified=False, now=112)
|
||||
store.publish(replacement, chunks, vectors_verified=True, now=112)
|
||||
assert store.status(generation)["ready"] == 1
|
||||
assert store.active_chunks(generation) == ()
|
||||
store.activate_generation(generation)
|
||||
assert len(store.active_chunks(generation)) == 1
|
||||
assert store.claim(generation, "three", now=113) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", ["revision", "unrecorded", "revoke", "deleted"])
|
||||
def test_changed_inputs_cannot_publish_or_be_read(fixture, change):
|
||||
path, evidence, grants, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
row = put(evidence)
|
||||
store.consume_pending(generation)
|
||||
lease = store.claim(generation, "worker", now=100)
|
||||
if change == "revision":
|
||||
put(evidence, "updated text")
|
||||
elif change == "unrecorded":
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute(
|
||||
"UPDATE evidence SET verbatim='unrecorded' WHERE evidence_id=?", (row.evidence_id,)
|
||||
)
|
||||
elif change == "revoke":
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("UPDATE source_grants SET revoked_at='2026-09-19' WHERE source_id='codex'")
|
||||
else:
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("DELETE FROM evidence_origins WHERE evidence_id=?", (row.evidence_id,))
|
||||
db.execute("DELETE FROM evidence WHERE evidence_id=?", (row.evidence_id,))
|
||||
ingest_journal.record(db, row.evidence_id, cause="delete")
|
||||
with pytest.raises(ValueError):
|
||||
store.publish(
|
||||
lease,
|
||||
(TextChunk(0, 0, len(row.verbatim), 4, row.verbatim),),
|
||||
vectors_verified=True,
|
||||
now=101,
|
||||
)
|
||||
assert store.active_chunks(generation) == ()
|
||||
|
||||
|
||||
def test_current_read_rejects_revised_published_content_before_consumer_catches_up(fixture):
|
||||
_, evidence, _, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
row = put(evidence)
|
||||
store.consume_pending(generation)
|
||||
lease = store.claim(generation, "worker", now=100)
|
||||
store.publish(
|
||||
lease,
|
||||
(TextChunk(0, 0, len(row.verbatim), 4, row.verbatim),),
|
||||
vectors_verified=True,
|
||||
now=101,
|
||||
)
|
||||
store.activate_generation(generation)
|
||||
assert store.active_chunks(generation)
|
||||
put(evidence, "new version")
|
||||
assert store.active_chunks(generation) == ()
|
||||
store.consume_pending(generation)
|
||||
updated = store.claim(generation, "worker", now=102)
|
||||
assert updated.build_id != lease.build_id
|
||||
|
||||
|
||||
def test_pausing_harness_sync_does_not_hide_published_memory(fixture):
|
||||
path, evidence, _, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
row = put(evidence)
|
||||
store.consume_pending(generation)
|
||||
lease = store.claim(generation, "worker", now=100)
|
||||
store.publish(
|
||||
lease,
|
||||
(TextChunk(0, 0, len(row.verbatim), 4, row.verbatim),),
|
||||
vectors_verified=True,
|
||||
now=101,
|
||||
)
|
||||
store.activate_generation(generation)
|
||||
before = store.active_chunks(generation)
|
||||
assert before
|
||||
|
||||
config = ConfigStore(path.parent)
|
||||
config.set_sync("codex", True)
|
||||
config.set_sync("codex", False)
|
||||
|
||||
assert store.active_chunks(generation) == before
|
||||
|
||||
|
||||
def test_partial_or_fabricated_chunk_manifest_never_publishes(fixture):
|
||||
_, evidence, _, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
row = put(evidence)
|
||||
store.consume_pending(generation)
|
||||
lease = store.claim(generation, "worker", now=100)
|
||||
for chunks in [
|
||||
(),
|
||||
(TextChunk(0, 0, 3, 3, "ori"),),
|
||||
(TextChunk(0, 0, len(row.verbatim), 4, "not original"),),
|
||||
]:
|
||||
with pytest.raises(ValueError, match="chunk"):
|
||||
store.publish(lease, chunks, vectors_verified=True, now=101)
|
||||
assert store.status(generation)["ready"] == 0
|
||||
|
||||
|
||||
def test_incomplete_index_cannot_replace_existing_active_generation(fixture):
|
||||
path, evidence, _, store, identity = fixture
|
||||
first = initialize(fixture)
|
||||
row = put(evidence)
|
||||
store.consume_pending(first)
|
||||
lease = store.claim(first, "worker", now=100)
|
||||
chunks = (TextChunk(0, 0, len(row.verbatim), 4, row.verbatim),)
|
||||
store.publish(lease, chunks, vectors_verified=True, now=101)
|
||||
store.activate_generation(first)
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
second = store.create_generation(db, EncoderIdentity("second", "c" * 64, "d" * 64, 3))
|
||||
with pytest.raises(ValueError, match="caught_up"):
|
||||
store.activate_generation(second)
|
||||
assert store.active_chunks(first)
|
||||
store.consume_pending(second)
|
||||
with pytest.raises(ValueError, match="incomplete"):
|
||||
store.activate_generation(second)
|
||||
claim = store.claim(second, "worker", now=102)
|
||||
store.publish(claim, chunks, vectors_verified=True, now=103)
|
||||
store.activate_generation(second)
|
||||
assert store.active_chunks(first) == () and store.active_chunks(second)
|
||||
|
||||
|
||||
def test_block_and_explicit_retry_require_same_lease_and_authority(fixture):
|
||||
path, evidence, _, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
put(evidence)
|
||||
store.consume_pending(generation)
|
||||
first = store.claim(generation, "one", now=100, lease_seconds=1)
|
||||
second = store.claim(generation, "two", now=102)
|
||||
store.block(first)
|
||||
assert store.status(generation)["leased"] == 1
|
||||
store.block(second)
|
||||
assert store.status(generation)["blocked"] == 1
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("UPDATE source_grants SET revoked_at='revoked'")
|
||||
assert store.retry_blocked(generation) == 0
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("UPDATE source_grants SET revoked_at=NULL")
|
||||
assert store.retry_blocked(generation) == 1
|
||||
assert store.claim(generation, "three", now=103) is not None
|
||||
|
||||
|
||||
def test_scope_filter_and_revocation_apply_before_exposing_manifest(fixture):
|
||||
path, evidence, _, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
row = put(evidence)
|
||||
store.consume_pending(generation)
|
||||
lease = store.claim(generation, "worker", now=100)
|
||||
assert store.read_lease(lease, now=101).verbatim == row.verbatim
|
||||
store.publish(
|
||||
lease,
|
||||
(TextChunk(0, 0, len(row.verbatim), 4, row.verbatim),),
|
||||
vectors_verified=True,
|
||||
now=101,
|
||||
)
|
||||
store.activate_generation(generation)
|
||||
assert store.active_chunks(generation, source_ids=frozenset()) == ()
|
||||
assert store.active_chunks(generation, source_ids=frozenset({"codex"}))
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("UPDATE source_grants SET revision=revision+1")
|
||||
assert store.active_chunks(generation) == ()
|
||||
|
||||
|
||||
def test_source_and_time_eligibility_are_applied_before_manifest_validation(fixture):
|
||||
path, evidence, _, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
row = put(evidence)
|
||||
store.consume_pending(generation)
|
||||
lease = store.claim(generation, "worker", now=100)
|
||||
chunks = (TextChunk(0, 0, len(row.verbatim), 4, row.verbatim),)
|
||||
store.publish(lease, chunks, vectors_verified=True, now=101)
|
||||
store.activate_generation(generation)
|
||||
|
||||
assert store.active_chunks(generation, source_ids=frozenset()) == ()
|
||||
assert store.active_chunks(generation, source_ids=frozenset({"other"})) == ()
|
||||
assert len(store.active_chunks(generation, source_ids=frozenset({"codex"}))) == 1
|
||||
assert (
|
||||
store.active_chunks(
|
||||
generation,
|
||||
time_range=TimeRange(
|
||||
start=datetime(2030, 1, 1, tzinfo=UTC),
|
||||
end=datetime(2030, 1, 2, tzinfo=UTC),
|
||||
),
|
||||
)
|
||||
== ()
|
||||
)
|
||||
|
||||
|
||||
def test_generation_and_metadata_survive_reopen_without_body_copy(fixture):
|
||||
path, evidence, _, store, identity = fixture
|
||||
generation = initialize(fixture)
|
||||
put(evidence, "sensitive-body-must-not-copy")
|
||||
store.consume_pending(generation)
|
||||
reopened = RetrievalIndexStore(path)
|
||||
assert reopened.status(generation) == store.status(generation)
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
reopened.initialize(db)
|
||||
assert reopened.create_generation(db, identity) == generation
|
||||
rows = db.execute("SELECT * FROM retrieval_documents").fetchall()
|
||||
assert "sensitive-body-must-not-copy" not in str(rows)
|
||||
|
||||
|
||||
def test_partial_schema_is_not_blessed_as_new_instance(fixture):
|
||||
path, _, _, store, _ = fixture
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("CREATE TABLE retrieval_chunks (unexpected TEXT)")
|
||||
with closing(sqlite3.connect(path)) as db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
with pytest.raises(ValueError, match="partial_initialization"):
|
||||
store.initialize(db)
|
||||
assert not db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE name='retrieval_index_state'"
|
||||
).fetchone()
|
||||
|
||||
|
||||
def test_lease_renewal_rechecks_authority_and_rejects_old_expiry(fixture):
|
||||
path, evidence, _, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
put(evidence)
|
||||
store.consume_pending(generation)
|
||||
lease = store.claim(generation, "worker", now=100)
|
||||
renewed = store.renew(lease, now=125)
|
||||
assert renewed.expires_at == 185
|
||||
with pytest.raises(ValueError, match="lease_changed"):
|
||||
store.read_lease(lease, now=126)
|
||||
assert store.read_lease(renewed, now=126)
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("UPDATE source_grants SET revision=revision+1")
|
||||
with pytest.raises(ValueError, match="authority"):
|
||||
store.renew(renewed, now=127)
|
||||
|
||||
|
||||
def test_regrant_requeues_and_reuses_immutable_manifest(fixture):
|
||||
path, evidence, grants, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
row = put(evidence)
|
||||
store.consume_pending(generation)
|
||||
lease = store.claim(generation, "worker", now=100)
|
||||
chunks = (TextChunk(0, 0, len(row.verbatim), 4, row.verbatim),)
|
||||
store.publish(lease, chunks, vectors_verified=True, now=101)
|
||||
store.activate_generation(generation)
|
||||
assert len(store.active_chunks(generation)) == 1
|
||||
grants.revoke(user_id="local", source_id="codex")
|
||||
assert store.active_chunks(generation) == ()
|
||||
grants.grant(
|
||||
user_id="local",
|
||||
source_id="codex",
|
||||
root=path.parent / "source",
|
||||
candidate_kind=RootCandidateKind.DEFAULT,
|
||||
)
|
||||
assert store.active_chunks(generation) == ()
|
||||
assert store.retry_blocked(generation) == 1
|
||||
renewed = store.claim(generation, "worker", now=102)
|
||||
assert renewed.build_id == lease.build_id
|
||||
store.publish(renewed, chunks, vectors_verified=True, now=103)
|
||||
assert len(store.active_chunks(generation)) == 1
|
||||
|
||||
|
||||
def test_retry_cursor_passes_ineligible_prefix_across_reopen(fixture):
|
||||
path, evidence, grants, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
for index in range(257):
|
||||
put(evidence, message=str(index))
|
||||
store.consume_pending(generation)
|
||||
store.consume_pending(generation)
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("UPDATE retrieval_documents SET state='blocked',source_id='revoked'")
|
||||
last_id = db.execute(
|
||||
"SELECT evidence_id FROM retrieval_documents ORDER BY event_seq DESC LIMIT 1"
|
||||
).fetchone()[0]
|
||||
db.execute(
|
||||
"UPDATE retrieval_documents SET source_id='codex' WHERE evidence_id=?", (last_id,)
|
||||
)
|
||||
assert store.retry_blocked(generation) == 0
|
||||
reopened = RetrievalIndexStore(path)
|
||||
assert reopened.retry_blocked(generation) == 1
|
||||
assert reopened.claim(generation, "worker", now=100).evidence_id == last_id
|
||||
|
||||
|
||||
def test_same_build_replay_requires_exact_complete_manifest(fixture):
|
||||
path, evidence, _, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
row = put(evidence)
|
||||
store.consume_pending(generation)
|
||||
lease = store.claim(generation, "worker", now=100)
|
||||
chunks = (TextChunk(0, 0, len(row.verbatim), 4, row.verbatim),)
|
||||
store.publish(lease, chunks, vectors_verified=True, now=101)
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("UPDATE retrieval_documents SET state='pending',active_build_id=NULL")
|
||||
replay = store.claim(generation, "worker", now=102)
|
||||
with pytest.raises(ValueError, match="manifest"):
|
||||
store.publish(
|
||||
replay,
|
||||
(TextChunk(0, 0, len(row.verbatim), 5, row.verbatim),),
|
||||
vectors_verified=True,
|
||||
now=103,
|
||||
)
|
||||
assert store.status(generation)["ready"] == 0
|
||||
store.publish(replay, chunks, vectors_verified=True, now=103)
|
||||
assert store.status(generation)["ready"] == 1
|
||||
|
||||
|
||||
def test_origin_metadata_roundtrip_reuses_old_build_without_losing_new_event(fixture):
|
||||
_, evidence, _, store, _ = fixture
|
||||
generation = initialize(fixture)
|
||||
builds = []
|
||||
seqs = []
|
||||
for version in ("a", "b", "a"):
|
||||
row = evidence.put(
|
||||
EvidenceInput(
|
||||
source_id="codex",
|
||||
conversation_id="thread",
|
||||
message_id="one",
|
||||
role="user",
|
||||
written_by=WrittenBy.PERSON,
|
||||
verbatim="unchanged original",
|
||||
origin_meta={EvidenceOrigin.FILE: OriginMetadata(source_version=version)},
|
||||
)
|
||||
).record
|
||||
assert row.revision == 1
|
||||
assert store.consume_pending(generation) == 1
|
||||
lease = store.claim(generation, "worker", now=100)
|
||||
builds.append(lease.build_id)
|
||||
seqs.append(lease.event_seq)
|
||||
store.publish(
|
||||
lease,
|
||||
(TextChunk(0, 0, len(row.verbatim), 5, row.verbatim),),
|
||||
vectors_verified=True,
|
||||
now=101,
|
||||
)
|
||||
assert builds[0] == builds[2] and builds[0] != builds[1]
|
||||
assert seqs[0] < seqs[1] < seqs[2]
|
||||
store.activate_generation(generation)
|
||||
assert len(store.active_chunks(generation)) == 1
|
||||
@@ -1,357 +0,0 @@
|
||||
import asyncio
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import FrozenInstanceError
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.retrieval.index_supervisor import RetrievalIndexSupervisor
|
||||
from one_memory.retrieval.index_worker import IndexWorkResult
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_network(monkeypatch):
|
||||
def reject(*args, **kwargs):
|
||||
raise AssertionError("supervisor tests must remain offline")
|
||||
|
||||
monkeypatch.setattr(socket.socket, "connect", reject)
|
||||
monkeypatch.setattr(socket.socket, "connect_ex", reject)
|
||||
|
||||
|
||||
class Worker:
|
||||
def __init__(self, state="idle", blocking=False):
|
||||
self.state = state
|
||||
self.entered = threading.Event()
|
||||
self.release = threading.Event()
|
||||
self.exited = threading.Event()
|
||||
self.cancel = None
|
||||
self.blocking = blocking
|
||||
self.calls = []
|
||||
self.thread_ids = set()
|
||||
self.active = 0
|
||||
self.maximum = 0
|
||||
self.closed = False
|
||||
|
||||
def run_once(self, *, cancel=None):
|
||||
self.cancel = cancel
|
||||
self.active += 1
|
||||
self.maximum = max(self.maximum, self.active)
|
||||
self.calls.append(time.monotonic())
|
||||
self.thread_ids.add(threading.get_ident())
|
||||
self.entered.set()
|
||||
try:
|
||||
if self.blocking:
|
||||
assert self.release.wait(3), "test cleanup failed"
|
||||
if self.state == "exception":
|
||||
raise RuntimeError("SECRET_RAW_ERROR")
|
||||
if self.state == "invalid":
|
||||
return object()
|
||||
return IndexWorkResult(self.state, "SECRET_REASON")
|
||||
finally:
|
||||
self.active -= 1
|
||||
self.exited.set()
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
async def until(predicate):
|
||||
deadline = time.monotonic() + 2
|
||||
while not predicate():
|
||||
assert time.monotonic() < deadline
|
||||
await asyncio.sleep(0.002)
|
||||
|
||||
|
||||
def test_stop_waits_for_native_exit_and_does_not_close_resources():
|
||||
async def run():
|
||||
worker = Worker(blocking=True)
|
||||
supervisor = RetrievalIndexSupervisor(worker)
|
||||
supervisor.start()
|
||||
await until(worker.entered.is_set)
|
||||
stopping = asyncio.create_task(supervisor.stop())
|
||||
try:
|
||||
await until(lambda: worker.cancel.is_set())
|
||||
await asyncio.sleep(0.03)
|
||||
assert not stopping.done() and not worker.exited.is_set()
|
||||
assert supervisor.snapshot().thread_alive
|
||||
finally:
|
||||
worker.release.set()
|
||||
await stopping
|
||||
assert worker.exited.is_set()
|
||||
assert not supervisor.snapshot().thread_alive
|
||||
assert not worker.closed
|
||||
assert supervisor.snapshot().state == "stopped"
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["stop", "wait"])
|
||||
def test_async_cancellation_repeatedly_waits_for_native_thread(method):
|
||||
async def run():
|
||||
worker = Worker(blocking=True)
|
||||
supervisor = RetrievalIndexSupervisor(worker)
|
||||
supervisor.start()
|
||||
await until(worker.entered.is_set)
|
||||
task = asyncio.create_task(getattr(supervisor, method)())
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
try:
|
||||
await until(lambda: worker.cancel.is_set())
|
||||
task.cancel()
|
||||
await asyncio.sleep(0.03)
|
||||
assert not task.done()
|
||||
assert not worker.exited.is_set()
|
||||
finally:
|
||||
worker.release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
await supervisor.stop()
|
||||
assert worker.exited.is_set() and not supervisor.snapshot().thread_alive
|
||||
assert not worker.closed
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_idle_no_trigger_no_spin_and_duplicate_start_rejected():
|
||||
async def run():
|
||||
worker = Worker()
|
||||
supervisor = RetrievalIndexSupervisor(worker, idle_wait_seconds=0.3)
|
||||
assert not supervisor.trigger()
|
||||
assert worker.calls == []
|
||||
supervisor.start()
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="index_supervisor_already_started"):
|
||||
supervisor.start()
|
||||
await until(lambda: len(worker.calls) == 1)
|
||||
await asyncio.sleep(0.06)
|
||||
assert len(worker.calls) == 1
|
||||
assert supervisor.snapshot().last_result == "idle"
|
||||
finally:
|
||||
await supervisor.stop()
|
||||
with pytest.raises(RuntimeError, match="index_supervisor_already_started"):
|
||||
supervisor.start()
|
||||
assert not supervisor.trigger()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_trigger_storm_coalesces_during_native_call_without_parallel_workers():
|
||||
async def run():
|
||||
worker = Worker(blocking=True)
|
||||
supervisor = RetrievalIndexSupervisor(worker, idle_wait_seconds=0.5)
|
||||
supervisor.start()
|
||||
await until(worker.entered.is_set)
|
||||
try:
|
||||
for _ in range(200):
|
||||
assert supervisor.trigger()
|
||||
worker.release.set()
|
||||
await until(lambda: len(worker.calls) >= 2)
|
||||
await asyncio.sleep(0.05)
|
||||
assert len(worker.calls) == 2
|
||||
assert worker.maximum == 1 and len(worker.thread_ids) == 1
|
||||
finally:
|
||||
worker.release.set()
|
||||
await supervisor.stop()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_interactive_activity_pauses_index_slices_then_resumes() -> None:
|
||||
async def run() -> None:
|
||||
worker = Worker()
|
||||
interactive = True
|
||||
supervisor = RetrievalIndexSupervisor(
|
||||
worker,
|
||||
yield_wait_seconds=0.02,
|
||||
background_should_yield=lambda: interactive,
|
||||
)
|
||||
supervisor.start()
|
||||
try:
|
||||
await asyncio.sleep(0.06)
|
||||
assert worker.calls == []
|
||||
interactive = False
|
||||
supervisor.trigger()
|
||||
await until(lambda: len(worker.calls) == 1)
|
||||
finally:
|
||||
await supervisor.stop()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"state,code",
|
||||
[
|
||||
("exception", "index_worker_failed"),
|
||||
("invalid", "index_worker_invalid_result"),
|
||||
("blocked", "index_worker_blocked"),
|
||||
],
|
||||
)
|
||||
def test_failures_have_fixed_error_and_backoff_even_with_triggers(state, code):
|
||||
async def run():
|
||||
worker = Worker(state)
|
||||
supervisor = RetrievalIndexSupervisor(worker, failure_wait_seconds=0.2)
|
||||
supervisor.start()
|
||||
try:
|
||||
await until(lambda: supervisor.snapshot().error_code is not None)
|
||||
for _ in range(100):
|
||||
supervisor.trigger()
|
||||
await asyncio.sleep(0.05)
|
||||
assert len(worker.calls) == 1
|
||||
snapshot = supervisor.snapshot()
|
||||
assert snapshot.error_code == code
|
||||
assert "SECRET" not in repr(snapshot)
|
||||
await until(lambda: len(worker.calls) >= 2)
|
||||
assert worker.calls[1] - worker.calls[0] >= 0.18
|
||||
finally:
|
||||
await supervisor.stop()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("state", ["ready", "excluded"])
|
||||
def test_backlog_yields_after_maximum_slices(state):
|
||||
async def run():
|
||||
worker = Worker(state)
|
||||
supervisor = RetrievalIndexSupervisor(
|
||||
worker, max_consecutive_slices=3, yield_wait_seconds=0.15
|
||||
)
|
||||
supervisor.start()
|
||||
try:
|
||||
await until(lambda: len(worker.calls) >= 3)
|
||||
for _ in range(50):
|
||||
supervisor.trigger()
|
||||
await asyncio.sleep(0.04)
|
||||
assert len(worker.calls) == 3
|
||||
await until(lambda: len(worker.calls) >= 4)
|
||||
assert worker.calls[3] - worker.calls[2] >= 0.13
|
||||
finally:
|
||||
await supervisor.stop()
|
||||
assert worker.maximum == 1
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_concurrent_stop_is_idempotent_and_stop_before_start_is_terminal():
|
||||
async def run():
|
||||
worker = Worker()
|
||||
supervisor = RetrievalIndexSupervisor(worker)
|
||||
supervisor.start()
|
||||
await asyncio.gather(supervisor.stop(), supervisor.stop())
|
||||
assert not supervisor.snapshot().thread_alive
|
||||
unused = RetrievalIndexSupervisor(worker)
|
||||
await unused.stop()
|
||||
with pytest.raises(RuntimeError, match="index_supervisor_already_started"):
|
||||
unused.start()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs",
|
||||
[
|
||||
{"idle_wait_seconds": 0},
|
||||
{"failure_wait_seconds": float("nan")},
|
||||
{"yield_wait_seconds": 0},
|
||||
{"max_consecutive_slices": True},
|
||||
{"max_consecutive_slices": 257},
|
||||
],
|
||||
)
|
||||
def test_invalid_configuration_never_starts_worker(kwargs):
|
||||
worker = Worker()
|
||||
with pytest.raises(ValueError, match="index_supervisor_invalid_config"):
|
||||
RetrievalIndexSupervisor(worker, **kwargs)
|
||||
assert worker.calls == []
|
||||
|
||||
|
||||
def test_idle_timeout_polls_again_and_stop_interrupts_long_wait():
|
||||
async def run():
|
||||
worker = Worker()
|
||||
supervisor = RetrievalIndexSupervisor(worker, idle_wait_seconds=0.04)
|
||||
supervisor.start()
|
||||
try:
|
||||
await until(lambda: len(worker.calls) >= 2)
|
||||
assert worker.calls[1] - worker.calls[0] >= 0.035
|
||||
finally:
|
||||
await supervisor.stop()
|
||||
long_idle = RetrievalIndexSupervisor(Worker(), idle_wait_seconds=300)
|
||||
long_idle.start()
|
||||
await until(lambda: long_idle.snapshot().slices == 1)
|
||||
started = time.monotonic()
|
||||
await long_idle.stop()
|
||||
assert time.monotonic() - started < 0.5
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_failure_clears_on_recovery_without_leaking_worker_reason():
|
||||
async def run():
|
||||
worker = Worker("exception")
|
||||
supervisor = RetrievalIndexSupervisor(worker, failure_wait_seconds=0.06)
|
||||
supervisor.start()
|
||||
try:
|
||||
await until(lambda: supervisor.snapshot().failures == 1)
|
||||
worker.state = "idle"
|
||||
await until(lambda: supervisor.snapshot().slices == 2)
|
||||
current = supervisor.snapshot()
|
||||
assert current.error_code is None and current.last_result == "idle"
|
||||
assert current.failures == 1
|
||||
assert "SECRET" not in repr(current)
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
current.slices = 0
|
||||
finally:
|
||||
await supervisor.stop()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_thread_start_failure_is_fixed_and_terminal(monkeypatch):
|
||||
def fail(_self):
|
||||
raise OSError("SECRET_THREAD_ERROR")
|
||||
|
||||
monkeypatch.setattr(threading.Thread, "start", fail)
|
||||
supervisor = RetrievalIndexSupervisor(Worker())
|
||||
with pytest.raises(RuntimeError, match="^index_supervisor_start_failed$"):
|
||||
supervisor.start()
|
||||
assert not supervisor.snapshot().thread_alive
|
||||
assert supervisor.snapshot().state == "stopped"
|
||||
assert "SECRET" not in repr(supervisor.snapshot())
|
||||
asyncio.run(supervisor.stop())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("state", ["deferred", "cancelled"])
|
||||
def test_nonprogress_results_do_not_drain_in_a_tight_loop(state):
|
||||
async def run():
|
||||
worker = Worker(state)
|
||||
supervisor = RetrievalIndexSupervisor(worker, idle_wait_seconds=0.3)
|
||||
supervisor.start()
|
||||
try:
|
||||
await until(lambda: supervisor.snapshot().slices == 1)
|
||||
await asyncio.sleep(0.05)
|
||||
assert len(worker.calls) == 1
|
||||
finally:
|
||||
await supervisor.stop()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_thread_is_non_daemon_and_joined_without_default_executor_threads():
|
||||
async def run():
|
||||
before = set(threading.enumerate())
|
||||
worker = Worker(blocking=True)
|
||||
supervisor = RetrievalIndexSupervisor(worker)
|
||||
supervisor.start()
|
||||
await until(worker.entered.is_set)
|
||||
own = [t for t in threading.enumerate() if t.ident in worker.thread_ids]
|
||||
assert len(own) == 1 and not own[0].daemon
|
||||
try:
|
||||
task = asyncio.create_task(supervisor.stop())
|
||||
await until(lambda: worker.cancel.is_set())
|
||||
assert set(threading.enumerate()) - before == set(own)
|
||||
finally:
|
||||
worker.release.set()
|
||||
await task
|
||||
assert not own[0].is_alive()
|
||||
assert set(threading.enumerate()) == before
|
||||
|
||||
asyncio.run(run())
|
||||
@@ -1,222 +0,0 @@
|
||||
import socket
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
from contextlib import closing
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.embedding_assets import assets_present
|
||||
from one_memory.local_embedding import TokenizedText, TokenOffset
|
||||
from one_memory.retrieval.index_store import RetrievalIndexStore
|
||||
from one_memory.retrieval.index_worker import LocalRetrievalEncoder, RetrievalIndexWorker
|
||||
from one_memory.retrieval.vector_store import RawVectorStore, VectorStoreError
|
||||
from tests.test_retrieval_index_store import fixture, initialize, put # noqa: F401
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def forbid_network(monkeypatch):
|
||||
def reject(*args, **kwargs):
|
||||
raise AssertionError("index worker must stay offline")
|
||||
|
||||
monkeypatch.setattr(socket.socket, "connect", reject)
|
||||
monkeypatch.setattr(socket.socket, "connect_ex", reject)
|
||||
|
||||
|
||||
class Encoder:
|
||||
def __init__(self, identity):
|
||||
self.identity = identity
|
||||
self.calls = []
|
||||
self.after_encode = lambda: None
|
||||
|
||||
def tokenize_with_offsets(self, text):
|
||||
return TokenizedText(
|
||||
tuple(range(len(text))), tuple(TokenOffset(i, i + 1) for i in range(len(text)))
|
||||
)
|
||||
|
||||
def encode_documents(self, texts):
|
||||
self.calls.append(texts)
|
||||
self.after_encode()
|
||||
return [[1.0, 0.5, 0.25] for _ in texts]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def setup(fixture, tmp_path): # noqa: F811
|
||||
generation = initialize(fixture)
|
||||
store = fixture[3]
|
||||
encoder = Encoder(fixture[4])
|
||||
vectors = RawVectorStore.create(tmp_path / "vectors", generation_id=generation, dimensions=3)
|
||||
worker = RetrievalIndexWorker(store, vectors, encoder, generation_id=generation, owner="test")
|
||||
try:
|
||||
yield fixture, generation, encoder, vectors, worker
|
||||
finally:
|
||||
vectors.close()
|
||||
|
||||
|
||||
def test_tick_indexes_complete_original_and_reopen_can_search(setup):
|
||||
f, generation, encoder, vectors, worker = setup
|
||||
record = put(f[1], "start " + "middle " * 180 + " end")
|
||||
assert worker.run_once().state == "ready"
|
||||
assert worker.run_once().state == "idle"
|
||||
assert all(len(batch) <= 8 for batch in encoder.calls)
|
||||
assert "".join(encoder.calls[0]) # Nonempty positive encoding control.
|
||||
f[3].activate_generation(generation)
|
||||
chunks = RetrievalIndexStore(f[0]).active_chunks(generation)
|
||||
assert len(chunks) > 1 and chunks[0].start_cp == 0
|
||||
assert chunks[-1].end_cp == len(record.verbatim)
|
||||
result = vectors.search(
|
||||
[1.0, 0.5, 0.25], [c.chunk_id for c in chunks], 8, time.monotonic() + 10
|
||||
)
|
||||
assert result.complete and result.hits
|
||||
with closing(sqlite3.connect(f[0])) as db:
|
||||
assert not db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE name='continuous_policy_history'"
|
||||
).fetchone()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", ["revision", "revoke"])
|
||||
def test_change_during_native_encoding_never_publishes(setup, change):
|
||||
f, generation, encoder, vectors, worker = setup
|
||||
put(f[1])
|
||||
|
||||
def mutate():
|
||||
if change == "revision":
|
||||
put(f[1], "changed original")
|
||||
else:
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute("UPDATE source_grants SET revoked_at='revoked'")
|
||||
|
||||
encoder.after_encode = mutate
|
||||
result = worker.run_once()
|
||||
assert result.state == "blocked"
|
||||
assert f[3].status(generation)["ready"] == 0
|
||||
assert f[3].active_chunks(generation) == ()
|
||||
|
||||
|
||||
def test_cancel_after_encode_does_not_publish_and_is_reclaimable(setup):
|
||||
f, generation, encoder, _, worker = setup
|
||||
put(f[1])
|
||||
cancel = threading.Event()
|
||||
encoder.after_encode = cancel.set
|
||||
assert worker.run_once(cancel=cancel).state == "cancelled"
|
||||
assert f[3].status(generation)["leased"] == 1
|
||||
assert f[3].claim(generation, "replacement", now=time.time() + 61)
|
||||
|
||||
|
||||
def test_cancel_before_tick_does_not_consume_or_encode(setup):
|
||||
f, generation, encoder, _, worker = setup
|
||||
put(f[1])
|
||||
cancel = threading.Event()
|
||||
cancel.set()
|
||||
assert worker.run_once(cancel=cancel).state == "cancelled"
|
||||
assert f[3].status(generation)["after_seq"] == 0
|
||||
assert not encoder.calls
|
||||
|
||||
|
||||
def test_partial_vector_write_can_retry_same_build_without_duplicates(setup, monkeypatch):
|
||||
f, generation, encoder, vectors, worker = setup
|
||||
put(f[1], "long message " * 90)
|
||||
original = vectors.upsert
|
||||
|
||||
def partial(records):
|
||||
original(records[:1])
|
||||
raise VectorStoreError("vector_write_unavailable")
|
||||
|
||||
monkeypatch.setattr(vectors, "upsert", partial)
|
||||
assert worker.run_once().state == "blocked"
|
||||
assert f[3].status(generation)["ready"] == 0
|
||||
monkeypatch.setattr(vectors, "upsert", original)
|
||||
assert f[3].retry_blocked(generation) == 1
|
||||
assert worker.run_once().state == "ready"
|
||||
f[3].activate_generation(generation)
|
||||
assert len(f[3].active_chunks(generation)) > 1
|
||||
|
||||
|
||||
def test_bad_vector_count_or_failed_verification_never_publishes(setup, monkeypatch):
|
||||
f, generation, encoder, vectors, worker = setup
|
||||
put(f[1])
|
||||
monkeypatch.setattr(encoder, "encode_documents", lambda _: [])
|
||||
assert worker.run_once().state == "blocked"
|
||||
assert f[3].status(generation)["ready"] == 0
|
||||
assert f[3].retry_blocked(generation) == 1
|
||||
monkeypatch.setattr(
|
||||
encoder, "encode_documents", lambda texts: [[1.0, 0.5, 0.25] for _ in texts]
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
vectors, "verify", lambda _: (_ for _ in ()).throw(VectorStoreError("failure"))
|
||||
)
|
||||
assert worker.run_once().state == "blocked"
|
||||
assert f[3].status(generation)["ready"] == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
["token=example-secret-value", "<environment_context>system wrapper</environment_context>"],
|
||||
)
|
||||
def test_content_requiring_source_mapping_is_not_encoded_or_marked_ready(setup, text):
|
||||
f, generation, encoder, _, worker = setup
|
||||
put(f[1], text)
|
||||
result = worker.run_once()
|
||||
assert result.state == "blocked" and result.reason == "content_mapping_required"
|
||||
assert not encoder.calls
|
||||
assert f[3].status(generation)["ready"] == 0
|
||||
|
||||
|
||||
def test_deadline_after_native_encode_discards_late_work(setup):
|
||||
f, generation, encoder, vectors, _ = setup
|
||||
put(f[1])
|
||||
clock = [1.0]
|
||||
encoder.after_encode = lambda: clock.__setitem__(0, 50.0)
|
||||
worker = RetrievalIndexWorker(
|
||||
f[3],
|
||||
vectors,
|
||||
encoder,
|
||||
generation_id=generation,
|
||||
owner="test",
|
||||
clock=lambda: clock[0],
|
||||
max_job_seconds=10,
|
||||
)
|
||||
assert worker.run_once().state == "deferred"
|
||||
assert f[3].status(generation)["ready"] == 0
|
||||
|
||||
|
||||
def test_generation_encoder_mismatch_rejected_before_claim(setup):
|
||||
f, generation, encoder, vectors, _ = setup
|
||||
from dataclasses import replace
|
||||
|
||||
encoder.identity = replace(encoder.identity, encoder_digest="f" * 64)
|
||||
with pytest.raises(ValueError, match="encoder_identity"):
|
||||
RetrievalIndexWorker(f[3], vectors, encoder, generation_id=generation, owner="other")
|
||||
assert f[3].status(generation)["leased"] == 0
|
||||
|
||||
|
||||
@pytest.mark.skipif(not assets_present(), reason="locked local assets absent")
|
||||
def test_real_local_encoder_chroma_full_text_chain(fixture, tmp_path): # noqa: F811
|
||||
path, evidence, _, store, _ = fixture
|
||||
encoder = LocalRetrievalEncoder(content_version=1)
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
store.initialize(db)
|
||||
generation = store.create_generation(db, encoder.identity)
|
||||
text = (
|
||||
"A long source message about source control and releases. " * 80
|
||||
+ " Final artifact at /project/design/final.png"
|
||||
)
|
||||
record = put(evidence, text)
|
||||
vectors = RawVectorStore.create(
|
||||
tmp_path / "real-vectors", generation_id=generation, dimensions=encoder.identity.dimensions
|
||||
)
|
||||
try:
|
||||
worker = RetrievalIndexWorker(
|
||||
store, vectors, encoder, generation_id=generation, owner="real"
|
||||
)
|
||||
assert worker.run_once().state == "ready"
|
||||
store.activate_generation(generation)
|
||||
chunks = store.active_chunks(generation)
|
||||
assert len(chunks) > 1 and chunks[-1].end_cp == len(text)
|
||||
query = encoder.encode_documents([text[chunks[-1].start_cp : chunks[-1].end_cp]])[0]
|
||||
result = vectors.search(query, [c.chunk_id for c in chunks], 8, time.monotonic() + 10)
|
||||
assert result.complete and result.hits[0].chunk_id == chunks[-1].chunk_id
|
||||
assert record.verbatim.endswith("final.png")
|
||||
finally:
|
||||
vectors.close()
|
||||
@@ -1,23 +1,14 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import urllib.request
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.api import create_app
|
||||
from one_memory.config import Settings
|
||||
from one_memory.secrets import InMemorySecretStore
|
||||
|
||||
|
||||
def test_retrieval_defaults_off_and_does_not_create_index(tmp_path, monkeypatch):
|
||||
def test_startup_does_not_create_a_private_retrieval_index(tmp_path, monkeypatch):
|
||||
"""Raw reads go through the SDK and lexical FTS; MemHub keeps no vector copy."""
|
||||
monkeypatch.setattr("one_memory.model_presets._preset_data", lambda: None)
|
||||
|
||||
async def scenario():
|
||||
@@ -29,10 +20,10 @@ def test_retrieval_defaults_off_and_does_not_create_index(tmp_path, monkeypatch)
|
||||
automatic_memory_processing=False,
|
||||
live_sync_enabled=False,
|
||||
)
|
||||
assert settings.retrieval_index_enabled is False
|
||||
app = create_app(settings, secret_store=InMemorySecretStore())
|
||||
async with app.router.lifespan_context(app):
|
||||
assert app.state.retrieval_runtime is None
|
||||
assert not hasattr(app.state, "retrieval_runtime")
|
||||
assert not (tmp_path / "retrieval-vectors").exists()
|
||||
with closing(sqlite3.connect(tmp_path / "one-memory.sqlite3")) as db:
|
||||
assert (
|
||||
db.execute(
|
||||
@@ -42,188 +33,3 @@ def test_retrieval_defaults_off_and_does_not_create_index(tmp_path, monkeypatch)
|
||||
)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_enabling_retrieval_requires_runtime_ownership(tmp_path):
|
||||
with pytest.raises(ValueError, match="retrieval_runtime_ownership_required"):
|
||||
create_app(
|
||||
Settings(_env_file=None, data_dir=tmp_path, retrieval_index_enabled=True),
|
||||
secret_store=InMemorySecretStore(),
|
||||
)
|
||||
|
||||
|
||||
def test_index_enabled_serve_imports_without_building_an_unmanaged_app(tmp_path, monkeypatch):
|
||||
"""The real serve entry must not be killed by an import-time default app.
|
||||
|
||||
`cli.serve()` imports `create_app` and then builds the app with
|
||||
`manage_runtime=True`. If this module also builds a default app at import
|
||||
time, enabling the index makes that throw before `serve()` ever binds.
|
||||
Run the public entry in a subprocess, not through `create_app()` directly,
|
||||
because the regression exists only in the import/entry-point combination.
|
||||
"""
|
||||
monkeypatch.setenv("ONE_MEMORY_DATA_DIR", str(tmp_path))
|
||||
monkeypatch.setenv("ONE_MEMORY_RETRIEVAL_INDEX_ENABLED", "1")
|
||||
monkeypatch.setenv("ONE_MEMORY_RETRIEVAL_MINIMUM_FREE_BYTES", "0")
|
||||
with socket.socket() as probe:
|
||||
probe.bind(("127.0.0.1", 0))
|
||||
port = probe.getsockname()[1]
|
||||
monkeypatch.setenv("ONE_MEMORY_HOST", "127.0.0.1")
|
||||
monkeypatch.setenv("ONE_MEMORY_PORT", str(port))
|
||||
monkeypatch.setenv("HOME", str(tmp_path / "home"))
|
||||
monkeypatch.setenv("ONE_MEMORY_SOURCE_HOME", str(tmp_path / "source-home"))
|
||||
monkeypatch.setenv("ONE_MEMORY_IMPORTS_ENABLED", "0")
|
||||
monkeypatch.setenv("ONE_MEMORY_LIVE_SYNC_ENABLED", "0")
|
||||
monkeypatch.setenv("ONE_MEMORY_AUTOMATIC_MEMORY_PROCESSING", "0")
|
||||
|
||||
process = subprocess.Popen(
|
||||
[sys.executable, "-m", "one_memory.cli", "serve"],
|
||||
env={
|
||||
**os.environ,
|
||||
"PYTHONPATH": str(Path(__file__).resolve().parents[1] / "src"),
|
||||
},
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
stdout = ""
|
||||
stderr = ""
|
||||
try:
|
||||
deadline = time.monotonic() + 15
|
||||
health = None
|
||||
while time.monotonic() < deadline and process.poll() is None:
|
||||
try:
|
||||
with urllib.request.urlopen(
|
||||
f"http://127.0.0.1:{port}/api/v1/health", timeout=0.5
|
||||
) as response:
|
||||
health = json.loads(response.read())
|
||||
break
|
||||
except OSError:
|
||||
time.sleep(0.1)
|
||||
if health is None:
|
||||
process.terminate()
|
||||
try:
|
||||
stdout, stderr = process.communicate(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
stdout, stderr = process.communicate(timeout=5)
|
||||
pytest.fail(
|
||||
"serve did not become healthy within 15s; "
|
||||
f"exit={process.returncode} stdout={stdout!r} stderr={stderr!r}"
|
||||
)
|
||||
assert health["status"] == "ok"
|
||||
assert health["service"] == "one-memory-core"
|
||||
finally:
|
||||
if process.poll() is None:
|
||||
process.terminate()
|
||||
try:
|
||||
process.communicate(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.communicate(timeout=5)
|
||||
|
||||
|
||||
def managed(tmp_path):
|
||||
return create_app(
|
||||
Settings(
|
||||
_env_file=None,
|
||||
data_dir=tmp_path,
|
||||
source_home=tmp_path / "home",
|
||||
imports_enabled=False,
|
||||
automatic_memory_processing=False,
|
||||
live_sync_enabled=False,
|
||||
retrieval_index_enabled=True,
|
||||
),
|
||||
manage_runtime=True,
|
||||
secret_store=InMemorySecretStore(),
|
||||
)
|
||||
|
||||
|
||||
def test_managed_start_without_an_encoder_preserves_normal_core_startup(tmp_path, monkeypatch):
|
||||
"""No built-in weights is the real reason an index cannot be admitted.
|
||||
|
||||
Core must still come up and the lexical lane must still answer; the runtime
|
||||
says not_initialized rather than taking the process down with it.
|
||||
"""
|
||||
monkeypatch.setattr("one_memory.model_presets._preset_data", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
"one_memory.retrieval.index_provisioning.assets_present", lambda *_args: False
|
||||
)
|
||||
|
||||
async def scenario():
|
||||
app = managed(tmp_path)
|
||||
async with app.router.lifespan_context(app):
|
||||
assert app.state.retrieval_runtime.snapshot().state == "not_initialized"
|
||||
assert (tmp_path / "runtime.json").exists()
|
||||
assert not (tmp_path / "retrieval-vectors").exists()
|
||||
assert not (tmp_path / "runtime.json").exists()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_enabling_retrieval_provisions_the_index_instead_of_waiting_to_be_handed_one(
|
||||
tmp_path, monkeypatch
|
||||
):
|
||||
"""The product entry point, which until now only a test helper could reach."""
|
||||
monkeypatch.setattr("one_memory.model_presets._preset_data", lambda: None)
|
||||
|
||||
async def scenario():
|
||||
app = managed(tmp_path)
|
||||
async with app.router.lifespan_context(app):
|
||||
assert app.state.retrieval_runtime.snapshot().state == "running"
|
||||
assert (tmp_path / "retrieval-vectors").is_dir()
|
||||
with closing(sqlite3.connect(tmp_path / "one-memory.sqlite3")) as db:
|
||||
assert [
|
||||
row[0]
|
||||
for row in db.execute(
|
||||
"SELECT state FROM retrieval_index_generations ORDER BY rowid"
|
||||
)
|
||||
] == ["active"]
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", [RuntimeError, asyncio.CancelledError])
|
||||
def test_retrieval_stop_failure_still_drains_other_owned_services(tmp_path, monkeypatch, failure):
|
||||
monkeypatch.setattr("one_memory.model_presets._preset_data", lambda: None)
|
||||
events = []
|
||||
|
||||
class StubRuntime:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
async def start(self):
|
||||
events.append("start")
|
||||
|
||||
async def stop(self):
|
||||
events.append("retrieval_drained")
|
||||
raise failure("after drain")
|
||||
|
||||
monkeypatch.setattr("one_memory.retrieval.runtime.ExistingRetrievalRuntime", StubRuntime)
|
||||
|
||||
async def scenario():
|
||||
app = create_app(
|
||||
Settings(
|
||||
_env_file=None,
|
||||
data_dir=tmp_path,
|
||||
source_home=tmp_path / "home",
|
||||
imports_enabled=False,
|
||||
automatic_memory_processing=False,
|
||||
live_sync_enabled=False,
|
||||
retrieval_index_enabled=True,
|
||||
),
|
||||
manage_runtime=True,
|
||||
secret_store=InMemorySecretStore(),
|
||||
)
|
||||
with pytest.raises(failure):
|
||||
async with app.router.lifespan_context(app):
|
||||
original = app.state.sdk_memory_runtime.shutdown
|
||||
|
||||
async def shutdown():
|
||||
await original()
|
||||
events.append("processing_drained")
|
||||
|
||||
monkeypatch.setattr(app.state.sdk_memory_runtime, "shutdown", shutdown)
|
||||
assert events == ["start", "retrieval_drained", "processing_drained"]
|
||||
assert not (tmp_path / "runtime.json").exists()
|
||||
|
||||
asyncio.run(scenario())
|
||||
@@ -1,422 +0,0 @@
|
||||
import hashlib
|
||||
import json
|
||||
import socket
|
||||
from dataclasses import FrozenInstanceError, replace
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.imports.models import EvidenceContentKind, EvidenceRecord, WrittenBy
|
||||
from one_memory.input_sanitization import SanitizedCopySpan, SanitizedInput, sanitize_model_input
|
||||
from one_memory.local_embedding import LocalEmbedder, TokenizedText, TokenOffset
|
||||
from one_memory.retrieval import mapped_content as mapped
|
||||
from one_memory.retrieval.chunking import ChunkingError
|
||||
from one_memory.retrieval.mapped_content import (
|
||||
MappedContentError,
|
||||
chunk_mapped_content,
|
||||
prepare_mapped_content,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def offline(monkeypatch):
|
||||
def reject(*args, **kwargs):
|
||||
raise AssertionError("mapped content must remain offline")
|
||||
|
||||
monkeypatch.setattr(socket.socket, "connect", reject)
|
||||
monkeypatch.setattr(socket.socket, "connect_ex", reject)
|
||||
|
||||
|
||||
def record(text, **changes):
|
||||
values = dict(
|
||||
evidence_id="synthetic-evidence",
|
||||
user_id="local",
|
||||
source_id="synthetic-source",
|
||||
workspace=None,
|
||||
project=None,
|
||||
conversation_id="synthetic-conversation",
|
||||
message_id="synthetic-message",
|
||||
correlation_id=None,
|
||||
source_ordinal=None,
|
||||
content_kind=EvidenceContentKind.RAW_CONVERSATION,
|
||||
role="user",
|
||||
written_by=WrittenBy.PERSON,
|
||||
created_at=None,
|
||||
verbatim=text,
|
||||
origin_meta={},
|
||||
content_hash="outer-manifest-owned",
|
||||
revision=4,
|
||||
imported_at=datetime(2026, 9, 19, tzinfo=UTC),
|
||||
updated_at=datetime(2026, 9, 19, tzinfo=UTC),
|
||||
)
|
||||
return EvidenceRecord(**(values | changes))
|
||||
|
||||
|
||||
class Characters:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def tokenize_with_offsets(self, text):
|
||||
self.calls.append(text)
|
||||
offsets = (
|
||||
TokenOffset(0, 0, True),
|
||||
*(TokenOffset(i, i + 1) for i in range(len(text))),
|
||||
TokenOffset(0, 0, True),
|
||||
)
|
||||
return TokenizedText(tuple(range(len(offsets))), offsets)
|
||||
|
||||
|
||||
def assert_map(original, content):
|
||||
meta = content.metadata
|
||||
assert meta.source_length_cp == len(original)
|
||||
assert meta.sanitized_length_cp == len(content.actual_text)
|
||||
assert meta.sanitized_sha256 == hashlib.sha256(content.actual_text.encode()).hexdigest()
|
||||
output_end = 0
|
||||
copied = []
|
||||
for span in meta.segments:
|
||||
assert span.start_cp == output_end < span.end_cp
|
||||
output_end = span.end_cp
|
||||
if span.kind == "copy":
|
||||
assert (
|
||||
content.actual_text[span.start_cp : span.end_cp]
|
||||
== original[span.source_start_cp : span.source_end_cp]
|
||||
)
|
||||
copied.append((span.source_start_cp, span.source_end_cp))
|
||||
else:
|
||||
assert span.kind == "synthetic"
|
||||
assert span.source_start_cp is span.source_end_cp is None
|
||||
assert output_end == len(content.actual_text)
|
||||
ranges = sorted(copied + [(s.start_cp, s.end_cp) for s in meta.excluded])
|
||||
end = 0
|
||||
for start, right in ranges:
|
||||
assert start == end < right
|
||||
end = right
|
||||
assert end == len(original)
|
||||
|
||||
|
||||
def assert_chunks(original, prepared, chunks):
|
||||
assert [c.ordinal for c in chunks] == list(range(len(chunks)))
|
||||
coverage = set()
|
||||
for c in chunks:
|
||||
assert c.text == original[c.start_cp : c.end_cp]
|
||||
assert c.token_count <= 512
|
||||
assert any(
|
||||
s.kind == "copy" and s.source_start_cp <= c.start_cp < c.end_cp <= s.source_end_cp
|
||||
for s in prepared.metadata.segments
|
||||
)
|
||||
coverage.update(range(c.start_cp, c.end_cp))
|
||||
for s in prepared.metadata.segments:
|
||||
if s.kind == "copy":
|
||||
assert all(
|
||||
i in coverage
|
||||
for i in range(s.source_start_cp, s.source_end_cp)
|
||||
if not original[i].isspace()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"role,author", [("user", WrittenBy.PERSON), ("assistant", WrittenBy.AGENT)]
|
||||
)
|
||||
def test_raw_authors_keep_exact_body_and_do_not_mutate_evidence(role, author):
|
||||
original = " \t\n\u4e2d\U0001f680e\u0301 same\n same\t "
|
||||
evidence = record(original, role=role, written_by=author)
|
||||
before = evidence.model_dump()
|
||||
content = prepare_mapped_content(evidence)
|
||||
assert content.state == "ready"
|
||||
assert content.actual_text == original.strip()
|
||||
assert evidence.model_dump() == before
|
||||
assert_map(original, content)
|
||||
assert_chunks(original, content, chunk_mapped_content(content, Characters()))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"change",
|
||||
[
|
||||
{"role": "tool", "written_by": WrittenBy.TOOL},
|
||||
{"role": "system"},
|
||||
{"content_kind": EvidenceContentKind.MEMORY_CONTENT},
|
||||
],
|
||||
)
|
||||
def test_raw_role_qualification_rejects_without_tokenizing(change):
|
||||
with pytest.raises(MappedContentError, match="mapped_content_ineligible"):
|
||||
prepare_mapped_content(record("ordinary body", **change))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("author", list(WrittenBy))
|
||||
def test_host_author_guess_does_not_change_protocol_role_qualification(author):
|
||||
content = prepare_mapped_content(record("ordinary body", role="user", written_by=author))
|
||||
assert content.state == "ready"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text,reason",
|
||||
[
|
||||
("<user_query>one</user_query><user_query>two</user_query>", "ambiguous"),
|
||||
("<user_info>unclosed", "ambiguous"),
|
||||
("<rules>context</rules> unexplained trailing body", "ambiguous"),
|
||||
("<dynamic_tool_catalog>generated</dynamic_tool_catalog>", "generated"),
|
||||
('{"request_id":"fixture","type":"write","steps":[]}', "generated"),
|
||||
(
|
||||
"<user_query><dynamic_tool_catalog>generated</dynamic_tool_catalog></user_query>",
|
||||
"generated",
|
||||
),
|
||||
("<user_query><environment_context>harness</environment_context></user_query>", "hygiene"),
|
||||
],
|
||||
)
|
||||
def test_ambiguous_generated_and_body_hygiene_are_not_admitted(text, reason):
|
||||
with pytest.raises(MappedContentError, match=reason):
|
||||
prepare_mapped_content(record(text))
|
||||
|
||||
|
||||
def test_wrapper_offsets_compose_by_position_despite_repeated_text_and_unicode():
|
||||
prefix = " \t<rules>same \U0001f680</rules>\n<user_query> \n"
|
||||
body = "same \U0001f680 token=same same e\u0301"
|
||||
original = prefix + body + " \n</user_query><image_local_path>same</image_local_path> "
|
||||
prepared = prepare_mapped_content(record(original))
|
||||
assert prepared.actual_text == sanitize_model_input(body)
|
||||
assert prepared.metadata.body_start_cp == len(prefix)
|
||||
assert prepared.metadata.body_end_cp == len(prefix) + len(body)
|
||||
assert_map(original, prepared)
|
||||
tokenizer = Characters()
|
||||
chunks = chunk_mapped_content(prepared, tokenizer, target_tokens=10, overlap_tokens=3)
|
||||
assert_chunks(original, prepared, chunks)
|
||||
assert chunks[0].start_cp == len(prefix)
|
||||
assert all("<rules>" not in text and "redacted-secret" not in text for text in tokenizer.calls)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"secret",
|
||||
[
|
||||
"token=private-value",
|
||||
"Authorization: Bearer abcdefghijklmnop",
|
||||
'--password "two words"',
|
||||
"sk-abcdefghijklmnopqrstuvwxyz",
|
||||
"https://alice:private-password@example.test/path",
|
||||
"-----BEGIN PRIVATE KEY-----\nprivate\n-----END PRIVATE KEY-----",
|
||||
],
|
||||
)
|
||||
def test_secrets_and_synthetic_replacements_never_reach_tokenizer(secret):
|
||||
original = "left \u4e2d\U0001f680 " + secret + " right e\u0301 /final/file"
|
||||
content = prepare_mapped_content(record(original))
|
||||
tokenizer = Characters()
|
||||
chunks = chunk_mapped_content(content, tokenizer, target_tokens=12, overlap_tokens=4)
|
||||
assert_map(original, content)
|
||||
assert_chunks(original, content, chunks)
|
||||
assert all(secret not in text and "redacted-secret" not in text for text in tokenizer.calls)
|
||||
assert any("left" in c.text for c in chunks) and any("right" in c.text for c in chunks)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text", ["", " \n\t ", "<rules>only wrapper</rules>", "token=private", " token=x \t token=y "]
|
||||
)
|
||||
def test_empty_wrapper_or_all_excluded_is_explicit_nonindexable(text):
|
||||
content = prepare_mapped_content(record(text))
|
||||
assert content.state == "no_indexable_content"
|
||||
assert_map(text, content)
|
||||
tokenizer = Characters()
|
||||
assert chunk_mapped_content(content, tokenizer) == ()
|
||||
assert tokenizer.calls == []
|
||||
|
||||
|
||||
def test_metadata_is_immutable_canonical_and_never_contains_text_or_secret_hashes():
|
||||
original = "<user_query>prefix token=private-value suffix</user_query>"
|
||||
content = prepare_mapped_content(record(original))
|
||||
encoded = content.metadata.canonical_json()
|
||||
assert encoded == json.dumps(json.loads(encoded), sort_keys=True, separators=(",", ":"))
|
||||
assert content.metadata.digest == hashlib.sha256(encoded.encode()).hexdigest()
|
||||
for text in [
|
||||
"prefix",
|
||||
"suffix",
|
||||
"private-value",
|
||||
"redacted-secret",
|
||||
hashlib.sha256(b"private-value").hexdigest(),
|
||||
]:
|
||||
assert text not in encoded
|
||||
assert content.metadata == prepare_mapped_content(record(original)).metadata
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
content.metadata.source_length_cp = 1
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
content.metadata.segments[0].start_cp = 1
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
content.actual_text = "other"
|
||||
|
||||
|
||||
def test_digest_changes_for_safe_text_or_coordinate_version_not_equal_length_excluded_secret():
|
||||
first = prepare_mapped_content(record("left token=secret-one right"))
|
||||
second = prepare_mapped_content(record("left token=secret-two right"))
|
||||
assert first.metadata.digest == second.metadata.digest
|
||||
assert (
|
||||
first.metadata.digest
|
||||
!= prepare_mapped_content(record("LEFT token=secret-one right")).metadata.digest
|
||||
)
|
||||
assert first.metadata.digest != replace(first.metadata, map_version="future-map").digest
|
||||
|
||||
|
||||
def test_overlapping_sanitizer_passes_and_literal_placeholder_are_distinct():
|
||||
text = "<redacted-secret> token=sk-abcdefghijklmnopqrstuvwxyz /tail"
|
||||
content = prepare_mapped_content(record(text))
|
||||
assert_map(text, content)
|
||||
chunks = chunk_mapped_content(content, Characters())
|
||||
assert chunks[0].text == "<redacted-secret> " # Literal source text is copy, not synthetic.
|
||||
assert chunks[-1].text == " /tail"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text", ["\u4e2d\x00\u6587", "\u4e2d\u200b\u6587", "\u4e2d\u200d\u6587"])
|
||||
def test_control_characters_are_not_silently_sanitized_or_dropped(text):
|
||||
content = prepare_mapped_content(record(text))
|
||||
assert content.actual_text == text
|
||||
with pytest.raises(ChunkingError, match="uncovered_nonwhitespace"):
|
||||
chunk_mapped_content(content, LocalEmbedder())
|
||||
|
||||
|
||||
def test_locked_tokenizer_unicode_coverage_and_prefix_budget():
|
||||
body = ("\u4e2d\u6587\U0001f680e\u0301\n" * 400) + " token=private " + ("\u5c3e\u6587\n" * 400)
|
||||
original = "<user_query>" + body + "</user_query>"
|
||||
content = prepare_mapped_content(record(original))
|
||||
tokenizer = LocalEmbedder()
|
||||
prefix = "\u68c0\u7d22\u524d\u7f00" * 30
|
||||
chunks = chunk_mapped_content(content, tokenizer, prefix=prefix)
|
||||
assert len(chunks) > 4
|
||||
assert_chunks(original, content, chunks)
|
||||
assert all(
|
||||
c.token_count == len(tokenizer.tokenize_with_offsets(prefix + c.text).token_ids)
|
||||
for c in chunks
|
||||
)
|
||||
|
||||
|
||||
def test_more_than_8192_chunks_aborts_one_copy_span_without_returning_partial_content():
|
||||
tokenizer = Characters()
|
||||
content = prepare_mapped_content(record("x" * 9000))
|
||||
with pytest.raises(MappedContentError, match="mapped_chunk_limit_exceeded"):
|
||||
chunk_mapped_content(content, tokenizer, target_tokens=1, overlap_tokens=0)
|
||||
assert len(tokenizer.calls) <= 8194
|
||||
|
||||
|
||||
def test_global_limit_across_copy_spans_and_exact_limit_succeeds():
|
||||
content = prepare_mapped_content(record("a token=private b"))
|
||||
with pytest.raises(MappedContentError, match="mapped_chunk_limit_exceeded"):
|
||||
chunk_mapped_content(content, Characters(), max_chunks=1)
|
||||
chunks = chunk_mapped_content(content, Characters(), max_chunks=2)
|
||||
assert [c.ordinal for c in chunks] == [0, 1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("limit", [0, -1, True, 8193])
|
||||
def test_chunk_limit_is_not_raiseable_or_coercible(limit):
|
||||
with pytest.raises(MappedContentError, match="mapped_invalid_chunk_limit"):
|
||||
chunk_mapped_content(prepare_mapped_content(record("body")), Characters(), max_chunks=limit)
|
||||
|
||||
|
||||
def test_forged_text_and_incomplete_sanitizer_mapping_fail_closed(monkeypatch):
|
||||
prepared = prepare_mapped_content(record("body"))
|
||||
with pytest.raises(MappedContentError, match="mapped_invalid"):
|
||||
chunk_mapped_content(replace(prepared, actual_text="BODY"), Characters())
|
||||
monkeypatch.setattr(
|
||||
mapped,
|
||||
"sanitize_model_input_with_spans",
|
||||
lambda text: SanitizedInput(text, (SanitizedCopySpan(0, 2, 0, 1),), ()),
|
||||
)
|
||||
with pytest.raises(MappedContentError, match="mapped_invalid"):
|
||||
prepare_mapped_content(record("body"))
|
||||
|
||||
|
||||
def test_exact_8192_chunks_and_shrink_retries_do_not_spend_extra_chunk_budget():
|
||||
content = prepare_mapped_content(record("x" * 8192))
|
||||
chunks = chunk_mapped_content(content, Characters(), target_tokens=1, overlap_tokens=0)
|
||||
assert len(chunks) == 8192 and chunks[-1].ordinal == 8191
|
||||
content = prepare_mapped_content(record("x" * 10))
|
||||
chunks = chunk_mapped_content(
|
||||
content,
|
||||
Characters(),
|
||||
target_tokens=8,
|
||||
overlap_tokens=0,
|
||||
max_tokens=8,
|
||||
prefix="PP",
|
||||
max_chunks=3,
|
||||
)
|
||||
assert len(chunks) == 3
|
||||
assert_chunks("x" * 10, content, chunks)
|
||||
|
||||
|
||||
def test_fragmented_content_limit_is_global_and_rejects_before_next_copy_tokenization():
|
||||
text = " token=private ".join(["x"] * 8193)
|
||||
prepared = prepare_mapped_content(record(text))
|
||||
tokenizer = Characters()
|
||||
with pytest.raises(MappedContentError, match="mapped_chunk_limit_exceeded"):
|
||||
chunk_mapped_content(prepared, tokenizer)
|
||||
assert len(tokenizer.calls) == 2 * 8192
|
||||
|
||||
|
||||
def test_random_repeated_wrapper_and_sanitizer_composition():
|
||||
import random
|
||||
|
||||
randomizer = random.Random(20260919)
|
||||
for _ in range(80):
|
||||
body = " ".join(
|
||||
randomizer.choices(
|
||||
[
|
||||
"same",
|
||||
"\u4e2d\U0001f680e\u0301",
|
||||
"token=same",
|
||||
"--password same",
|
||||
"Authorization: Bearer abcdefghijklmnop",
|
||||
"\t",
|
||||
"https://u:p@host/path",
|
||||
],
|
||||
k=10,
|
||||
)
|
||||
)
|
||||
original = "<rules>same</rules><user_query> " + body + " </user_query>"
|
||||
content = prepare_mapped_content(record(original))
|
||||
assert_map(original, content)
|
||||
assert_chunks(original, content, chunk_mapped_content(content, Characters()))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"case", ["output_gap", "source_overlap", "wrong_copy", "missing_exclusion"]
|
||||
)
|
||||
def test_bad_upstream_maps_do_not_authorize_original_slices(monkeypatch, case):
|
||||
samples = {
|
||||
"output_gap": SanitizedInput("body", (SanitizedCopySpan(2, 4, 2, 4),), ()),
|
||||
"source_overlap": SanitizedInput(
|
||||
"bobo", (SanitizedCopySpan(0, 2, 0, 2), SanitizedCopySpan(2, 4, 0, 2)), ((2, 4),)
|
||||
),
|
||||
"wrong_copy": SanitizedInput("BODY", (SanitizedCopySpan(0, 4, 0, 4),), ()),
|
||||
"missing_exclusion": SanitizedInput("bo", (SanitizedCopySpan(0, 2, 0, 2),), ()),
|
||||
}
|
||||
monkeypatch.setattr(mapped, "sanitize_model_input_with_spans", lambda text: samples[case])
|
||||
with pytest.raises(MappedContentError, match="mapped_invalid"):
|
||||
prepare_mapped_content(record("body"))
|
||||
|
||||
|
||||
def test_assistant_context_keeps_existing_hygiene_semantics_but_not_generated_payloads():
|
||||
text = "Discuss <environment_context>literal</environment_context> as context"
|
||||
content = prepare_mapped_content(record(text, role="assistant", written_by=WrittenBy.AGENT))
|
||||
assert content.state == "ready" and content.actual_text == text
|
||||
with pytest.raises(MappedContentError, match="generated"):
|
||||
prepare_mapped_content(
|
||||
record(
|
||||
"<dynamic_tool_catalog>generated</dynamic_tool_catalog>",
|
||||
role="assistant",
|
||||
written_by=WrittenBy.AGENT,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_copy_boundaries_cannot_reconstruct_a_url_across_secret_gap():
|
||||
original = "https://example.test/left token=private /right"
|
||||
content = prepare_mapped_content(record(original))
|
||||
chunks = chunk_mapped_content(content, Characters())
|
||||
assert len(chunks) == 2
|
||||
assert chunks[0].text == "https://example.test/left "
|
||||
assert chunks[1].text == " /right"
|
||||
assert chunks[0].end_cp < chunks[1].start_cp
|
||||
assert_chunks(original, content, chunks)
|
||||
|
||||
|
||||
def test_source_size_limit_rejects_before_parser(monkeypatch):
|
||||
def unexpected(_):
|
||||
raise AssertionError("oversize input reached parser")
|
||||
|
||||
monkeypatch.setattr(mapped, "unwrap_host_envelope_with_positions", unexpected)
|
||||
with pytest.raises(MappedContentError, match="mapped_source_limit_exceeded"):
|
||||
prepare_mapped_content(record("x" * 1_000_001))
|
||||
@@ -1,549 +0,0 @@
|
||||
import socket
|
||||
import sqlite3
|
||||
import time
|
||||
from contextlib import closing
|
||||
from dataclasses import replace
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.retrieval.chunking import TextChunk
|
||||
from one_memory.retrieval.index_store import RetrievalIndexStore
|
||||
from one_memory.retrieval.index_worker import RetrievalIndexWorker
|
||||
from one_memory.retrieval.mapped_content import chunk_mapped_content, prepare_mapped_content
|
||||
from one_memory.retrieval.raw_candidates import read_raw_semantic_candidates
|
||||
from one_memory.retrieval.raw_evidence import materialize_raw_candidates
|
||||
from one_memory.retrieval.vector_store import RawVectorStore
|
||||
from tests.test_retrieval_index_store import fixture, initialize, put # noqa: F401
|
||||
from tests.test_retrieval_index_worker import Encoder
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def offline(monkeypatch):
|
||||
def fail(*args, **kwargs):
|
||||
raise AssertionError("no network permitted")
|
||||
|
||||
monkeypatch.setattr(socket.socket, "connect", fail)
|
||||
monkeypatch.setattr(socket.socket, "connect_ex", fail)
|
||||
|
||||
|
||||
def migrate(f):
|
||||
initialize(f)
|
||||
identity = replace(f[4], normalizer_version=2, chunker_version=2)
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
f[3].migrate_v2(db)
|
||||
generation = f[3].create_generation(db, identity)
|
||||
return generation, identity
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mapped_setup(fixture, tmp_path): # noqa: F811
|
||||
generation, identity = migrate(fixture)
|
||||
encoder = Encoder(identity)
|
||||
vectors = RawVectorStore.create(
|
||||
tmp_path / "mapped-vectors", generation_id=generation, dimensions=3
|
||||
)
|
||||
worker = RetrievalIndexWorker(
|
||||
fixture[3], vectors, encoder, generation_id=generation, owner="mapped"
|
||||
)
|
||||
try:
|
||||
yield fixture, generation, encoder, vectors, worker
|
||||
finally:
|
||||
vectors.close()
|
||||
|
||||
|
||||
def candidates(f, generation, vectors):
|
||||
return read_raw_semantic_candidates(
|
||||
f[3],
|
||||
vectors,
|
||||
[1.0, 0.5, 0.25],
|
||||
generation_id=generation,
|
||||
source_ids=frozenset({"codex"}),
|
||||
deadline=time.monotonic() + 10,
|
||||
)
|
||||
|
||||
|
||||
def materialize(f, generation, hits):
|
||||
return materialize_raw_candidates(
|
||||
f[3],
|
||||
hits,
|
||||
generation_id=generation,
|
||||
source_ids=frozenset({"codex"}),
|
||||
deadline=time.monotonic() + 10,
|
||||
)
|
||||
|
||||
|
||||
def test_empty_materialization_rechecks_excluded_map(mapped_setup):
|
||||
f, generation, _, vectors, worker = mapped_setup
|
||||
put(f[1], "token=private-value")
|
||||
assert worker.run_once().state == "excluded"
|
||||
f[3].activate_generation(generation)
|
||||
empty = candidates(f, generation, vectors)
|
||||
assert empty.complete and not empty.hits
|
||||
assert materialize(f, generation, empty).complete
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute("DELETE FROM retrieval_build_maps")
|
||||
assert not materialize(f, generation, empty).complete
|
||||
|
||||
|
||||
def test_completed_excluded_missing_active_binding_is_not_empty_success(mapped_setup):
|
||||
f, generation, _, vectors, worker = mapped_setup
|
||||
put(f[1], "token=private-value")
|
||||
assert worker.run_once().state == "excluded"
|
||||
f[3].activate_generation(generation)
|
||||
assert candidates(f, generation, vectors).complete
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute("UPDATE retrieval_documents SET active_build_id=NULL")
|
||||
result = candidates(f, generation, vectors)
|
||||
assert not result.complete and result.reason == "index_unavailable"
|
||||
|
||||
|
||||
def test_old_empty_completion_is_invalidated_when_a_new_indexed_record_arrives(mapped_setup):
|
||||
f, generation, _, vectors, worker = mapped_setup
|
||||
put(f[1], "token=private-value")
|
||||
assert worker.run_once().state == "excluded"
|
||||
f[3].activate_generation(generation)
|
||||
empty = candidates(f, generation, vectors)
|
||||
assert empty.complete and not empty.hits
|
||||
put(f[1], "new actual content", message="new")
|
||||
assert worker.run_once().state == "ready"
|
||||
result = materialize(f, generation, empty)
|
||||
assert not result.complete and result.reason == "evidence_changed"
|
||||
|
||||
|
||||
def test_schema2_roundtrip_only_encodes_copy_and_separates_original_spans(mapped_setup):
|
||||
f, generation, encoder, vectors, worker = mapped_setup
|
||||
original = (
|
||||
"<rules>private wrapper</rules><user_query>before token=private-value after</user_query>"
|
||||
)
|
||||
record = put(f[1], original)
|
||||
result = worker.run_once()
|
||||
assert result.state == "ready" and result.chunk_count == 2
|
||||
assert encoder.calls == [["before ", " after"]]
|
||||
f[3].activate_generation(generation)
|
||||
hits = candidates(f, generation, vectors)
|
||||
assert hits.complete and len(hits.hits) == 2
|
||||
output = materialize(f, generation, hits)
|
||||
assert output.complete and len(output.items) == 1
|
||||
assert [s.text for s in output.items[0].spans] == ["before ", " after"]
|
||||
assert all(original[s.start_cp : s.end_cp] == s.text for s in output.items[0].spans)
|
||||
with closing(sqlite3.connect(f[0])) as db:
|
||||
rows = db.execute(
|
||||
"SELECT source_map_json,map_digest,coverage_state FROM retrieval_build_maps"
|
||||
).fetchall()
|
||||
assert len(rows) == 1 and rows[0][2] == "indexable"
|
||||
assert "before" not in rows[0][0] and "private" not in rows[0][0]
|
||||
assert rows[0][1] == prepare_mapped_content(record).metadata.digest
|
||||
assert RetrievalIndexStore(f[0]).active_chunks(generation)
|
||||
|
||||
|
||||
def test_explicit_migration_rolls_back_and_v1_remains_readable(fixture): # noqa: F811
|
||||
f = fixture
|
||||
generation = initialize(f)
|
||||
row = put(f[1], "legacy text")
|
||||
f[3].consume_pending(generation)
|
||||
lease = f[3].claim(generation, "legacy")
|
||||
f[3].publish(
|
||||
lease, (TextChunk(0, 0, len(row.verbatim), 4, row.verbatim),), vectors_verified=True
|
||||
)
|
||||
f[3].activate_generation(generation)
|
||||
before = f[3].active_chunks(generation)
|
||||
with closing(sqlite3.connect(f[0])) as db:
|
||||
with pytest.raises(RuntimeError):
|
||||
f[3].migrate_v2(db)
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
f[3].migrate_v2(db)
|
||||
db.rollback()
|
||||
assert db.execute("SELECT schema_version FROM retrieval_index_state").fetchone()[0] == 1
|
||||
assert not db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE name='retrieval_build_maps'"
|
||||
).fetchone()
|
||||
assert f[3].active_chunks(generation) == before
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
f[3].migrate_v2(db)
|
||||
f[3].migrate_v2(db)
|
||||
assert f[3].active_chunks(generation) == before
|
||||
from one_memory.retrieval.raw_candidates import RawCandidateBatch, RawSemanticCandidate
|
||||
|
||||
output = materialize(
|
||||
f, generation, RawCandidateBatch(tuple(RawSemanticCandidate(c, 0.0) for c in before), True)
|
||||
)
|
||||
assert output.complete and output.items[0].spans[0].text == "legacy text"
|
||||
|
||||
|
||||
def test_version2_requires_explicit_schema_upgrade(fixture): # noqa: F811
|
||||
f = fixture
|
||||
initialize(f)
|
||||
with closing(sqlite3.connect(f[0])) as db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
with pytest.raises(ValueError, match="schema"):
|
||||
f[3].create_generation(db, replace(f[4], normalizer_version=2, chunker_version=2))
|
||||
assert db.execute("SELECT schema_version FROM retrieval_index_state").fetchone()[0] == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text", ["token=private", "<rules>only wrapper</rules>"])
|
||||
def test_excluded_completion_has_no_vectors_and_is_not_ready(mapped_setup, text):
|
||||
f, generation, encoder, vectors, worker = mapped_setup
|
||||
put(f[1], text)
|
||||
result = worker.run_once()
|
||||
assert result.state == "excluded" and result.chunk_count == 0
|
||||
assert encoder.calls == []
|
||||
assert f[3].status(generation)["excluded"] == 1
|
||||
assert f[3].status(generation)["ready"] == 0
|
||||
f[3].activate_generation(generation)
|
||||
assert candidates(f, generation, vectors).hits == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tamper", ["json", "digest", "offset", "missing_map", "missing_chunk"])
|
||||
def test_persisted_tampering_rejected_at_candidate_and_final_read(mapped_setup, tamper):
|
||||
f, generation, _, vectors, worker = mapped_setup
|
||||
put(f[1], "before token=private after")
|
||||
assert worker.run_once().state == "ready"
|
||||
f[3].activate_generation(generation)
|
||||
hits = candidates(f, generation, vectors)
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
if tamper == "json":
|
||||
db.execute("UPDATE retrieval_build_maps SET source_map_json='{}'")
|
||||
elif tamper == "digest":
|
||||
db.execute("UPDATE retrieval_build_maps SET map_digest=?", ("0" * 64,))
|
||||
elif tamper == "offset":
|
||||
db.execute("UPDATE retrieval_chunks SET end_cp=end_cp+1 WHERE ordinal=0")
|
||||
elif tamper == "missing_map":
|
||||
db.execute("DELETE FROM retrieval_build_maps")
|
||||
else:
|
||||
db.execute("DELETE FROM retrieval_chunks WHERE ordinal=1")
|
||||
assert not candidates(f, generation, vectors).complete
|
||||
final = materialize(f, generation, hits)
|
||||
assert not final.complete and not final.items
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", ["secret", "revoke", "policy"])
|
||||
def test_current_source_secret_revision_and_grant_changes_invalidate_hits(mapped_setup, change):
|
||||
f, generation, _, vectors, worker = mapped_setup
|
||||
put(f[1], "before token=private-one after")
|
||||
assert worker.run_once().state == "ready"
|
||||
f[3].activate_generation(generation)
|
||||
hits = candidates(f, generation, vectors)
|
||||
if change == "secret":
|
||||
put(f[1], "before token=private-two after")
|
||||
else:
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute(
|
||||
"UPDATE source_grants SET revoked_at='revoked'"
|
||||
if change == "revoke"
|
||||
else "UPDATE source_grants SET revision=revision+1"
|
||||
)
|
||||
assert not materialize(f, generation, hits).items
|
||||
assert not f[3].active_chunks(generation)
|
||||
|
||||
|
||||
def test_publish_recomputes_map_and_rejects_cross_gap_or_incomplete_manifest(mapped_setup):
|
||||
f, generation, encoder, _, _ = mapped_setup
|
||||
row = put(f[1], "before token=private after")
|
||||
f[3].consume_pending(generation)
|
||||
lease = f[3].claim(generation, "publisher")
|
||||
prepared = prepare_mapped_content(row)
|
||||
chunks = chunk_mapped_content(prepared, encoder)
|
||||
for bad in [chunks[:1], (TextChunk(0, 0, len(row.verbatim), 10, row.verbatim),)]:
|
||||
with pytest.raises(ValueError, match="manifest"):
|
||||
f[3].publish(lease, bad, vectors_verified=True, source_map=prepared.metadata)
|
||||
with pytest.raises(ValueError, match="map"):
|
||||
f[3].publish(
|
||||
lease,
|
||||
chunks,
|
||||
vectors_verified=True,
|
||||
source_map=replace(prepared.metadata, sanitized_sha256="0" * 64),
|
||||
)
|
||||
assert f[3].status(generation)["ready"] == 0
|
||||
|
||||
|
||||
def test_many_copy_regions_remain_partial_not_a_cross_gap_quote(mapped_setup):
|
||||
f, generation, _, vectors, worker = mapped_setup
|
||||
put(f[1], " token=private ".join([f"part-{i}" for i in range(6)]))
|
||||
assert worker.run_once().chunk_count == 6
|
||||
f[3].activate_generation(generation)
|
||||
final = materialize(f, generation, candidates(f, generation, vectors))
|
||||
assert not final.complete and len(final.items[0].spans) == 4
|
||||
assert all("token=" not in s.text for s in final.items[0].spans)
|
||||
|
||||
|
||||
def test_local_onnx_chroma_mapped_roundtrip_without_download(fixture, tmp_path): # noqa: F811
|
||||
from one_memory.retrieval.index_worker import LocalRetrievalEncoder
|
||||
|
||||
f = fixture
|
||||
encoder = LocalRetrievalEncoder()
|
||||
assert encoder.identity.normalizer_version == encoder.identity.chunker_version == 2
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
f[3].initialize(db)
|
||||
f[3].migrate_v2(db)
|
||||
generation = f[3].create_generation(db, encoder.identity)
|
||||
row = put(
|
||||
f[1],
|
||||
"<user_query>\u4e2d\u6587\U0001f680e\u0301\n"
|
||||
+ "\u524d\u6587" * 300
|
||||
+ " token=private-value "
|
||||
+ "\u540e\u6587" * 300
|
||||
+ "</user_query>",
|
||||
)
|
||||
vectors = RawVectorStore.create(
|
||||
tmp_path / "native-mapped", generation_id=generation, dimensions=encoder.identity.dimensions
|
||||
)
|
||||
try:
|
||||
worker = RetrievalIndexWorker(
|
||||
f[3], vectors, encoder, generation_id=generation, owner="local"
|
||||
)
|
||||
result = worker.run_once()
|
||||
assert result.state == "ready" and result.chunk_count >= 4
|
||||
f[3].activate_generation(generation)
|
||||
hits = read_raw_semantic_candidates(
|
||||
f[3],
|
||||
vectors,
|
||||
encoder.encode_documents(["\u540e\u6587"])[0],
|
||||
generation_id=generation,
|
||||
source_ids=frozenset({"codex"}),
|
||||
deadline=time.monotonic() + 15,
|
||||
)
|
||||
final = materialize(f, generation, hits)
|
||||
assert final.complete and len(final.items[0].spans) == 2
|
||||
for span in final.items[0].spans:
|
||||
assert span.text == row.verbatim[span.start_cp : span.end_cp]
|
||||
assert "private-value" not in span.text and "redacted-secret" not in span.text
|
||||
finally:
|
||||
vectors.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
"<user_query>one</user_query><user_query>two</user_query>",
|
||||
"<dynamic_tool_catalog>generated</dynamic_tool_catalog>",
|
||||
],
|
||||
)
|
||||
def test_ambiguous_generated_stay_blocked(mapped_setup, text):
|
||||
f, generation, encoder, _, worker = mapped_setup
|
||||
put(f[1], text)
|
||||
assert worker.run_once().state == "blocked"
|
||||
assert not encoder.calls
|
||||
with pytest.raises(ValueError, match="incomplete"):
|
||||
f[3].activate_generation(generation)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", ["revoke", "secret"])
|
||||
def test_change_during_encoding_never_publishes(mapped_setup, change):
|
||||
f, generation, encoder, _, worker = mapped_setup
|
||||
put(f[1], "before token=private-one after")
|
||||
original = encoder.encode_documents
|
||||
|
||||
def changed(texts):
|
||||
result = original(texts)
|
||||
if change == "secret":
|
||||
put(f[1], "before token=private-two after")
|
||||
else:
|
||||
f[2].revoke(user_id="local", source_id="codex")
|
||||
return result
|
||||
|
||||
encoder.encode_documents = changed
|
||||
assert worker.run_once().state == "blocked"
|
||||
assert f[3].status(generation)["ready"] == 0
|
||||
with closing(sqlite3.connect(f[0])) as db:
|
||||
assert not db.execute("SELECT * FROM retrieval_build_maps").fetchall()
|
||||
|
||||
|
||||
def test_cancellation_inside_publish_rolls_back_map_chunks_and_completion(
|
||||
mapped_setup, monkeypatch
|
||||
):
|
||||
import threading
|
||||
|
||||
f, generation, _, _, worker = mapped_setup
|
||||
put(f[1], "before token=private after")
|
||||
cancel = threading.Event()
|
||||
original = f[3]._prepare
|
||||
|
||||
def cancelled(*args):
|
||||
prepared = original(*args)
|
||||
cancel.set()
|
||||
return prepared
|
||||
|
||||
monkeypatch.setattr(f[3], "_prepare", cancelled)
|
||||
assert worker.run_once(cancel=cancel).state == "cancelled"
|
||||
assert f[3].status(generation)["leased"] == 1
|
||||
with closing(sqlite3.connect(f[0])) as db:
|
||||
assert not db.execute("SELECT * FROM retrieval_build_maps").fetchall()
|
||||
assert not db.execute("SELECT * FROM retrieval_chunks").fetchall()
|
||||
|
||||
|
||||
def test_excluded_metadata_is_validated_on_activation_and_candidate_read(mapped_setup):
|
||||
f, generation, _, vectors, worker = mapped_setup
|
||||
put(f[1], "token=private")
|
||||
assert worker.run_once().state == "excluded"
|
||||
f[3].activate_generation(generation)
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute("DELETE FROM retrieval_build_maps")
|
||||
assert not candidates(f, generation, vectors).complete
|
||||
with pytest.raises(ValueError, match="map"):
|
||||
f[3].activate_generation(generation)
|
||||
|
||||
|
||||
def test_migration_preserves_an_existing_v1_lease_and_store_identity(fixture): # noqa: F811
|
||||
f = fixture
|
||||
generation = initialize(f)
|
||||
row = put(f[1], "legacy leased")
|
||||
f[3].consume_pending(generation)
|
||||
lease = f[3].claim(generation, "legacy")
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
f[3].migrate_v2(db)
|
||||
assert f[3].read_lease(lease) == row
|
||||
f[3].publish(
|
||||
lease, (TextChunk(0, 0, len(row.verbatim), 4, row.verbatim),), vectors_verified=True
|
||||
)
|
||||
f[3].activate_generation(generation)
|
||||
assert f[3].active_chunks(generation)
|
||||
|
||||
|
||||
def test_bad_token_offsets_leave_no_partial_mapped_manifest(mapped_setup):
|
||||
from one_memory.local_embedding import TokenizedText, TokenOffset
|
||||
|
||||
f, generation, encoder, _, worker = mapped_setup
|
||||
put(f[1], "before token=private after")
|
||||
encoder.tokenize_with_offsets = lambda text: TokenizedText((1,), (TokenOffset(0, 1),))
|
||||
assert worker.run_once().state == "blocked"
|
||||
assert not encoder.calls
|
||||
assert f[3].status(generation)["ready"] == 0
|
||||
|
||||
|
||||
def test_failed_migration_keeps_original_table_lease_and_schema(fixture): # noqa: F811
|
||||
f = fixture
|
||||
generation = initialize(f)
|
||||
put(f[1], "preserved")
|
||||
f[3].consume_pending(generation)
|
||||
lease = f[3].claim(generation, "legacy")
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute("CREATE TABLE retrieval_build_maps (foreign_value TEXT)")
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
with pytest.raises(sqlite3.Error):
|
||||
f[3].migrate_v2(db)
|
||||
assert db.execute("SELECT schema_version FROM retrieval_index_state").fetchone()[0] == 1
|
||||
assert not db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE name='retrieval_documents_v2'"
|
||||
).fetchone()
|
||||
assert f[3].read_lease(lease).verbatim == "preserved"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text,expected", [("before token=private after", "ready"), ("token=private", "excluded")]
|
||||
)
|
||||
def test_regrant_reuses_immutable_map_and_manifest_only_after_explicit_retry(
|
||||
mapped_setup, text, expected
|
||||
):
|
||||
from one_memory.imports.sources.registry import RootCandidateKind
|
||||
|
||||
f, generation, _, _, worker = mapped_setup
|
||||
put(f[1], text)
|
||||
assert worker.run_once().state == expected
|
||||
f[3].activate_generation(generation)
|
||||
with closing(sqlite3.connect(f[0])) as db:
|
||||
original = db.execute("SELECT * FROM retrieval_build_maps").fetchall()
|
||||
f[2].revoke(user_id="local", source_id="codex")
|
||||
assert not f[3].active_chunks(generation)
|
||||
f[2].grant(
|
||||
user_id="local",
|
||||
source_id="codex",
|
||||
root=f[0].parent / "source",
|
||||
candidate_kind=RootCandidateKind.DEFAULT,
|
||||
)
|
||||
assert not f[3].active_chunks(generation)
|
||||
assert f[3].retry_blocked(generation) == 1
|
||||
assert worker.run_once().state == expected
|
||||
with closing(sqlite3.connect(f[0])) as db:
|
||||
assert db.execute("SELECT * FROM retrieval_build_maps").fetchall() == original
|
||||
|
||||
|
||||
def test_map_replay_conflict_rolls_back_new_chunks(mapped_setup):
|
||||
f, generation, encoder, _, _ = mapped_setup
|
||||
row = put(f[1], "before token=private after")
|
||||
f[3].consume_pending(generation)
|
||||
lease = f[3].claim(generation, "publisher")
|
||||
prepared = prepare_mapped_content(row)
|
||||
chunks = chunk_mapped_content(prepared, encoder)
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute(
|
||||
"INSERT INTO retrieval_build_maps VALUES (?,?,?,?)",
|
||||
(lease.build_id, "{}", "0" * 64, "indexable"),
|
||||
)
|
||||
with pytest.raises(ValueError, match="immutable_source_map"):
|
||||
f[3].publish(lease, chunks, vectors_verified=True, source_map=prepared.metadata)
|
||||
with closing(sqlite3.connect(f[0])) as db:
|
||||
assert not db.execute("SELECT * FROM retrieval_chunks").fetchall()
|
||||
assert f[3].status(generation)["leased"] == 1
|
||||
|
||||
|
||||
def test_deadline_expiring_during_publish_rolls_back_completion(mapped_setup, monkeypatch):
|
||||
f, generation, _, _, worker = mapped_setup
|
||||
put(f[1], "before token=private after")
|
||||
current = [0.0]
|
||||
worker.clock = lambda: current[0]
|
||||
worker.max_job_seconds = 10
|
||||
original = f[3]._prepare
|
||||
|
||||
def expires(*args):
|
||||
result = original(*args)
|
||||
current[0] = 11
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(f[3], "_prepare", expires)
|
||||
result = worker.run_once()
|
||||
assert result.state == "deferred" and result.reason == "deadline"
|
||||
assert f[3].status(generation)["leased"] == 1
|
||||
with closing(sqlite3.connect(f[0])) as db:
|
||||
assert not db.execute("SELECT * FROM retrieval_chunks").fetchall()
|
||||
assert not db.execute("SELECT * FROM retrieval_build_maps").fetchall()
|
||||
|
||||
|
||||
def test_deadline_and_character_budget_never_release_cross_gap_spans(mapped_setup):
|
||||
f, generation, _, vectors, worker = mapped_setup
|
||||
put(f[1], "before token=private after")
|
||||
assert worker.run_once().state == "ready"
|
||||
f[3].activate_generation(generation)
|
||||
hits = candidates(f, generation, vectors)
|
||||
expired = materialize_raw_candidates(
|
||||
f[3],
|
||||
hits,
|
||||
generation_id=generation,
|
||||
source_ids=frozenset({"codex"}),
|
||||
deadline=time.monotonic() - 1,
|
||||
)
|
||||
assert not expired.items and expired.reason == "deadline"
|
||||
bounded = materialize_raw_candidates(
|
||||
f[3],
|
||||
hits,
|
||||
generation_id=generation,
|
||||
source_ids=frozenset({"codex"}),
|
||||
deadline=time.monotonic() + 10,
|
||||
character_budget=7,
|
||||
)
|
||||
assert not bounded.complete and bounded.reason == "character_budget"
|
||||
assert bounded.used_characters <= 7
|
||||
assert sum(len(item.spans) for item in bounded.items) == 1
|
||||
|
||||
|
||||
def test_lease_expiring_inside_publish_cannot_commit(mapped_setup, monkeypatch):
|
||||
f, generation, _, _, worker = mapped_setup
|
||||
put(f[1], "before token=private after")
|
||||
current = [100.0]
|
||||
worker.wall_clock = lambda: current[0]
|
||||
original = f[3]._prepare
|
||||
|
||||
def expires(*args):
|
||||
result = original(*args)
|
||||
current[0] = 1000.0
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(f[3], "_prepare", expires)
|
||||
assert worker.run_once().state == "blocked"
|
||||
with closing(sqlite3.connect(f[0])) as db:
|
||||
assert not db.execute("SELECT * FROM retrieval_chunks").fetchall()
|
||||
assert not db.execute("SELECT * FROM retrieval_build_maps").fetchall()
|
||||
assert f[3].status(generation)["ready"] == 0
|
||||
@@ -1,125 +0,0 @@
|
||||
"""The index has to come up through the product, not through a test helper.
|
||||
|
||||
Every existing retrieval test builds its generation by hand, so none of them
|
||||
can tell whether a user who enables the feature ever gets an index at all.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.imports.storage import EvidenceStore
|
||||
from one_memory.retrieval.index_provisioning import provision_retrieval_index
|
||||
from one_memory.retrieval.index_store import RetrievalIndexStore
|
||||
from one_memory.retrieval.index_worker import LocalRetrievalEncoder
|
||||
|
||||
|
||||
def database(tmp_path: Path) -> Path:
|
||||
path = tmp_path / "one-memory.sqlite3"
|
||||
EvidenceStore(path)
|
||||
return path
|
||||
|
||||
|
||||
def active(path: Path) -> list[str]:
|
||||
with closing(sqlite3.connect(path)) as db:
|
||||
if not db.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE name='retrieval_index_generations'"
|
||||
).fetchone():
|
||||
return []
|
||||
return [
|
||||
row[0]
|
||||
for row in db.execute(
|
||||
"SELECT generation_id FROM retrieval_index_generations WHERE state='active'"
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_first_enable_publishes_one_active_generation_and_its_vector_store(tmp_path):
|
||||
database(tmp_path)
|
||||
|
||||
generation = provision_retrieval_index(tmp_path)
|
||||
|
||||
assert generation is not None
|
||||
assert active(tmp_path / "one-memory.sqlite3") == [generation]
|
||||
assert (tmp_path / "retrieval-vectors").is_dir()
|
||||
|
||||
|
||||
def test_second_start_reuses_the_same_generation_rather_than_stacking_one(tmp_path):
|
||||
database(tmp_path)
|
||||
|
||||
first = provision_retrieval_index(tmp_path)
|
||||
second = provision_retrieval_index(tmp_path)
|
||||
|
||||
assert first == second
|
||||
assert active(tmp_path / "one-memory.sqlite3") == [first]
|
||||
|
||||
|
||||
def test_a_v1_index_is_left_alone_instead_of_being_migrated_on_startup(tmp_path, caplog):
|
||||
"""Falsifiable against the first test, where provisioning does publish one.
|
||||
|
||||
Migrating someone's existing index while they were only starting the app is
|
||||
the one thing the schema note rules out, so this must stay a no-op. The
|
||||
absence of a warning is asserted too: crashing into the same `None` would
|
||||
otherwise read as a deliberate decision.
|
||||
"""
|
||||
path = database(tmp_path)
|
||||
store = RetrievalIndexStore(path)
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
store.initialize(db)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
assert provision_retrieval_index(tmp_path) is None
|
||||
assert active(path) == []
|
||||
assert "retrieval_index_provision_failed" not in caplog.text
|
||||
with closing(sqlite3.connect(path)) as db:
|
||||
version = db.execute("SELECT schema_version FROM retrieval_index_state").fetchone()[0]
|
||||
assert version == 1
|
||||
|
||||
|
||||
def test_a_generation_for_another_encoder_is_not_swapped_out_silently(tmp_path):
|
||||
path = database(tmp_path)
|
||||
store = RetrievalIndexStore(path)
|
||||
other = LocalRetrievalEncoder(content_version=1).identity
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
store.initialize(db)
|
||||
store.migrate_v2(db)
|
||||
foreign = store.create_generation(db, other)
|
||||
store.activate_generation(foreign)
|
||||
|
||||
assert provision_retrieval_index(tmp_path) is None
|
||||
assert active(path) == [foreign]
|
||||
|
||||
|
||||
def test_without_the_built_in_encoder_no_index_is_created(tmp_path, monkeypatch):
|
||||
path = database(tmp_path)
|
||||
monkeypatch.setattr(
|
||||
"one_memory.retrieval.index_provisioning.assets_present", lambda *_args: False
|
||||
)
|
||||
|
||||
assert provision_retrieval_index(tmp_path) is None
|
||||
assert active(path) == []
|
||||
assert not (tmp_path / "retrieval-vectors").exists()
|
||||
|
||||
|
||||
def test_a_missing_database_is_not_created_by_provisioning(tmp_path):
|
||||
assert provision_retrieval_index(tmp_path) is None
|
||||
assert not (tmp_path / "one-memory.sqlite3").exists()
|
||||
assert not (tmp_path / "retrieval-vectors").exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["initialize_v2", "create_generation"])
|
||||
def test_a_failed_provision_reports_none_instead_of_half_an_index(tmp_path, monkeypatch, failure):
|
||||
path = database(tmp_path)
|
||||
|
||||
def explode(*_args, **_kwargs):
|
||||
raise sqlite3.OperationalError("synthetic")
|
||||
|
||||
monkeypatch.setattr(RetrievalIndexStore, failure, explode)
|
||||
|
||||
assert provision_retrieval_index(tmp_path) is None
|
||||
assert active(path) == []
|
||||
@@ -1,224 +0,0 @@
|
||||
import sqlite3
|
||||
import time
|
||||
from contextlib import closing
|
||||
from datetime import datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.imports.models import EvidenceInput
|
||||
from one_memory.retrieval.index_worker import RetrievalIndexWorker
|
||||
from one_memory.retrieval.models import TimeRange
|
||||
from one_memory.retrieval.raw_candidates import (
|
||||
compose_raw_retrieval_path,
|
||||
read_raw_semantic_candidates,
|
||||
)
|
||||
from one_memory.retrieval.vector_store import RawVectorStore
|
||||
from tests.test_retrieval_index_store import fixture, initialize, put # noqa: F401
|
||||
from tests.test_retrieval_index_worker import Encoder, forbid_network # noqa: F401
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def indexed(fixture, tmp_path): # noqa: F811
|
||||
generation = initialize(fixture)
|
||||
record = put(fixture[1])
|
||||
encoder = Encoder(fixture[4])
|
||||
vectors = RawVectorStore.create(tmp_path / "vectors", generation_id=generation, dimensions=3)
|
||||
worker = RetrievalIndexWorker(
|
||||
fixture[3], vectors, encoder, generation_id=generation, owner="test"
|
||||
)
|
||||
assert worker.run_once().state == "ready"
|
||||
fixture[3].activate_generation(generation)
|
||||
try:
|
||||
yield fixture, generation, record, vectors
|
||||
finally:
|
||||
vectors.close()
|
||||
|
||||
|
||||
def search(indexed, **kwargs):
|
||||
f, generation, _, vectors = indexed
|
||||
options = dict(source_ids=frozenset({"codex"}), deadline=time.monotonic() + 10)
|
||||
options.update(kwargs)
|
||||
return read_raw_semantic_candidates(
|
||||
f[3], vectors, [1.0, 0.5, 0.25], generation_id=generation, **options
|
||||
)
|
||||
|
||||
|
||||
def test_positive_only_current_authorized_published_chunks_are_searched(indexed):
|
||||
result = search(indexed)
|
||||
assert result.complete and len(result.hits) == 1
|
||||
assert result.hits[0].chunk.evidence_id == indexed[2].evidence_id
|
||||
assert result.coverage == "indexed_only"
|
||||
|
||||
|
||||
def test_empty_authority_never_reads_database_or_vector(indexed, monkeypatch):
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError("unauthorized access")
|
||||
|
||||
monkeypatch.setattr(indexed[0][3], "active_chunks", forbidden)
|
||||
monkeypatch.setattr(indexed[3], "search", forbidden)
|
||||
result = search(indexed, source_ids=frozenset())
|
||||
assert not result.hits and result.reason == "not_authorized"
|
||||
|
||||
|
||||
def test_unapproved_source_never_enters_vector_query(indexed, monkeypatch):
|
||||
monkeypatch.setattr(indexed[3], "search", lambda *a, **k: pytest.fail("not eligible"))
|
||||
assert search(indexed, source_ids=frozenset({"cursor"})).hits == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", ["revoke", "revise", "policy"])
|
||||
def test_mutation_during_vector_query_discards_stale_candidate(indexed, monkeypatch, change):
|
||||
f, _, _, vectors = indexed
|
||||
original = vectors.search
|
||||
|
||||
def query(*args, **kwargs):
|
||||
result = original(*args, **kwargs)
|
||||
if change == "revise":
|
||||
put(f[1], "changed text")
|
||||
elif change == "revoke":
|
||||
f[2].revoke(user_id="local", source_id="codex")
|
||||
else:
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute("UPDATE source_grants SET revision=revision+1")
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(vectors, "search", query)
|
||||
result = search(indexed)
|
||||
assert not result.hits and not result.complete and result.reason == "evidence_changed"
|
||||
|
||||
|
||||
def test_expired_deadline_does_not_read_or_query(indexed, monkeypatch):
|
||||
monkeypatch.setattr(indexed[3], "search", lambda *a, **k: pytest.fail("late query"))
|
||||
result = search(indexed, deadline=time.monotonic() - 1)
|
||||
assert not result.hits and not result.complete and result.reason == "deadline"
|
||||
|
||||
|
||||
def test_missing_published_vector_is_error_not_empty(indexed):
|
||||
indexed[3]._collection.delete(ids=[c.chunk_id for c in indexed[0][3].active_chunks(indexed[1])])
|
||||
result = search(indexed)
|
||||
assert not result.hits and not result.complete and result.reason == "vector_unavailable"
|
||||
|
||||
|
||||
def test_small_allowlist_uses_query_as_integrity_check_without_duplicate_get(indexed, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
indexed[3],
|
||||
"verify_candidates",
|
||||
lambda *args, **kwargs: pytest.fail("small allowlist must not be read twice"),
|
||||
)
|
||||
result = search(indexed)
|
||||
assert result.complete and len(result.hits) == 1
|
||||
|
||||
|
||||
def test_missing_event_time_never_matches_hard_range(indexed, monkeypatch):
|
||||
window = TimeRange(start="2026-09-18T00:00:00+08:00", end="2026-09-19T00:00:00+08:00")
|
||||
monkeypatch.setattr(indexed[3], "search", lambda *a, **k: pytest.fail("unknown event time"))
|
||||
assert search(indexed, time_range=window).hits == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sources", [None, ["codex"], "codex"])
|
||||
def test_caller_must_provide_explicit_resolved_authority(indexed, sources):
|
||||
with pytest.raises(ValueError, match="source"):
|
||||
search(indexed, source_ids=sources)
|
||||
|
||||
|
||||
def test_event_time_uses_utc_half_open_range_not_import_time(indexed):
|
||||
f, generation, record, vectors = indexed
|
||||
f[1].put(
|
||||
EvidenceInput(
|
||||
source_id=record.source_id,
|
||||
conversation_id=record.conversation_id,
|
||||
message_id=record.message_id,
|
||||
role=record.role,
|
||||
written_by=record.written_by,
|
||||
verbatim=record.verbatim,
|
||||
created_at=datetime.fromisoformat("2026-09-18T16:00:00+00:00"),
|
||||
origin_meta=record.origin_meta,
|
||||
)
|
||||
)
|
||||
worker = RetrievalIndexWorker(
|
||||
f[3], vectors, Encoder(f[4]), generation_id=generation, owner="test"
|
||||
)
|
||||
assert worker.run_once().state == "ready"
|
||||
before = TimeRange(start="2026-09-18T00:00:00+08:00", end="2026-09-19T00:00:00+08:00")
|
||||
after = TimeRange(start="2026-09-19T00:00:00+08:00", end="2026-09-20T00:00:00+08:00")
|
||||
assert not search(indexed, time_range=before).hits
|
||||
result = search(indexed, time_range=after)
|
||||
assert result.complete and len(result.hits) == 1
|
||||
assert result.hits[0].chunk.event_time == after.start
|
||||
|
||||
|
||||
def test_final_deadline_failure_discards_previously_found_hits(indexed, monkeypatch):
|
||||
from one_memory.retrieval import raw_candidates
|
||||
|
||||
original = raw_candidates._check
|
||||
calls = []
|
||||
|
||||
def check(deadline):
|
||||
calls.append(deadline)
|
||||
if len(calls) == 3:
|
||||
raise TimeoutError("late final revalidation")
|
||||
return original(deadline)
|
||||
|
||||
monkeypatch.setattr(raw_candidates, "_check", check)
|
||||
result = search(indexed)
|
||||
assert len(calls) == 3 and not result.hits and result.reason == "deadline"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("corruption", ["all_missing", "tail_missing", "offset", "hash"])
|
||||
def test_corrupted_published_manifest_is_unavailable_not_complete_empty(indexed, corruption):
|
||||
f, generation, _, vectors = indexed
|
||||
put(f[1], "long source text " * 90)
|
||||
worker = RetrievalIndexWorker(
|
||||
f[3], vectors, Encoder(f[4]), generation_id=generation, owner="test"
|
||||
)
|
||||
assert worker.run_once().state == "ready"
|
||||
assert len(search(indexed).hits) > 1
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
if corruption == "all_missing":
|
||||
db.execute("DELETE FROM retrieval_chunks")
|
||||
elif corruption == "tail_missing":
|
||||
db.execute("DELETE FROM retrieval_chunks WHERE ordinal>0")
|
||||
elif corruption == "offset":
|
||||
db.execute("UPDATE retrieval_chunks SET start_cp=3 WHERE ordinal=0")
|
||||
else:
|
||||
db.execute("UPDATE retrieval_chunks SET text_hash='broken'")
|
||||
result = search(indexed)
|
||||
assert not result.complete and not result.hits and result.reason == "index_unavailable"
|
||||
|
||||
|
||||
def test_null_stored_grant_with_revocation_never_queries_vectors(indexed, monkeypatch):
|
||||
f, _, _, vectors = indexed
|
||||
assert search(indexed).hits
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute("UPDATE retrieval_documents SET grant_fingerprint=NULL")
|
||||
f[2].revoke(user_id="local", source_id="codex")
|
||||
monkeypatch.setattr(vectors, "search", lambda *a, **k: pytest.fail("revoked vector query"))
|
||||
assert not search(indexed).hits
|
||||
|
||||
|
||||
def test_unpublished_authorized_evidence_cannot_claim_complete_empty(indexed):
|
||||
"""A partial index with no hits is incomplete, not an honest empty library."""
|
||||
f, generation, _, _vectors = indexed
|
||||
put(f[1], "用户决定 One Memory 采用 MIT 协议发布。", message="unindexed")
|
||||
assert f[3].consume_pending(generation) >= 1
|
||||
pending = f[3].status(generation)
|
||||
assert pending["pending"] >= 1
|
||||
# Published chunks may still match a probe vector. The unpublished
|
||||
# authorized document is what forbids complete-empty, so the candidate
|
||||
# batch is forced empty: that is the lying-empty case, not a miss on
|
||||
# already-indexed text.
|
||||
empty = type(search(indexed))(hits=(), complete=True, reason="no_eligible_chunks")
|
||||
path = compose_raw_retrieval_path(
|
||||
empty, unpublished_authorized=pending["pending"] > 0, lexical_hit_count=0
|
||||
)
|
||||
assert path.coverage == "partial"
|
||||
assert path.status != "empty"
|
||||
assert "semantic" in path.searched_channels
|
||||
|
||||
|
||||
def test_complete_index_with_no_hits_may_be_empty(indexed):
|
||||
batch = search(indexed, source_ids=frozenset({"codex"}))
|
||||
# Query vector still matches the only indexed chunk; force a no-hit batch.
|
||||
empty = type(batch)(hits=(), complete=True, reason="no_eligible_chunks")
|
||||
path = compose_raw_retrieval_path(empty, unpublished_authorized=False, lexical_hit_count=0)
|
||||
assert path.status == "empty"
|
||||
assert path.coverage == "complete"
|
||||
@@ -1,88 +0,0 @@
|
||||
import sqlite3
|
||||
import time
|
||||
from contextlib import closing
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.retrieval.index_worker import RetrievalIndexWorker
|
||||
from one_memory.retrieval.raw_evidence import materialize_raw_candidates
|
||||
from tests.test_retrieval_index_store import fixture, put # noqa: F401
|
||||
from tests.test_retrieval_index_worker import Encoder, forbid_network # noqa: F401
|
||||
from tests.test_retrieval_raw_candidates import indexed as indexed_fixture
|
||||
from tests.test_retrieval_raw_candidates import search
|
||||
|
||||
indexed = indexed_fixture
|
||||
|
||||
|
||||
def materialize(indexed, candidates=None, **kwargs):
|
||||
options = dict(
|
||||
source_ids=frozenset({"codex"}), deadline=time.monotonic() + 10, character_budget=6000
|
||||
)
|
||||
options.update(kwargs)
|
||||
return materialize_raw_candidates(
|
||||
indexed[0][3], candidates or search(indexed), generation_id=indexed[1], **options
|
||||
)
|
||||
|
||||
|
||||
def test_materialized_text_is_exact_current_original(indexed):
|
||||
result = materialize(indexed)
|
||||
assert result.complete and len(result.items) == 1
|
||||
item = result.items[0]
|
||||
assert item.evidence_id == indexed[2].evidence_id
|
||||
assert item.spans[0].text == indexed[2].verbatim
|
||||
assert item.written_by == "person" and item.matched_by == ("semantic",)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", ["revision", "revoke", "unrecorded"])
|
||||
def test_stale_candidates_do_not_authorize_original_read(indexed, change):
|
||||
candidates = search(indexed)
|
||||
assert candidates.hits
|
||||
f = indexed[0]
|
||||
if change == "revision":
|
||||
put(f[1], "updated original")
|
||||
elif change == "revoke":
|
||||
f[2].revoke(user_id="local", source_id="codex")
|
||||
else:
|
||||
with closing(sqlite3.connect(f[0])) as db, db:
|
||||
db.execute(
|
||||
"UPDATE evidence SET verbatim='unrecorded' WHERE evidence_id=?",
|
||||
(indexed[2].evidence_id,),
|
||||
)
|
||||
result = materialize(indexed, candidates)
|
||||
assert not result.items and not result.complete and result.reason == "evidence_changed"
|
||||
|
||||
|
||||
def test_budget_does_not_cut_a_chunk_into_fake_partial_object(indexed):
|
||||
result = materialize(indexed, character_budget=5)
|
||||
assert not result.items and not result.complete and result.reason == "character_budget"
|
||||
|
||||
|
||||
def test_forged_offsets_and_modified_manifest_do_not_release_body(indexed):
|
||||
candidates = search(indexed)
|
||||
with closing(sqlite3.connect(indexed[0][0])) as db, db:
|
||||
db.execute("UPDATE retrieval_chunks SET start_cp=2")
|
||||
result = materialize(indexed, candidates)
|
||||
assert not result.items and not result.complete
|
||||
|
||||
|
||||
def test_empty_authority_and_expired_deadline_return_no_body(indexed):
|
||||
candidates = search(indexed)
|
||||
assert not materialize(indexed, candidates, source_ids=frozenset()).items
|
||||
assert not materialize(indexed, candidates, deadline=time.monotonic() - 1).items
|
||||
|
||||
|
||||
def test_overlapping_unicode_chunks_merge_without_duplicate_characters(indexed):
|
||||
f, generation, _, vectors = indexed
|
||||
text = "\u6c49\u5b57\U0001f680e\u0301 " * 180 + " /project/final-design.png"
|
||||
put(f[1], text)
|
||||
worker = RetrievalIndexWorker(
|
||||
f[3], vectors, Encoder(f[4]), generation_id=generation, owner="test"
|
||||
)
|
||||
assert worker.run_once().state == "ready"
|
||||
candidates = search(indexed)
|
||||
assert len(candidates.hits) > 1
|
||||
result = materialize(indexed, candidates)
|
||||
assert result.complete and len(result.items) == 1
|
||||
assert len(result.items[0].spans) == 1
|
||||
assert result.items[0].spans[0].text == text
|
||||
assert result.used_characters == len(text)
|
||||
@@ -1,356 +0,0 @@
|
||||
"""v1 raw lane must not claim complete empty when the derived index is partial."""
|
||||
|
||||
import asyncio
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.context_models import ContextLaneResult, ContextSearchRequest
|
||||
from one_memory.context_raw import read_raw_context
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.embedding_assets import assets_present
|
||||
from one_memory.retrieval.raw_candidates import RawCandidateBatch
|
||||
from one_memory.retrieval.raw_lane import fuse_raw_lane, read_raw_indexed_context
|
||||
from tests.test_context_raw import database as database # noqa: F401
|
||||
from tests.test_context_raw import put
|
||||
|
||||
|
||||
def _require_assets() -> None:
|
||||
"""Make an absent local fixture explicit instead of a lease-style hang."""
|
||||
if not assets_present():
|
||||
pytest.skip("_assets_unavailable")
|
||||
|
||||
|
||||
def test_unpublished_authorized_evidence_degrades_empty_lexical_lane():
|
||||
lexical = ContextLaneResult(status="empty", reason="lexical_only")
|
||||
result = fuse_raw_lane(
|
||||
lexical,
|
||||
RawCandidateBatch(complete=True, reason="no_eligible_chunks"),
|
||||
unpublished_authorized=True,
|
||||
)
|
||||
assert result.status == "degraded"
|
||||
assert result.reason == "lexical_only"
|
||||
assert result.items == ()
|
||||
|
||||
|
||||
def test_complete_index_with_no_hits_stays_empty():
|
||||
lexical = ContextLaneResult(status="empty", reason="lexical_only")
|
||||
result = fuse_raw_lane(
|
||||
lexical,
|
||||
RawCandidateBatch(complete=True, reason="no_eligible_chunks"),
|
||||
unpublished_authorized=False,
|
||||
)
|
||||
assert result.status == "empty"
|
||||
assert result.items == ()
|
||||
|
||||
|
||||
def test_interactive_raw_lane_reads_one_snapshot_while_a_writer_holds(database):
|
||||
"""The read lane must not wait for an import writer, nor see its uncommitted row.
|
||||
|
||||
This is the C4-owned half of the import-contention BLOCK: WAL readers are
|
||||
expected to proceed from the last committed snapshot even while the single
|
||||
writer holds a transaction. A regression that lets context reads block on
|
||||
the writer - or read its uncommitted row - makes this fail.
|
||||
"""
|
||||
with sqlite3.connect(database) as db, db:
|
||||
db.execute("PRAGMA journal_mode=WAL")
|
||||
put(db, "committed", "needle committed row")
|
||||
writer_holds = threading.Event()
|
||||
release_writer = threading.Event()
|
||||
|
||||
def hold_write_transaction() -> None:
|
||||
with main_database(database).write(autocommit=True) as writer:
|
||||
writer.execute("BEGIN IMMEDIATE")
|
||||
writer.execute(
|
||||
"""INSERT INTO evidence
|
||||
(evidence_id,idempotency_key,user_id,source_id,conversation_id,role,
|
||||
written_by,created_at,verbatim,content_hash,revision,imported_at,updated_at)
|
||||
VALUES ('uncommitted','uncommitted','local','cursor','c','user',
|
||||
'person','2026-09-20','needle uncommitted row','synthetic',
|
||||
1,'2026-09-20','2026-09-20')"""
|
||||
)
|
||||
writer_holds.set()
|
||||
assert release_writer.wait(5)
|
||||
writer.execute("ROLLBACK")
|
||||
|
||||
thread = threading.Thread(target=hold_write_transaction)
|
||||
thread.start()
|
||||
assert writer_holds.wait(5)
|
||||
try:
|
||||
before = main_database(database).stats()
|
||||
started = time.monotonic()
|
||||
result = read_raw_context(
|
||||
database,
|
||||
ContextSearchRequest(query="needle", kinds=("raw",)),
|
||||
source_ids=frozenset({"cursor"}),
|
||||
deadline=time.monotonic() + 2,
|
||||
)
|
||||
elapsed = time.monotonic() - started
|
||||
after = main_database(database).stats()
|
||||
finally:
|
||||
release_writer.set()
|
||||
thread.join(timeout=5)
|
||||
assert not thread.is_alive()
|
||||
assert elapsed < 1.0
|
||||
assert after.reader_waits == before.reader_waits
|
||||
assert [item.anchor_evidence_id for item in result.items] == ["committed"]
|
||||
assert all(
|
||||
message.evidence_id != "uncommitted" for item in result.items for message in item.messages
|
||||
)
|
||||
|
||||
|
||||
def test_authorization_and_unavailability_are_not_rewritten():
|
||||
denied = ContextLaneResult(status="not_authorized", reason="permission_required")
|
||||
assert (
|
||||
fuse_raw_lane(denied, RawCandidateBatch(), unpublished_authorized=True).status
|
||||
== "not_authorized"
|
||||
)
|
||||
missing = ContextLaneResult(status="unavailable", reason="read_unavailable")
|
||||
assert fuse_raw_lane(missing, RawCandidateBatch(), unpublished_authorized=True).status == (
|
||||
"unavailable"
|
||||
)
|
||||
|
||||
|
||||
def test_indexed_reader_without_runtime_matches_lexical_empty(database):
|
||||
with sqlite3.connect(database) as db, db:
|
||||
put(db, "e-unrelated", "alpha transcript only")
|
||||
request = ContextSearchRequest(query="no-matching-zzzzzz", kinds=("raw",))
|
||||
deadline = time.monotonic() + 5
|
||||
lexical = read_raw_context(
|
||||
database, request, source_ids=frozenset({"cursor"}), deadline=deadline
|
||||
)
|
||||
indexed = read_raw_indexed_context(
|
||||
database, request, source_ids=frozenset({"cursor"}), deadline=deadline, index=None
|
||||
)
|
||||
assert lexical.status == "empty"
|
||||
assert indexed == lexical
|
||||
assert indexed.searched_channels == ("lexical",)
|
||||
assert indexed.semantic_status == "not_searched"
|
||||
assert indexed.index_complete is False
|
||||
|
||||
|
||||
def test_indexed_reader_with_unpublished_index_is_not_empty(database):
|
||||
with sqlite3.connect(database) as db, db:
|
||||
put(db, "e-unrelated", "alpha transcript only")
|
||||
request = ContextSearchRequest(query="no-matching-zzzzzz", kinds=("raw",))
|
||||
index = SimpleNamespace(
|
||||
lookup_raw=lambda query, *, source_ids, deadline: (
|
||||
RawCandidateBatch(complete=True, reason="no_eligible_chunks"),
|
||||
True,
|
||||
)
|
||||
)
|
||||
result = read_raw_indexed_context(
|
||||
database,
|
||||
request,
|
||||
source_ids=frozenset({"cursor"}),
|
||||
deadline=time.monotonic() + 5,
|
||||
index=index,
|
||||
)
|
||||
assert result.status == "degraded"
|
||||
assert result.items == ()
|
||||
assert result.searched_channels == ("lexical", "semantic")
|
||||
assert result.semantic_status == "completed"
|
||||
assert result.index_complete is False
|
||||
|
||||
|
||||
def test_semantic_hit_without_lexical_overlap_is_not_released(database):
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from one_memory.retrieval.index_store import PublishedChunk
|
||||
from one_memory.retrieval.raw_candidates import RawSemanticCandidate
|
||||
|
||||
with sqlite3.connect(database) as db, db:
|
||||
put(db, "e-mit", "用户决定 One Memory 采用 MIT 协议发布。")
|
||||
chunk = PublishedChunk(
|
||||
chunk_id="c1",
|
||||
build_id="b1",
|
||||
evidence_id="e-mit",
|
||||
revision=1,
|
||||
start_cp=0,
|
||||
end_cp=8,
|
||||
ordinal=0,
|
||||
source_id="cursor",
|
||||
event_time=datetime(2026, 9, 19, tzinfo=UTC),
|
||||
)
|
||||
index = SimpleNamespace(
|
||||
lookup_raw=lambda query, *, source_ids, deadline: (
|
||||
RawCandidateBatch(
|
||||
hits=(RawSemanticCandidate(chunk, 0.1),),
|
||||
complete=True,
|
||||
),
|
||||
False,
|
||||
)
|
||||
)
|
||||
result = read_raw_indexed_context(
|
||||
database,
|
||||
ContextSearchRequest(query="开源许可证怎么选", kinds=("raw",)),
|
||||
source_ids=frozenset({"cursor"}),
|
||||
deadline=time.monotonic() + 5,
|
||||
index=index,
|
||||
)
|
||||
assert result.status == "degraded"
|
||||
assert result.reason == "lexical_only"
|
||||
assert result.items == ()
|
||||
assert result.searched_channels == ("lexical", "semantic")
|
||||
assert result.semantic_status == "completed"
|
||||
assert result.index_complete is True
|
||||
|
||||
|
||||
def test_semantic_candidate_with_lexical_corroboration_is_released(database):
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from one_memory.retrieval.index_store import PublishedChunk
|
||||
from one_memory.retrieval.raw_candidates import RawSemanticCandidate
|
||||
|
||||
with sqlite3.connect(database) as db, db:
|
||||
put(db, "e-mit", "用户决定 One Memory 采用 MIT 协议发布。")
|
||||
chunk = PublishedChunk(
|
||||
chunk_id="c1",
|
||||
build_id="b1",
|
||||
evidence_id="e-mit",
|
||||
revision=1,
|
||||
start_cp=0,
|
||||
end_cp=8,
|
||||
ordinal=0,
|
||||
source_id="cursor",
|
||||
event_time=datetime(2026, 9, 19, tzinfo=UTC),
|
||||
)
|
||||
index = SimpleNamespace(
|
||||
lookup_raw=lambda query, *, source_ids, deadline: (
|
||||
RawCandidateBatch(
|
||||
hits=(RawSemanticCandidate(chunk, 0.1),),
|
||||
complete=True,
|
||||
),
|
||||
False,
|
||||
)
|
||||
)
|
||||
result = read_raw_indexed_context(
|
||||
database,
|
||||
ContextSearchRequest(query="MIT 协议", kinds=("raw",)),
|
||||
source_ids=frozenset({"cursor"}),
|
||||
deadline=time.monotonic() + 5,
|
||||
index=index,
|
||||
)
|
||||
assert result.status == "ok"
|
||||
assert [item.anchor_evidence_id for item in result.items] == ["e-mit"]
|
||||
|
||||
|
||||
def test_raw_semantic_timeout_is_not_reported_as_complete_empty(database):
|
||||
with sqlite3.connect(database) as db, db:
|
||||
put(db, "e-unrelated", "alpha transcript only")
|
||||
index = SimpleNamespace(
|
||||
lookup_raw=lambda query, *, source_ids, deadline: (
|
||||
RawCandidateBatch(reason="deadline"),
|
||||
False,
|
||||
)
|
||||
)
|
||||
result = read_raw_indexed_context(
|
||||
database,
|
||||
ContextSearchRequest(query="no-matching-zzzzzz", kinds=("raw",)),
|
||||
source_ids=frozenset({"cursor"}),
|
||||
deadline=time.monotonic() + 5,
|
||||
index=index,
|
||||
)
|
||||
assert result.status == "degraded"
|
||||
assert result.semantic_status == "timed_out"
|
||||
assert result.index_complete is False
|
||||
assert result.timed_out
|
||||
|
||||
|
||||
def test_real_built_index_semantic_only_hit_is_not_released_without_confirmation(tmp_path):
|
||||
"""The whole chain for real: ONNX encode, Chroma, published chunks, authority.
|
||||
|
||||
Every other test here hands the lane a hand-written candidate, so none of
|
||||
them can tell whether a real build would ever produce one. The query is
|
||||
chosen to share no term with the stored sentence, and the lexical-only read
|
||||
is asserted empty first, so an `ok` here can only have come from the index.
|
||||
"""
|
||||
_require_assets()
|
||||
from one_memory.imports.grants import SourceGrantStore
|
||||
from one_memory.imports.sources.registry import RootCandidateKind
|
||||
from one_memory.retrieval.runtime import ExistingRetrievalRuntime
|
||||
from tests.test_retrieval_index_store import put as put_evidence
|
||||
from tests.test_retrieval_runtime import configured, until
|
||||
|
||||
store, evidence, generation = configured(tmp_path)
|
||||
source = tmp_path / "source"
|
||||
source.mkdir()
|
||||
SourceGrantStore(store.database_path, home=tmp_path).grant(
|
||||
user_id="local",
|
||||
source_id="codex",
|
||||
root=source,
|
||||
candidate_kind=RootCandidateKind.DEFAULT,
|
||||
)
|
||||
put_evidence(evidence, "我们最后决定 One Memory 用 MIT 协议发布。")
|
||||
request = ContextSearchRequest(query="开源软件的授权条款如何挑选", kinds=("raw",))
|
||||
sources = frozenset({"codex"})
|
||||
|
||||
lexical = read_raw_context(
|
||||
store.database_path, request, source_ids=sources, deadline=time.monotonic() + 5
|
||||
)
|
||||
assert lexical.status == "empty", "词法已经命中,这条测试证明不了语义"
|
||||
|
||||
async def scenario():
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert await handle.start()
|
||||
try:
|
||||
await until(lambda: store.status(generation)["ready"] >= 1)
|
||||
return await asyncio.to_thread(
|
||||
read_raw_indexed_context,
|
||||
store.database_path,
|
||||
request,
|
||||
source_ids=sources,
|
||||
deadline=time.monotonic() + 10,
|
||||
index=handle,
|
||||
)
|
||||
finally:
|
||||
await handle.stop()
|
||||
|
||||
result = asyncio.run(scenario())
|
||||
# The semantic lane found the row, but lexical-only evidence is not enough
|
||||
# confirmation to release it to a host. Keep the candidate visible as a
|
||||
# degraded coverage receipt until a relevance/corroboration gate exists.
|
||||
assert result.status == "degraded"
|
||||
assert result.reason == "lexical_only"
|
||||
assert result.items == ()
|
||||
|
||||
|
||||
def test_real_built_index_skips_with_a_named_asset_reason(tmp_path, monkeypatch):
|
||||
"""A missing untracked fixture must not look like a lease or worker hang."""
|
||||
monkeypatch.setattr("test_retrieval_raw_lane.assets_present", lambda: False)
|
||||
with pytest.raises(pytest.skip.Exception, match="_assets_unavailable"):
|
||||
test_real_built_index_semantic_only_hit_is_not_released_without_confirmation(tmp_path)
|
||||
|
||||
|
||||
def test_revoked_source_semantic_hit_does_not_release_body(database):
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from one_memory.retrieval.index_store import PublishedChunk
|
||||
from one_memory.retrieval.raw_candidates import RawSemanticCandidate
|
||||
from one_memory.retrieval.raw_lane import load_raw_items_from_candidates
|
||||
from tests.test_context_raw import revoke
|
||||
|
||||
with sqlite3.connect(database) as db, db:
|
||||
put(db, "e-revoked", "用户决定 One Memory 采用 MIT 协议发布。")
|
||||
revoke(db, "cursor")
|
||||
chunk = PublishedChunk(
|
||||
chunk_id="c1",
|
||||
build_id="b1",
|
||||
evidence_id="e-revoked",
|
||||
revision=1,
|
||||
start_cp=0,
|
||||
end_cp=8,
|
||||
ordinal=0,
|
||||
source_id="cursor",
|
||||
event_time=datetime(2026, 9, 19, tzinfo=UTC),
|
||||
)
|
||||
items = load_raw_items_from_candidates(
|
||||
database,
|
||||
RawCandidateBatch(hits=(RawSemanticCandidate(chunk, 0.1),), complete=True),
|
||||
source_ids=frozenset({"cursor"}),
|
||||
deadline=time.monotonic() + 5,
|
||||
)
|
||||
assert items == ()
|
||||
@@ -1,522 +0,0 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
from contextlib import closing
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.imports.grants import SourceGrantStore
|
||||
from one_memory.imports.sources.registry import RootCandidateKind
|
||||
from one_memory.imports.storage import EvidenceStore
|
||||
from one_memory.retrieval import runtime
|
||||
from one_memory.retrieval.index_store import RetrievalIndexStore
|
||||
from one_memory.retrieval.index_worker import IndexWorkResult, LocalRetrievalEncoder
|
||||
from one_memory.retrieval.runtime import ExistingRetrievalRuntime
|
||||
from one_memory.retrieval.vector_store import RawVectorStore
|
||||
from tests.test_retrieval_index_store import put
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def offline(monkeypatch):
|
||||
def reject(*args, **kwargs):
|
||||
raise AssertionError("offline runtime test")
|
||||
|
||||
monkeypatch.setattr(socket.socket, "connect", reject)
|
||||
monkeypatch.setattr(socket.socket, "connect_ex", reject)
|
||||
|
||||
|
||||
def configured(tmp_path, *, version=2, collection=True):
|
||||
path = tmp_path / "one-memory.sqlite3"
|
||||
evidence = EvidenceStore(path)
|
||||
identity = LocalRetrievalEncoder(content_version=version).identity
|
||||
store = RetrievalIndexStore(path)
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
store.initialize(db)
|
||||
store.migrate_v2(db)
|
||||
generation = store.create_generation(db, identity)
|
||||
store.activate_generation(generation)
|
||||
if collection:
|
||||
vectors = RawVectorStore.create(
|
||||
tmp_path / "retrieval-vectors", generation_id=generation, dimensions=identity.dimensions
|
||||
)
|
||||
vectors.close()
|
||||
return store, evidence, generation
|
||||
|
||||
|
||||
async def until(predicate):
|
||||
end = time.monotonic() + 10
|
||||
while not predicate():
|
||||
assert time.monotonic() < end
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
|
||||
def test_constructor_and_missing_start_create_nothing(tmp_path):
|
||||
target = tmp_path / "missing"
|
||||
handle = ExistingRetrievalRuntime(target)
|
||||
assert not target.exists()
|
||||
assert not asyncio.run(handle.start())
|
||||
assert handle.snapshot().state == "not_initialized"
|
||||
assert not target.exists()
|
||||
asyncio.run(handle.stop())
|
||||
|
||||
|
||||
def test_stop_releases_supervisor_worker_but_preserves_immutable_snapshot(tmp_path):
|
||||
configured(tmp_path)
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
|
||||
async def scenario():
|
||||
assert await handle.start()
|
||||
assert handle._supervisor is not None
|
||||
await handle.stop()
|
||||
|
||||
asyncio.run(scenario())
|
||||
assert handle._supervisor is None
|
||||
assert handle.snapshot().supervisor is not None
|
||||
assert not handle.snapshot().supervisor.thread_alive
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kind",
|
||||
["schema1", "no_active", "missing_collection", "bad_identity", "extra_identity", "two_active"],
|
||||
)
|
||||
def test_not_initialized_or_invalid_configuration_never_creates_index(tmp_path, kind):
|
||||
store, _, gen = configured(tmp_path, collection=False)
|
||||
with closing(sqlite3.connect(store.database_path)) as db, db:
|
||||
if kind == "schema1":
|
||||
db.execute("UPDATE retrieval_index_state SET schema_version=1")
|
||||
elif kind == "no_active":
|
||||
db.execute("UPDATE retrieval_index_generations SET state='building'")
|
||||
elif kind in ("bad_identity", "extra_identity"):
|
||||
raw = json.loads(
|
||||
db.execute("SELECT identity_json FROM retrieval_index_generations").fetchone()[0]
|
||||
)
|
||||
if kind == "bad_identity":
|
||||
raw["encoder_digest"] = "0" * 64
|
||||
else:
|
||||
raw["unexpected"] = "nope"
|
||||
db.execute("UPDATE retrieval_index_generations SET identity_json=?", (json.dumps(raw),))
|
||||
elif kind == "two_active":
|
||||
db.execute("DROP INDEX retrieval_one_active")
|
||||
db.execute(
|
||||
"INSERT INTO retrieval_index_generations SELECT 'other',identity_json,state,"
|
||||
"created_at FROM retrieval_index_generations"
|
||||
)
|
||||
before = db.execute("SELECT * FROM retrieval_index_generations").fetchall()
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert not asyncio.run(handle.start())
|
||||
assert handle.snapshot().state in {"not_initialized", "failed"}
|
||||
assert not (tmp_path / "retrieval-vectors").exists()
|
||||
with closing(sqlite3.connect(store.database_path)) as db:
|
||||
assert db.execute("SELECT * FROM retrieval_index_generations").fetchall() == before
|
||||
assert str(tmp_path) not in repr(handle.snapshot())
|
||||
|
||||
|
||||
def test_low_disk_rejects_before_open_writer(tmp_path, monkeypatch):
|
||||
configured(tmp_path)
|
||||
monkeypatch.setattr(
|
||||
runtime.shutil, "disk_usage", lambda _: shutil._ntuple_diskusage(1000, 999, 1)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
runtime.RawVectorStore,
|
||||
"open_writer_existing",
|
||||
lambda *a, **k: pytest.fail("opened on low disk"),
|
||||
)
|
||||
handle = ExistingRetrievalRuntime(tmp_path)
|
||||
assert not asyncio.run(handle.start())
|
||||
assert handle.snapshot().error_code == "retrieval_runtime_low_disk"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", ["root", "database", "vectors"])
|
||||
def test_symlink_rejected(tmp_path, kind):
|
||||
actual = tmp_path / "actual"
|
||||
actual.mkdir()
|
||||
configured(actual)
|
||||
selected = actual
|
||||
if kind == "root":
|
||||
selected = tmp_path / "link"
|
||||
selected.symlink_to(actual, target_is_directory=True)
|
||||
else:
|
||||
path = actual / ("one-memory.sqlite3" if kind == "database" else "retrieval-vectors")
|
||||
moved = tmp_path / "moved"
|
||||
path.rename(moved)
|
||||
path.symlink_to(moved, target_is_directory=kind == "vectors")
|
||||
handle = ExistingRetrievalRuntime(selected, minimum_free_bytes=0)
|
||||
assert not asyncio.run(handle.start())
|
||||
assert handle.snapshot().error_code == "retrieval_runtime_unsafe_path"
|
||||
|
||||
|
||||
def test_local_existing_integration_indexes_only_after_explicit_start(tmp_path):
|
||||
store, evidence, gen = configured(tmp_path)
|
||||
source = tmp_path / "source"
|
||||
source.mkdir()
|
||||
SourceGrantStore(store.database_path, home=tmp_path).grant(
|
||||
user_id="local", source_id="codex", root=source, candidate_kind=RootCandidateKind.DEFAULT
|
||||
)
|
||||
row = put(evidence, "before token=private after")
|
||||
|
||||
async def run():
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert store.status(gen)["ready"] == 0
|
||||
assert await handle.start()
|
||||
try:
|
||||
await until(lambda: store.status(gen)["ready"] == 1)
|
||||
assert len(store.active_chunks(gen)) == 2
|
||||
assert handle.snapshot().generation_id == gen
|
||||
with pytest.raises(RuntimeError, match="already_started"):
|
||||
await handle.start()
|
||||
finally:
|
||||
await handle.stop()
|
||||
assert handle.snapshot().state == "stopped"
|
||||
assert not handle.snapshot().supervisor.thread_alive
|
||||
assert evidence.get(row.evidence_id).verbatim == row.verbatim
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def fake_resources(tmp_path, monkeypatch, *, fail_worker=False, block_open=False, version=2):
|
||||
configured(tmp_path, version=version)
|
||||
entered, release, exited, opened = (threading.Event() for _ in range(4))
|
||||
events = []
|
||||
|
||||
class Vectors:
|
||||
def close(self):
|
||||
events.append("close")
|
||||
|
||||
def open_vectors(*args, **kwargs):
|
||||
opened.set()
|
||||
if block_open:
|
||||
assert release.wait(5)
|
||||
return Vectors()
|
||||
|
||||
class Worker:
|
||||
def __init__(self, *args, **kwargs):
|
||||
if fail_worker:
|
||||
raise RuntimeError("SECRET failure")
|
||||
|
||||
def run_once(self, *, cancel):
|
||||
entered.set()
|
||||
while not cancel.is_set():
|
||||
time.sleep(0.001)
|
||||
events.append("cancel")
|
||||
assert release.wait(5)
|
||||
events.append("exit")
|
||||
exited.set()
|
||||
return IndexWorkResult("cancelled")
|
||||
|
||||
monkeypatch.setattr(runtime.RawVectorStore, "open_writer_existing", open_vectors)
|
||||
monkeypatch.setattr(runtime, "RetrievalIndexWorker", Worker)
|
||||
return entered, release, exited, opened, events
|
||||
|
||||
|
||||
def test_partial_start_failure_closes_owned_vectors(tmp_path, monkeypatch):
|
||||
*_, events = fake_resources(tmp_path, monkeypatch, fail_worker=True)
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert not asyncio.run(handle.start())
|
||||
assert events == ["close"]
|
||||
assert handle.snapshot().error_code == "retrieval_runtime_start_failed"
|
||||
assert "SECRET" not in repr(handle.snapshot())
|
||||
|
||||
|
||||
def test_cancelled_stop_drains_before_close(tmp_path, monkeypatch):
|
||||
entered, release, exited, _, events = fake_resources(tmp_path, monkeypatch)
|
||||
|
||||
async def run():
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert await handle.start()
|
||||
await until(entered.is_set)
|
||||
task = asyncio.create_task(handle.stop())
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
try:
|
||||
await until(lambda: "cancel" in events)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0.03)
|
||||
assert not task.done() and "close" not in events
|
||||
finally:
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert exited.is_set() and events == ["cancel", "exit", "close"]
|
||||
await handle.stop()
|
||||
assert events.count("close") == 1
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_cancelled_start_drains_open_and_cleans_resources(tmp_path, monkeypatch):
|
||||
_, release, _, opened, events = fake_resources(
|
||||
tmp_path, monkeypatch, fail_worker=True, block_open=True
|
||||
)
|
||||
|
||||
async def run():
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
task = asyncio.create_task(handle.start())
|
||||
await until(opened.is_set)
|
||||
task.cancel()
|
||||
try:
|
||||
await asyncio.sleep(0.03)
|
||||
assert not task.done() and not events
|
||||
finally:
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert events == ["close"]
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_cancelled_successful_open_does_not_start_worker_and_closes_once(tmp_path, monkeypatch):
|
||||
entered, release, _, opened, events = fake_resources(tmp_path, monkeypatch, block_open=True)
|
||||
|
||||
async def run():
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
starting = asyncio.create_task(handle.start())
|
||||
await until(opened.is_set)
|
||||
starting.cancel()
|
||||
await asyncio.sleep(0.02)
|
||||
starting.cancel()
|
||||
try:
|
||||
await asyncio.sleep(0.02)
|
||||
assert not starting.done() and not events and not entered.is_set()
|
||||
finally:
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await starting
|
||||
assert events == ["close"] and not entered.is_set()
|
||||
assert handle.snapshot().state == "stopped"
|
||||
await handle.stop()
|
||||
assert events == ["close"]
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_old_identity_one_recovers_without_generation_upgrade(tmp_path, monkeypatch):
|
||||
entered, release, _, _, events = fake_resources(tmp_path, monkeypatch, version=1)
|
||||
path = tmp_path / "one-memory.sqlite3"
|
||||
with closing(sqlite3.connect(path)) as db:
|
||||
before = db.execute("SELECT * FROM retrieval_index_generations").fetchall()
|
||||
|
||||
async def run():
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert await handle.start()
|
||||
await until(entered.is_set)
|
||||
release.set()
|
||||
await handle.stop()
|
||||
|
||||
asyncio.run(run())
|
||||
assert events == ["cancel", "exit", "close"]
|
||||
with closing(sqlite3.connect(path)) as db:
|
||||
assert db.execute("SELECT * FROM retrieval_index_generations").fetchall() == before
|
||||
|
||||
|
||||
def test_missing_identity_fields_fail_even_when_dataclass_defaults_match(tmp_path):
|
||||
store, _, _ = configured(tmp_path, version=1, collection=False)
|
||||
with closing(sqlite3.connect(store.database_path)) as db, db:
|
||||
raw = json.loads(
|
||||
db.execute("SELECT identity_json FROM retrieval_index_generations").fetchone()[0]
|
||||
)
|
||||
raw.pop("normalizer_version")
|
||||
db.execute("UPDATE retrieval_index_generations SET identity_json=?", (json.dumps(raw),))
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert not asyncio.run(handle.start())
|
||||
assert handle.snapshot().error_code == "retrieval_runtime_invalid_index"
|
||||
|
||||
|
||||
def test_valid_hash_but_foreign_encoder_never_opens_vectors(tmp_path, monkeypatch):
|
||||
from dataclasses import asdict
|
||||
|
||||
from one_memory.retrieval.index_store import EncoderIdentity, digest
|
||||
|
||||
store, _, _ = configured(tmp_path)
|
||||
with closing(sqlite3.connect(store.database_path)) as db, db:
|
||||
raw = json.loads(
|
||||
db.execute("SELECT identity_json FROM retrieval_index_generations").fetchone()[0]
|
||||
)
|
||||
raw["encoder_digest"] = "0" * 64
|
||||
identity = EncoderIdentity(**raw)
|
||||
store_id = db.execute("SELECT data_store_id FROM retrieval_index_state").fetchone()[0]
|
||||
db.execute(
|
||||
"UPDATE retrieval_index_generations SET generation_id=?,identity_json=?",
|
||||
(digest([store_id, asdict(identity)]), json.dumps(raw)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
runtime.RawVectorStore,
|
||||
"open_writer_existing",
|
||||
lambda *a, **k: pytest.fail("foreign encoder opened"),
|
||||
)
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert not asyncio.run(handle.start())
|
||||
assert handle.snapshot().error_code == "retrieval_runtime_encoder_mismatch"
|
||||
|
||||
|
||||
def test_concurrent_stop_waits_start_and_runs_cleanup_once(tmp_path, monkeypatch):
|
||||
entered, release, _, opened, events = fake_resources(tmp_path, monkeypatch, block_open=True)
|
||||
|
||||
async def run():
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
starting = asyncio.create_task(handle.start())
|
||||
await until(opened.is_set)
|
||||
stops = [asyncio.create_task(handle.stop()), asyncio.create_task(handle.stop())]
|
||||
try:
|
||||
await asyncio.sleep(0.03)
|
||||
assert all(not task.done() for task in stops)
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(*stops)
|
||||
assert not await starting
|
||||
assert events == ["close"] and not entered.is_set()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_stop_before_start_never_opens_resources(tmp_path):
|
||||
handle = ExistingRetrievalRuntime(tmp_path)
|
||||
asyncio.run(handle.stop())
|
||||
with pytest.raises(RuntimeError, match="already_started"):
|
||||
asyncio.run(handle.start())
|
||||
assert not list(tmp_path.iterdir())
|
||||
|
||||
|
||||
def test_missing_collection_in_existing_chroma_does_not_create_collection(tmp_path):
|
||||
store, _, gen = configured(tmp_path)
|
||||
with closing(sqlite3.connect(tmp_path / "retrieval-vectors" / "chroma.sqlite3")) as db, db:
|
||||
db.execute("UPDATE collections SET name='foreign_collection'")
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert not asyncio.run(handle.start())
|
||||
assert handle.snapshot().state == "not_initialized"
|
||||
with closing(sqlite3.connect(tmp_path / "retrieval-vectors" / "chroma.sqlite3")) as db:
|
||||
assert db.execute("SELECT name FROM collections").fetchall() == [("foreign_collection",)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target", ["root", "database", "vectors", "vector_database"])
|
||||
def test_group_writable_paths_rejected_without_open(tmp_path, monkeypatch, target):
|
||||
configured(tmp_path)
|
||||
paths = {
|
||||
"root": tmp_path,
|
||||
"database": tmp_path / "one-memory.sqlite3",
|
||||
"vectors": tmp_path / "retrieval-vectors",
|
||||
"vector_database": tmp_path / "retrieval-vectors/chroma.sqlite3",
|
||||
}
|
||||
path = paths[target]
|
||||
original_mode = path.stat().st_mode & 0o777
|
||||
path.chmod(original_mode | 0o020)
|
||||
monkeypatch.setattr(
|
||||
runtime.RawVectorStore,
|
||||
"open_writer_existing",
|
||||
lambda *a, **k: pytest.fail("unsafe resource opened"),
|
||||
)
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
try:
|
||||
assert not asyncio.run(handle.start())
|
||||
assert handle.snapshot().error_code == "retrieval_runtime_unsafe_path"
|
||||
finally:
|
||||
path.chmod(original_mode)
|
||||
|
||||
|
||||
def test_foreign_uid_and_windows_are_failclosed(tmp_path, monkeypatch):
|
||||
configured(tmp_path)
|
||||
uid = os.getuid()
|
||||
monkeypatch.setattr(runtime.os, "getuid", lambda: uid + 1)
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert not asyncio.run(handle.start())
|
||||
assert handle.snapshot().error_code == "retrieval_runtime_unsafe_path"
|
||||
monkeypatch.setattr(runtime.os, "getuid", lambda: uid)
|
||||
monkeypatch.setattr(runtime.sys, "platform", "win32")
|
||||
windows = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert not asyncio.run(windows.start())
|
||||
assert windows.snapshot().error_code == "retrieval_runtime_platform_unsupported"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target", ["root", "database"])
|
||||
def test_identical_store_id_copy_replacement_during_open_is_rejected(tmp_path, monkeypatch, target):
|
||||
root = tmp_path / "profile"
|
||||
root.mkdir()
|
||||
configured(root)
|
||||
closed = []
|
||||
|
||||
class Vectors:
|
||||
def close(self):
|
||||
closed.append(True)
|
||||
|
||||
def replace_path(*args, **kwargs):
|
||||
if target == "database":
|
||||
path = root / "one-memory.sqlite3"
|
||||
copy = root / "replacement.sqlite3"
|
||||
shutil.copy2(path, copy)
|
||||
os.replace(copy, path)
|
||||
else:
|
||||
old = tmp_path / "old-profile"
|
||||
root.rename(old)
|
||||
shutil.copytree(old, root)
|
||||
return Vectors()
|
||||
|
||||
monkeypatch.setattr(runtime.RawVectorStore, "open_writer_existing", replace_path)
|
||||
handle = ExistingRetrievalRuntime(root, minimum_free_bytes=0)
|
||||
assert not asyncio.run(handle.start())
|
||||
assert handle.snapshot().error_code == "retrieval_runtime_index_changed"
|
||||
assert closed == [True]
|
||||
|
||||
|
||||
def test_failed_drain_keeps_live_worker_vectors_open_and_raises_fixed_code(tmp_path, monkeypatch):
|
||||
entered, release, _, _, events = fake_resources(tmp_path, monkeypatch)
|
||||
actual = runtime.RetrievalIndexSupervisor
|
||||
instances = []
|
||||
|
||||
class BadStop(actual):
|
||||
def __init__(self, worker, **kwargs):
|
||||
super().__init__(worker, **kwargs)
|
||||
instances.append(self)
|
||||
|
||||
async def stop(self):
|
||||
raise RuntimeError("SECRET_STOP_ERROR")
|
||||
|
||||
monkeypatch.setattr(runtime, "RetrievalIndexSupervisor", BadStop)
|
||||
|
||||
async def run():
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert await handle.start()
|
||||
await until(entered.is_set)
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="^retrieval_runtime_stop_failed$"):
|
||||
await handle.stop()
|
||||
assert "close" not in events
|
||||
assert handle.snapshot().supervisor.thread_alive
|
||||
assert handle.snapshot().state == "failed"
|
||||
assert "SECRET" not in repr(handle.snapshot())
|
||||
finally:
|
||||
release.set()
|
||||
await actual.stop(instances[0])
|
||||
# Test alone owns cleanup of the intentionally broken supervisor.
|
||||
handle._vectors.close()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_stop_error_after_confirmed_exit_still_closes_vectors(tmp_path, monkeypatch):
|
||||
entered, release, _, _, events = fake_resources(tmp_path, monkeypatch)
|
||||
actual = runtime.RetrievalIndexSupervisor
|
||||
|
||||
class BadStop(actual):
|
||||
async def stop(self):
|
||||
await super().stop()
|
||||
raise RuntimeError("SECRET_AFTER_EXIT")
|
||||
|
||||
monkeypatch.setattr(runtime, "RetrievalIndexSupervisor", BadStop)
|
||||
|
||||
async def run():
|
||||
handle = ExistingRetrievalRuntime(tmp_path, minimum_free_bytes=0)
|
||||
assert await handle.start()
|
||||
await until(entered.is_set)
|
||||
release.set()
|
||||
with pytest.raises(RuntimeError, match="^retrieval_runtime_stop_failed$"):
|
||||
await handle.stop()
|
||||
assert events == ["cancel", "exit", "close"]
|
||||
assert not handle.snapshot().supervisor.thread_alive
|
||||
|
||||
asyncio.run(run())
|
||||
@@ -1,528 +0,0 @@
|
||||
import math
|
||||
import socket
|
||||
from time import monotonic
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from one_memory.retrieval.vector_store import (
|
||||
RawVectorRecord,
|
||||
RawVectorStore,
|
||||
VectorStoreError,
|
||||
_unit_vector,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_network(monkeypatch):
|
||||
def refused(*args, **kwargs):
|
||||
pytest.fail("vector adapter attempted network")
|
||||
|
||||
monkeypatch.setattr(socket.socket, "connect", refused)
|
||||
monkeypatch.setattr(socket, "create_connection", refused)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(tmp_path):
|
||||
value = RawVectorStore.create(
|
||||
tmp_path / "retrieval-vectors", generation_id="generation-a", dimensions=3
|
||||
)
|
||||
try:
|
||||
yield value
|
||||
finally:
|
||||
value.close()
|
||||
|
||||
|
||||
def record(chunk="chunk-a", *, build="build-a", generation="generation-a", vector=(1.0, 0.0, 0.0)):
|
||||
return RawVectorRecord(chunk, build, generation, vector)
|
||||
|
||||
|
||||
def test_constructor_does_no_io_and_unopened_calls_do_not_create(tmp_path):
|
||||
directory = tmp_path / "never-created"
|
||||
value = RawVectorStore(directory, generation_id="generation-a", dimensions=3)
|
||||
assert not directory.exists()
|
||||
with pytest.raises(VectorStoreError, match="not_open"):
|
||||
value.search((1.0, 0.0, 0.0), ["chunk-a"], 1, monotonic() + 1)
|
||||
assert not directory.exists()
|
||||
value.close()
|
||||
|
||||
|
||||
def test_open_existing_missing_store_or_collection_never_creates(tmp_path, store):
|
||||
missing = tmp_path / "missing"
|
||||
with pytest.raises(VectorStoreError, match="not_found"):
|
||||
RawVectorStore.open_existing(missing, generation_id="generation-a", dimensions=3)
|
||||
assert not missing.exists()
|
||||
before = store._client.list_collections()
|
||||
with pytest.raises(VectorStoreError, match="not_found"):
|
||||
RawVectorStore.open_existing(store.directory, generation_id="other", dimensions=3)
|
||||
assert store._client.list_collections() == before
|
||||
|
||||
|
||||
def test_real_chroma_collection_has_no_embedding_function_or_documents(store):
|
||||
store.upsert([record()])
|
||||
collection = store._collection
|
||||
assert collection._embedding_function is None
|
||||
result = collection.get(ids=["chunk-a"], include=["embeddings", "metadatas", "documents"])
|
||||
assert result["documents"] == [None]
|
||||
assert result["metadatas"][0]["chunk_id"] == "chunk-a"
|
||||
assert result["metadatas"][0]["build_id"] == "build-a"
|
||||
assert result["metadatas"][0]["generation_id"] == "generation-a"
|
||||
assert len(result["embeddings"][0]) == 3
|
||||
store.verify([record()])
|
||||
|
||||
|
||||
def test_existing_read_handle_cannot_write_or_repair_and_releases_shared_client(store):
|
||||
store.upsert([record()])
|
||||
opened = RawVectorStore.open_existing(
|
||||
store.directory, generation_id="generation-a", dimensions=3
|
||||
)
|
||||
try:
|
||||
result = opened.search((1.0, 0.0, 0.0), ["chunk-a"], 1, monotonic() + 10)
|
||||
assert result.complete and result.hits[0].chunk_id == "chunk-a"
|
||||
with pytest.raises(VectorStoreError, match="read_only"):
|
||||
opened.upsert([record("new")])
|
||||
opened.verify([record()])
|
||||
finally:
|
||||
opened.close()
|
||||
store.verify([record()])
|
||||
with pytest.raises(VectorStoreError, match="closed"):
|
||||
opened.search((1.0, 0.0, 0.0), ["chunk-a"], 1, monotonic() + 1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("changed", [record(build="different"), record(vector=(0.0, 1.0, 0.0))])
|
||||
def test_immutable_id_rejects_different_metadata_or_vector_without_overwriting(store, changed):
|
||||
store.upsert([record()])
|
||||
store.upsert([record()])
|
||||
with pytest.raises(VectorStoreError, match="immutable_conflict"):
|
||||
store.upsert([record("new"), changed])
|
||||
store.verify([record()])
|
||||
assert store._collection.get(ids=["new"])["ids"] == []
|
||||
|
||||
|
||||
def test_verify_detects_missing_and_tampered_records(store):
|
||||
with pytest.raises(VectorStoreError, match="missing"):
|
||||
store.verify([record()])
|
||||
store.upsert([record()])
|
||||
store._collection.update(ids=["chunk-a"], embeddings=[[0.0, 1.0, 0.0]])
|
||||
with pytest.raises(VectorStoreError):
|
||||
store.verify([record()])
|
||||
with pytest.raises(VectorStoreError):
|
||||
store.search((1.0, 0.0, 0.0), ["chunk-a"], 1, monotonic() + 10)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vector",
|
||||
[
|
||||
(1.0, 0.0),
|
||||
(1.0, 0.0, math.nan),
|
||||
(math.inf, 0.0, 1.0),
|
||||
(0.0, 0.0, 0.0),
|
||||
(True, 0.0, 1.0),
|
||||
(1e100, 0.0, 1.0),
|
||||
],
|
||||
)
|
||||
def test_invalid_vectors_refused_before_writes_or_search(store, vector):
|
||||
with pytest.raises(VectorStoreError):
|
||||
store.upsert([record(vector=vector)])
|
||||
with pytest.raises(VectorStoreError):
|
||||
store.search(vector, [], 1, monotonic() + 1)
|
||||
assert store._collection.count() == 0
|
||||
|
||||
|
||||
def test_wrong_generation_and_duplicate_ids_rejected(store):
|
||||
with pytest.raises(VectorStoreError):
|
||||
store.upsert([record(generation="other")])
|
||||
with pytest.raises(VectorStoreError):
|
||||
store.upsert([record(), record()])
|
||||
assert store._collection.count() == 0
|
||||
|
||||
|
||||
def test_empty_eligible_ids_never_queries_all_collection(store, monkeypatch):
|
||||
store.upsert([record()])
|
||||
monkeypatch.setattr(
|
||||
store._collection, "query", lambda **kwargs: pytest.fail("unfiltered query")
|
||||
)
|
||||
result = store.search((1.0, 0.0, 0.0), [], 3, monotonic() - 1)
|
||||
assert result.complete and result.hits == ()
|
||||
assert result.eligible_count == result.searched_count == 0
|
||||
|
||||
|
||||
def test_real_in_filter_excludes_closer_unapproved_id_and_matches_small_cosine_oracle(store):
|
||||
rows = [
|
||||
record("denied", vector=(1.0, 0.0, 0.0)),
|
||||
record("b", vector=(0.9, 0.1, 0.0)),
|
||||
record("c", vector=(0.0, 1.0, 0.0)),
|
||||
record("d", vector=(-1.0, 0.0, 0.0)),
|
||||
]
|
||||
store.upsert(rows)
|
||||
result = store.search((1.0, 0.0, 0.0), ["b", "c", "d"], 3, monotonic() + 10)
|
||||
assert result.complete and result.searched_count == 3
|
||||
assert [hit.chunk_id for hit in result.hits] == ["b", "c", "d"]
|
||||
for hit, row in zip(result.hits, rows[1:], strict=True):
|
||||
vector = np.array(row.vector, dtype=np.float32)
|
||||
expected = 1 - float(vector[0] / np.linalg.norm(vector))
|
||||
assert hit.distance == pytest.approx(expected, abs=1e-5)
|
||||
|
||||
|
||||
def test_eligible_batches_cover_more_than_2048_and_merge_topk(store, monkeypatch):
|
||||
ids = [f"chunk-{i:04d}" for i in range(2053)]
|
||||
rows = [record(chunk, vector=(float(i + 1), 1.0, 0.0)) for i, chunk in enumerate(ids)]
|
||||
store.upsert(rows)
|
||||
calls = []
|
||||
get = store._collection.get
|
||||
|
||||
def tracked(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return get(**kwargs)
|
||||
|
||||
monkeypatch.setattr(store._collection, "get", tracked)
|
||||
result = store.search((1.0, 0.0, 0.0), ids, 3, monotonic() + 30)
|
||||
assert result.complete and result.searched_count == 2053
|
||||
assert len(calls) == 2
|
||||
assert [len(call["ids"]) for call in calls] == [2048, 5]
|
||||
assert set().union(*(set(call["ids"]) for call in calls)) == set(ids)
|
||||
assert all(call["include"] == ["embeddings", "metadatas"] for call in calls)
|
||||
assert len(result.hits) == 3
|
||||
# The exact ranking over every eligible chunk, not one batch's local top-K:
|
||||
# a per-batch top-K would have kept the second batch's steepest vectors out.
|
||||
# Distances come from the stored unit vectors, so float32 rounding decides
|
||||
# the top of this ranking just as it does for the adapter.
|
||||
stored = {row.chunk_id: np.array(_unit_vector(row.vector), dtype=np.float32) for row in rows}
|
||||
expected = {
|
||||
chunk_id: float(1 - np.dot(unit, np.array([1.0, 0.0, 0.0], dtype=np.float32)))
|
||||
for chunk_id, unit in stored.items()
|
||||
}
|
||||
ranked = sorted(expected.items(), key=lambda pair: (pair[1], pair[0]))[:3]
|
||||
assert [hit.chunk_id for hit in result.hits] == [chunk for chunk, _ in ranked]
|
||||
for hit in result.hits:
|
||||
assert hit.distance == pytest.approx(expected[hit.chunk_id], abs=1e-6)
|
||||
|
||||
|
||||
def test_deadline_before_query_and_late_batch_return_are_partial(store, monkeypatch):
|
||||
store.upsert([record()])
|
||||
calls = []
|
||||
now = [0.0]
|
||||
store._clock = lambda: now[0]
|
||||
get = store._collection.get
|
||||
|
||||
def late(**kwargs):
|
||||
calls.append(kwargs)
|
||||
result = get(**kwargs)
|
||||
now[0] = 2.0
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(store._collection, "get", late)
|
||||
first = store.search((1.0, 0.0, 0.0), ["chunk-a"], 1, 0.0)
|
||||
assert not first.complete and first.reason == "deadline" and first.hits == ()
|
||||
assert calls == []
|
||||
second = store.search((1.0, 0.0, 0.0), ["chunk-a"], 1, 1.0)
|
||||
assert not second.complete and second.searched_count == 0 and second.hits == ()
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
def test_search_never_scores_a_vector_the_caller_did_not_declare_eligible(store, monkeypatch):
|
||||
"""Chroma's own filtered query is no longer used at all.
|
||||
|
||||
It does not score every ID in its `where $in` allowlist, so a query could
|
||||
come back short while claiming completion, and a global query filtered
|
||||
afterwards would let unauthorized vectors take the eligible ones' places.
|
||||
"""
|
||||
store.upsert([record(), record("denied", vector=(1.0, 0.0, 0.0))])
|
||||
monkeypatch.setattr(
|
||||
store._collection,
|
||||
"query",
|
||||
lambda **kwargs: pytest.fail("global or filtered ANN query is not used"),
|
||||
)
|
||||
result = store.search((1.0, 0.0, 0.0), ["chunk-a"], 1, monotonic() + 10)
|
||||
assert result.complete and [hit.chunk_id for hit in result.hits] == ["chunk-a"]
|
||||
|
||||
|
||||
def test_an_eligible_vector_missing_from_the_backend_is_unavailable_not_a_miss(store):
|
||||
"""The caller published this chunk, so an absent vector is a broken build."""
|
||||
store.upsert([record("present")])
|
||||
with pytest.raises(VectorStoreError, match="^vector_query_unavailable$"):
|
||||
store.search((1.0, 0.0, 0.0), ["present", "absent"], 2, monotonic() + 10)
|
||||
|
||||
|
||||
def test_open_existing_validates_generation_dimensions_and_collection_config(store):
|
||||
with pytest.raises(VectorStoreError):
|
||||
RawVectorStore.open_existing(store.directory, generation_id="generation-a", dimensions=4)
|
||||
store._collection.modify(metadata={"generation_id": "other"})
|
||||
with pytest.raises(VectorStoreError):
|
||||
RawVectorStore.open_existing(store.directory, generation_id="generation-a", dimensions=3)
|
||||
|
||||
|
||||
def test_open_existing_refuses_symlink_database_before_chroma(tmp_path):
|
||||
directory = tmp_path / "retrieval-vectors"
|
||||
directory.mkdir()
|
||||
target = tmp_path / "not-a-chroma-db"
|
||||
target.write_bytes(b"do not open")
|
||||
(directory / "chroma.sqlite3").symlink_to(target)
|
||||
with pytest.raises(VectorStoreError):
|
||||
RawVectorStore.open_existing(directory, generation_id="generation-a", dimensions=3)
|
||||
assert target.read_bytes() == b"do not open"
|
||||
|
||||
|
||||
def test_writer_reopen_resumes_idempotently_but_query_handle_still_cannot_write(tmp_path):
|
||||
directory = tmp_path / "retrieval-vectors"
|
||||
original = RawVectorStore.create(directory, generation_id="generation-a", dimensions=3)
|
||||
try:
|
||||
original.upsert([record(vector=(1.0, 1.0, 0.0))])
|
||||
finally:
|
||||
original.close()
|
||||
writer = RawVectorStore.open_writer_existing(
|
||||
directory, generation_id="generation-a", dimensions=3
|
||||
)
|
||||
try:
|
||||
writer.upsert([record(vector=(1.0, 1.0, 0.0))])
|
||||
writer.verify([record(vector=(1.0, 1.0, 0.0))])
|
||||
with pytest.raises(VectorStoreError, match="immutable_conflict"):
|
||||
writer.upsert([record(build="other", vector=(1.0, 1.0, 0.0))])
|
||||
with pytest.raises(VectorStoreError, match="immutable_conflict"):
|
||||
writer.upsert([record(vector=(2.0, 2.0, 0.0))])
|
||||
assert writer._collection.count() == 1
|
||||
finally:
|
||||
writer.close()
|
||||
reader = RawVectorStore.open_existing(directory, generation_id="generation-a", dimensions=3)
|
||||
try:
|
||||
with pytest.raises(VectorStoreError, match="read_only"):
|
||||
reader.upsert([record("new")])
|
||||
result = reader.search((1.0, 1.0, 0.0), ["chunk-a"], 1, monotonic() + 10)
|
||||
assert result.complete and result.hits[0].distance == pytest.approx(0, abs=1e-5)
|
||||
finally:
|
||||
reader.close()
|
||||
|
||||
|
||||
def test_writer_existing_does_not_create_missing_store_or_collection(tmp_path, store):
|
||||
missing = tmp_path / "writer-missing"
|
||||
with pytest.raises(VectorStoreError, match="not_found"):
|
||||
RawVectorStore.open_writer_existing(missing, generation_id="generation-a", dimensions=3)
|
||||
assert not missing.exists()
|
||||
with pytest.raises(VectorStoreError, match="not_found"):
|
||||
RawVectorStore.open_writer_existing(store.directory, generation_id="other", dimensions=3)
|
||||
|
||||
|
||||
def test_partial_keeps_only_completed_eligible_batches(store, monkeypatch):
|
||||
ids = [f"present-{i}" for i in range(2048)] + ["last"]
|
||||
store.upsert([*(record(chunk, vector=(1.0, float(i + 1), 0.0)) for i, chunk in enumerate(ids))])
|
||||
now = [0.0]
|
||||
store._clock = lambda: now[0]
|
||||
get = store._collection.get
|
||||
calls = []
|
||||
|
||||
def timed(**kwargs):
|
||||
calls.append(kwargs)
|
||||
result = get(**kwargs)
|
||||
if len(calls) == 2:
|
||||
now[0] = 2.0
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(store._collection, "get", timed)
|
||||
result = store.search((1.0, 0.0, 0.0), ids, 2, 1.0)
|
||||
assert not result.complete and result.reason == "deadline"
|
||||
assert result.searched_count == 2048 and result.eligible_count == 2049
|
||||
assert [hit.chunk_id for hit in result.hits] == ["present-0", "present-1"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", ["metadata", "dimension", "hash", "extra_id"])
|
||||
def test_malformed_search_response_is_rejected_not_partial_success(store, monkeypatch, mutation):
|
||||
store.upsert([record()])
|
||||
get = store._collection.get
|
||||
|
||||
def malformed(**kwargs):
|
||||
result = get(**kwargs)
|
||||
if mutation == "metadata":
|
||||
result["metadatas"][0][0]["generation_id"] = "wrong"
|
||||
elif mutation == "dimension":
|
||||
result["embeddings"] = [[1.0, 0.0]]
|
||||
elif mutation == "hash":
|
||||
result["metadatas"][0]["vector_sha256"] = "0" * 64
|
||||
else:
|
||||
result["ids"] = result["ids"] * 2
|
||||
result["metadatas"] = result["metadatas"] * 2
|
||||
result["embeddings"] = [result["embeddings"][0], result["embeddings"][0]]
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(store._collection, "get", malformed)
|
||||
with pytest.raises(VectorStoreError):
|
||||
store.search((1.0, 0.0, 0.0), ["chunk-a"], 1, monotonic() + 10)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ids,limit,deadline",
|
||||
[
|
||||
(["chunk-a", "chunk-a"], 1, 100.0),
|
||||
("chunk-a", 1, 100.0),
|
||||
(["chunk-a"], 0, 100.0),
|
||||
(["chunk-a"], True, 100.0),
|
||||
(["chunk-a"], 1, math.inf),
|
||||
(["chunk-a"], 1, math.nan),
|
||||
],
|
||||
)
|
||||
def test_invalid_query_parameters_do_not_reach_chroma(store, monkeypatch, ids, limit, deadline):
|
||||
monkeypatch.setattr(store._collection, "query", lambda **kwargs: pytest.fail("invalid query"))
|
||||
with pytest.raises(VectorStoreError):
|
||||
store.search((1.0, 0.0, 0.0), ids, limit, deadline)
|
||||
|
||||
|
||||
def test_query_backend_error_is_unavailable_not_complete_empty(store, monkeypatch):
|
||||
def broken(**kwargs):
|
||||
raise RuntimeError("backend request payload")
|
||||
|
||||
monkeypatch.setattr(store._collection, "query", broken)
|
||||
with pytest.raises(VectorStoreError, match="^vector_query_unavailable$"):
|
||||
store.search((1.0, 0.0, 0.0), ["chunk-a"], 1, monotonic() + 10)
|
||||
|
||||
|
||||
def test_last_owned_client_close_releases_chroma_system(tmp_path):
|
||||
from chromadb.api.shared_system_client import SharedSystemClient
|
||||
|
||||
directory = tmp_path / "retrieval-vectors"
|
||||
writer = RawVectorStore.create(directory, generation_id="close-test", dimensions=3)
|
||||
reader = RawVectorStore.open_existing(directory, generation_id="close-test", dimensions=3)
|
||||
identifier = writer._client._identifier
|
||||
try:
|
||||
assert identifier in SharedSystemClient._identifier_to_system
|
||||
writer.close()
|
||||
assert identifier in SharedSystemClient._identifier_to_system
|
||||
reader.close()
|
||||
assert identifier not in SharedSystemClient._identifier_to_system
|
||||
reader.close()
|
||||
finally:
|
||||
reader.close()
|
||||
writer.close()
|
||||
|
||||
|
||||
def test_open_existing_missing_collection_does_not_construct_a_client(store, monkeypatch):
|
||||
import chromadb
|
||||
|
||||
monkeypatch.setattr(
|
||||
chromadb, "PersistentClient", lambda *args, **kwargs: pytest.fail("must preflight first")
|
||||
)
|
||||
with pytest.raises(VectorStoreError, match="collection_not_found"):
|
||||
RawVectorStore.open_existing(store.directory, generation_id="absent", dimensions=3)
|
||||
|
||||
|
||||
def test_existing_dimension_mismatch_does_not_modify_collection(store):
|
||||
before = store._collection.metadata.copy()
|
||||
with pytest.raises(VectorStoreError):
|
||||
RawVectorStore.open_writer_existing(
|
||||
store.directory, generation_id="generation-a", dimensions=4
|
||||
)
|
||||
assert store._collection.metadata == before
|
||||
store.upsert([record()])
|
||||
store.verify([record()])
|
||||
|
||||
|
||||
def test_verify_candidates_reader_checks_every_batch_without_writing(store, monkeypatch):
|
||||
rows = [record(f"candidate-{i}") for i in range(2050)]
|
||||
store.upsert(rows)
|
||||
reader = RawVectorStore.open_existing(
|
||||
store.directory, generation_id="generation-a", dimensions=3
|
||||
)
|
||||
try:
|
||||
calls = []
|
||||
get = reader._collection.get
|
||||
|
||||
def tracked(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return get(**kwargs)
|
||||
|
||||
monkeypatch.setattr(reader._collection, "get", tracked)
|
||||
for name in ("add", "upsert", "update", "delete", "query"):
|
||||
monkeypatch.setattr(
|
||||
reader._collection, name, lambda **kwargs: pytest.fail("unexpected backend call")
|
||||
)
|
||||
assert (
|
||||
reader.verify_candidates(
|
||||
[(row.chunk_id, row.build_id) for row in rows], deadline=monotonic() + 30
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert [len(call["ids"]) for call in calls] == [2048, 2]
|
||||
assert all(call["include"] == ["embeddings", "metadatas"] for call in calls)
|
||||
finally:
|
||||
reader.close()
|
||||
|
||||
|
||||
def test_verify_candidates_empty_and_expired_do_not_read(store, monkeypatch):
|
||||
monkeypatch.setattr(store._collection, "get", lambda **kwargs: pytest.fail("unexpected read"))
|
||||
store._clock = lambda: 10.0
|
||||
assert store.verify_candidates([], deadline=0.0) is True
|
||||
assert store.verify_candidates([("chunk-a", "build-a")], deadline=10.0) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expires", ["after_read", "between_batches"])
|
||||
def test_verify_candidates_deadline_stops_before_another_batch(store, monkeypatch, expires):
|
||||
rows = [record(f"candidate-{i}") for i in range(2049)]
|
||||
store.upsert(rows)
|
||||
calls = []
|
||||
get = store._collection.get
|
||||
|
||||
def tracked(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return get(**kwargs)
|
||||
|
||||
ticks = iter([0.0, 2.0] if expires == "after_read" else [0.0, 0.0, 2.0])
|
||||
store._clock = lambda: next(ticks)
|
||||
monkeypatch.setattr(store._collection, "get", tracked)
|
||||
assert (
|
||||
store.verify_candidates([(row.chunk_id, row.build_id) for row in rows], deadline=1.0)
|
||||
is False
|
||||
)
|
||||
assert len(calls) == 1 and len(calls[0]["ids"]) == 2048
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pair", [("missing", "build-a"), ("chunk-a", "other-build")])
|
||||
def test_verify_candidates_missing_or_wrong_build_refuses(store, pair):
|
||||
store.upsert([record()])
|
||||
with pytest.raises(VectorStoreError, match="missing|identity_mismatch"):
|
||||
store.verify_candidates([pair], deadline=monotonic() + 10)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", ["generation", "dimension", "hash", "extra_id"])
|
||||
def test_verify_candidates_corrupt_backend_record_refuses(store, monkeypatch, mutation):
|
||||
store.upsert([record()])
|
||||
get = store._collection.get
|
||||
|
||||
def corrupt(**kwargs):
|
||||
data = get(**kwargs)
|
||||
if mutation == "generation":
|
||||
data["metadatas"][0]["generation_id"] = "other-generation"
|
||||
elif mutation == "dimension":
|
||||
data["embeddings"] = [[1.0, 0.0]]
|
||||
elif mutation == "hash":
|
||||
data["metadatas"][0]["vector_sha256"] = "0" * 64
|
||||
else:
|
||||
data["ids"][0] = "not-eligible"
|
||||
return data
|
||||
|
||||
monkeypatch.setattr(store._collection, "get", corrupt)
|
||||
with pytest.raises(VectorStoreError):
|
||||
store.verify_candidates([("chunk-a", "build-a")], deadline=monotonic() + 10)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"candidates",
|
||||
[
|
||||
[("chunk-a", "build-a"), ("chunk-a", "build-a")],
|
||||
[("chunk-a", "build-a"), ("chunk-a", "other")],
|
||||
[("../escape", "build-a")],
|
||||
[("chunk-a", "")],
|
||||
[("chunk-a", 123)],
|
||||
[("chunk-a",)],
|
||||
["ab"],
|
||||
"ab",
|
||||
None,
|
||||
],
|
||||
)
|
||||
def test_verify_candidates_invalid_input_refuses_before_read(store, monkeypatch, candidates):
|
||||
monkeypatch.setattr(store._collection, "get", lambda **kwargs: pytest.fail("unexpected read"))
|
||||
with pytest.raises(VectorStoreError):
|
||||
store.verify_candidates(candidates, deadline=monotonic() + 10)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("deadline", [math.nan, math.inf, True, "soon"])
|
||||
def test_verify_candidates_invalid_deadline_refuses_even_empty(store, deadline):
|
||||
with pytest.raises(VectorStoreError, match="invalid_deadline"):
|
||||
store.verify_candidates([], deadline=deadline)
|
||||
@@ -4,15 +4,22 @@ import asyncio
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from pydantic import SecretStr
|
||||
|
||||
from one_memory import cli
|
||||
from one_memory.api import create_app
|
||||
from one_memory.config import Settings
|
||||
from one_memory.database_access import close_main_databases, main_database
|
||||
from one_memory.hy_sdk_migration import repair_sdk_metadata_from_bundle
|
||||
from one_memory.hy_sdk_types import HySdkRawPage
|
||||
from one_memory.imports.storage import EvidenceStore
|
||||
from one_memory.memory_server.contracts import BatchIdentity, WriteRequest
|
||||
from one_memory.memory_server.durable_queue import DurableWriteQueue
|
||||
from one_memory.memory_server.migration_runner import SqliteMigrationStore, migrate
|
||||
@@ -35,6 +42,7 @@ from one_memory.sdk_migration_runner import (
|
||||
resume_legacy_cleanup,
|
||||
run_startup_migration,
|
||||
)
|
||||
from one_memory.secrets import InMemorySecretStore
|
||||
from one_memory.target_schema import legacy_tables_present
|
||||
|
||||
|
||||
@@ -62,6 +70,260 @@ def _start(store, migration_id="m1"):
|
||||
)
|
||||
|
||||
|
||||
def test_post_cutover_metadata_repair_restores_source_and_original_times(store):
|
||||
data_dir = store.database_path.parent
|
||||
migration_id = "sdk-authority-v2-source-repair"
|
||||
with main_database(store.database_path).write() as database, database:
|
||||
database.execute(
|
||||
"UPDATE memory_authority_state SET authority='sdk',migration_id=? WHERE singleton=1",
|
||||
(migration_id,),
|
||||
)
|
||||
legacy = (
|
||||
data_dir
|
||||
/ "backups"
|
||||
/ f"{migration_id}-pre-cutover.sdk-bundle"
|
||||
/ "management"
|
||||
/ "one-memory.sqlite3"
|
||||
)
|
||||
legacy.parent.mkdir(parents=True)
|
||||
with sqlite3.connect(legacy) as database:
|
||||
database.executescript(
|
||||
"""
|
||||
CREATE TABLE memory_task_outputs (body TEXT);
|
||||
CREATE TABLE governed_memories (
|
||||
memory_id TEXT, created_at TEXT, updated_at TEXT
|
||||
);
|
||||
CREATE TABLE evidence (
|
||||
evidence_id TEXT, source_id TEXT, created_at TEXT, imported_at TEXT
|
||||
);
|
||||
"""
|
||||
)
|
||||
database.execute(
|
||||
"INSERT INTO governed_memories VALUES (?,?,?)",
|
||||
("legacy-1", "2026-10-18T00:00:00+00:00", "2026-10-18T01:00:00+00:00"),
|
||||
)
|
||||
database.execute(
|
||||
"INSERT INTO evidence VALUES (?,?,?,?)",
|
||||
("e1", "workbuddy", "2254-01-01T00:00:00+00:00", "2026-10-17T00:00:00+00:00"),
|
||||
)
|
||||
database.execute(
|
||||
"INSERT INTO memory_task_outputs VALUES (?)",
|
||||
(
|
||||
json.dumps(
|
||||
{
|
||||
"memory_id": "legacy-1",
|
||||
"hy_memory_id": "sdk-1",
|
||||
"support_refs": [{"evidence_id": "e1"}],
|
||||
}
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
import chromadb
|
||||
from chromadb.config import Settings as ChromaSettings
|
||||
|
||||
vector = data_dir / "hy-memory" / "governed-v2" / "vector"
|
||||
collection = chromadb.PersistentClient(
|
||||
path=str(vector),
|
||||
settings=ChromaSettings(anonymized_telemetry=False),
|
||||
).get_or_create_collection("memory")
|
||||
collection.add(
|
||||
ids=["sdk-1"],
|
||||
documents=["remember"],
|
||||
embeddings=[[0.1, 0.2]],
|
||||
metadatas=[
|
||||
{
|
||||
"source": "default",
|
||||
"gmt_created": 1,
|
||||
"memory_at": 8_962_051_200,
|
||||
}
|
||||
],
|
||||
)
|
||||
coding = data_dir / "hy-memory" / "governed-v2" / "coding_cv.sqlite3"
|
||||
with sqlite3.connect(coding) as database:
|
||||
database.execute(
|
||||
"CREATE TABLE coding_memory_meta ("
|
||||
"experience_id TEXT,user_id TEXT,agent_id TEXT,solution TEXT,"
|
||||
"source TEXT,created_at TEXT,updated_at TEXT)"
|
||||
)
|
||||
database.execute(
|
||||
"CREATE TABLE working_fact ("
|
||||
"fact_id TEXT,user_id TEXT,agent_id TEXT,statement TEXT,"
|
||||
"source TEXT,created_at TEXT)"
|
||||
)
|
||||
database.execute(
|
||||
"INSERT INTO coding_memory_meta VALUES "
|
||||
"('sdk-1','local','scope','remember','default','now','now')"
|
||||
)
|
||||
|
||||
result = repair_sdk_metadata_from_bundle(
|
||||
data_dir=data_dir,
|
||||
database_path=store.database_path,
|
||||
now=datetime(2026, 10, 25, tzinfo=UTC),
|
||||
)
|
||||
metadata = collection.get(ids=["sdk-1"], include=["metadatas"])["metadatas"][0]
|
||||
assert result.vector_rows == 1
|
||||
assert result.coding_rows == 1
|
||||
assert metadata["source"] == "workbuddy"
|
||||
assert datetime.fromtimestamp(metadata["memory_at"], tz=UTC) == datetime(
|
||||
2026, 10, 17, tzinfo=UTC
|
||||
)
|
||||
assert datetime.fromtimestamp(metadata["gmt_created"], tz=UTC) == datetime(
|
||||
2026, 10, 18, tzinfo=UTC
|
||||
)
|
||||
with sqlite3.connect(coding) as database:
|
||||
source, created_at = database.execute(
|
||||
"SELECT source,created_at FROM coding_memory_meta WHERE experience_id='sdk-1'"
|
||||
).fetchone()
|
||||
assert source == "workbuddy"
|
||||
assert created_at.startswith("2026-10-18T00:00:00")
|
||||
|
||||
repeated = repair_sdk_metadata_from_bundle(
|
||||
data_dir=data_dir,
|
||||
database_path=store.database_path,
|
||||
now=datetime(2026, 10, 25, tzinfo=UTC),
|
||||
)
|
||||
assert repeated == result
|
||||
|
||||
|
||||
def test_metadata_repair_reaches_memories_the_cutover_backfilled_under_legacy_ids(store):
|
||||
"""The backfill wrote governed memories with no source; their evidence links them."""
|
||||
data_dir = store.database_path.parent
|
||||
migration_id = "sdk-authority-v2-backfill"
|
||||
with main_database(store.database_path).write() as database, database:
|
||||
database.execute(
|
||||
"UPDATE memory_authority_state SET authority='sdk',migration_id=? WHERE singleton=1",
|
||||
(migration_id,),
|
||||
)
|
||||
# A v1 receipt from the earlier, narrower repair must not block this one.
|
||||
(data_dir / f".sdk-metadata-repair-v1-{migration_id}.json").write_text("{}\n")
|
||||
# The later bundle was taken after legacy tables were retired; the earlier
|
||||
# cutover's bundle still carries them.
|
||||
for bundle, populated in (("sdk-authority-v1-old", True), (migration_id, False)):
|
||||
legacy = data_dir / "backups" / f"{bundle}-pre-cutover.sdk-bundle" / "management"
|
||||
legacy.mkdir(parents=True)
|
||||
with closing(sqlite3.connect(legacy / "one-memory.sqlite3")) as database, database:
|
||||
database.executescript(
|
||||
"""
|
||||
CREATE TABLE evidence (
|
||||
evidence_id TEXT, source_id TEXT, created_at TEXT, imported_at TEXT
|
||||
);
|
||||
"""
|
||||
)
|
||||
if not populated:
|
||||
continue
|
||||
database.executescript(
|
||||
"""
|
||||
CREATE TABLE governed_memories (
|
||||
memory_id TEXT, created_at TEXT, updated_at TEXT
|
||||
);
|
||||
CREATE TABLE governed_memory_evidence (
|
||||
memory_id TEXT, evidence_id TEXT, active INTEGER
|
||||
);
|
||||
"""
|
||||
)
|
||||
database.execute(
|
||||
"INSERT INTO governed_memories VALUES ('legacy-7','2026-09-02T00:00:00+00:00',NULL)"
|
||||
)
|
||||
database.execute("INSERT INTO governed_memory_evidence VALUES ('legacy-7','e7',1)")
|
||||
database.execute(
|
||||
"INSERT INTO evidence VALUES ('e7','workbuddy','2026-09-01T00:00:00+00:00',NULL)"
|
||||
)
|
||||
|
||||
import chromadb
|
||||
from chromadb.config import Settings as ChromaSettings
|
||||
|
||||
vector = data_dir / "hy-memory" / "governed-v2" / "vector"
|
||||
client = chromadb.PersistentClient(
|
||||
path=str(vector), settings=ChromaSettings(anonymized_telemetry=False)
|
||||
)
|
||||
collection = client.get_or_create_collection("memory")
|
||||
imported = int(datetime(2026, 10, 20, tzinfo=UTC).timestamp())
|
||||
collection.add(
|
||||
ids=["legacy-7", "unlinked"],
|
||||
documents=["backfilled", "hash-dated"],
|
||||
embeddings=[[0.1, 0.2], [0.2, 0.1]],
|
||||
metadatas=[
|
||||
{"source": "default", "gmt_created": imported, "memory_at": imported},
|
||||
{"source": "cursor", "gmt_created": imported, "memory_at": 8_962_051_200},
|
||||
],
|
||||
)
|
||||
client.close()
|
||||
|
||||
result = repair_sdk_metadata_from_bundle(
|
||||
data_dir=data_dir,
|
||||
database_path=store.database_path,
|
||||
now=datetime(2026, 10, 25, tzinfo=UTC),
|
||||
)
|
||||
|
||||
client = chromadb.PersistentClient(
|
||||
path=str(vector), settings=ChromaSettings(anonymized_telemetry=False)
|
||||
)
|
||||
rows = client.get_collection("memory").get(ids=["legacy-7", "unlinked"], include=["metadatas"])
|
||||
client.close()
|
||||
metadata = dict(zip(rows["ids"], rows["metadatas"], strict=True))
|
||||
assert result.vector_rows == 2
|
||||
assert metadata["legacy-7"]["source"] == "workbuddy"
|
||||
assert metadata["legacy-7"]["memory_at"] == int(datetime(2026, 9, 1, tzinfo=UTC).timestamp())
|
||||
assert metadata["legacy-7"]["gmt_created"] == int(datetime(2026, 9, 2, tzinfo=UTC).timestamp())
|
||||
assert metadata["unlinked"]["source"] == "cursor"
|
||||
assert metadata["unlinked"]["memory_at"] == imported
|
||||
|
||||
|
||||
def test_metadata_repair_keeps_source_unknown_when_outputs_disagree(tmp_path):
|
||||
"""Two legacy memories folded into one SDK memory pool their evidence."""
|
||||
from one_memory.hy_sdk_migration import _legacy_metadata
|
||||
|
||||
legacy = tmp_path / "one-memory.sqlite3"
|
||||
with closing(sqlite3.connect(legacy)) as database, database:
|
||||
database.executescript(
|
||||
"""
|
||||
CREATE TABLE memory_task_outputs (body TEXT);
|
||||
CREATE TABLE governed_memories (
|
||||
memory_id TEXT, created_at TEXT, updated_at TEXT
|
||||
);
|
||||
CREATE TABLE evidence (
|
||||
evidence_id TEXT, source_id TEXT, created_at TEXT, imported_at TEXT
|
||||
);
|
||||
"""
|
||||
)
|
||||
database.executemany(
|
||||
"INSERT INTO governed_memories VALUES (?,?,?)",
|
||||
[
|
||||
("legacy-1", "2026-10-18T00:00:00+00:00", None),
|
||||
("legacy-2", "2026-10-12T00:00:00+00:00", None),
|
||||
],
|
||||
)
|
||||
database.executemany(
|
||||
"INSERT INTO evidence VALUES (?,?,?,NULL)",
|
||||
[
|
||||
("e1", "workbuddy", "2026-10-17T00:00:00+00:00"),
|
||||
("e2", "cursor", "2026-10-11T00:00:00+00:00"),
|
||||
],
|
||||
)
|
||||
database.executemany(
|
||||
"INSERT INTO memory_task_outputs VALUES (?)",
|
||||
[
|
||||
(
|
||||
json.dumps(
|
||||
{
|
||||
"memory_id": legacy_id,
|
||||
"hy_memory_id": "sdk-1",
|
||||
"support_refs": [{"evidence_id": evidence_id}],
|
||||
}
|
||||
),
|
||||
)
|
||||
for legacy_id, evidence_id in (("legacy-1", "e1"), ("legacy-2", "e2"))
|
||||
],
|
||||
)
|
||||
|
||||
repaired = _legacy_metadata(legacy, now=datetime(2026, 10, 25, tzinfo=UTC))["sdk-1"]
|
||||
|
||||
assert repaired.source is None
|
||||
assert repaired.memory_at == int(datetime(2026, 10, 17, tzinfo=UTC).timestamp())
|
||||
assert repaired.created_at == int(datetime(2026, 10, 12, tzinfo=UTC).timestamp())
|
||||
|
||||
|
||||
def _insert_evidence(store, evidence_id, *, conversation="session-1", role="user", seq=1):
|
||||
verbatim = f"text of {evidence_id}"
|
||||
with main_database(store.database_path).write() as database, database:
|
||||
@@ -283,6 +545,235 @@ def test_finished_jobs_without_input_do_not_stop_the_migration(store):
|
||||
prepare_legacy_write_queue_for_cutover(store.database_path)
|
||||
|
||||
|
||||
def _chroma(home: Path):
|
||||
import chromadb
|
||||
from chromadb.config import Settings as ChromaSettings
|
||||
|
||||
return chromadb.PersistentClient(
|
||||
path=str(home / "hy-memory" / "governed-v2" / "vector"),
|
||||
settings=ChromaSettings(anonymized_telemetry=False),
|
||||
)
|
||||
|
||||
|
||||
def _fill(client, name: str, ids: list[str], *, text: str = "old") -> None:
|
||||
collection = client.get_or_create_collection(name, metadata={"hnsw:space": "cosine"})
|
||||
collection.upsert(
|
||||
ids=ids,
|
||||
embeddings=[[float(index + 1)] * 4 for index, _ in enumerate(ids)],
|
||||
documents=[f"{text} {identifier}" for identifier in ids],
|
||||
metadatas=[{"layer": "l2_fact"} for _ in ids],
|
||||
)
|
||||
|
||||
|
||||
def test_legacy_vectors_are_renamed_and_newer_writes_are_kept(tmp_path) -> None:
|
||||
"""Build 20: 3317 memories in `one_memory_*`, the Runtime reading an empty `memhub_*`."""
|
||||
from one_memory.hy_sdk_runtime import builtin_retrieval_snapshot, vector_collection_bases
|
||||
from one_memory.sdk_vector_reattach import (
|
||||
legacy_vector_collections,
|
||||
reattach_legacy_vectors,
|
||||
)
|
||||
|
||||
profile = builtin_retrieval_snapshot()
|
||||
current, legacy = vector_collection_bases(profile)
|
||||
client = _chroma(tmp_path)
|
||||
_fill(client, f"{legacy[0]}_4", ["a", "b", "both"])
|
||||
_fill(client, f"{legacy[0]}_coding_keys_cv_4", ["k1"])
|
||||
_fill(client, f"{current}_4", ["new", "both"], text="newer")
|
||||
del client
|
||||
|
||||
assert len(legacy_vector_collections(tmp_path, profile)) == 2
|
||||
done = reattach_legacy_vectors(tmp_path, profile)
|
||||
|
||||
assert {(item.current, item.final_count) for item in done} == {
|
||||
(f"{current}_4", 4),
|
||||
(f"{current}_coding_keys_cv_4", 1),
|
||||
}
|
||||
assert legacy_vector_collections(tmp_path, profile) == ()
|
||||
client = _chroma(tmp_path)
|
||||
assert sorted(collection.name for collection in client.list_collections()) == [
|
||||
f"{current}_4",
|
||||
f"{current}_coding_keys_cv_4",
|
||||
]
|
||||
chat = client.get_collection(f"{current}_4").get(ids=["both"], include=["documents"])
|
||||
assert chat["documents"] == ["newer both"]
|
||||
assert reattach_legacy_vectors(tmp_path, profile) == ()
|
||||
|
||||
|
||||
def _confirmed(store, *memory_ids: str, kind: str = "chat") -> None:
|
||||
with main_database(store.database_path).write() as database, database:
|
||||
for memory_id in memory_ids:
|
||||
database.execute(
|
||||
"INSERT INTO governed_memories (memory_id,user_id,kind,status,content,"
|
||||
"scopes_json,revision,created_at,updated_at) "
|
||||
"VALUES (?,'local',?,'confirmed','x','[]',1,'now','now')",
|
||||
(memory_id, kind),
|
||||
)
|
||||
|
||||
|
||||
class _ReadingClient:
|
||||
"""Public-API shaped SDK client that holds only `readable` ids."""
|
||||
|
||||
def __init__(self, readable: set[str]) -> None:
|
||||
self.readable = readable
|
||||
|
||||
def list_memories(self, *, user_id, statuses, limit, offset, order):
|
||||
del user_id, statuses, order
|
||||
ids = sorted(self.readable)
|
||||
rows = [{"memory_id": memory_id} for memory_id in ids[offset : offset + limit]]
|
||||
return {"vdb": {"memories": rows, "total": len(ids)}}
|
||||
|
||||
def get(self, memory_id):
|
||||
return {"memory_id": memory_id} if memory_id in self.readable else None
|
||||
|
||||
|
||||
def _reading(readable: set[str]):
|
||||
from contextlib import contextmanager
|
||||
|
||||
@contextmanager
|
||||
def open_client(_profile):
|
||||
yield _ReadingClient(readable)
|
||||
|
||||
return open_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def active_profile(monkeypatch):
|
||||
from one_memory.hy_sdk_runtime import builtin_retrieval_snapshot
|
||||
|
||||
profile = builtin_retrieval_snapshot()
|
||||
monkeypatch.setattr("one_memory.sdk_migration_runner.read_active_profile", lambda _db: profile)
|
||||
monkeypatch.setattr("one_memory.sdk_migration_verify.read_active_profile", lambda _db: profile)
|
||||
return profile
|
||||
|
||||
|
||||
def test_cutover_waits_until_the_new_runtime_reads_every_chat_memory(store, active_profile):
|
||||
"""Build 20 reported success on 37 of 3317. Now the legacy tables stay."""
|
||||
from one_memory.sdk_migration_verify import MemoryReadbackFailed
|
||||
|
||||
_confirmed(store, "m1", "m2", "m3")
|
||||
home = store.database_path.parent
|
||||
|
||||
with pytest.raises(MemoryReadbackFailed) as caught:
|
||||
run_startup_migration(
|
||||
data_dir=home,
|
||||
runtime_build="21",
|
||||
sdk_version="2.0.5",
|
||||
namespace="governed-v2",
|
||||
open_client=_reading({"m1"}),
|
||||
)
|
||||
|
||||
assert caught.value.code == "sdk_migration_memory_unreadable"
|
||||
assert caught.value.report.as_document() == {
|
||||
"chat_expected": 3,
|
||||
"chat_readable": 1,
|
||||
"productivity_expected": 0,
|
||||
"productivity_readable": 0,
|
||||
"passed": False,
|
||||
}
|
||||
assert store.authority() is MemoryAuthority.LEGACY
|
||||
assert "governed_memories" in legacy_tables_present(store.database_path)
|
||||
|
||||
result = run_startup_migration(
|
||||
data_dir=home,
|
||||
runtime_build="21",
|
||||
sdk_version="2.0.5",
|
||||
namespace="governed-v2",
|
||||
open_client=_reading({"m1", "m2", "m3"}),
|
||||
)
|
||||
|
||||
assert result.state == "cutover"
|
||||
assert result.readback is not None and result.readback.passed
|
||||
assert store.authority() is MemoryAuthority.SDK
|
||||
|
||||
|
||||
def test_retired_coding_experiences_do_not_block_but_an_empty_store_does(store, active_profile):
|
||||
_confirmed(store, "chat-1")
|
||||
_confirmed(store, "exp-kept", "exp-retired", kind="coding")
|
||||
home = store.database_path.parent
|
||||
from one_memory.sdk_migration_verify import MemoryReadbackFailed
|
||||
|
||||
with pytest.raises(MemoryReadbackFailed):
|
||||
run_startup_migration(
|
||||
data_dir=home,
|
||||
runtime_build="21",
|
||||
sdk_version="2.0.5",
|
||||
namespace="governed-v2",
|
||||
open_client=_reading({"chat-1"}),
|
||||
)
|
||||
result = run_startup_migration(
|
||||
data_dir=home,
|
||||
runtime_build="21",
|
||||
sdk_version="2.0.5",
|
||||
namespace="governed-v2",
|
||||
open_client=_reading({"chat-1", "exp-kept"}),
|
||||
)
|
||||
assert result.readback is not None
|
||||
assert (result.readback.productivity_expected, result.readback.productivity_readable) == (2, 1)
|
||||
|
||||
|
||||
def test_memories_without_a_model_profile_keep_legacy_authority(store, monkeypatch):
|
||||
monkeypatch.setattr("one_memory.sdk_migration_runner.read_active_profile", lambda _db: None)
|
||||
_confirmed(store, "m1")
|
||||
|
||||
with pytest.raises(DataPlaneMigrationError, match="sdk_migration_model_profile_required"):
|
||||
run_startup_migration(
|
||||
data_dir=store.database_path.parent,
|
||||
runtime_build="21",
|
||||
sdk_version="2.0.5",
|
||||
namespace="governed-v2",
|
||||
)
|
||||
assert store.authority() is MemoryAuthority.LEGACY
|
||||
|
||||
|
||||
def test_a_build_20_home_gets_its_vectors_back_and_is_checked_against_its_backup(
|
||||
store, active_profile, monkeypatch
|
||||
):
|
||||
from one_memory.hy_sdk_runtime import vector_collection_bases
|
||||
|
||||
home = store.database_path.parent
|
||||
current, legacy_bases = vector_collection_bases(active_profile)
|
||||
client = _chroma(home)
|
||||
_fill(client, f"{legacy_bases[0]}_512", ["m1", "m2"])
|
||||
_fill(client, f"{current}_512", ["after-upgrade"], text="newer")
|
||||
del client
|
||||
_confirmed(store, "m1", "m2")
|
||||
first = run_startup_migration(
|
||||
data_dir=home,
|
||||
runtime_build="20",
|
||||
sdk_version="2.0.5",
|
||||
namespace="governed-v2",
|
||||
open_client=_reading({"m1", "m2"}),
|
||||
)
|
||||
assert first.state == "cutover"
|
||||
# Reproduce build 20: cut over while the legacy collection stayed stranded.
|
||||
client = _chroma(home)
|
||||
client.get_collection(f"{current}_512").modify(name=f"{legacy_bases[0]}_512")
|
||||
_fill(client, f"{current}_512", ["after-upgrade"], text="newer")
|
||||
del client
|
||||
monkeypatch.setattr(cli, "load_settings", lambda: Settings(_env_file=None, data_dir=home))
|
||||
status = io.StringIO()
|
||||
assert cli.run(["sdk-migration-status", "--json"], output=status) == 0
|
||||
assert json.loads(status.getvalue())["state"] == "required"
|
||||
|
||||
repaired = run_startup_migration(
|
||||
data_dir=home,
|
||||
runtime_build="21",
|
||||
sdk_version="2.0.5",
|
||||
namespace="governed-v2",
|
||||
open_client=_reading({"m1", "m2", "after-upgrade"}),
|
||||
)
|
||||
|
||||
assert repaired.state == "complete"
|
||||
assert [(item.current, item.final_count) for item in repaired.reattached] == [
|
||||
(f"{current}_512", 3)
|
||||
]
|
||||
assert repaired.readback is not None and repaired.readback.passed
|
||||
assert repaired.readback.chat_expected == 2
|
||||
status = io.StringIO()
|
||||
assert cli.run(["sdk-migration-status", "--json"], output=status) == 0
|
||||
assert json.loads(status.getvalue())["state"] == "not_required"
|
||||
|
||||
|
||||
def test_pending_dispatch_becomes_one_write_job_per_conversation(store):
|
||||
_insert_evidence(store, "a1", conversation="a", seq=1)
|
||||
_insert_evidence(store, "a2", conversation="a", role="assistant", seq=2)
|
||||
@@ -560,3 +1051,223 @@ def test_sdk_raw_retention_keeps_only_the_latest_three_days():
|
||||
)
|
||||
assert asyncio.run(runtime.prune_expired_raw(now=now)) == 1
|
||||
assert [row["memory_id"] for row in client.rows] == ["recent"]
|
||||
|
||||
|
||||
# The UI asks which memory data plane is authoritative and must get the
|
||||
# answer the API gates use.
|
||||
|
||||
|
||||
def _app_settings(tmp_path) -> Settings:
|
||||
return Settings(
|
||||
_env_file=None,
|
||||
data_dir=tmp_path / "data",
|
||||
source_home=tmp_path,
|
||||
api_token=SecretStr("authority-test"),
|
||||
imports_enabled=False,
|
||||
automatic_memory_processing=False,
|
||||
live_sync_enabled=False,
|
||||
)
|
||||
|
||||
|
||||
async def _authority(tmp_path) -> tuple[dict[str, object], str]:
|
||||
app = create_app(_app_settings(tmp_path), secret_store=InMemorySecretStore())
|
||||
async with (
|
||||
app.router.lifespan_context(app),
|
||||
AsyncClient(
|
||||
transport=ASGITransport(app=app),
|
||||
base_url="http://test",
|
||||
headers={"Authorization": "Bearer authority-test"},
|
||||
) as client,
|
||||
):
|
||||
response = await client.get("/api/v2/memory-authority")
|
||||
assert response.status_code == 200
|
||||
return response.json(), app.state.memory_authority.value
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.fresh_sdk
|
||||
def test_a_fresh_home_reports_sdk_authority(tmp_path) -> None:
|
||||
document, gate = asyncio.run(_authority(tmp_path))
|
||||
|
||||
assert document == {"authority": "sdk", "migration_id": None}
|
||||
assert gate == "sdk"
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
def test_a_migrated_home_still_reports_sdk_after_its_receipt_is_dropped(
|
||||
tmp_path, monkeypatch
|
||||
) -> None:
|
||||
"""Build 20: the cleanup dropped the cutover receipt, the UI fell back to legacy."""
|
||||
data_dir = tmp_path / "data"
|
||||
EvidenceStore(data_dir / "one-memory.sqlite3")
|
||||
monkeypatch.setattr(cli, "load_settings", lambda: _app_settings(tmp_path))
|
||||
monkeypatch.setattr(
|
||||
cli,
|
||||
"inspect_runtime",
|
||||
lambda _data_dir: RuntimeStatus("stopped", "runtime_receipt_missing"),
|
||||
)
|
||||
output = io.StringIO()
|
||||
assert cli.run(["sdk-migrate", "--runtime-build", "21", "--yes", "--json"], output=output) == 0
|
||||
migration_id = json.loads(output.getvalue().splitlines()[-1])["migration_id"]
|
||||
assert legacy_tables_present(data_dir / "one-memory.sqlite3") == ()
|
||||
|
||||
document, gate = asyncio.run(_authority(tmp_path))
|
||||
|
||||
assert document == {"authority": "sdk", "migration_id": migration_id}
|
||||
assert gate == "sdk"
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
def test_a_legacy_home_reports_legacy_authority(tmp_path) -> None:
|
||||
document, gate = asyncio.run(_authority(tmp_path))
|
||||
|
||||
assert document == {"authority": "legacy", "migration_id": None}
|
||||
assert gate == "legacy"
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.fresh_sdk
|
||||
def test_the_sdk_memory_map_draws_sdk_memories_without_legacy_tables(tmp_path) -> None:
|
||||
"""Build 21: the memory map had no SDK read and waited forever."""
|
||||
from one_memory.hy_graph_projection import HyGraphEdge, HyGraphNode, HyGraphSnapshot
|
||||
|
||||
rows = [
|
||||
{
|
||||
"memory_id": memory_id,
|
||||
"content": f"memory {memory_id}",
|
||||
"layer": "l2_fact",
|
||||
"status": "active",
|
||||
"is_latest": True,
|
||||
"user_id": "local",
|
||||
"agent_id": "all-apps",
|
||||
"session_id": "s1",
|
||||
"source": "cursor",
|
||||
}
|
||||
for memory_id in ("a", "b")
|
||||
] + [
|
||||
{
|
||||
"memory_id": memory_id,
|
||||
"content": f"backfilled {memory_id}",
|
||||
"layer": "l2_fact",
|
||||
"status": "active",
|
||||
"user_id": "local",
|
||||
"agent_id": "one-memory",
|
||||
"session_id": "default_session",
|
||||
"source": "default",
|
||||
}
|
||||
for memory_id in ("c", "d")
|
||||
]
|
||||
listed = []
|
||||
|
||||
class Client:
|
||||
def list_memories(self, **kwargs):
|
||||
listed.append(kwargs)
|
||||
return {"vdb": {"memories": rows, "total": 7}}
|
||||
|
||||
class Handle:
|
||||
async def client(self):
|
||||
return Client()
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
async def scenario():
|
||||
app = create_app(_app_settings(tmp_path), secret_store=InMemorySecretStore())
|
||||
async with (
|
||||
app.router.lifespan_context(app),
|
||||
AsyncClient(
|
||||
transport=ASGITransport(app=app),
|
||||
base_url="http://test",
|
||||
headers={"Authorization": "Bearer authority-test"},
|
||||
) as client,
|
||||
):
|
||||
original = app.state.sdk_client_handle
|
||||
app.state.sdk_client_handle = Handle()
|
||||
|
||||
async def graph_snapshot(*, limit):
|
||||
del limit
|
||||
return HyGraphSnapshot(
|
||||
nodes=(HyGraphNode(node_id="t1", content="topic", layer="TOPIC"),),
|
||||
edges=(
|
||||
HyGraphEdge(
|
||||
source_node_id="t1", target_node_id="a", relation="derived_from"
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
app.state.sdk_memory_runtime.graph_snapshot = graph_snapshot
|
||||
assert (await client.get("/api/v1/memories/graph")).status_code == 409
|
||||
response = await client.get(
|
||||
"/api/v2/memories/graph",
|
||||
params={"source": "cursor", "offset": 100, "limit": 100},
|
||||
)
|
||||
app.state.sdk_client_handle = original
|
||||
assert response.status_code == 200
|
||||
return response.json()
|
||||
|
||||
graph = asyncio.run(scenario())
|
||||
assert listed[0]["include_sources"] == ["cursor"]
|
||||
assert (listed[0]["offset"], listed[0]["limit"]) == (100, 100)
|
||||
assert (graph["total_memories"], graph["projected_memories"]) == (7, 4)
|
||||
assert {node["node_id"] for node in graph["nodes"]} == {
|
||||
"memory:a",
|
||||
"memory:b",
|
||||
"memory:c",
|
||||
"memory:d",
|
||||
"source:cursor",
|
||||
"hy:t1",
|
||||
}
|
||||
# "default_session" is the SDK's filler, not a conversation.
|
||||
assert [
|
||||
(edge["source_node_id"], edge["target_node_id"])
|
||||
for edge in graph["edges"]
|
||||
if edge["relation"] == "same_conversation"
|
||||
] == [("memory:a", "memory:b")]
|
||||
assert {edge["relation"] for edge in graph["edges"]} == {
|
||||
"evidence_source",
|
||||
"same_conversation",
|
||||
"derived_from",
|
||||
}
|
||||
assert graph["hy_graph_available"] is True
|
||||
|
||||
|
||||
def test_a_completed_sdk_write_does_not_touch_retired_evidence(store) -> None:
|
||||
"""Build 21 logged RetiredAuthorityAccess after every successful write."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
from one_memory.imports.storage import EvidenceStore as Evidence
|
||||
from one_memory.memory_server.durable_queue import DurableWriteDrain, execute_job
|
||||
from one_memory.retired_authority import RetiredEvidenceStore
|
||||
|
||||
path = store.database_path
|
||||
_insert_evidence(store, "e1")
|
||||
record = Evidence(path).get("e1")
|
||||
assert record is not None
|
||||
queue = DurableWriteQueue(path, sdk_authority=True)
|
||||
job = queue.submit(
|
||||
WriteRequest(
|
||||
internal_session_id="s",
|
||||
source_id="cursor",
|
||||
batch_identity=BatchIdentity(internal_session_id="s", first_round=1, last_round=1),
|
||||
evidence_ids=["e1"],
|
||||
input_records=(record,),
|
||||
idempotency_key="k",
|
||||
)
|
||||
)
|
||||
DurableWriteDrain(queue, execute=lambda _job: asyncio.sleep(0)).claim(limit=1)
|
||||
|
||||
class Supervisor:
|
||||
async def process_sdk_job(self, records, *, request_id, before_hy_invocation):
|
||||
del records, request_id
|
||||
await before_hy_invocation()
|
||||
return SimpleNamespace(raw_ids=("raw-1",), memory_ids=("m-1",))
|
||||
|
||||
asyncio.run(
|
||||
execute_job(
|
||||
job.job_id,
|
||||
queue=queue,
|
||||
supervisor=Supervisor(),
|
||||
evidence=RetiredEvidenceStore(path),
|
||||
)
|
||||
)
|
||||
assert queue.get("k").state == "completed"
|
||||
@@ -112,9 +112,11 @@ import {
|
||||
installBundledCliIfNeeded,
|
||||
installServiceWithDisabledFallback,
|
||||
isInstalledApplicationsApp,
|
||||
requireCurrentRuntimeAtLeast,
|
||||
stableRuntimeErrorCode,
|
||||
stableRuntimeFailureDetail,
|
||||
startStableRuntimeRepair,
|
||||
startupFailureCause,
|
||||
} from "./stable-runtime";
|
||||
import { STARTUP_SCREEN_URL, startupScreenHtml } from "./startup-screen";
|
||||
import {
|
||||
@@ -1093,7 +1095,7 @@ async function initialize(isQuitting: () => boolean): Promise<void> {
|
||||
component: "desktop.sdk_migration",
|
||||
event: "sdk_migration.startup",
|
||||
outcome:
|
||||
migration.state === "cutover"
|
||||
migration.state === "cutover" || migration.state === "complete"
|
||||
? "succeeded"
|
||||
: migration.state === "deferred"
|
||||
? "blocked"
|
||||
@@ -1158,30 +1160,52 @@ async function initialize(isQuitting: () => boolean): Promise<void> {
|
||||
const previousOwner = owner;
|
||||
let runtimeChanged = false;
|
||||
let transientRuntime = false;
|
||||
if (owner?.owner !== "cli") {
|
||||
const installed = await installPackagedRuntime({
|
||||
resourcesPath: process.resourcesPath,
|
||||
executablePath: process.execPath,
|
||||
installRoot,
|
||||
platform: "darwin",
|
||||
arch: process.arch,
|
||||
...(packagedBuildNumber === undefined
|
||||
? {}
|
||||
: { buildNumber: packagedBuildNumber }),
|
||||
...(packagedCommit === undefined
|
||||
? {}
|
||||
: { gitCommit: packagedCommit }),
|
||||
// Only the packaged Desktop path opts into a valid same-build,
|
||||
// different-commit takeover. CLI installs remain strict.
|
||||
replaceSameBuild: true,
|
||||
serviceEnableMarker: runtimeServiceEnableMarkerPath(dataDir),
|
||||
});
|
||||
runtimeChanged = !installed.reused_current;
|
||||
await installBundledCliIfNeeded({
|
||||
home: app.getPath("home"),
|
||||
installRoot,
|
||||
platform: "darwin",
|
||||
});
|
||||
// Install whoever owns the service. Owner decides who starts and stops
|
||||
// it, not which build runs: skipping the install for a CLI owner (also
|
||||
// the App's own guess when its receipt is missing) left an upgraded App
|
||||
// on the previous Runtime and, with `current` gone, on none at all.
|
||||
{
|
||||
let installed: Awaited<
|
||||
ReturnType<typeof installPackagedRuntime>
|
||||
> | null = null;
|
||||
try {
|
||||
installed = await installPackagedRuntime({
|
||||
resourcesPath: process.resourcesPath,
|
||||
executablePath: process.execPath,
|
||||
installRoot,
|
||||
platform: "darwin",
|
||||
arch: process.arch,
|
||||
...(packagedBuildNumber === undefined
|
||||
? {}
|
||||
: { buildNumber: packagedBuildNumber }),
|
||||
...(packagedCommit === undefined
|
||||
? {}
|
||||
: { gitCommit: packagedCommit }),
|
||||
// Only the packaged Desktop path opts into a valid same-build,
|
||||
// different-commit takeover. CLI installs remain strict.
|
||||
replaceSameBuild: true,
|
||||
serviceEnableMarker: runtimeServiceEnableMarkerPath(dataDir),
|
||||
});
|
||||
} catch (error) {
|
||||
// `memhub upgrade` may already have installed a newer build.
|
||||
if (
|
||||
!(
|
||||
owner?.owner === "cli" &&
|
||||
error instanceof Error &&
|
||||
error.message === "distribution_downgrade_refused"
|
||||
)
|
||||
)
|
||||
throw error;
|
||||
}
|
||||
if (installed !== null) {
|
||||
runtimeChanged = !installed.reused_current;
|
||||
await installBundledCliIfNeeded({
|
||||
home: app.getPath("home"),
|
||||
installRoot,
|
||||
platform: "darwin",
|
||||
});
|
||||
}
|
||||
await requireCurrentRuntimeAtLeast(installRoot, packagedBuildNumber);
|
||||
}
|
||||
if (
|
||||
!desktopBackgroundLaunch &&
|
||||
@@ -1485,6 +1509,9 @@ async function initialize(isQuitting: () => boolean): Promise<void> {
|
||||
outcome: "failed",
|
||||
error_code: stableRuntimeErrorCode(error),
|
||||
...diagnosticError(error),
|
||||
// The code alone is often `desktop_start_failed`; the cause is what
|
||||
// support needs. Home is shortened and the text bounded.
|
||||
details: { cause: startupFailureCause(error, app.getPath("home")) },
|
||||
});
|
||||
throw error;
|
||||
}
|
||||
|
||||
@@ -15,10 +15,12 @@ import {
|
||||
installServiceWithDisabledFallback,
|
||||
isDisabledServiceInstallError,
|
||||
isInstalledApplicationsApp,
|
||||
requireCurrentRuntimeAtLeast,
|
||||
stableRuntimeErrorCode,
|
||||
stableRuntimeFailureDetail,
|
||||
stableRuntimeUserMessage,
|
||||
startStableRuntimeRepair,
|
||||
startupFailureCause,
|
||||
} from "./stable-runtime";
|
||||
|
||||
it("only treats a packaged macOS executable under Applications as installed", () => {
|
||||
@@ -285,6 +287,12 @@ describe("startup install errors", () => {
|
||||
expect(noSpace).toContain("重新启动");
|
||||
expect(noSpace).toContain("不会丢失");
|
||||
expect(noSpace).not.toContain("sdk_migration_disk_space_insufficient");
|
||||
|
||||
const unreadable = stableRuntimeUserMessage(
|
||||
new Error("sdk_migration_memory_unreadable"),
|
||||
);
|
||||
expect(unreadable).toContain("没有切换数据源");
|
||||
expect(unreadable).toContain("不会丢失");
|
||||
});
|
||||
});
|
||||
|
||||
@@ -305,4 +313,53 @@ describe("startup dialog body", () => {
|
||||
expect(detail).toContain("磁盘空间不足");
|
||||
expect(detail).not.toContain("sdk_migration_disk_space_insufficient");
|
||||
});
|
||||
|
||||
it("turns runtime_service_not_ready into a user-facing sentence", () => {
|
||||
const error = new Error("runtime_service_not_ready");
|
||||
expect(stableRuntimeErrorCode(error)).toBe("runtime_service_not_ready");
|
||||
expect(stableRuntimeFailureDetail(error)).toContain("首次启动");
|
||||
expect(stableRuntimeFailureDetail(error)).not.toContain(
|
||||
"runtime_service_not_ready",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("packaged runtime build after an App upgrade", () => {
|
||||
it("refuses to run on a Runtime older than the App, or none", async () => {
|
||||
const root = await mkdtemp(path.join(os.tmpdir(), "memhub-current-"));
|
||||
await expect(requireCurrentRuntimeAtLeast(root, 21)).rejects.toThrow(
|
||||
"packaged_runtime_not_current",
|
||||
);
|
||||
await mkdir(path.join(root, "current"));
|
||||
const manifest = path.join(root, "current", "manifest.json");
|
||||
await writeFile(manifest, JSON.stringify({ build_number: 19 }));
|
||||
await expect(requireCurrentRuntimeAtLeast(root, 21)).rejects.toThrow(
|
||||
"packaged_runtime_not_current",
|
||||
);
|
||||
await writeFile(manifest, JSON.stringify({ build_number: 21 }));
|
||||
await expect(
|
||||
requireCurrentRuntimeAtLeast(root, 21),
|
||||
).resolves.toBeUndefined();
|
||||
await writeFile(manifest, JSON.stringify({ build_number: 22 }));
|
||||
await expect(
|
||||
requireCurrentRuntimeAtLeast(root, 21),
|
||||
).resolves.toBeUndefined();
|
||||
await expect(
|
||||
requireCurrentRuntimeAtLeast(root, undefined),
|
||||
).resolves.toBeUndefined();
|
||||
expect(
|
||||
stableRuntimeUserMessage(new Error("packaged_runtime_not_current")),
|
||||
).toContain("重新安装");
|
||||
});
|
||||
|
||||
it("keeps the cause of a startup failure, without the home path", () => {
|
||||
const cause = startupFailureCause(
|
||||
new Error("ENOENT: no such file /Users/someone/Library/x\nmore"),
|
||||
"/Users/someone",
|
||||
);
|
||||
expect(cause).toBe("Error: ENOENT: no such file ~/Library/x more");
|
||||
expect(startupFailureCause(new Error("x".repeat(500)), "/h").length).toBe(
|
||||
300,
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -129,13 +129,19 @@ const SAFE_STABLE_RUNTIME_ERROR_CODES = new Set([
|
||||
// Migration runs before the app window exists, so these are the only names
|
||||
// a first launch can be given. They stay codes here and become sentences in
|
||||
// STABLE_RUNTIME_USER_MESSAGES: a person reading the dialog never sees them.
|
||||
"packaged_runtime_not_current",
|
||||
"sdk_migration_preflight_failed",
|
||||
"sdk_migration_memory_unreadable",
|
||||
"sdk_migration_model_profile_required",
|
||||
"sdk_migration_model_secret_unavailable",
|
||||
"sdk_migration_secure_storage_unavailable",
|
||||
"sdk_migration_disk_space_insufficient",
|
||||
"sdk_migration_cutover_failed",
|
||||
"sdk_migration_bundle_copy_failed",
|
||||
"sdk_migration_response_invalid",
|
||||
"sdk_migration_status_unavailable",
|
||||
"sdk_migration_incomplete",
|
||||
"runtime_service_not_ready",
|
||||
]);
|
||||
|
||||
export function stableRuntimeErrorCode(error: unknown): string {
|
||||
@@ -201,6 +207,16 @@ export const STABLE_RUNTIME_USER_MESSAGES: Record<string, string> = {
|
||||
"已安装的运行环境比这个安装包更新,无法用旧安装包启动。请安装最新版本的 MemHub;你的记忆数据不会受到影响。",
|
||||
install_current_unmanaged:
|
||||
"运行环境的安装目录被改动过,无法安全接管。请重新安装 MemHub,你的记忆数据不会受到影响。",
|
||||
packaged_runtime_not_current:
|
||||
"新版本的运行环境没有装上,MemHub 不会用旧版本继续运行。请重新安装 MemHub;如果仍然失败,请把这个提示反馈给我们。你的记忆数据不会受到影响。",
|
||||
sdk_migration_memory_unreadable:
|
||||
"升级前核对发现,新版本还读不到你的全部记忆,所以没有切换数据源,旧数据保持原样可用。请把这个提示反馈给我们;原有记忆数据不会丢失。",
|
||||
sdk_migration_model_profile_required:
|
||||
"升级记忆数据格式需要先有可用的模型配置,用来核对新版本能否读回你的记忆。请在旧版本中完成模型配置后重启;你的记忆数据不会受到影响。",
|
||||
sdk_migration_model_secret_unavailable:
|
||||
"模型密钥读不到(可能被系统钥匙串拒绝),无法核对新版本能否读回你的记忆。请允许 MemHub 访问钥匙串后重新启动;你的记忆数据不会受到影响。",
|
||||
sdk_migration_secure_storage_unavailable:
|
||||
"系统钥匙串暂时不可用,无法核对新版本能否读回你的记忆。请解锁钥匙串后重新启动;你的记忆数据不会受到影响。",
|
||||
sdk_migration_disk_space_insufficient:
|
||||
"磁盘空间不足,升级记忆数据格式时中断了。请至少腾出与记忆数据同等大小的可用空间(可在「关于本机 → 储存空间」查看),然后重新启动 MemHub;升级会自动继续,原有记忆数据不会丢失。",
|
||||
sdk_migration_bundle_copy_failed:
|
||||
@@ -215,6 +231,8 @@ export const STABLE_RUNTIME_USER_MESSAGES: Record<string, string> = {
|
||||
"升级过程返回了无法识别的内容,已停止。原有记忆数据不会丢失;请重新启动重试,如果仍然失败,请把这个提示反馈给我们。",
|
||||
sdk_migration_incomplete:
|
||||
"记忆数据格式升级没有完成,MemHub 暂时无法打开。请重新启动重试;如果仍然失败,请把这个提示反馈给我们,原有记忆数据不会丢失。",
|
||||
runtime_service_not_ready:
|
||||
"记忆服务仍在首次启动,请稍候后重试。你的记忆数据不会受到影响。",
|
||||
desktop_start_failed:
|
||||
"MemHub 启动失败。请重新启动;如果仍然打不开,请重新安装,你的记忆数据不会受到影响。",
|
||||
};
|
||||
@@ -240,3 +258,31 @@ export function stableRuntimeFailureDetail(error: unknown): string {
|
||||
}
|
||||
return STABLE_RUNTIME_USER_MESSAGES[error.message] ?? error.message;
|
||||
}
|
||||
|
||||
/** The installed Runtime must be at least the build this App carries. */
|
||||
export async function requireCurrentRuntimeAtLeast(
|
||||
installRoot: string,
|
||||
buildNumber: number | undefined,
|
||||
): Promise<void> {
|
||||
if (buildNumber === undefined) return;
|
||||
const manifest = await readFile(
|
||||
path.join(installRoot, "current", "manifest.json"),
|
||||
"utf8",
|
||||
)
|
||||
.then((text) => JSON.parse(text) as { build_number?: unknown })
|
||||
.catch(() => null);
|
||||
const installed = manifest?.build_number;
|
||||
if (typeof installed !== "number" || installed < buildNumber) {
|
||||
throw new Error("packaged_runtime_not_current");
|
||||
}
|
||||
}
|
||||
|
||||
export function startupFailureCause(error: unknown, home: string): string {
|
||||
const text =
|
||||
error instanceof Error ? `${error.name}: ${error.message}` : String(error);
|
||||
return text
|
||||
.split(home)
|
||||
.join("~")
|
||||
.replace(/[\r\n]+/g, " ")
|
||||
.slice(0, 300);
|
||||
}
|
||||
@@ -177,12 +177,13 @@ it("OBS-22: the graph inspector has no editing entry and never says changes cann
|
||||
await act(async () => {});
|
||||
// The detail pane is open and the page is being re-read, which used to be
|
||||
// the exact condition that printed the warning.
|
||||
expect(screen.getByText("共同原文")).toBeTruthy();
|
||||
expect(screen.getByRole("heading", { name: "第一条" })).toBeTruthy();
|
||||
expect(screen.queryByText(/暂不能保存更改/)).toBeNull();
|
||||
});
|
||||
|
||||
it("OBS-19: Enter in the memory filter re-reads, exactly like the refresh button", async () => {
|
||||
it("OBS-19: Enter in the memory filter submits it; submitting it again re-reads", async () => {
|
||||
const onRetry = vi.fn();
|
||||
const onQuery = vi.fn();
|
||||
const page = {
|
||||
items: memories,
|
||||
total: memories.length,
|
||||
@@ -200,7 +201,7 @@ it("OBS-19: Enter in the memory filter re-reads, exactly like the refresh button
|
||||
loading: false,
|
||||
failed: false,
|
||||
stale: false,
|
||||
onQuery: vi.fn(),
|
||||
onQuery,
|
||||
onSource: vi.fn(),
|
||||
resetPage: vi.fn(),
|
||||
onPage: vi.fn(),
|
||||
@@ -211,11 +212,39 @@ it("OBS-19: Enter in the memory filter re-reads, exactly like the refresh button
|
||||
await act(async () => {});
|
||||
await userEvent.click(screen.getByRole("button", { name: "刷新地图记忆" }));
|
||||
expect(onRetry).toHaveBeenCalledTimes(1);
|
||||
// Typing alone does not search; Enter submits the words.
|
||||
const input = screen.getByRole("textbox", { name: "筛选地图记忆" });
|
||||
await userEvent.type(input, "第一");
|
||||
expect(onQuery).not.toHaveBeenCalled();
|
||||
await userEvent.type(input, "{Enter}");
|
||||
expect(onQuery).toHaveBeenCalledWith("第一");
|
||||
// Submitting the words already applied re-reads, like the refresh button.
|
||||
await userEvent.clear(input);
|
||||
await userEvent.type(input, "{Enter}");
|
||||
expect(onRetry).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("searches only when 搜索 is pressed, not on every keystroke", async () => {
|
||||
const onQuery = vi.fn();
|
||||
const resetPage = vi.fn();
|
||||
render(
|
||||
<Host
|
||||
simplified
|
||||
libraryView={{
|
||||
...pagedView({ items: memories, total: 2, offset: 0, limit: 100 }),
|
||||
onQuery,
|
||||
resetPage,
|
||||
}}
|
||||
/>,
|
||||
);
|
||||
await userEvent.type(
|
||||
screen.getByRole("textbox", { name: "筛选地图记忆" }),
|
||||
"第一{Enter}",
|
||||
"版本号",
|
||||
);
|
||||
expect(onRetry).toHaveBeenCalledTimes(2);
|
||||
expect(onQuery).not.toHaveBeenCalled();
|
||||
await userEvent.click(screen.getByRole("button", { name: "搜索" }));
|
||||
expect(onQuery).toHaveBeenCalledWith("版本号");
|
||||
expect(resetPage).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("OBS-15: switching memory shows a read in progress instead of the previous memory", async () => {
|
||||
@@ -224,11 +253,11 @@ it("OBS-15: switching memory shows a read in progress instead of the previous me
|
||||
// screen while it is open.
|
||||
render(<Host simplified deferMemoryId="m1" />);
|
||||
const list = screen.getByRole("list", { name: "地图记忆列表" });
|
||||
expect(screen.getByText("共同原文")).toBeTruthy();
|
||||
expect(screen.getByRole("heading", { name: "第一条" })).toBeTruthy();
|
||||
await user.click(within(list).getByRole("button", { name: /第二条/ }));
|
||||
// The first memory's evidence must not still be presented as if it
|
||||
// belonged to the selection that is now the second one.
|
||||
expect(screen.queryByText("共同原文")).toBeNull();
|
||||
expect(screen.queryByRole("heading", { name: "第一条" })).toBeNull();
|
||||
expect(screen.getByRole("status").textContent).toContain("正在读取这条记忆");
|
||||
});
|
||||
|
||||
@@ -247,16 +276,26 @@ it("labels the right-hand empty state as memory details", () => {
|
||||
expect(screen.queryByText("原始证据")).toBeNull();
|
||||
});
|
||||
|
||||
it("renders SDK-projected memory tags with the existing Workbench chips", () => {
|
||||
const { container } = render(<Host simplified />);
|
||||
expect(
|
||||
Array.from(
|
||||
container.querySelectorAll(".map-memory-tags .wb-tag"),
|
||||
(node) => node.textContent,
|
||||
),
|
||||
).toEqual(["产品", "产品"]);
|
||||
});
|
||||
|
||||
it("opens the newest listed memory in the detail pane and keeps it closed after dismiss", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<Host simplified />);
|
||||
const list = screen.getByRole("list", { name: "地图记忆列表" });
|
||||
const first = within(list).getByRole("button", { name: /第一条/ });
|
||||
expect(first.getAttribute("aria-pressed")).toBe("true");
|
||||
expect(screen.getByText("共同原文")).toBeTruthy();
|
||||
expect(screen.getByText(/原对话/)).toBeTruthy();
|
||||
expect(screen.getByRole("heading", { name: "第一条" })).toBeTruthy();
|
||||
expect(screen.queryByText(/原对话/)).toBeNull();
|
||||
await user.click(screen.getByRole("button", { name: "关闭记忆详情" }));
|
||||
expect(screen.queryByText("共同原文")).toBeNull();
|
||||
expect(screen.queryByRole("heading", { name: "第一条" })).toBeNull();
|
||||
});
|
||||
|
||||
it("OBS-14: lists what the conversation settled instead of asking to settle first", async () => {
|
||||
@@ -501,6 +540,65 @@ it("D96: lists a memory by when its conversation happened, not when it was gener
|
||||
).toBeTruthy();
|
||||
});
|
||||
|
||||
it("names the SDK's no-source bucket 未识别来源, apart from a source under review", () => {
|
||||
const items = ["有来源", "无来源"].map((content, i) => ({
|
||||
...memories[0],
|
||||
memory_id: `s${i}`,
|
||||
content,
|
||||
})) as GovernedMemory[];
|
||||
render(
|
||||
<Host
|
||||
simplified
|
||||
items={items}
|
||||
sources={[
|
||||
{
|
||||
source_id: "codex",
|
||||
display_name: "Codex",
|
||||
authorized: true,
|
||||
imported_messages: 1,
|
||||
imported_conversations: 1,
|
||||
} as SourceResponse,
|
||||
]}
|
||||
libraryView={{
|
||||
...pagedView({
|
||||
items,
|
||||
total: 2,
|
||||
offset: 0,
|
||||
limit: 100,
|
||||
source_ids_by_memory: { s0: ["codex"], s1: ["default"] },
|
||||
search_window: 500,
|
||||
}),
|
||||
query: "版本",
|
||||
unattributedSource: "default",
|
||||
}}
|
||||
/>,
|
||||
);
|
||||
const list = screen.getByRole("list", { name: "地图记忆列表" });
|
||||
expect(
|
||||
within(list).getByRole("button", { name: /有来源/ }).textContent,
|
||||
).toContain("Codex");
|
||||
const unattributed = within(list).getByRole("button", { name: /无来源/ });
|
||||
expect(unattributed.textContent).toContain("未识别来源");
|
||||
expect(unattributed.textContent).not.toContain("待核对");
|
||||
expect(screen.getByText("只在最相关的 500 条记忆里查找关键词")).toBeTruthy();
|
||||
});
|
||||
|
||||
it("says a filtered source is empty rather than that nothing matched", () => {
|
||||
render(
|
||||
<Host
|
||||
simplified
|
||||
items={[]}
|
||||
libraryView={{
|
||||
...pagedView({ items: [], total: 0, offset: 0, limit: 100 }),
|
||||
source: "default",
|
||||
unattributedSource: "default",
|
||||
}}
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByText("这个来源下还没有记忆")).toBeTruthy();
|
||||
expect(screen.queryByText("没有匹配的记忆")).toBeNull();
|
||||
});
|
||||
|
||||
it("D96: with conversations imported and no memory yet, says how many were imported", () => {
|
||||
render(
|
||||
<Host
|
||||
|
||||
@@ -63,6 +63,8 @@ type Props = Parameters<typeof Workbench>[0] & {
|
||||
loading: boolean;
|
||||
failed: boolean;
|
||||
stale: boolean;
|
||||
/** The source id the SDK uses for "no source given" (shown as 未识别来源). */
|
||||
unattributedSource?: string;
|
||||
onQuery(value: string): void;
|
||||
onSource(value: string): void;
|
||||
resetPage(): void;
|
||||
@@ -168,18 +170,40 @@ export function MemoryMap(props: Props) {
|
||||
} = props;
|
||||
const library = props.libraryView;
|
||||
const serverPaged = library !== undefined;
|
||||
const unattributedLabel = text("未识别来源");
|
||||
const graph = useMemo(
|
||||
() =>
|
||||
serverPaged
|
||||
? projectPagedSourceMap(baseGraph, library?.page ?? null, sources)
|
||||
? projectPagedSourceMap(
|
||||
baseGraph,
|
||||
library?.page ?? null,
|
||||
sources,
|
||||
library?.unattributedSource
|
||||
? { [library.unattributedSource]: unattributedLabel }
|
||||
: undefined,
|
||||
)
|
||||
: baseGraph,
|
||||
[baseGraph, serverPaged, library?.page, sources],
|
||||
[
|
||||
baseGraph,
|
||||
serverPaged,
|
||||
library?.page,
|
||||
library?.unattributedSource,
|
||||
sources,
|
||||
unattributedLabel,
|
||||
],
|
||||
);
|
||||
const refreshing = !!library?.page && library.loading;
|
||||
const overlaysSuspended = hidden || !!props.overlayBlocked;
|
||||
const [axis, setAxis] = useState<MapAxis>("source");
|
||||
const [groupKey, setGroupKey] = useState("");
|
||||
const [search, setSearch] = useState("");
|
||||
// On a server-paged list, typing only edits the words; the search runs when
|
||||
// the reader submits it (the 搜索 button or Enter), not on every keystroke.
|
||||
const [searchDraft, setSearchDraft] = useState(library?.query ?? "");
|
||||
const submittedQuery = library?.query;
|
||||
useEffect(() => {
|
||||
if (submittedQuery !== undefined) setSearchDraft(submittedQuery);
|
||||
}, [submittedQuery]);
|
||||
const [selectedId, setSelectedId] = useState<string | null>(null);
|
||||
const [groupNode, setGroupNode] = useState<MemoryGraphNode | null>(null);
|
||||
const [question, setQuestion] = useState("");
|
||||
@@ -338,7 +362,11 @@ export function MemoryMap(props: Props) {
|
||||
const ids = library.page?.source_ids_by_memory?.[memoryId];
|
||||
return ids?.length
|
||||
? ids
|
||||
.map((id) => sourceNameOr(sources, id, text("来源待核对")))
|
||||
.map((id) =>
|
||||
id === library.unattributedSource
|
||||
? text("未识别来源")
|
||||
: sourceNameOr(sources, id, text("来源待核对")),
|
||||
)
|
||||
.join("、")
|
||||
: text("来源待核对");
|
||||
}
|
||||
@@ -556,6 +584,8 @@ export function MemoryMap(props: Props) {
|
||||
setGroupNode(null);
|
||||
if (library) {
|
||||
clear();
|
||||
// A source picked on the map sets this too; keep both on the one choice.
|
||||
setGroupKey(key);
|
||||
library.onSource(key.startsWith("source:") ? key.slice(7) : "");
|
||||
library.resetPage();
|
||||
return;
|
||||
@@ -611,6 +641,10 @@ export function MemoryMap(props: Props) {
|
||||
(source) =>
|
||||
source.authorized && groupKey === `source:${source.source_id}`,
|
||||
)?.display_name ??
|
||||
(library?.unattributedSource &&
|
||||
groupKey === `source:${library.unattributedSource}`
|
||||
? text("未识别来源")
|
||||
: null) ??
|
||||
text("全部记忆");
|
||||
if (
|
||||
props.simplified &&
|
||||
@@ -740,6 +774,7 @@ export function MemoryMap(props: Props) {
|
||||
options={[
|
||||
{ value: "", label: text("全部应用") },
|
||||
...(library?.source &&
|
||||
library.source !== library.unattributedSource &&
|
||||
!sources.some(
|
||||
(source) =>
|
||||
source.source_id === library.source && source.authorized,
|
||||
@@ -752,12 +787,22 @@ export function MemoryMap(props: Props) {
|
||||
]
|
||||
: []),
|
||||
...(library
|
||||
? sources
|
||||
.filter((source) => source.authorized)
|
||||
.map((source) => ({
|
||||
value: `source:${source.source_id}`,
|
||||
label: source.display_name,
|
||||
}))
|
||||
? [
|
||||
...sources
|
||||
.filter((source) => source.authorized)
|
||||
.map((source) => ({
|
||||
value: `source:${source.source_id}`,
|
||||
label: source.display_name,
|
||||
})),
|
||||
...(library.unattributedSource
|
||||
? [
|
||||
{
|
||||
value: `source:${library.unattributedSource}`,
|
||||
label: text("未识别来源"),
|
||||
},
|
||||
]
|
||||
: []),
|
||||
]
|
||||
: groups.map((group) => ({
|
||||
value: group.key,
|
||||
label: group.label,
|
||||
@@ -789,7 +834,23 @@ export function MemoryMap(props: Props) {
|
||||
))}
|
||||
</fieldset>
|
||||
)}
|
||||
<div className="map-search">
|
||||
<form
|
||||
className="map-search"
|
||||
role="search"
|
||||
onSubmit={(e) => {
|
||||
e.preventDefault();
|
||||
if (!library || libraryPane === "dialogues") return;
|
||||
const words = searchDraft.trim();
|
||||
setFilterNotice("");
|
||||
clear();
|
||||
if (words === library.query.trim()) {
|
||||
library.onRetry();
|
||||
return;
|
||||
}
|
||||
library.onQuery(words);
|
||||
library.resetPage();
|
||||
}}
|
||||
>
|
||||
<Search size={16} aria-hidden="true" />
|
||||
<input
|
||||
aria-label={text("筛选地图记忆")}
|
||||
@@ -805,31 +866,39 @@ export function MemoryMap(props: Props) {
|
||||
value={
|
||||
libraryPane === "dialogues"
|
||||
? dialogueFilter
|
||||
: (library?.query ?? search)
|
||||
: library
|
||||
? searchDraft
|
||||
: search
|
||||
}
|
||||
maxLength={library ? 512 : undefined}
|
||||
onKeyDown={(e) => {
|
||||
// The refresh button next to this list re-reads the page;
|
||||
// Enter has to do the same thing, not sit there doing nothing.
|
||||
if (e.key !== "Enter" || e.nativeEvent.isComposing) return;
|
||||
e.preventDefault();
|
||||
library?.onRetry();
|
||||
// Enter that confirms an IME candidate is not a search.
|
||||
if (e.key === "Enter" && e.nativeEvent.isComposing)
|
||||
e.preventDefault();
|
||||
}}
|
||||
onChange={(e) => {
|
||||
if (libraryPane === "dialogues") {
|
||||
setDialogueFilter(e.target.value);
|
||||
return;
|
||||
}
|
||||
if (library) {
|
||||
setSearchDraft(e.target.value);
|
||||
return;
|
||||
}
|
||||
setFilterNotice("");
|
||||
setSearch(e.target.value);
|
||||
if (library) {
|
||||
clear();
|
||||
library.onQuery(e.target.value);
|
||||
library.resetPage();
|
||||
}
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
{library && libraryPane === "memories" ? (
|
||||
<button
|
||||
type="submit"
|
||||
className="map-search-submit"
|
||||
disabled={library.loading || hidden}
|
||||
>
|
||||
{text("搜索")}
|
||||
</button>
|
||||
) : null}
|
||||
</form>
|
||||
{library && libraryPane === "memories" && (
|
||||
<div className="map-list-heading" aria-live="polite">
|
||||
<span>
|
||||
@@ -988,7 +1057,16 @@ export function MemoryMap(props: Props) {
|
||||
onClick={() => pick(m.memory_id)}
|
||||
>
|
||||
<span className={`map-state ${m.status}`} aria-hidden="true" />
|
||||
<span>{m.content}</span>
|
||||
<span className="map-memory-content">{m.content}</span>
|
||||
{m.tags.length ? (
|
||||
<span className="map-memory-tags wb-tags">
|
||||
{m.tags.map((tag) => (
|
||||
<span key={tag} className="wb-tag">
|
||||
{tag}
|
||||
</span>
|
||||
))}
|
||||
</span>
|
||||
) : null}
|
||||
<small>
|
||||
{props.simplified ? (
|
||||
<>
|
||||
@@ -1029,11 +1107,23 @@ export function MemoryMap(props: Props) {
|
||||
!library.stale &&
|
||||
!filtered.length && (
|
||||
<p className="map-empty">
|
||||
{library.query || library.source
|
||||
{library.query
|
||||
? text("没有匹配的记忆")
|
||||
: groupKey
|
||||
? text("本页没有匹配的主题记忆")
|
||||
: text("还没有已生成的记忆")}
|
||||
: library.source
|
||||
? text("这个来源下还没有记忆")
|
||||
: groupKey
|
||||
? text("本页没有匹配的主题记忆")
|
||||
: text("还没有已生成的记忆")}
|
||||
</p>
|
||||
)}
|
||||
{libraryPane === "memories" &&
|
||||
library?.page?.search_window &&
|
||||
!library.loading &&
|
||||
!library.failed && (
|
||||
<p className="map-empty">
|
||||
{text("只在最相关的 {count} 条记忆里查找关键词", {
|
||||
count: count(library.page.search_window),
|
||||
})}
|
||||
</p>
|
||||
)}
|
||||
{libraryPane === "memories" &&
|
||||
|
||||
@@ -218,7 +218,7 @@ it("default map search finds an existing 501st memory rather than declaring no m
|
||||
fireEvent.change(screen.getByLabelText("筛选地图记忆"), {
|
||||
target: { value: "第501条独有目标" },
|
||||
});
|
||||
fireEvent.keyDown(screen.getByLabelText("筛选地图记忆"), { key: "Enter" });
|
||||
fireEvent.submit(screen.getByRole("search"));
|
||||
await waitFor(() =>
|
||||
expect(
|
||||
within(screen.getByRole("list", { name: "地图记忆列表" })).queryByText(
|
||||
@@ -310,6 +310,7 @@ it("OBS-13: the settled-count headline names the filter it is counting", async (
|
||||
fireEvent.change(screen.getByLabelText("筛选地图记忆"), {
|
||||
target: { value: "第501条独有目标" },
|
||||
});
|
||||
fireEvent.submit(screen.getByRole("search"));
|
||||
await waitFor(() =>
|
||||
expect(
|
||||
screen.getByText(
|
||||
@@ -320,6 +321,7 @@ it("OBS-13: the settled-count headline names the filter it is counting", async (
|
||||
fireEvent.change(screen.getByLabelText("筛选地图记忆"), {
|
||||
target: { value: "" },
|
||||
});
|
||||
fireEvent.submit(screen.getByRole("search"));
|
||||
await chooseSource("codex");
|
||||
await waitFor(() =>
|
||||
expect(
|
||||
@@ -473,6 +475,7 @@ it("distinguishes failed search from empty, then retries the same query", async
|
||||
fireEvent.change(screen.getByLabelText("筛选地图记忆"), {
|
||||
target: { value: "目标" },
|
||||
});
|
||||
fireEvent.submit(screen.getByRole("search"));
|
||||
await screen.findByText("记忆读取失败,请重试。");
|
||||
expect(screen.queryByText("没有匹配的记忆")).toBeNull();
|
||||
expect(screen.queryByText("还没有已生成的记忆")).toBeNull();
|
||||
@@ -490,6 +493,7 @@ it("distinguishes failed search from empty, then retries the same query", async
|
||||
fireEvent.change(screen.getByLabelText("筛选地图记忆"), {
|
||||
target: { value: "确实没有这项" },
|
||||
});
|
||||
fireEvent.submit(screen.getByRole("search"));
|
||||
await screen.findByText("没有匹配的记忆");
|
||||
expect(screen.queryByText("还没有已生成的记忆")).toBeNull();
|
||||
expect(
|
||||
@@ -510,6 +514,7 @@ it("aborts an older query after a source change and never replaces the newer res
|
||||
fireEvent.change(screen.getByLabelText("筛选地图记忆"), {
|
||||
target: { value: "旧搜索" },
|
||||
});
|
||||
fireEvent.submit(screen.getByRole("search"));
|
||||
await waitFor(() => expect(resolve).toHaveLength(1));
|
||||
const signals = vi
|
||||
.mocked(api.getMemoryLibrary)
|
||||
@@ -521,6 +526,7 @@ it("aborts an older query after a source change and never replaces the newer res
|
||||
fireEvent.change(screen.getByLabelText("筛选地图记忆"), {
|
||||
target: { value: "目标" },
|
||||
});
|
||||
fireEvent.submit(screen.getByRole("search"));
|
||||
await within(screen.getByRole("list", { name: "地图记忆列表" })).findByText(
|
||||
"第501条独有目标",
|
||||
);
|
||||
@@ -553,6 +559,7 @@ it("hidden map aborts page reads and does not accept late body or poll until rea
|
||||
fireEvent.change(screen.getByLabelText("筛选地图记忆"), {
|
||||
target: { value: "迟到" },
|
||||
});
|
||||
fireEvent.submit(screen.getByRole("search"));
|
||||
await act(async () => {});
|
||||
const signals = vi
|
||||
.mocked(api.getMemoryLibrary)
|
||||
@@ -611,6 +618,7 @@ it("keeps the populated map while a search is in flight and does not treat it as
|
||||
fireEvent.change(screen.getByLabelText("筛选地图记忆"), {
|
||||
target: { value: "onememory" },
|
||||
});
|
||||
fireEvent.submit(screen.getByRole("search"));
|
||||
await act(async () => {});
|
||||
expect(screen.queryByText("正在读取…")).toBeNull();
|
||||
expect(screen.getByText("已沉淀 501 条 · 当前展示 100 条")).toBeTruthy();
|
||||
@@ -700,6 +708,7 @@ it("explicit refresh keeps the search scope and replaces the page with the lates
|
||||
fireEvent.change(screen.getByLabelText("筛选地图记忆"), {
|
||||
target: { value: "目标" },
|
||||
});
|
||||
fireEvent.submit(screen.getByRole("search"));
|
||||
const list = screen.getByRole("list", { name: "地图记忆列表" });
|
||||
await within(list).findByText("第501条独有目标");
|
||||
vi.mocked(api.getMemoryLibrary).mockResolvedValueOnce({
|
||||
|
||||
@@ -0,0 +1,333 @@
|
||||
// @vitest-environment jsdom
|
||||
import type {
|
||||
GovernedMemory,
|
||||
MemoryGraphResponse,
|
||||
SdkMemoryRecord,
|
||||
} from "@one-memory/contracts";
|
||||
import {
|
||||
cleanup,
|
||||
fireEvent,
|
||||
render,
|
||||
screen,
|
||||
waitFor,
|
||||
} from "@testing-library/react";
|
||||
import { afterEach, expect, it, vi } from "vitest";
|
||||
import { MemorySpace } from "./MemorySpace";
|
||||
import { sdkMapMemory } from "./sdk-memory-map";
|
||||
|
||||
const api = vi.hoisted(() => ({
|
||||
authority: vi.fn(),
|
||||
legacyLibrary: vi.fn(),
|
||||
legacyGraph: vi.fn(),
|
||||
sdkLibrary: vi.fn(),
|
||||
sdkGraph: vi.fn(),
|
||||
sdkMemory: vi.fn(),
|
||||
}));
|
||||
vi.mock("./api", () => ({
|
||||
getMemoryAuthority: api.authority,
|
||||
getMemoryLibrary: api.legacyLibrary,
|
||||
getMemories: async () => [],
|
||||
getMemoryGraph: api.legacyGraph,
|
||||
getSdkMemoryLibrary: api.sdkLibrary,
|
||||
getSdkMemoryGraph: api.sdkGraph,
|
||||
getSdkMemory: api.sdkMemory,
|
||||
updateSdkMemory: vi.fn(),
|
||||
deleteSdkMemory: vi.fn(),
|
||||
getMemoryProcessing: async () => null,
|
||||
getModelProfiles: async () => [],
|
||||
apiErrorCode: () => null,
|
||||
}));
|
||||
vi.mock("./GraphWidthPrompt", () => ({ GraphWidthPrompt: () => null }));
|
||||
vi.mock("./MemoryNoteComposer", () => ({ MemoryNoteComposer: () => null }));
|
||||
vi.mock("./MemoryMap", () => ({
|
||||
MemoryMap: ({
|
||||
hidden,
|
||||
memories,
|
||||
graph,
|
||||
loaded,
|
||||
onOpenMemory,
|
||||
libraryView,
|
||||
}: {
|
||||
hidden: boolean;
|
||||
memories: GovernedMemory[];
|
||||
graph: MemoryGraphResponse | null;
|
||||
loaded: boolean;
|
||||
onOpenMemory(id: string): void;
|
||||
libraryView?: {
|
||||
page: { source_ids_by_memory?: Record<string, string[]> } | null;
|
||||
unattributedSource?: string;
|
||||
onSource(value: string): void;
|
||||
resetPage(): void;
|
||||
};
|
||||
}) => (
|
||||
<div hidden={hidden} data-testid="map">
|
||||
<span data-testid="map-memories">
|
||||
{memories.map((m) => `${m.memory_id}:${m.status}`).join(",")}
|
||||
</span>
|
||||
<span data-testid="map-nodes">{graph?.nodes.length ?? "none"}</span>
|
||||
<span data-testid="map-loaded">{String(loaded)}</span>
|
||||
<span data-testid="map-sources">
|
||||
{memories
|
||||
.map((m) =>
|
||||
(libraryView?.page?.source_ids_by_memory?.[m.memory_id] ?? []).join(
|
||||
"+",
|
||||
),
|
||||
)
|
||||
.join(",")}
|
||||
</span>
|
||||
<span data-testid="map-unattributed">
|
||||
{libraryView?.unattributedSource ?? ""}
|
||||
</span>
|
||||
<button type="button" onClick={() => onOpenMemory("a")}>
|
||||
open a
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
libraryView?.onSource("workbuddy");
|
||||
libraryView?.resetPage();
|
||||
}}
|
||||
>
|
||||
pick workbuddy
|
||||
</button>
|
||||
</div>
|
||||
),
|
||||
}));
|
||||
vi.mock("./SdkMemoryLibrary", () => ({
|
||||
SdkMemoryLibrary: ({
|
||||
focusedMemoryId,
|
||||
}: {
|
||||
focusedMemoryId: string | null;
|
||||
}) => <div data-testid="sdk-library">{focusedMemoryId ?? ""}</div>,
|
||||
}));
|
||||
vi.mock("./MemoryRecallScope", () => ({
|
||||
MemoryRecallScope: () => null,
|
||||
memoryProvenanceLabel: () => "原话",
|
||||
}));
|
||||
|
||||
const record = (memory_id: string): SdkMemoryRecord => ({
|
||||
memory_id,
|
||||
content: `memory ${memory_id}`,
|
||||
layer: "l2_fact",
|
||||
status: "active",
|
||||
is_latest: true,
|
||||
supersedes: [],
|
||||
superseded_by: [],
|
||||
tags: ["工作"],
|
||||
owner: null,
|
||||
source_raw_ids: [],
|
||||
user_id: "local",
|
||||
agent_id: "all-apps",
|
||||
session_id: "s1",
|
||||
source: "cursor",
|
||||
memory_at: null,
|
||||
gmt_created: 1_761_000_000,
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
cleanup();
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
function sdkAuthority() {
|
||||
api.authority.mockResolvedValue({ authority: "sdk", migration_id: "m" });
|
||||
api.sdkLibrary.mockResolvedValue({
|
||||
memories: [record("a"), record("b")],
|
||||
total: 3018,
|
||||
limit: 300,
|
||||
offset: 0,
|
||||
});
|
||||
api.sdkGraph.mockResolvedValue({
|
||||
total_memories: 3018,
|
||||
projected_memories: 2,
|
||||
nodes: [
|
||||
{ node_id: "memory:a", node_type: "memory", label: "a", memory_id: "a" },
|
||||
{ node_id: "memory:b", node_type: "memory", label: "b", memory_id: "b" },
|
||||
{ node_id: "source:cursor", node_type: "source", label: "Cursor" },
|
||||
],
|
||||
edges: [],
|
||||
hy_graph_available: true,
|
||||
});
|
||||
}
|
||||
|
||||
it("draws the map from SDK memories on SDK authority, without legacy reads", async () => {
|
||||
sdkAuthority();
|
||||
|
||||
render(<MemorySpace simplified active />);
|
||||
|
||||
await waitFor(() =>
|
||||
expect(screen.getByTestId("map-loaded").textContent).toBe("true"),
|
||||
);
|
||||
expect(screen.getByTestId("map").hidden).toBe(false);
|
||||
expect(screen.getByTestId("map-memories").textContent).toBe(
|
||||
"a:confirmed,b:confirmed",
|
||||
);
|
||||
expect(screen.getByTestId("map-nodes").textContent).toBe("3");
|
||||
expect(await screen.findByText(/3,018/)).toBeTruthy();
|
||||
expect(api.sdkLibrary).toHaveBeenCalledWith(
|
||||
{
|
||||
statuses: ["active"],
|
||||
source: undefined,
|
||||
query: undefined,
|
||||
offset: 0,
|
||||
limit: 100,
|
||||
},
|
||||
expect.any(AbortSignal),
|
||||
);
|
||||
expect(api.sdkGraph).toHaveBeenCalledWith(
|
||||
{ source: undefined, offset: 0, limit: 100 },
|
||||
expect.any(AbortSignal),
|
||||
);
|
||||
expect(api.legacyLibrary).not.toHaveBeenCalled();
|
||||
expect(api.legacyGraph).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("keeps rows, sources, count and graph on the one filtered page", async () => {
|
||||
sdkAuthority();
|
||||
render(<MemorySpace simplified active />);
|
||||
await waitFor(() =>
|
||||
expect(screen.getByTestId("map-loaded").textContent).toBe("true"),
|
||||
);
|
||||
api.sdkLibrary.mockResolvedValue({
|
||||
memories: [
|
||||
{ ...record("w1"), source: "workbuddy" },
|
||||
{ ...record("w2"), source: "default" },
|
||||
],
|
||||
total: 566,
|
||||
limit: 100,
|
||||
offset: 0,
|
||||
});
|
||||
|
||||
fireEvent.click(screen.getByText("pick workbuddy"));
|
||||
|
||||
await waitFor(() =>
|
||||
expect(screen.getByTestId("map-memories").textContent).toBe(
|
||||
"w1:confirmed,w2:confirmed",
|
||||
),
|
||||
);
|
||||
expect(screen.getByTestId("map-sources").textContent).toBe(
|
||||
"workbuddy,default",
|
||||
);
|
||||
expect(screen.getByTestId("map-unattributed").textContent).toBe("default");
|
||||
expect(api.sdkLibrary).toHaveBeenLastCalledWith(
|
||||
expect.objectContaining({ source: "workbuddy", offset: 0 }),
|
||||
expect.any(AbortSignal),
|
||||
);
|
||||
expect(api.sdkGraph).toHaveBeenLastCalledWith(
|
||||
{ source: "workbuddy", offset: 0, limit: 100 },
|
||||
expect.any(AbortSignal),
|
||||
);
|
||||
expect(await screen.findByText(/566/)).toBeTruthy();
|
||||
});
|
||||
|
||||
it("opens an SDK memory in place instead of switching away from the map", async () => {
|
||||
sdkAuthority();
|
||||
api.sdkMemory.mockResolvedValue(record("a"));
|
||||
|
||||
render(<MemorySpace simplified active />);
|
||||
await waitFor(() =>
|
||||
expect(screen.getByTestId("map-loaded").textContent).toBe("true"),
|
||||
);
|
||||
fireEvent.click(screen.getByText("open a"));
|
||||
|
||||
await waitFor(() => expect(api.sdkMemory).toHaveBeenCalledWith("a"));
|
||||
// The inspector stays on the map; the reader is never moved to a second list.
|
||||
expect(screen.getByTestId("map").hidden).toBe(false);
|
||||
expect(screen.getByTestId("sdk-library").textContent).toBe("");
|
||||
});
|
||||
|
||||
it("keeps asking for the authority until Core answers", async () => {
|
||||
vi.useFakeTimers({ shouldAdvanceTime: true });
|
||||
try {
|
||||
sdkAuthority();
|
||||
api.authority
|
||||
.mockRejectedValueOnce(new Error("core_not_ready"))
|
||||
.mockRejectedValueOnce(new Error("core_not_ready"));
|
||||
|
||||
render(<MemorySpace simplified active />);
|
||||
await vi.advanceTimersByTimeAsync(2_000);
|
||||
|
||||
await waitFor(() => expect(api.authority).toHaveBeenCalledTimes(3));
|
||||
await waitFor(() =>
|
||||
expect(screen.getByTestId("map-loaded").textContent).toBe("true"),
|
||||
);
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
|
||||
it("projects SDK records into the fields the map groups by", () => {
|
||||
const coding = sdkMapMemory({ ...record("c"), layer: "coding_experience" });
|
||||
expect(coding).toMatchObject({
|
||||
memory_id: "c",
|
||||
kind: "coding",
|
||||
status: "confirmed",
|
||||
tags: ["工作"],
|
||||
deleted_at: null,
|
||||
});
|
||||
expect(sdkMapMemory({ ...record("s"), layer: "l6_schema" }).kind).toBe(
|
||||
"schema",
|
||||
);
|
||||
expect(sdkMapMemory({ ...record("x"), status: "superseded" }).status).toBe(
|
||||
"superseded",
|
||||
);
|
||||
});
|
||||
|
||||
const node = (
|
||||
node_id: string,
|
||||
node_type: MemoryGraphResponse["nodes"][number]["node_type"],
|
||||
extra: Partial<MemoryGraphResponse["nodes"][number]> = {},
|
||||
): MemoryGraphResponse["nodes"][number] => ({
|
||||
node_id,
|
||||
node_type,
|
||||
label: node_id,
|
||||
memory_id: null,
|
||||
source_id: null,
|
||||
kind: null,
|
||||
status: null,
|
||||
confidence: null,
|
||||
...extra,
|
||||
});
|
||||
const provenance = (
|
||||
from: string,
|
||||
to: string,
|
||||
): MemoryGraphResponse["edges"][number] => ({
|
||||
edge_id: `${from}->${to}`,
|
||||
source_node_id: from,
|
||||
target_node_id: to,
|
||||
relation: "evidence_source",
|
||||
weight: null,
|
||||
origin: "provenance",
|
||||
});
|
||||
|
||||
it("groups and draws SDK memories by source with the real map model", async () => {
|
||||
const { mapGroups, projectMap } = await import("./memory-map-model");
|
||||
const memories = [record("a"), record("b")].map(sdkMapMemory);
|
||||
const graph: MemoryGraphResponse = {
|
||||
total_memories: 2,
|
||||
projected_memories: 2,
|
||||
nodes: [
|
||||
node("memory:a", "memory", { memory_id: "a" }),
|
||||
node("memory:b", "memory", { memory_id: "b" }),
|
||||
node("source:cursor", "source", { source_id: "cursor", label: "Cursor" }),
|
||||
node("hy:t1", "topic"),
|
||||
],
|
||||
edges: [
|
||||
provenance("source:cursor", "memory:a"),
|
||||
provenance("source:cursor", "memory:b"),
|
||||
],
|
||||
hy_graph_available: true,
|
||||
};
|
||||
|
||||
const groups = mapGroups(memories, graph, "source");
|
||||
expect(groups.map((group) => [group.key, group.memoryIds])).toEqual([
|
||||
["source:cursor", ["a", "b"]],
|
||||
]);
|
||||
const drawn = projectMap(graph, memories, "source", ["a", "b"]);
|
||||
expect(
|
||||
drawn.nodes
|
||||
.filter((node) => node.node_type === "memory")
|
||||
.map((n) => n.memory_id),
|
||||
).toEqual(["a", "b"]);
|
||||
expect(drawn.edges.length).toBeGreaterThan(0);
|
||||
});
|
||||
@@ -5,6 +5,7 @@ import type {
|
||||
ImportJob,
|
||||
MemoryDetail,
|
||||
MemoryGraphResponse,
|
||||
MemoryLibraryPage,
|
||||
MemoryProcessingSummary,
|
||||
SdkAskMemory,
|
||||
SourceResponse,
|
||||
@@ -32,6 +33,7 @@ import {
|
||||
apiErrorCode,
|
||||
askMemoriesStream,
|
||||
correctMemory,
|
||||
deleteSdkMemory,
|
||||
digestMemoryGraph,
|
||||
getMemories,
|
||||
getMemory,
|
||||
@@ -39,12 +41,16 @@ import {
|
||||
getMemoryGraph,
|
||||
getMemoryProcessing,
|
||||
getModelProfiles,
|
||||
getSdkMemory,
|
||||
getSdkMemoryGraph,
|
||||
getSdkMemoryLibrary,
|
||||
manualExtract,
|
||||
mutateMemory,
|
||||
retryUnknownMemoryProcessing,
|
||||
saveMemoryNote,
|
||||
setExtractionWindow,
|
||||
syncMemoryIndex,
|
||||
updateSdkMemory,
|
||||
} from "./api";
|
||||
import { ConversationSearch } from "./ConversationSearch";
|
||||
import { clockStamp } from "./clock-stamp";
|
||||
@@ -56,7 +62,8 @@ import { subscribePoll } from "./poll-scheduler";
|
||||
import { R6Button } from "./R6Controls";
|
||||
import { SdkMemoryLibrary } from "./SdkMemoryLibrary";
|
||||
import { SdkRawHistory } from "./SdkRawHistory";
|
||||
import { sourceNameOr } from "./source-label";
|
||||
import { sdkDetail, sdkLibraryPage } from "./sdk-memory-map";
|
||||
import { sourceNameOr, UNATTRIBUTED_SOURCE_ID } from "./source-label";
|
||||
import { traceUiAction } from "./ui-trace";
|
||||
import { useMemoryLibrary } from "./use-memory-library";
|
||||
import { isDefiniteRefusal, writeTraceStatus } from "./write-outcome";
|
||||
@@ -168,13 +175,23 @@ export function MemorySpace({
|
||||
if ((!active && !preload) || typeof getMemoryAuthority !== "function")
|
||||
return;
|
||||
let current = true;
|
||||
void getMemoryAuthority()
|
||||
.then((result) => {
|
||||
if (current) setMemoryAuthority(result.authority);
|
||||
})
|
||||
.catch(() => {});
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
// The app opens before Core is ready, so the first ask usually fails.
|
||||
// Without a retry the authority stayed unknown and nothing ever loaded.
|
||||
const ask = (delay: number) => {
|
||||
void getMemoryAuthority()
|
||||
.then((result) => {
|
||||
if (current) setMemoryAuthority(result.authority);
|
||||
})
|
||||
.catch(() => {
|
||||
if (!current) return;
|
||||
timer = setTimeout(() => ask(Math.min(delay * 2, 10_000)), delay);
|
||||
});
|
||||
};
|
||||
ask(500);
|
||||
return () => {
|
||||
current = false;
|
||||
if (timer !== undefined) clearTimeout(timer);
|
||||
};
|
||||
}, [active, preload]);
|
||||
useEffect(() => {
|
||||
@@ -283,9 +300,125 @@ export function MemorySpace({
|
||||
text: string;
|
||||
memories: AskMemory[];
|
||||
} | null>(null);
|
||||
const [sdkFocusedMemoryId, setSdkFocusedMemoryId] = useState<string | null>(
|
||||
null,
|
||||
);
|
||||
// The SDK library is no longer where a memory opens from the map; the id is
|
||||
// kept so the panel still has a focus target for its own tests.
|
||||
const [sdkFocusedMemoryId] = useState<string | null>(null);
|
||||
// The map on SDK authority. Legacy reads (library page, governed graph) are
|
||||
// refused there. Like the governed map, the list, its count, its source and
|
||||
// time labels and the graph are all one page: the library page and the graph
|
||||
// for that same source and window, read together. Reading the graph from a
|
||||
// different, unfiltered window made rows, sources and totals disagree.
|
||||
const [sdkGraph, setSdkGraph] = useState<MemoryGraphResponse | null>(null);
|
||||
// A write on SDK authority has no governed refresh to ride on; bumping this
|
||||
// re-reads the SDK page and graph after an edit or a delete.
|
||||
const [sdkMapRevision, setSdkMapRevision] = useState(0);
|
||||
const bumpSdkMap = useCallback(() => setSdkMapRevision((v) => v + 1), []);
|
||||
const [sdkLibrary, setSdkLibrary] = useState<{
|
||||
page: MemoryLibraryPage | null;
|
||||
loading: boolean;
|
||||
failed: boolean;
|
||||
/** The filter the shown page was read with, for the count above it. */
|
||||
settledSource: string;
|
||||
settledQuery: string;
|
||||
}>({
|
||||
page: null,
|
||||
loading: false,
|
||||
failed: false,
|
||||
settledSource: "",
|
||||
settledQuery: "",
|
||||
});
|
||||
const [sdkLibraryQuery, setSdkLibraryQuery] = useState("");
|
||||
const [sdkLibrarySource, setSdkLibrarySource] = useState("");
|
||||
const [sdkLibraryOffset, setSdkLibraryOffset] = useState(0);
|
||||
const sdkMapVisible =
|
||||
memoryAuthority === "sdk" &&
|
||||
simplified &&
|
||||
active &&
|
||||
displayedMode === "map";
|
||||
useEffect(() => {
|
||||
if (!sdkMapVisible) return;
|
||||
// A write bumps the revision; reading it here re-runs the read at once
|
||||
// instead of waiting for the 30s poll.
|
||||
void sdkMapRevision;
|
||||
let cancelled = false;
|
||||
let controller: AbortController | null = null;
|
||||
const source = sdkLibrarySource || undefined;
|
||||
const query = sdkLibraryQuery || undefined;
|
||||
const read = async (first: boolean) => {
|
||||
controller?.abort();
|
||||
const request = new AbortController();
|
||||
controller = request;
|
||||
if (first) setSdkLibrary((current) => ({ ...current, loading: true }));
|
||||
try {
|
||||
const [page, graph] = await Promise.all([
|
||||
getSdkMemoryLibrary(
|
||||
{
|
||||
statuses: ["active"],
|
||||
source,
|
||||
query,
|
||||
offset: sdkLibraryOffset,
|
||||
limit: 100,
|
||||
},
|
||||
request.signal,
|
||||
),
|
||||
// The overlay is decoration: the list stands without it.
|
||||
getSdkMemoryGraph(
|
||||
{ source, offset: query ? 0 : sdkLibraryOffset, limit: 100 },
|
||||
request.signal,
|
||||
).catch(() => null),
|
||||
]);
|
||||
if (cancelled) return;
|
||||
setSdkLibrary({
|
||||
page: sdkLibraryPage(page),
|
||||
loading: false,
|
||||
failed: false,
|
||||
settledSource: sdkLibrarySource,
|
||||
settledQuery: sdkLibraryQuery,
|
||||
});
|
||||
setSdkGraph(graph);
|
||||
} catch {
|
||||
if (cancelled || request.signal.aborted) return;
|
||||
setSdkLibrary((current) => ({
|
||||
...current,
|
||||
loading: false,
|
||||
failed: true,
|
||||
}));
|
||||
}
|
||||
};
|
||||
void read(true);
|
||||
const timer = setInterval(() => void read(false), 30_000);
|
||||
return () => {
|
||||
cancelled = true;
|
||||
clearInterval(timer);
|
||||
controller?.abort();
|
||||
};
|
||||
}, [
|
||||
sdkMapVisible,
|
||||
sdkMapRevision,
|
||||
sdkLibraryOffset,
|
||||
sdkLibraryQuery,
|
||||
sdkLibrarySource,
|
||||
]);
|
||||
const sdkFilterLabel = [
|
||||
sdkLibrary.settledSource
|
||||
? `${t("memory:来源")} ${
|
||||
sdkLibrary.settledSource === UNATTRIBUTED_SOURCE_ID
|
||||
? t("memory:未识别来源")
|
||||
: sourceNameOr(
|
||||
sources,
|
||||
sdkLibrary.settledSource,
|
||||
sdkLibrary.settledSource,
|
||||
)
|
||||
}`
|
||||
: "",
|
||||
sdkLibrary.settledQuery.trim()
|
||||
? t("memory:关键词“{query}”", {
|
||||
query: sdkLibrary.settledQuery.trim(),
|
||||
})
|
||||
: "",
|
||||
]
|
||||
.filter(Boolean)
|
||||
.join(" · ");
|
||||
const [draft, setDraft] = useState("");
|
||||
const [selection, setSelection] = useState<Set<string>>(new Set());
|
||||
const bulkInFlight = useRef(false);
|
||||
@@ -969,8 +1102,9 @@ export function MemorySpace({
|
||||
open(memory.memory_id);
|
||||
return;
|
||||
}
|
||||
setSdkFocusedMemoryId(memory.memory_id);
|
||||
setMode("library");
|
||||
// An SDK answer cites SDK memories; they open in the inspector like any
|
||||
// other rather than switching the reader to a second list.
|
||||
void open(memory.memory_id);
|
||||
}
|
||||
|
||||
async function saveAnswer(draft: { content: string }) {
|
||||
@@ -1019,7 +1153,13 @@ export function MemorySpace({
|
||||
setDetailFreshness(null);
|
||||
setDetailLoad({ memoryId, status: "loading" });
|
||||
try {
|
||||
const result = await getMemory(memoryId);
|
||||
// Governed detail is refused on SDK authority, so the inspector reads the
|
||||
// same memory through v2 and renders it in place. Staying here keeps the
|
||||
// map's list, graph and detail where the design put them.
|
||||
const result =
|
||||
memoryAuthority === "sdk"
|
||||
? sdkDetail(await getSdkMemory(memoryId))
|
||||
: await getMemory(memoryId);
|
||||
if (epoch !== detailEpoch.current || result.memory.memory_id !== memoryId)
|
||||
return;
|
||||
setDetail(result);
|
||||
@@ -1067,6 +1207,28 @@ export function MemorySpace({
|
||||
return;
|
||||
setBusy(true);
|
||||
setError(null);
|
||||
if (memoryAuthority === "sdk") {
|
||||
// An SDK memory is removed outright: there is no revision to undo.
|
||||
if (actionName !== "delete") {
|
||||
setBusy(false);
|
||||
return;
|
||||
}
|
||||
const trace = traceUiAction("memory", "confirm");
|
||||
try {
|
||||
await deleteSdkMemory(detail.memory.memory_id);
|
||||
trace.done("succeeded");
|
||||
closeDetail();
|
||||
bumpSdkMap();
|
||||
} catch (reason) {
|
||||
trace.done(writeTraceStatus(reason), {
|
||||
reason: "memory_action_failed",
|
||||
});
|
||||
setError(message(reason, t));
|
||||
} finally {
|
||||
setBusy(false);
|
||||
}
|
||||
return;
|
||||
}
|
||||
const intent = `${detail.memory.memory_id}:${actionName}:${detail.memory.revision}`;
|
||||
const requestId = requestIdFor(requestIds.current, intent);
|
||||
const trace = traceUiAction("memory", "confirm");
|
||||
@@ -1273,6 +1435,28 @@ export function MemorySpace({
|
||||
return;
|
||||
setBusy(true);
|
||||
setError(null);
|
||||
if (memoryAuthority === "sdk") {
|
||||
// A rewrite replaces the SDK row instead of adding a revision, so there
|
||||
// is no revision to send and nothing to reconcile afterwards.
|
||||
const trace = traceUiAction("memory", "submit");
|
||||
try {
|
||||
const saved = await updateSdkMemory(detail.memory.memory_id, draft);
|
||||
trace.done("succeeded");
|
||||
setDetail(sdkDetail(saved));
|
||||
setDraft(saved.content);
|
||||
bumpSdkMap();
|
||||
} catch (reason) {
|
||||
trace.done(writeTraceStatus(reason), {
|
||||
reason: isDefiniteRefusal(reason)
|
||||
? "memory_correct_failed"
|
||||
: "memory_correct_result_unknown",
|
||||
});
|
||||
setError(message(reason, t));
|
||||
} finally {
|
||||
setBusy(false);
|
||||
}
|
||||
return;
|
||||
}
|
||||
const intent = `${detail.memory.memory_id}:correct:${detail.memory.revision}`;
|
||||
const requestId = requestIdFor(requestIds.current, intent);
|
||||
const trace = traceUiAction("memory", "submit");
|
||||
@@ -1444,39 +1628,62 @@ export function MemorySpace({
|
||||
{simplified ? t("memory:记忆") : t("memory:记忆空间")}
|
||||
</h1>
|
||||
<span className="memory-loaded-count">
|
||||
{simplified && !management
|
||||
? mapLibrary.failed
|
||||
? t("memory:记忆读取失败")
|
||||
: mapLibrary.stale
|
||||
? t("memory:正在导入历史会话,记忆将陆续出现")
|
||||
: mapLibrary.page
|
||||
? tn(
|
||||
"memory:已沉淀 {count} 条{filter} · 当前展示 {shown} 条",
|
||||
mapLibrary.page.total,
|
||||
{
|
||||
filter: settledFilterLabel
|
||||
? t("memory:({filter})", {
|
||||
filter: settledFilterLabel,
|
||||
})
|
||||
: "",
|
||||
shown: new Intl.NumberFormat(locale).format(
|
||||
mapLibrary.page.items.length,
|
||||
),
|
||||
},
|
||||
)
|
||||
: t("memory:正在读取…")
|
||||
: loaded
|
||||
{memoryAuthority === "sdk" && simplified && !management
|
||||
? sdkLibrary.page
|
||||
? tn(
|
||||
"memory:已加载 {count} 条记忆",
|
||||
memories.filter((m) => m.status !== "deleted").length,
|
||||
"memory:已沉淀 {count} 条{filter} · 当前展示 {shown} 条",
|
||||
sdkLibrary.page.total,
|
||||
{
|
||||
filter: sdkFilterLabel
|
||||
? t("memory:({filter})", { filter: sdkFilterLabel })
|
||||
: "",
|
||||
shown: new Intl.NumberFormat(locale).format(
|
||||
sdkLibrary.page.items.length,
|
||||
),
|
||||
},
|
||||
)
|
||||
: t("memory:正在读取…")}
|
||||
: // The list below says a failed read once, with its retry.
|
||||
sdkLibrary.failed
|
||||
? null
|
||||
: t("memory:正在读取…")
|
||||
: memoryAuthority === "sdk"
|
||||
? null
|
||||
: simplified && !management
|
||||
? mapLibrary.failed
|
||||
? t("memory:记忆读取失败")
|
||||
: mapLibrary.stale
|
||||
? t("memory:正在导入历史会话,记忆将陆续出现")
|
||||
: mapLibrary.page
|
||||
? tn(
|
||||
"memory:已沉淀 {count} 条{filter} · 当前展示 {shown} 条",
|
||||
mapLibrary.page.total,
|
||||
{
|
||||
filter: settledFilterLabel
|
||||
? t("memory:({filter})", {
|
||||
filter: settledFilterLabel,
|
||||
})
|
||||
: "",
|
||||
shown: new Intl.NumberFormat(locale).format(
|
||||
mapLibrary.page.items.length,
|
||||
),
|
||||
},
|
||||
)
|
||||
: t("memory:正在读取…")
|
||||
: loaded
|
||||
? tn(
|
||||
"memory:已加载 {count} 条记忆",
|
||||
memories.filter((m) => m.status !== "deleted").length,
|
||||
)
|
||||
: t("memory:正在读取…")}
|
||||
</span>
|
||||
</div>
|
||||
<MemoryNoteComposer
|
||||
suspended={!active}
|
||||
onPendingChange={setPendingNote}
|
||||
onSaved={() => {
|
||||
// A finished note is a new SDK memory: re-read the page now
|
||||
// rather than leave it for the next poll.
|
||||
if (memoryAuthority === "sdk") bumpSdkMap();
|
||||
void refresh(query).catch((reason) => setError(message(reason, t)));
|
||||
}}
|
||||
onOpenMemory={(id) => {
|
||||
@@ -1654,9 +1861,11 @@ export function MemorySpace({
|
||||
<ModelPrompt
|
||||
processing={processing}
|
||||
hasMemories={
|
||||
simplified
|
||||
? knownMapTotal.current > 0 || (mapLibrary.page?.total ?? 0) > 0
|
||||
: memories.length > 0
|
||||
memoryAuthority === "sdk"
|
||||
? (sdkLibrary.page?.total ?? 0) > 0
|
||||
: simplified
|
||||
? knownMapTotal.current > 0 || (mapLibrary.page?.total ?? 0) > 0
|
||||
: memories.length > 0
|
||||
}
|
||||
onConfigureModel={onConfigureModel}
|
||||
/>
|
||||
@@ -1814,23 +2023,38 @@ export function MemorySpace({
|
||||
focusDialogue={focusDialogue}
|
||||
onFocusDialogueConsumed={onFocusDialogueConsumed}
|
||||
libraryView={
|
||||
simplified
|
||||
simplified && memoryAuthority === "sdk"
|
||||
? {
|
||||
query: mapQuery,
|
||||
source: mapSource,
|
||||
page: mapLibrary.page,
|
||||
loading: mapLibrary.loading,
|
||||
failed: mapLibrary.failed,
|
||||
stale: sourceSnapshotStale,
|
||||
onQuery: setMapQuery,
|
||||
onSource: setMapSource,
|
||||
resetPage: () => mapLibrary.go(0),
|
||||
onPage: mapLibrary.go,
|
||||
onRetry: () => {
|
||||
void mapLibrary.refresh().catch(() => {});
|
||||
},
|
||||
query: sdkLibraryQuery,
|
||||
source: sdkLibrarySource,
|
||||
page: sdkLibrary.page,
|
||||
loading: sdkLibrary.loading,
|
||||
failed: sdkLibrary.failed,
|
||||
stale: false,
|
||||
unattributedSource: UNATTRIBUTED_SOURCE_ID,
|
||||
onQuery: setSdkLibraryQuery,
|
||||
onSource: setSdkLibrarySource,
|
||||
resetPage: () => setSdkLibraryOffset(0),
|
||||
onPage: (offset) => setSdkLibraryOffset(offset),
|
||||
onRetry: bumpSdkMap,
|
||||
}
|
||||
: undefined
|
||||
: simplified
|
||||
? {
|
||||
query: mapQuery,
|
||||
source: mapSource,
|
||||
page: mapLibrary.page,
|
||||
loading: mapLibrary.loading,
|
||||
failed: mapLibrary.failed,
|
||||
stale: sourceSnapshotStale,
|
||||
onQuery: setMapQuery,
|
||||
onSource: setMapSource,
|
||||
resetPage: () => mapLibrary.go(0),
|
||||
onPage: mapLibrary.go,
|
||||
onRetry: () => {
|
||||
void mapLibrary.refresh().catch(() => {});
|
||||
},
|
||||
}
|
||||
: undefined
|
||||
}
|
||||
onMutationComplete={(id) => {
|
||||
void refresh(query)
|
||||
@@ -1838,19 +2062,39 @@ export function MemorySpace({
|
||||
.catch((reason) => setError(message(reason, t)));
|
||||
}}
|
||||
hidden={!active || displayedMode !== "map"}
|
||||
memories={simplified ? (mapLibrary.page?.items ?? []) : memories}
|
||||
graph={simplified && sourceSnapshotStale ? null : graph}
|
||||
memories={
|
||||
memoryAuthority === "sdk"
|
||||
? (sdkLibrary.page?.items ?? [])
|
||||
: simplified
|
||||
? (mapLibrary.page?.items ?? [])
|
||||
: memories
|
||||
}
|
||||
graph={
|
||||
memoryAuthority === "sdk"
|
||||
? sdkGraph
|
||||
: simplified && sourceSnapshotStale
|
||||
? null
|
||||
: graph
|
||||
}
|
||||
sources={sources}
|
||||
processing={processing}
|
||||
arrival={memoryArrival}
|
||||
onViewHistory={simplified ? onConnect : () => setMode("history")}
|
||||
loaded={simplified ? mapLibrary.page !== null : loaded}
|
||||
loaded={
|
||||
memoryAuthority === "sdk"
|
||||
? sdkLibrary.page !== null
|
||||
: simplified
|
||||
? mapLibrary.page !== null
|
||||
: loaded
|
||||
}
|
||||
detail={detail}
|
||||
detailLoad={detailLoad}
|
||||
overlayBlocked={
|
||||
!active ||
|
||||
deleteCount > 0 ||
|
||||
(simplified && (!mapLibrary.page || sourceSnapshotStale))
|
||||
(memoryAuthority === "sdk"
|
||||
? sdkLibrary.page === null
|
||||
: simplified && (!mapLibrary.page || sourceSnapshotStale))
|
||||
}
|
||||
draft={draft}
|
||||
busy={busy || detailReadBlocked}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
// @vitest-environment jsdom
|
||||
|
||||
import type { MemoryDetail } from "@one-memory/contracts";
|
||||
import type { MemoryDetail, SourceResponse } from "@one-memory/contracts";
|
||||
import {
|
||||
act,
|
||||
cleanup,
|
||||
@@ -14,7 +14,7 @@ import { afterEach, beforeEach, expect, it, vi } from "vitest";
|
||||
import { LANGUAGE_STORAGE_KEY, LanguageProvider } from "./gui-language";
|
||||
import { SelectedMemoryDetail } from "./SelectedMemoryDetail";
|
||||
|
||||
it("localizes informational details and dates while preserving evidence and memory", async () => {
|
||||
it("localizes memory details without exposing the retired evidence surface", async () => {
|
||||
delete window.oneMemory;
|
||||
localStorage.setItem(LANGUAGE_STORAGE_KEY, "es");
|
||||
render(
|
||||
@@ -29,9 +29,9 @@ it("localizes informational details and dates while preserving evidence and memo
|
||||
);
|
||||
expect(await screen.findByText("Guardada · Versión 2")).toBeTruthy();
|
||||
expect(
|
||||
screen.getByRole("heading", { name: "Evidencia original" }),
|
||||
).toBeTruthy();
|
||||
expect(screen.getByText("这次先确认范围。")).toBeTruthy();
|
||||
screen.queryByRole("heading", { name: "Evidencia original" }),
|
||||
).toBeNull();
|
||||
expect(screen.queryByText("这次先确认范围。")).toBeNull();
|
||||
expect(correct).not.toHaveBeenCalled();
|
||||
expect(remove).not.toHaveBeenCalled();
|
||||
localStorage.clear();
|
||||
@@ -145,6 +145,72 @@ it.each([
|
||||
expect(screen.getByText(`${label} · 第 2 版`)).toBeTruthy();
|
||||
expect(screen.queryByText("待核对 · 第 2 版")).toBeNull();
|
||||
});
|
||||
it("shows the native SDK state without inventing a revision", () => {
|
||||
render(
|
||||
<SelectedMemoryDetail
|
||||
detail={{
|
||||
...detail,
|
||||
memory: { ...detail.memory, revision: 0 },
|
||||
evidence: [],
|
||||
sdk_status: "shadow",
|
||||
source_id: "cursor",
|
||||
sdk_memory_at: "2026-09-18T00:00:00Z",
|
||||
sdk_created_at: "2026-09-19T00:00:00Z",
|
||||
}}
|
||||
sources={[
|
||||
{ source_id: "cursor", display_name: "Cursor" } as SourceResponse,
|
||||
]}
|
||||
onChanged={vi.fn()}
|
||||
hideActions
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByText("已隐藏")).toBeTruthy();
|
||||
expect(screen.queryByText(/第 0 版/)).toBeNull();
|
||||
expect(screen.getByText(/^来源 Cursor · /)).toBeTruthy();
|
||||
});
|
||||
it("splits long memory text into a bounded heading and readable body", () => {
|
||||
const content = "长".repeat(100);
|
||||
render(
|
||||
<SelectedMemoryDetail
|
||||
detail={{
|
||||
...detail,
|
||||
memory: { ...detail.memory, content },
|
||||
evidence: [],
|
||||
sdk_status: "active",
|
||||
}}
|
||||
sources={[]}
|
||||
onChanged={vi.fn()}
|
||||
hideActions
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByRole("heading", { name: "长".repeat(72) })).toBeTruthy();
|
||||
expect(screen.getByText("长".repeat(28)).className).toContain(
|
||||
"selected-memory-body",
|
||||
);
|
||||
});
|
||||
it.each([
|
||||
[
|
||||
"a late sentence end",
|
||||
`${"长".repeat(80)}。${"尾".repeat(10)}`,
|
||||
"长".repeat(72),
|
||||
],
|
||||
["words", `${"word ".repeat(20)}end`, "word ".repeat(14).trim()],
|
||||
])("keeps the heading within the limit for %s", (_case, content, heading) => {
|
||||
render(
|
||||
<SelectedMemoryDetail
|
||||
detail={{
|
||||
...detail,
|
||||
memory: { ...detail.memory, content },
|
||||
evidence: [],
|
||||
sdk_status: "active",
|
||||
}}
|
||||
sources={[]}
|
||||
onChanged={vi.fn()}
|
||||
hideActions
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByRole("heading", { name: heading })).toBeTruthy();
|
||||
});
|
||||
it("opens full editor and saves content without rewriting scope", async () => {
|
||||
const changed = vi.fn();
|
||||
render(
|
||||
@@ -204,27 +270,7 @@ it("keeps draft after conflict until latest revision is explicitly acknowledged"
|
||||
),
|
||||
);
|
||||
});
|
||||
it("original page never substitutes unavailable evidence with stale text", () => {
|
||||
const view = render(
|
||||
<SelectedMemoryDetail detail={detail} sources={[]} onChanged={() => {}} />,
|
||||
);
|
||||
fireEvent.click(screen.getByRole("button", { name: "查看原文" }));
|
||||
expect(screen.getAllByText("这次先确认范围。").length).toBeGreaterThan(0);
|
||||
view.rerender(
|
||||
<SelectedMemoryDetail
|
||||
detail={{
|
||||
...detail,
|
||||
evidence: [{ ...detail.evidence[0], available: false }],
|
||||
}}
|
||||
sources={[]}
|
||||
onChanged={() => {}}
|
||||
/>,
|
||||
);
|
||||
expect(screen.queryByText("这次先确认范围。")).toBeNull();
|
||||
expect(screen.getByText("来源权限已变化或原文不可用。")).toBeTruthy();
|
||||
});
|
||||
|
||||
it.each(["编辑", "查看原文", "删除记忆"])(
|
||||
it.each(["编辑", "删除记忆"])(
|
||||
"returns focus to %s without scrolling the background",
|
||||
(name) => {
|
||||
render(
|
||||
@@ -275,7 +321,6 @@ it.each(["escape", "cancel"])(
|
||||
"value",
|
||||
"尚未保存的草稿",
|
||||
);
|
||||
expect(screen.getByText("这次先确认范围。")).toBeTruthy();
|
||||
expect(correct).not.toHaveBeenCalled();
|
||||
confirm.mockReturnValue(true);
|
||||
close();
|
||||
@@ -535,7 +580,7 @@ it("suspends the editor without losing its draft or leaving the background inert
|
||||
);
|
||||
});
|
||||
|
||||
it("stamps today's source time the same way as 06A and 06E", () => {
|
||||
it("stamps today's source and conversation time without an evidence page", () => {
|
||||
render(
|
||||
<SelectedMemoryDetail
|
||||
detail={{
|
||||
@@ -552,8 +597,6 @@ it("stamps today's source time the same way as 06A and 06E", () => {
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByText(/来源 workbuddy · 今天 /)).toBeTruthy();
|
||||
expect(screen.getByText(/原对话 今天 /)).toBeTruthy();
|
||||
fireEvent.click(screen.getByRole("button", { name: "查看原文" }));
|
||||
expect(screen.getByRole("heading", { name: "原始对话" })).toBeTruthy();
|
||||
expect(screen.getByText(/^workbuddy · 今天 /)).toBeTruthy();
|
||||
expect(screen.getByText(/对话时间 今天 /)).toBeTruthy();
|
||||
expect(screen.queryByRole("button", { name: "查看原文" })).toBeNull();
|
||||
});
|
||||
@@ -1,16 +1,22 @@
|
||||
import type { MemoryDetail, SourceResponse } from "@one-memory/contracts";
|
||||
import type {
|
||||
MemoryDetail,
|
||||
SdkMemoryStatus,
|
||||
SourceResponse,
|
||||
} from "@one-memory/contracts";
|
||||
import { useEffect, useRef, useState } from "react";
|
||||
import { createPortal } from "react-dom";
|
||||
import { apiErrorCode, correctMemory, getMemory, mutateMemory } from "./api";
|
||||
import { clockStamp } from "./clock-stamp";
|
||||
import { useGuiLanguage } from "./gui-language";
|
||||
import { R6Button } from "./R6Controls";
|
||||
import { sourceNameOr } from "./source-label";
|
||||
import { sourceNameOr, UNATTRIBUTED_SOURCE_ID } from "./source-label";
|
||||
import { traceUiAction } from "./ui-trace";
|
||||
import { isDefiniteRefusal, writeTraceStatus } from "./write-outcome";
|
||||
import "./selected-memory.css";
|
||||
|
||||
type Page = "read" | "edit" | "original" | "delete";
|
||||
type Page = "read" | "edit" | "delete";
|
||||
|
||||
const HEADING_LIMIT = 72;
|
||||
|
||||
export function SelectedMemoryDetail({
|
||||
detail,
|
||||
@@ -50,19 +56,28 @@ export function SelectedMemoryDetail({
|
||||
}, []);
|
||||
const memory = detail.memory;
|
||||
const evidence = detail.evidence[0];
|
||||
const statusLabel =
|
||||
memory.status === "confirmed"
|
||||
const sdkStatusLabels: Record<SdkMemoryStatus, Parameters<typeof t>[0]> = {
|
||||
active: "selectedMemory.status.active",
|
||||
superseded: "selectedMemory.status.superseded",
|
||||
shadow: "selectedMemory.status.shadow",
|
||||
negated: "selectedMemory.status.negated",
|
||||
conflicted: "selectedMemory.status.conflicted",
|
||||
archived: "selectedMemory.status.archived",
|
||||
};
|
||||
const statusLabel = detail.sdk_status
|
||||
? t(sdkStatusLabels[detail.sdk_status])
|
||||
: memory.status === "confirmed"
|
||||
? t("selectedMemory.status.confirmed")
|
||||
: memory.status === "deleted"
|
||||
? t("selectedMemory.status.deleted")
|
||||
: memory.status === "superseded"
|
||||
? t("selectedMemory.status.superseded")
|
||||
: t("selectedMemory.status.unknown");
|
||||
const source = sourceNameOr(
|
||||
sources,
|
||||
evidence?.source_id,
|
||||
t("selectedMemory.sourceUnknown"),
|
||||
);
|
||||
const sourceId = detail.source_id ?? evidence?.source_id;
|
||||
const source =
|
||||
detail.sdk_status && sourceId === UNATTRIBUTED_SOURCE_ID
|
||||
? t("selectedMemory.sourceUnattributed")
|
||||
: sourceNameOr(sources, sourceId, t("selectedMemory.sourceUnknown"));
|
||||
const dirty = page === "edit" && draft !== memory.content;
|
||||
useEffect(() => {
|
||||
if (page === "edit" && memory.revision > base) setConflictFound(true);
|
||||
@@ -221,20 +236,22 @@ export function SelectedMemoryDetail({
|
||||
if (alive.current) setBusy(false);
|
||||
}
|
||||
}
|
||||
const author =
|
||||
evidence?.written_by === "person" || evidence?.written_by === "user"
|
||||
? t("selectedMemory.author.you")
|
||||
: evidence?.written_by === "agent" || evidence?.written_by === "assistant"
|
||||
? t("selectedMemory.author.assistant")
|
||||
: t("selectedMemory.author.unknown");
|
||||
const timestamp = clockStamp(evidence?.created_at, locale);
|
||||
const split = memory.content.search(/[。\n]/);
|
||||
// The heading is the first sentence when it fits in HEADING_LIMIT
|
||||
// characters; otherwise the text breaks at the last space within the limit
|
||||
// (CJK text has none, so it breaks at the limit itself).
|
||||
const punctuation = memory.content.search(/[。!?!?\n]/);
|
||||
const space = memory.content.lastIndexOf(" ", HEADING_LIMIT);
|
||||
const split =
|
||||
punctuation > 0 && punctuation < HEADING_LIMIT
|
||||
? punctuation + 1
|
||||
: memory.content.length > HEADING_LIMIT
|
||||
? space > HEADING_LIMIT / 2
|
||||
? space
|
||||
: HEADING_LIMIT
|
||||
: -1;
|
||||
const heading =
|
||||
split > 0 && split < 85
|
||||
? memory.content.slice(0, split + 1)
|
||||
: memory.content;
|
||||
const body =
|
||||
split > 0 && split < 85 ? memory.content.slice(split + 1).trim() : "";
|
||||
split > 0 ? memory.content.slice(0, split).trim() : memory.content;
|
||||
const body = split > 0 ? memory.content.slice(split).trim() : "";
|
||||
return (
|
||||
<>
|
||||
<article
|
||||
@@ -242,23 +259,37 @@ export function SelectedMemoryDetail({
|
||||
data-hide-actions={hideActions ? "true" : "false"}
|
||||
>
|
||||
<p className="selected-memory-status">
|
||||
{statusLabel} ·{" "}
|
||||
{t("selectedMemory.version", {
|
||||
revision: new Intl.NumberFormat(locale).format(memory.revision),
|
||||
})}
|
||||
{statusLabel}
|
||||
{!detail.sdk_status &&
|
||||
` · ${t("selectedMemory.version", {
|
||||
revision: new Intl.NumberFormat(locale).format(memory.revision),
|
||||
})}`}
|
||||
</p>
|
||||
<h3>{heading}</h3>
|
||||
{body && <p className="selected-memory-body">{body}</p>}
|
||||
<p className="selected-memory-source">
|
||||
{t("selectedMemory.sourceLine", { source, time: timestamp })}
|
||||
</p>
|
||||
{detail.source_id || detail.evidence.length ? (
|
||||
<p className="selected-memory-source">
|
||||
{t("selectedMemory.sourceLine", {
|
||||
source,
|
||||
time: clockStamp(
|
||||
detail.sdk_memory_at ?? evidence?.created_at,
|
||||
locale,
|
||||
),
|
||||
})}
|
||||
</p>
|
||||
) : null}
|
||||
<p className="selected-memory-times">
|
||||
{t("selectedMemory.timestamps", {
|
||||
conversation: clockStamp(
|
||||
conversationAt?.at ?? evidence?.created_at,
|
||||
conversationAt?.at ??
|
||||
detail.sdk_memory_at ??
|
||||
evidence?.created_at,
|
||||
locale,
|
||||
),
|
||||
created: clockStamp(
|
||||
detail.sdk_created_at ?? memory.updated_at,
|
||||
locale,
|
||||
),
|
||||
created: clockStamp(memory.updated_at, locale),
|
||||
})}
|
||||
</p>
|
||||
{!hideActions && (
|
||||
@@ -276,27 +307,6 @@ export function SelectedMemoryDetail({
|
||||
{t("selectedMemory.edit")}
|
||||
</R6Button>
|
||||
)}
|
||||
<section className="selected-memory-evidence">
|
||||
<h4>{t("selectedMemory.evidence")}</h4>
|
||||
<small>
|
||||
{t("selectedMemory.evidenceLine", { author, time: timestamp })}
|
||||
</small>
|
||||
<blockquote>
|
||||
{evidence?.available
|
||||
? evidence.verbatim
|
||||
: t("selectedMemory.evidenceUnavailable")}
|
||||
</blockquote>
|
||||
{!hideActions && (
|
||||
<R6Button
|
||||
onClick={(event) => {
|
||||
trigger.current = event.currentTarget;
|
||||
setPage("original");
|
||||
}}
|
||||
>
|
||||
{t("selectedMemory.viewOriginal")}
|
||||
</R6Button>
|
||||
)}
|
||||
</section>
|
||||
{!hideActions && (
|
||||
<R6Button variant="quiet" onClick={() => setShowHistory((v) => !v)}>
|
||||
{t("selectedMemory.history")}
|
||||
@@ -345,9 +355,7 @@ export function SelectedMemoryDetail({
|
||||
aria-label={
|
||||
page === "edit"
|
||||
? t("selectedMemory.editTitle")
|
||||
: page === "original"
|
||||
? t("selectedMemory.originalTitle")
|
||||
: t("selectedMemory.delete")
|
||||
: t("selectedMemory.delete")
|
||||
}
|
||||
onKeyDown={(event) => {
|
||||
if (event.key === "Tab") {
|
||||
@@ -380,13 +388,8 @@ export function SelectedMemoryDetail({
|
||||
<h1 ref={focus} tabIndex={-1}>
|
||||
{page === "edit"
|
||||
? t("selectedMemory.editTitle")
|
||||
: page === "original"
|
||||
? t("selectedMemory.originalTitle")
|
||||
: t("selectedMemory.delete")}
|
||||
: t("selectedMemory.delete")}
|
||||
</h1>
|
||||
{page === "original" && (
|
||||
<R6Button onClick={close}>{t("selectedMemory.back")}</R6Button>
|
||||
)}
|
||||
</header>
|
||||
{page === "edit" ? (
|
||||
<div className="selected-editor">
|
||||
@@ -443,40 +446,6 @@ export function SelectedMemoryDetail({
|
||||
</R6Button>
|
||||
</div>
|
||||
</div>
|
||||
) : page === "original" ? (
|
||||
<div className="selected-original">
|
||||
<small>
|
||||
{source} · {timestamp}
|
||||
</small>
|
||||
<p>{author}</p>
|
||||
{detail.evidence.map((item) => (
|
||||
<section key={item.evidence_id}>
|
||||
<small>
|
||||
{sourceNameOr(
|
||||
sources,
|
||||
item.source_id,
|
||||
t("selectedMemory.sourceUnknown"),
|
||||
)}{" "}
|
||||
·{" "}
|
||||
{item.written_by === "person"
|
||||
? t("selectedMemory.author.user")
|
||||
: item.written_by === "agent"
|
||||
? t("selectedMemory.author.agent")
|
||||
: t("selectedMemory.author.unknown")}
|
||||
</small>
|
||||
<blockquote>
|
||||
{item.available
|
||||
? item.verbatim
|
||||
: t("selectedMemory.sourceUnavailable")}
|
||||
</blockquote>
|
||||
</section>
|
||||
))}
|
||||
<small>{t("selectedMemory.related")}</small>
|
||||
<h2>{heading}</h2>
|
||||
<R6Button onClick={close}>
|
||||
{t("selectedMemory.viewMemory")}
|
||||
</R6Button>
|
||||
</div>
|
||||
) : (
|
||||
<div className="selected-delete">
|
||||
<h2>{t("selectedMemory.deleteTitle")}</h2>
|
||||
|
||||
@@ -68,6 +68,8 @@ describe("paged memory library client", () => {
|
||||
expect(
|
||||
await getSdkMemoryLibrary({
|
||||
statuses: ["active", "superseded"],
|
||||
source: "cursor",
|
||||
query: "release",
|
||||
offset: 0,
|
||||
limit: 25,
|
||||
}),
|
||||
@@ -76,5 +78,7 @@ describe("paged memory library client", () => {
|
||||
const params = get.mock.calls[0][1].params as URLSearchParams;
|
||||
expect(params.getAll("statuses")).toEqual(["active", "superseded"]);
|
||||
expect(params.get("limit")).toBe("25");
|
||||
expect(params.get("source")).toBe("cursor");
|
||||
expect(params.get("query")).toBe("release");
|
||||
});
|
||||
});
|
||||
@@ -891,10 +891,14 @@ export async function getMemoryLibrary(
|
||||
export async function getSdkMemoryLibrary(
|
||||
{
|
||||
statuses = ["active"],
|
||||
source,
|
||||
query,
|
||||
offset = 0,
|
||||
limit = 100,
|
||||
}: {
|
||||
statuses?: SdkMemoryStatus[];
|
||||
source?: string;
|
||||
query?: string;
|
||||
offset?: number;
|
||||
limit?: number;
|
||||
} = {},
|
||||
@@ -905,6 +909,8 @@ export async function getSdkMemoryLibrary(
|
||||
limit: String(limit),
|
||||
});
|
||||
for (const status of statuses) params.append("statuses", status);
|
||||
if (source) params.set("source", source);
|
||||
if (query?.trim()) params.set("query", query.trim());
|
||||
const response = await client.get<SdkMemoryPage>("/api/v2/memory-library", {
|
||||
params,
|
||||
signal,
|
||||
@@ -1162,6 +1168,26 @@ export async function getMemory(memoryId: string): Promise<MemoryDetail> {
|
||||
*/
|
||||
let pendingMemoryGraph: Promise<MemoryGraphResponse> | null = null;
|
||||
|
||||
/** The memory map on SDK authority: SDK memories plus the graph overlay. */
|
||||
/** The SDK map for one list page: the same source and window the list reads. */
|
||||
export async function getSdkMemoryGraph(
|
||||
{
|
||||
source,
|
||||
offset = 0,
|
||||
limit = 100,
|
||||
}: { source?: string; offset?: number; limit?: number } = {},
|
||||
signal?: AbortSignal,
|
||||
): Promise<MemoryGraphResponse> {
|
||||
const response = await client.get<MemoryGraphResponse>(
|
||||
"/api/v2/memories/graph",
|
||||
{
|
||||
params: { limit, offset, ...(source ? { source } : {}) },
|
||||
signal,
|
||||
},
|
||||
);
|
||||
return response.data;
|
||||
}
|
||||
|
||||
export async function getMemoryGraphWith(
|
||||
transport: { get: typeof client.get } = client,
|
||||
): Promise<MemoryGraphResponse> {
|
||||
|
||||
@@ -54,6 +54,12 @@ describe("clockStamp", () => {
|
||||
expect(clockStamp(value, "es")).toContain("19");
|
||||
});
|
||||
|
||||
it("names the year of a date in another year", () => {
|
||||
expect(clockStamp("2254-01-13T13:54:00Z", "zh-CN")).toContain("2254");
|
||||
expect(clockStamp("2025-12-31T12:00:00Z", "es")).toContain("2025");
|
||||
expect(clockStamp("2026-08-03T15:07:00Z", "zh-CN")).not.toContain("2026");
|
||||
});
|
||||
|
||||
it("formats another date with month and day, without a today word", () => {
|
||||
const value = "2026-08-03T15:07:00Z";
|
||||
const es = clockStamp(value, "es");
|
||||
|
||||
@@ -34,7 +34,10 @@ export function clockStamp(
|
||||
date.getDate() === now.getDate()
|
||||
)
|
||||
return translate(locale, "clock.today", { clock });
|
||||
// Another year carries its year, as the memory list's day stamp does: a
|
||||
// bare "01/13" for 2254-01-13 reads as this January.
|
||||
return new Intl.DateTimeFormat(locale, {
|
||||
...(date.getFullYear() !== now.getFullYear() ? { year: "numeric" } : {}),
|
||||
month: "2-digit",
|
||||
day: "2-digit",
|
||||
hour: "2-digit",
|
||||
|
||||
@@ -689,6 +689,23 @@ button.rail-button.active::before {
|
||||
background: transparent;
|
||||
font-size: 12px;
|
||||
}
|
||||
.map-search input {
|
||||
flex: 1;
|
||||
}
|
||||
.map-search-submit {
|
||||
flex: none;
|
||||
padding: 4px 10px;
|
||||
border: 0;
|
||||
border-radius: 6px;
|
||||
color: #fff;
|
||||
background: #3d5a70;
|
||||
font-size: 12px;
|
||||
cursor: pointer;
|
||||
}
|
||||
.map-search-submit:disabled {
|
||||
opacity: 0.5;
|
||||
cursor: default;
|
||||
}
|
||||
.map-search:focus-within,
|
||||
.map-question:focus-within {
|
||||
outline: 2px solid #7892a6;
|
||||
|
||||
@@ -23,6 +23,8 @@ export function projectPagedSourceMap(
|
||||
graph: MemoryGraphResponse | null,
|
||||
page: MemoryLibraryPage | null,
|
||||
sources: SourceResponse[],
|
||||
/** Names for source ids that are not in the catalog (e.g. 未识别来源). */
|
||||
labels: Record<string, string> = {},
|
||||
): MemoryGraphResponse {
|
||||
const nodes = new Map<string, MemoryGraphNode>();
|
||||
const edges = new Map<string, MemoryGraphEdge>();
|
||||
@@ -93,7 +95,7 @@ export function projectPagedSourceMap(
|
||||
node_type: "source",
|
||||
memory_id: null,
|
||||
source_id: id,
|
||||
label: catalog.get(id)?.display_name || id,
|
||||
label: catalog.get(id)?.display_name || labels[id] || id,
|
||||
kind: null,
|
||||
status: null,
|
||||
confidence: null,
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
import type { SdkMemoryRecord } from "@one-memory/contracts";
|
||||
import { afterEach, beforeEach, expect, it, vi } from "vitest";
|
||||
import { projectPagedSourceMap } from "./memory-map-model";
|
||||
import { sdkDetail, sdkLibraryPage } from "./sdk-memory-map";
|
||||
|
||||
const NOW = Date.UTC(2026, 9, 25) / 1_000;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.useFakeTimers();
|
||||
vi.setSystemTime(NOW * 1_000);
|
||||
});
|
||||
afterEach(() => vi.useRealTimers());
|
||||
|
||||
const record = (extra: Partial<SdkMemoryRecord>): SdkMemoryRecord => ({
|
||||
memory_id: "m",
|
||||
content: "memory",
|
||||
layer: "l2_fact",
|
||||
status: "active",
|
||||
is_latest: true,
|
||||
supersedes: [],
|
||||
superseded_by: [],
|
||||
tags: [],
|
||||
owner: null,
|
||||
source_raw_ids: [],
|
||||
user_id: "local",
|
||||
agent_id: "all-apps",
|
||||
session_id: "s1",
|
||||
source: "cursor",
|
||||
memory_at: null,
|
||||
gmt_created: NOW - 86_400,
|
||||
...extra,
|
||||
});
|
||||
|
||||
it("keeps the SDK's no-source bucket as a source the list and map can name", () => {
|
||||
const page = sdkLibraryPage({
|
||||
memories: [record({ memory_id: "d", source: "default" })],
|
||||
total: 1,
|
||||
limit: 100,
|
||||
offset: 0,
|
||||
});
|
||||
expect(page.source_ids_by_memory).toEqual({ d: ["default"] });
|
||||
expect(sdkDetail(record({ source: "default" })).source_id).toBe("default");
|
||||
|
||||
const drawn = projectPagedSourceMap(null, page, [], {
|
||||
default: "未识别来源",
|
||||
});
|
||||
expect(
|
||||
drawn.nodes.find((node) => node.node_id === "source:default")?.label,
|
||||
).toBe("未识别来源");
|
||||
});
|
||||
|
||||
it("does not present an impossible conversation time as when it happened", () => {
|
||||
const future = Date.UTC(2254, 0, 13) / 1_000;
|
||||
const page = sdkLibraryPage({
|
||||
memories: [record({ memory_id: "f", memory_at: future })],
|
||||
total: 1,
|
||||
limit: 100,
|
||||
offset: 0,
|
||||
});
|
||||
expect(page.source_time_by_memory?.f).toEqual({
|
||||
at: new Date((NOW - 86_400) * 1_000).toISOString(),
|
||||
basis: "import",
|
||||
});
|
||||
expect(sdkDetail(record({ memory_at: future })).sdk_memory_at).toBeNull();
|
||||
|
||||
const real = NOW - 3 * 86_400;
|
||||
expect(
|
||||
sdkLibraryPage({
|
||||
memories: [record({ memory_id: "r", memory_at: real })],
|
||||
total: 1,
|
||||
limit: 100,
|
||||
offset: 0,
|
||||
}).source_time_by_memory?.r.basis,
|
||||
).toBe("conversation");
|
||||
});
|
||||
|
||||
it("carries a full search window through to the list", () => {
|
||||
expect(
|
||||
sdkLibraryPage({
|
||||
memories: [],
|
||||
total: 0,
|
||||
limit: 100,
|
||||
offset: 0,
|
||||
search_window: 500,
|
||||
}).search_window,
|
||||
).toBe(500);
|
||||
});
|
||||
@@ -0,0 +1,136 @@
|
||||
import type {
|
||||
GovernedMemory,
|
||||
MemoryDetail,
|
||||
MemoryKind,
|
||||
MemoryLibraryPage,
|
||||
SdkMemoryPage,
|
||||
SdkMemoryRecord,
|
||||
} from "@one-memory/contracts";
|
||||
|
||||
function kindOf(layer: string): MemoryKind {
|
||||
if (layer === "l6_schema") return "schema";
|
||||
if (layer === "l7_intention") return "intention";
|
||||
if (layer.startsWith("coding")) return "coding";
|
||||
if (layer.startsWith("working")) return "working";
|
||||
return "chat";
|
||||
}
|
||||
|
||||
function stamp(seconds: number | null | undefined): string {
|
||||
return typeof seconds === "number" && Number.isFinite(seconds)
|
||||
? new Date(seconds * 1_000).toISOString()
|
||||
: "";
|
||||
}
|
||||
|
||||
const EARLIEST_PLAUSIBLE = Date.UTC(2000, 0, 1) / 1_000;
|
||||
|
||||
/** A conversation time that can be true: not before 2000, not tomorrow or later.
|
||||
*
|
||||
* A historical Cursor bug stored a file hash as the message time (years like
|
||||
* 2254). Such a value is not shown as when the conversation happened; the
|
||||
* import time, or "time unknown", says what is actually known.
|
||||
*/
|
||||
function plausible(seconds: number | null | undefined): number | null {
|
||||
return typeof seconds === "number" &&
|
||||
Number.isFinite(seconds) &&
|
||||
seconds >= EARLIEST_PLAUSIBLE &&
|
||||
seconds <= Date.now() / 1_000 + 86_400
|
||||
? seconds
|
||||
: null;
|
||||
}
|
||||
|
||||
/** The SDK writes "default" when no source was given; keep it as that bucket. */
|
||||
function sourceOf(record: SdkMemoryRecord): string {
|
||||
return record.source || "default";
|
||||
}
|
||||
|
||||
/** An SDK memory in the shape the memory map draws.
|
||||
*
|
||||
* The map groups and labels by id, content, kind, status and tags only, so the
|
||||
* governed-only fields are inert placeholders here. Opening, editing and
|
||||
* deleting still go through the SDK (v2) endpoints — `sdkDetail()` below wraps
|
||||
* the same projection for the map's inspector, which is why the inspector can
|
||||
* stay where the design put it instead of sending the reader to another view.
|
||||
*/
|
||||
export function sdkMapMemory(record: SdkMemoryRecord): GovernedMemory {
|
||||
const created = stamp(record.gmt_created ?? record.memory_at);
|
||||
return {
|
||||
memory_id: record.memory_id,
|
||||
user_id: record.user_id,
|
||||
kind: kindOf(record.layer),
|
||||
status:
|
||||
record.status === "active"
|
||||
? "confirmed"
|
||||
: record.status === "superseded"
|
||||
? "superseded"
|
||||
: "deleted",
|
||||
content: record.content,
|
||||
scopes: [],
|
||||
evidence_ids: [],
|
||||
origin: "extracted",
|
||||
tags: [...record.tags],
|
||||
supersedes: [...record.supersedes],
|
||||
superseded_by: [...record.superseded_by],
|
||||
// GovernedMemory still requires a revision for the shared map model. SDK
|
||||
// has no revision contract; the inspector recognises sdk_status and never
|
||||
// renders this placeholder.
|
||||
revision: 0,
|
||||
created_at: created,
|
||||
updated_at: created,
|
||||
deleted_at: null,
|
||||
};
|
||||
}
|
||||
|
||||
/** One SDK memory as the map's inspector reads it.
|
||||
*
|
||||
* The inspector was built for a governed detail. SDK-native memory has no
|
||||
* product revision or "original evidence" section; source_raw_ids remain an
|
||||
* internal SDK provenance relation and are deliberately not projected as UI
|
||||
* evidence.
|
||||
*/
|
||||
export function sdkDetail(record: SdkMemoryRecord): MemoryDetail {
|
||||
return {
|
||||
memory: sdkMapMemory(record),
|
||||
evidence: [],
|
||||
history: [],
|
||||
index_sync: null,
|
||||
sdk_status: record.status,
|
||||
source_id: sourceOf(record),
|
||||
sdk_memory_at: stamp(plausible(record.memory_at)) || null,
|
||||
sdk_created_at: stamp(record.gmt_created) || null,
|
||||
};
|
||||
}
|
||||
|
||||
/** One SDK page as the map's left list reads it.
|
||||
*
|
||||
* The list asks a library page for its rows, its total and — per row — when
|
||||
* the cited conversation happened. The SDK page carries the row itself, so the
|
||||
* time is taken from the memory: `memory_at` is the moment the conversation
|
||||
* happened, `gmt_created` only the moment it was imported. Naming which one it
|
||||
* is keeps "对话于" from silently becoming "沉淀于".
|
||||
*/
|
||||
export function sdkLibraryPage(page: SdkMemoryPage): MemoryLibraryPage {
|
||||
const sourceTime: Record<
|
||||
string,
|
||||
{ at: string; basis: "conversation" | "import" }
|
||||
> = {};
|
||||
const sourceIds: Record<string, string[]> = {};
|
||||
for (const record of page.memories) {
|
||||
sourceIds[record.memory_id] = [sourceOf(record)];
|
||||
const conversation = plausible(record.memory_at);
|
||||
const at = conversation ?? plausible(record.gmt_created);
|
||||
if (at === null) continue;
|
||||
sourceTime[record.memory_id] = {
|
||||
at: new Date(at * 1_000).toISOString(),
|
||||
basis: conversation !== null ? "conversation" : "import",
|
||||
};
|
||||
}
|
||||
return {
|
||||
items: page.memories.map(sdkMapMemory),
|
||||
total: page.total,
|
||||
offset: page.offset,
|
||||
limit: page.limit,
|
||||
source_ids_by_memory: sourceIds,
|
||||
source_time_by_memory: sourceTime,
|
||||
search_window: page.search_window ?? null,
|
||||
};
|
||||
}
|
||||
@@ -138,13 +138,19 @@
|
||||
gap: 8px;
|
||||
font-family: "Noto Sans SC", "PingFang SC", sans-serif;
|
||||
}
|
||||
.r6-shell .r6-memory .map-memory-item > span:not(.map-state) {
|
||||
.r6-shell .r6-memory .map-memory-item > .map-memory-content {
|
||||
min-height: 40px;
|
||||
display: -webkit-box;
|
||||
-webkit-line-clamp: 2;
|
||||
-webkit-box-orient: vertical;
|
||||
overflow: hidden;
|
||||
}
|
||||
.r6-shell .r6-memory .map-memory-tags.wb-tags {
|
||||
grid-column: 2;
|
||||
min-height: 0;
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
}
|
||||
.r6-shell .r6-memory .map-memory-item small {
|
||||
font-size: 11px;
|
||||
line-height: 19px;
|
||||
@@ -212,15 +218,19 @@
|
||||
margin: 0 0 32px;
|
||||
}
|
||||
.selected-memory-detail > h3 {
|
||||
display: block;
|
||||
font-size: 22px;
|
||||
line-height: 38px;
|
||||
font-weight: 500;
|
||||
letter-spacing: normal;
|
||||
margin: 0 0 18px;
|
||||
overflow-wrap: anywhere;
|
||||
}
|
||||
.selected-memory-body {
|
||||
display: block;
|
||||
font-size: 17px;
|
||||
line-height: 29px;
|
||||
letter-spacing: normal;
|
||||
margin: 0;
|
||||
white-space: pre-wrap;
|
||||
overflow-wrap: anywhere;
|
||||
@@ -237,36 +247,6 @@
|
||||
color: #5e6e7a;
|
||||
line-height: 21px;
|
||||
}
|
||||
.selected-memory-evidence {
|
||||
border-top: 1px solid #c7d6e0;
|
||||
margin: 24px 0 16px;
|
||||
padding-top: 22px;
|
||||
}
|
||||
.selected-memory-evidence h4 {
|
||||
font-size: 16px;
|
||||
line-height: 28px;
|
||||
margin: 0 0 18px;
|
||||
font-weight: 400;
|
||||
}
|
||||
.selected-memory-evidence small {
|
||||
font-size: 12px;
|
||||
color: #5e6e7a;
|
||||
}
|
||||
.selected-memory-evidence blockquote {
|
||||
font-size: 15px;
|
||||
line-height: 26px;
|
||||
margin: 18px 0 44px;
|
||||
border: 0;
|
||||
padding: 0;
|
||||
background: none;
|
||||
white-space: pre-wrap;
|
||||
overflow-wrap: anywhere;
|
||||
}
|
||||
.selected-memory-detail[data-hide-actions="true"]
|
||||
.selected-memory-evidence
|
||||
blockquote {
|
||||
margin-bottom: 0;
|
||||
}
|
||||
.selected-memory-more {
|
||||
font-size: 12px;
|
||||
padding: 16px 0;
|
||||
@@ -304,8 +284,7 @@
|
||||
align-items: center;
|
||||
gap: 24px;
|
||||
}
|
||||
.selected-memory-page > header .r6-button,
|
||||
.selected-original > .r6-button {
|
||||
.selected-memory-page > header .r6-button {
|
||||
min-width: 150px;
|
||||
}
|
||||
.selected-memory-page h1 {
|
||||
@@ -394,40 +373,6 @@
|
||||
border-left: 2px solid #8da4b3;
|
||||
white-space: pre-wrap;
|
||||
}
|
||||
.selected-original {
|
||||
padding: 60px 46px;
|
||||
min-height: 0;
|
||||
overflow-y: auto;
|
||||
overscroll-behavior: contain;
|
||||
scrollbar-width: thin;
|
||||
}
|
||||
.selected-original > small {
|
||||
font-size: 14px;
|
||||
color: #5e6e7a;
|
||||
}
|
||||
.selected-original > p {
|
||||
margin-top: 70px;
|
||||
font-size: 14px;
|
||||
}
|
||||
.selected-original blockquote {
|
||||
margin: 24px 0 96px;
|
||||
font-size: 24px;
|
||||
line-height: 40px;
|
||||
white-space: pre-wrap;
|
||||
overflow-wrap: anywhere;
|
||||
}
|
||||
.selected-original section {
|
||||
margin-top: 24px;
|
||||
}
|
||||
.selected-original section > small {
|
||||
font-size: 13px;
|
||||
color: #5e6e7a;
|
||||
}
|
||||
.selected-original > h2 {
|
||||
font-size: 21px;
|
||||
font-weight: 400;
|
||||
margin: 24px 0 32px;
|
||||
}
|
||||
.selected-delete {
|
||||
padding: 130px 66px 0;
|
||||
min-height: 0;
|
||||
@@ -499,10 +444,6 @@
|
||||
.selected-delete > p {
|
||||
margin: 12px 0;
|
||||
}
|
||||
.selected-original {
|
||||
padding-top: 24px;
|
||||
padding-bottom: 24px;
|
||||
}
|
||||
}
|
||||
@media (max-width: 1250px) {
|
||||
.r6-shell .view.r6-memory {
|
||||
@@ -536,7 +477,6 @@
|
||||
.selected-memory-page {
|
||||
padding: 24px;
|
||||
}
|
||||
.selected-original,
|
||||
.selected-delete {
|
||||
padding-left: 0;
|
||||
padding-right: 0;
|
||||
|
||||
@@ -9,6 +9,13 @@ const INTERNAL_SOURCE_NAMES: Record<string, string> = {
|
||||
"memory-content-export": "导入的已有记忆",
|
||||
};
|
||||
|
||||
/**
|
||||
* The source the SDK records when a writer named none. The legacy-to-SDK
|
||||
* backfill wrote every legacy memory this way, so it is a real, filterable
|
||||
* bucket ("未识别来源"), not a source still being checked.
|
||||
*/
|
||||
export const UNATTRIBUTED_SOURCE_ID = "default";
|
||||
|
||||
type NamedSource = Pick<SourceResponse, "source_id" | "display_name">;
|
||||
|
||||
/** The name to show for a source id, falling back to the id itself. */
|
||||
|
||||
@@ -135,7 +135,7 @@ This exception does not extend to VDB, Coding, prompt, reconcile, or lifecycle.
|
||||
- Context wire:cutover marker 提交前只使用 legacy Context v1;提交后 Host Bridge 只使用 SDK-native Context v2。两种 schema 原样校验、原样返回,不把 SDK item 补造成 v1 的 `revision`、`evidence_ids`、`paths` 或 `adoption`,已选择的 v2 路径失败时也不回退 legacy 权威。
|
||||
- Ask wire:marker 前保留 Governance Ask v1;marker 后 Web / Host Bridge 只调用 SDK Ask v2。v2 的引用只含 SDK 原生 `memory_id/content/layer/status/source/gmt_created`,生成过程中重复 `get` 校验同一 SDK 快照,不补造 revision、Evidence 或 scopes,SDK 失败不回退 v1。
|
||||
- Matter:未形成已发布产品合同,不迁移占位数据;所有 `matter*` 表在升级清理时物理删除,fresh install 不创建,Runtime 不注册 Matter API。
|
||||
- 旧用户迁移:由 App 在可见的启动 loading 阶段统一执行;此时 Core 尚未对外提供产品读写。所有已发布版本都先经 SDK 写入,Governance/Evidence 只是镜像,所以迁移**不复制、不比对任何 memory 或 raw**:只创建并校验完整 bundle、让写入队列脱离旧表、提交 marker、删除旧表。由一个 Core 命令 `sdk-migrate` 完成;Desktop 启动与 CLI `memhub install`/`upgrade`(macOS、Linux)在启动新 Runtime 前调用它,普通用户不需要另跑命令。App 不得在 migration `deferred` 状态继续启动 Runtime。
|
||||
- 旧用户迁移:由 App 在可见的启动 loading 阶段统一执行;此时 Core 尚未对外提供产品读写。所有已发布版本都先经 SDK 写入,Governance/Evidence 只是镜像,所以迁移**不复制任何 memory 或 raw**:创建并校验完整 bundle、让写入队列脱离旧表、把旧版命名的 SDK 向量集合接回当前名字、用 Runtime 同样配置的公开 SDK client 按 ID 回读核对,通过后才提交 marker、删除旧表。由一个 Core 命令 `sdk-migrate` 完成;Desktop 启动与 CLI `memhub install`/`upgrade`(macOS、Linux)在启动新 Runtime 前调用它,普通用户不需要另跑命令。App 不得在 migration `deferred` 状态继续启动 Runtime。
|
||||
- 迁移期间宿主 capture hook 由 Desktop 活跃 PID marker 静默挡住,不再追加
|
||||
Evidence;cutover 后先从宿主 History checkpoint 补齐遗漏 turn,再恢复 live sync。
|
||||
- Matter 若未来重新开发,需以新的 SDK-native 合同重新设计,不继承当前占位 schema。
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# SDK Data Plane Migration v1
|
||||
|
||||
Updated: 2026-10-24.
|
||||
Updated: 2026-10-24 (vector reattach, readback).
|
||||
|
||||
This contract moves released MemHub installations from the historical
|
||||
Evidence/Governance authority to the Hy Memory SDK data plane defined by
|
||||
@@ -29,6 +29,14 @@ and no raw turn, and does not compare the mirror with the SDK:
|
||||
- SDK L1 raw is retained for three days by the Runtime's own
|
||||
`prune_expired_raw`, not by the migration.
|
||||
|
||||
"Already complete" holds only if the new Runtime opens the same stores. Builds
|
||||
before #145 named a profile's vector collections `one_memory_<12 hex>…`; #145
|
||||
renamed them `memhub_<16 hex>…` without carrying the old ones over, so later
|
||||
Runtimes read an empty collection (build 20 cut over showing 37 of 3317
|
||||
memories). The migration therefore reattaches those collections (§3, step 3)
|
||||
and proves the new Runtime reads the memories back before it cuts over
|
||||
(step 4).
|
||||
|
||||
## 2. Authority before and after cutover
|
||||
|
||||
Before the cutover marker is committed, legacy tables remain the product read
|
||||
@@ -68,9 +76,32 @@ then:
|
||||
restore; an existing verified bundle of the same run is reused;
|
||||
2. **prepare**: applies the management schema migrations, discards SDK staging
|
||||
left by builds through 19 (§5) and makes legacy queue work self-contained
|
||||
(§4); it needs no model profile and no SDK client;
|
||||
3. **cutover**: re-verifies the bundle, repeats §4, commits SDK authority;
|
||||
4. **cleanup**: drops the legacy tables (including the `evidence_search` index).
|
||||
(§4);
|
||||
3. **vectors**: for the active embedding profile, renames every
|
||||
`one_memory_<old hash>…` collection to the current `memhub_<hash>…` name
|
||||
through Chroma's public client. A current collection that already holds
|
||||
writes is merged in first (newer record wins on a shared id); counts are
|
||||
checked after each step;
|
||||
4. **verify**: opens the public SDK client exactly as the Runtime will and reads
|
||||
the legacy confirmed memory ids back (`list_memories`, then `get`). Every
|
||||
chat memory must read back. Coding/Working may be fewer, because the SDK's
|
||||
Coding reconciler deletes experiences it retires, but not zero when the
|
||||
legacy mirror has some. Otherwise the run stops with
|
||||
`sdk_migration_memory_unreadable` and the counts, before anything
|
||||
irreversible; legacy authority and tables stay intact;
|
||||
5. **cutover**: re-verifies the bundle, repeats §4, commits SDK authority;
|
||||
6. **cleanup**: drops the legacy tables (including the `evidence_search`
|
||||
index) and runs `VACUUM` (skipped if the disk is too full).
|
||||
|
||||
A home with legacy memories but no active model profile is not migrated
|
||||
(`sdk_migration_model_profile_required`): its collections can be neither named
|
||||
nor read back.
|
||||
|
||||
A home already on SDK authority still reports `required` while legacy-named
|
||||
vector collections exist for its profile (build 20 homes). `sdk-migrate` then
|
||||
reattaches them and reads back against the memory ids in the pre-cutover
|
||||
bundle; that report is recorded but does not block, as nothing irreversible is
|
||||
left. Runtime code never renames or reads Chroma directly.
|
||||
|
||||
With `--json` each step prints `{"state":"running","step":…}` and the last line
|
||||
is the result: `cutover`, `complete` (only cleanup was left), `not_required`, or
|
||||
@@ -148,6 +179,12 @@ The release gate covers:
|
||||
- CLI upgrade migrating before restart, rolling back on a failed migration,
|
||||
and keeping the new version when a migrated home is slow to come up;
|
||||
- host capture fenced for the whole command;
|
||||
- legacy-named vector collections reattached (with merge) before verifying,
|
||||
and on a home that already cut over;
|
||||
- readback refusing a cutover the new Runtime cannot read, and tolerating
|
||||
retired Coding experiences;
|
||||
- `/api/v2/memory-authority` agreeing with the API gates on fresh, migrated
|
||||
and legacy homes;
|
||||
- pending job freezing, legacy task and pending dispatch conversion;
|
||||
- invoked/unknown legacy tasks and pending deletions not blocking cutover;
|
||||
- interrupted directory swap recovery before and after authority;
|
||||
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
"name": "one-memory",
|
||||
"version": "1.0.5",
|
||||
"release": {
|
||||
"build_number": 17,
|
||||
"build_number": 26,
|
||||
"macos_min_version": "13.0",
|
||||
"app_update": "manual",
|
||||
"runtime_manifest_schema": 1,
|
||||
|
||||
@@ -470,6 +470,9 @@ export interface MemoryLibraryPage {
|
||||
* import): a partial snapshot, not an authoritative view. Show "importing"
|
||||
* and retry; never render this state as a failure. */
|
||||
stale?: boolean;
|
||||
/** Set on a search whose relevance window was full: matches were looked for
|
||||
* only among this many most relevant memories, so `total` may undercount. */
|
||||
search_window?: number | null;
|
||||
}
|
||||
|
||||
export interface MemoryLibraryQuery {
|
||||
@@ -518,6 +521,11 @@ export interface MemoryDetail {
|
||||
memory: GovernedMemory;
|
||||
evidence: MemoryEvidenceExcerpt[];
|
||||
history: MemoryVersion[];
|
||||
/** SDK-native display metadata. Legacy governed details leave these absent. */
|
||||
sdk_status?: SdkMemoryStatus;
|
||||
source_id?: string | null;
|
||||
sdk_memory_at?: string | null;
|
||||
sdk_created_at?: string | null;
|
||||
// Null for a memory nobody has rewritten: there is no index to be behind,
|
||||
// and the interface should not explain a state the user never caused.
|
||||
index_sync: MemoryIndexSync | null;
|
||||
@@ -555,6 +563,7 @@ export interface SdkMemoryPage {
|
||||
total: number;
|
||||
limit: number;
|
||||
offset: number;
|
||||
search_window?: number | null;
|
||||
}
|
||||
|
||||
export interface SdkUserProfileResponse {
|
||||
|
||||
@@ -317,6 +317,39 @@ it("selects SDK Context v2 after cutover and passes native fields unchanged", as
|
||||
);
|
||||
});
|
||||
|
||||
it.each([
|
||||
[10, 200],
|
||||
[11, 502],
|
||||
])(
|
||||
"accepts up to Core's v2 default of 10 items when no limit is sent (%i items)",
|
||||
async (count, status) => {
|
||||
const f = await fixture();
|
||||
const items = Array.from({ length: count }, (_, index) => ({
|
||||
kind: "memory",
|
||||
memory_id: `memory-${index}`,
|
||||
content: "m",
|
||||
layer: "l2_fact",
|
||||
status: "active",
|
||||
source: "codex",
|
||||
updated_at: "2026-10-21T00:00:00Z",
|
||||
score: 0.5,
|
||||
truncated: false,
|
||||
}));
|
||||
f.setAuthority("sdk");
|
||||
f.respond({
|
||||
schema_version: 2,
|
||||
request_id: "sdk-request",
|
||||
state: "grounded",
|
||||
items,
|
||||
used_characters: count,
|
||||
truncated: false,
|
||||
total_ms: 1,
|
||||
});
|
||||
|
||||
expect((await f.post({ query: "original" })).status).toBe(status);
|
||||
},
|
||||
);
|
||||
|
||||
it("does not fall back to legacy context when the selected SDK response is invalid", async () => {
|
||||
const f = await fixture();
|
||||
f.setAuthority("sdk");
|
||||
|
||||
@@ -74,7 +74,7 @@ export const CONTEXT_INPUT_SCHEMA = {
|
||||
description:
|
||||
"Filters raw only, within already granted sources; does not change memory scope.",
|
||||
},
|
||||
limit: { type: "integer", minimum: 1, maximum: 20, default: 8 },
|
||||
limit: { type: "integer", minimum: 1, maximum: 20, default: 10 },
|
||||
character_budget: {
|
||||
type: "integer",
|
||||
minimum: 200,
|
||||
@@ -285,8 +285,10 @@ export function isContextResponseForRequest(
|
||||
): value is HostContextSearchResponse {
|
||||
if (!isHostContextSearchResponse(value)) return false;
|
||||
const kinds = request.kinds ?? ["raw", "memory"];
|
||||
// Core's default item cap: 8 on v1, 10 on SDK-native v2.
|
||||
const limit = request.limit ?? (isSdkContextSearchResponse(value) ? 10 : 8);
|
||||
return (
|
||||
value.items.length <= (request.limit ?? 8) &&
|
||||
value.items.length <= limit &&
|
||||
value.used_characters <= (request.character_budget ?? 4000) &&
|
||||
value.items.every(
|
||||
(item) =>
|
||||
|
||||
@@ -74,10 +74,16 @@ export const currentChineseMessages = {
|
||||
"memorySpace.tabs.workbench": "工作台",
|
||||
|
||||
"selectedMemory.status.confirmed": "已沉淀",
|
||||
"selectedMemory.status.active": "已生效",
|
||||
"selectedMemory.status.deleted": "已删除",
|
||||
"selectedMemory.status.superseded": "已被取代",
|
||||
"selectedMemory.status.shadow": "已隐藏",
|
||||
"selectedMemory.status.negated": "已否定",
|
||||
"selectedMemory.status.conflicted": "有冲突",
|
||||
"selectedMemory.status.archived": "已归档",
|
||||
"selectedMemory.status.unknown": "状态待核对",
|
||||
"selectedMemory.sourceUnknown": "来源待核对",
|
||||
"selectedMemory.sourceUnattributed": "未识别来源",
|
||||
"selectedMemory.discardDraft": "放弃尚未保存的修改?",
|
||||
"selectedMemory.newVersionReadFailed": "暂时无法读取新版本,草稿已保留。",
|
||||
"selectedMemory.updateRejected":
|
||||
@@ -208,10 +214,16 @@ export const currentEnglishMessages: Record<
|
||||
"memorySpace.tabs.workbench": "Workbench",
|
||||
|
||||
"selectedMemory.status.confirmed": "Created",
|
||||
"selectedMemory.status.active": "Active",
|
||||
"selectedMemory.status.deleted": "Deleted",
|
||||
"selectedMemory.status.superseded": "Superseded",
|
||||
"selectedMemory.status.shadow": "Hidden",
|
||||
"selectedMemory.status.negated": "Negated",
|
||||
"selectedMemory.status.conflicted": "Conflicted",
|
||||
"selectedMemory.status.archived": "Archived",
|
||||
"selectedMemory.status.unknown": "Status needs review",
|
||||
"selectedMemory.sourceUnknown": "Source needs review",
|
||||
"selectedMemory.sourceUnattributed": "Unidentified source",
|
||||
"selectedMemory.discardDraft": "Discard unsaved changes?",
|
||||
"selectedMemory.newVersionReadFailed":
|
||||
"Could not load the new version. Your draft is retained.",
|
||||
@@ -356,10 +368,16 @@ export const currentSpanishMessages: Record<
|
||||
"memorySpace.tabs.workbench": "Espacio de trabajo",
|
||||
|
||||
"selectedMemory.status.confirmed": "Guardada",
|
||||
"selectedMemory.status.active": "Activa",
|
||||
"selectedMemory.status.deleted": "Eliminada",
|
||||
"selectedMemory.status.superseded": "Sustituida",
|
||||
"selectedMemory.status.shadow": "Oculta",
|
||||
"selectedMemory.status.negated": "Negada",
|
||||
"selectedMemory.status.conflicted": "En conflicto",
|
||||
"selectedMemory.status.archived": "Archivada",
|
||||
"selectedMemory.status.unknown": "Estado pendiente de comprobar",
|
||||
"selectedMemory.sourceUnknown": "Origen sin verificar",
|
||||
"selectedMemory.sourceUnattributed": "Origen no identificado",
|
||||
"selectedMemory.discardDraft": "¿Descartar los cambios sin guardar?",
|
||||
"selectedMemory.newVersionReadFailed":
|
||||
"No se pudo leer la nueva versión. Tu borrador se conserva.",
|
||||
|
||||
@@ -1571,6 +1571,10 @@ const baseEnglishMessages = {
|
||||
"memory:全部记忆": "All memories",
|
||||
"memory:来源待核对": "Source needs review",
|
||||
"memory:所选来源待核对": "Selected source needs review",
|
||||
"memory:未识别来源": "Unidentified source",
|
||||
"memory:这个来源下还没有记忆": "No memories from this source yet",
|
||||
"memory:只在最相关的 {count} 条记忆里查找关键词":
|
||||
"Keywords are matched only among the {count} most relevant memories",
|
||||
"memory:未知来源": "Unknown source",
|
||||
"memory:记忆读取失败": "Could not load memories",
|
||||
"memory:正在读取…": "Loading…",
|
||||
|
||||
@@ -146,6 +146,10 @@ export const memoryChineseMessages = {
|
||||
"memory:全部记忆": "全部记忆",
|
||||
"memory:来源待核对": "来源待核对",
|
||||
"memory:所选来源待核对": "所选来源待核对",
|
||||
"memory:未识别来源": "未识别来源",
|
||||
"memory:这个来源下还没有记忆": "这个来源下还没有记忆",
|
||||
"memory:只在最相关的 {count} 条记忆里查找关键词":
|
||||
"只在最相关的 {count} 条记忆里查找关键词",
|
||||
"memory:未知来源": "未知来源",
|
||||
"memory:记忆读取失败": "记忆读取失败",
|
||||
"memory:正在读取…": "正在读取…",
|
||||
@@ -538,6 +542,10 @@ export const memorySpanishMessages: Record<
|
||||
"memory:全部记忆": "Todas las memorias",
|
||||
"memory:来源待核对": "Origen sin verificar",
|
||||
"memory:所选来源待核对": "Origen seleccionado sin verificar",
|
||||
"memory:未识别来源": "Origen no identificado",
|
||||
"memory:这个来源下还没有记忆": "Aún no hay recuerdos de este origen",
|
||||
"memory:只在最相关的 {count} 条记忆里查找关键词":
|
||||
"Las palabras clave solo se buscan entre los {count} recuerdos más relevantes",
|
||||
"memory:未知来源": "Origen desconocido",
|
||||
"memory:记忆读取失败": "Error al leer las memorias",
|
||||
"memory:正在读取…": "Leyendo…",
|
||||
|
||||
@@ -711,7 +711,7 @@ describe("local Runtime installation", () => {
|
||||
expect(replacement.version_directory).not.toBe(firstPlan.version_directory);
|
||||
});
|
||||
|
||||
it("refuses different bytes from the same commit even with the opt-in", async () => {
|
||||
it("repairs different bytes from the same commit with the packaged opt-in", async () => {
|
||||
if (process.platform !== "darwin") return;
|
||||
const commit = "a".repeat(40);
|
||||
const first = await distribution("1.0.0", {
|
||||
@@ -738,18 +738,17 @@ describe("local Runtime installation", () => {
|
||||
.digest("hex");
|
||||
await writeFile(manifestPath, JSON.stringify(manifest));
|
||||
|
||||
await expect(
|
||||
planLocalInstall({
|
||||
source: rebuilt.source,
|
||||
installRoot: first.installRoot,
|
||||
platform: process.platform,
|
||||
arch: process.arch,
|
||||
allowSameBuildReplace: true,
|
||||
}),
|
||||
).rejects.toThrow("installed_version_conflict");
|
||||
const plan = await planLocalInstall({
|
||||
source: rebuilt.source,
|
||||
installRoot: first.installRoot,
|
||||
platform: process.platform,
|
||||
arch: process.arch,
|
||||
allowSameBuildReplace: true,
|
||||
});
|
||||
expect(plan.same_build_replace).toBe(true);
|
||||
});
|
||||
|
||||
it("refuses a tampered selected runtime instead of taking it over", async () => {
|
||||
it("replaces a dirty selected runtime from the trusted packaged payload", async () => {
|
||||
if (process.platform !== "darwin") return;
|
||||
const first = await distribution("1.0.0", {
|
||||
buildNumber: 5,
|
||||
@@ -771,15 +770,15 @@ describe("local Runtime installation", () => {
|
||||
gitCommit: "b".repeat(40),
|
||||
});
|
||||
|
||||
await expect(
|
||||
planLocalInstall({
|
||||
source: rebuilt.source,
|
||||
installRoot: first.installRoot,
|
||||
platform: process.platform,
|
||||
arch: process.arch,
|
||||
allowSameBuildReplace: true,
|
||||
}),
|
||||
).rejects.toThrow("distribution_hash_mismatch");
|
||||
const plan = await planLocalInstall({
|
||||
source: rebuilt.source,
|
||||
installRoot: first.installRoot,
|
||||
platform: process.platform,
|
||||
arch: process.arch,
|
||||
allowSameBuildReplace: true,
|
||||
});
|
||||
expect(plan.same_build_replace).toBe(true);
|
||||
expect(plan.version_directory).not.toBe(firstPlan.version_directory);
|
||||
});
|
||||
|
||||
it("does not let the same-build opt-in weaken downgrade refusal", async () => {
|
||||
@@ -805,6 +804,32 @@ describe("local Runtime installation", () => {
|
||||
).rejects.toThrow("distribution_downgrade_refused");
|
||||
});
|
||||
|
||||
it("does not let the same-build opt-in replace a tree whose build it cannot read", async () => {
|
||||
if (process.platform !== "darwin") return;
|
||||
const first = await distribution("1.0.0", { buildNumber: 5 });
|
||||
const firstPlan = await planLocalInstall({
|
||||
source: first.source,
|
||||
installRoot: first.installRoot,
|
||||
platform: process.platform,
|
||||
arch: process.arch,
|
||||
});
|
||||
await applyLocalInstall(firstPlan);
|
||||
await writeFile(
|
||||
path.join(firstPlan.version_directory, "manifest.json"),
|
||||
"{",
|
||||
);
|
||||
const older = await distribution("0.9.0", { buildNumber: 4 });
|
||||
await expect(
|
||||
planLocalInstall({
|
||||
source: older.source,
|
||||
installRoot: first.installRoot,
|
||||
platform: process.platform,
|
||||
arch: process.arch,
|
||||
allowSameBuildReplace: true,
|
||||
}),
|
||||
).rejects.toThrow("distribution_manifest_invalid");
|
||||
});
|
||||
|
||||
it("accepts an equal build_number when the bytes are identical (idempotent reinstall)", async () => {
|
||||
if (process.platform !== "darwin") return;
|
||||
const first = await distribution("1.0.0", { buildNumber: 5 });
|
||||
|
||||
@@ -367,13 +367,23 @@ async function inspectBuildNumberTransition(options: {
|
||||
if (relative.startsWith("..") || path.isAbsolute(relative)) {
|
||||
throw new Error("install_current_unmanaged");
|
||||
}
|
||||
const selected = options.allowSameBuildReplace
|
||||
? await verifyDistribution(
|
||||
let selectedVerified = true;
|
||||
let selected = await readManifest(target).catch(() => null);
|
||||
if (options.allowSameBuildReplace) {
|
||||
// The opt-in may replace a dirty tree, but only one whose build it can
|
||||
// still read: without a manifest the downgrade check below has nothing to
|
||||
// compare, and a newer runtime would be overwritten by an older one.
|
||||
if (selected === null) throw new Error("distribution_manifest_invalid");
|
||||
try {
|
||||
selected = await verifyDistribution(
|
||||
target,
|
||||
options.incoming.platform,
|
||||
options.incoming.arch,
|
||||
)
|
||||
: await readManifest(target).catch(() => null);
|
||||
);
|
||||
} catch {
|
||||
selectedVerified = false;
|
||||
}
|
||||
}
|
||||
if (selected === null) return { sameBuildReplace: false };
|
||||
// Pre-contract packages have no `build_number`. Comparing them against an
|
||||
// invented ordinal would either break the hourly chain or manufacture a
|
||||
@@ -385,6 +395,9 @@ async function inspectBuildNumberTransition(options: {
|
||||
if (options.incoming.build_number > selected.build_number)
|
||||
return { sameBuildReplace: false };
|
||||
if (options.incoming.build_number === selected.build_number) {
|
||||
if (!selectedVerified && options.allowSameBuildReplace) {
|
||||
return { sameBuildReplace: true };
|
||||
}
|
||||
// Same build: an idempotent reinstall is fine only if the bytes match. The
|
||||
// incoming source is verified, so compare it against the selected tree.
|
||||
try {
|
||||
@@ -395,11 +408,8 @@ async function inspectBuildNumberTransition(options: {
|
||||
}
|
||||
if (
|
||||
!isReleaseCommit(selected.git_commit) ||
|
||||
!isReleaseCommit(options.incoming.git_commit) ||
|
||||
selected.git_commit === options.incoming.git_commit
|
||||
!isReleaseCommit(options.incoming.git_commit)
|
||||
) {
|
||||
// Same commit with different bytes is not a new build: it is a
|
||||
// non-reproducible or corrupted artifact and must stay rejected.
|
||||
throw new Error("installed_version_conflict");
|
||||
}
|
||||
return { sameBuildReplace: true };
|
||||
|
||||
@@ -422,7 +422,7 @@ describe("installPackagedRuntime", () => {
|
||||
).toBe("rebuilt-body");
|
||||
});
|
||||
|
||||
it("refuses same-commit different bytes even when the bundle opts in", async () => {
|
||||
it("repairs same-commit drift when the trusted bundle opts in", async () => {
|
||||
const { root, resources, executable, installRoot } = await fixture();
|
||||
const { selected } = await installFormalCurrent({
|
||||
root,
|
||||
@@ -434,18 +434,43 @@ describe("installPackagedRuntime", () => {
|
||||
await writeFile(path.join(selected, "web", "index.html"), "unexpected");
|
||||
await rewriteManifestDigest(selected, "web/index.html");
|
||||
|
||||
await expect(
|
||||
installPackagedRuntime({
|
||||
resourcesPath: resources,
|
||||
executablePath: executable,
|
||||
installRoot,
|
||||
platform: "darwin",
|
||||
arch: "arm64",
|
||||
buildNumber: 3757,
|
||||
gitCommit: "a".repeat(40),
|
||||
replaceSameBuild: true,
|
||||
}),
|
||||
).rejects.toThrow("installed_version_conflict");
|
||||
const receipt = await installPackagedRuntime({
|
||||
resourcesPath: resources,
|
||||
executablePath: executable,
|
||||
installRoot,
|
||||
platform: "darwin",
|
||||
arch: "arm64",
|
||||
buildNumber: 3757,
|
||||
gitCommit: "a".repeat(40),
|
||||
replaceSameBuild: true,
|
||||
});
|
||||
expect(receipt.reused_current).toBe(false);
|
||||
expect(receipt.version_directory).not.toBe(selected);
|
||||
});
|
||||
|
||||
it("repairs an extra file left in the selected runtime tree", async () => {
|
||||
const { root, resources, executable, installRoot } = await fixture();
|
||||
const { selected } = await installFormalCurrent({
|
||||
root,
|
||||
resources,
|
||||
executable,
|
||||
installRoot,
|
||||
gitCommit: "a".repeat(40),
|
||||
});
|
||||
await writeFile(path.join(selected, "left-by-ditto"), "dirty");
|
||||
|
||||
const receipt = await installPackagedRuntime({
|
||||
resourcesPath: resources,
|
||||
executablePath: executable,
|
||||
installRoot,
|
||||
platform: "darwin",
|
||||
arch: "arm64",
|
||||
buildNumber: 3757,
|
||||
gitCommit: "a".repeat(40),
|
||||
replaceSameBuild: true,
|
||||
});
|
||||
expect(receipt.reused_current).toBe(false);
|
||||
expect(receipt.version_directory).not.toBe(selected);
|
||||
});
|
||||
|
||||
it.each(["missing"])(
|
||||
@@ -483,6 +508,31 @@ describe("installPackagedRuntime", () => {
|
||||
},
|
||||
);
|
||||
|
||||
it("refuses to repair a selected tree whose manifest is unreadable", async () => {
|
||||
const { root, resources, executable, installRoot } = await fixture();
|
||||
const { selected } = await installFormalCurrent({
|
||||
root,
|
||||
resources,
|
||||
executable,
|
||||
installRoot,
|
||||
buildNumber: 3758,
|
||||
});
|
||||
await writeFile(path.join(selected, "manifest.json"), "{");
|
||||
|
||||
await expect(
|
||||
installPackagedRuntime({
|
||||
resourcesPath: resources,
|
||||
executablePath: executable,
|
||||
installRoot,
|
||||
platform: "darwin",
|
||||
arch: "arm64",
|
||||
buildNumber: 3757,
|
||||
gitCommit: "a".repeat(40),
|
||||
replaceSameBuild: true,
|
||||
}),
|
||||
).rejects.toThrow("distribution_manifest_invalid");
|
||||
});
|
||||
|
||||
it("preserves downgrade refusal when selected current is newer", async () => {
|
||||
const { root, resources, executable, installRoot } = await fixture();
|
||||
await installFormalCurrent({
|
||||
|
||||
@@ -123,9 +123,9 @@ export async function installPackagedRuntime(options: {
|
||||
*
|
||||
* Off by default, which preserves `installed_version_conflict`. Desktop turns
|
||||
* it on: the bundle it carries has passed notarization, so its payload is
|
||||
* already trusted, and the same build number is re-released whenever a
|
||||
* release is rebuilt from a new commit. Refusing that case leaves an affected
|
||||
* machine with an app that exits on every launch, which is a worse outcome
|
||||
* already trusted. It may need to replace a rebuilt release or repair a
|
||||
* selected tree polluted by an interrupted copy. Refusing either case leaves
|
||||
* an affected machine with an app that exits on every launch, which is worse
|
||||
* than replacing the runtime. The superseded tree is kept, not edited. */
|
||||
replaceSameBuild?: boolean;
|
||||
gitCommit?: string;
|
||||
@@ -307,11 +307,36 @@ async function reuseSelectedSameBuild(options: {
|
||||
) {
|
||||
throw new Error("install_current_unmanaged");
|
||||
}
|
||||
const selectedManifest = await verifyInstalledDistribution(
|
||||
selected,
|
||||
options.platform,
|
||||
options.arch,
|
||||
);
|
||||
let selectedManifest: DistributionManifest;
|
||||
try {
|
||||
selectedManifest = await verifyInstalledDistribution(
|
||||
selected,
|
||||
options.platform,
|
||||
options.arch,
|
||||
);
|
||||
} catch (error) {
|
||||
if (!options.replaceSameBuild) throw error;
|
||||
// A notarized App bundle is the repair source. A dirty selected tree must
|
||||
// reach the normal staged install path instead of failing before the
|
||||
// same-build replacement policy can run.
|
||||
const declared = await readFile(
|
||||
path.join(selected, "manifest.json"),
|
||||
"utf8",
|
||||
)
|
||||
.then((value) => JSON.parse(value) as DistributionManifest)
|
||||
.catch(() => null);
|
||||
// Without a readable build the downgrade check cannot run; fail closed
|
||||
// rather than let an older bundle replace a newer runtime.
|
||||
if (declared === null) throw new Error("distribution_manifest_invalid");
|
||||
if (
|
||||
declared.build_number !== undefined &&
|
||||
options.buildNumber !== undefined &&
|
||||
declared.build_number > options.buildNumber
|
||||
) {
|
||||
throw new Error("distribution_downgrade_refused");
|
||||
}
|
||||
return null;
|
||||
}
|
||||
if (selectedManifest.build_number !== options.buildNumber) return null;
|
||||
// The bundle's own payload, read in place. It has to agree with the selected
|
||||
// tree file for file; there is no cheaper signal that is also exact, and a
|
||||
|
||||
@@ -83,7 +83,7 @@ it("does not stop Runtime when nothing needs migrating", async () => {
|
||||
expect(fenced).toBe(false);
|
||||
});
|
||||
|
||||
it("treats a finished cleanup as a completed migration", async () => {
|
||||
it("reports a repair of an already migrated home as complete, not a second cutover", async () => {
|
||||
const result = await runSdkMigration({
|
||||
dataDir: "/data",
|
||||
core,
|
||||
@@ -95,7 +95,7 @@ it("treats a finished cleanup as a completed migration", async () => {
|
||||
fence: noFence,
|
||||
});
|
||||
|
||||
expect(result).toEqual({ state: "cutover" });
|
||||
expect(result).toEqual({ state: "complete" });
|
||||
});
|
||||
|
||||
it("reports an unreadable migration status as deferred", async () => {
|
||||
|
||||
@@ -4,7 +4,8 @@ import { withRuntimeMaintenanceFence } from "./runtime-maintenance.js";
|
||||
const MAX_OUTPUT_BYTES = 1024 * 1024;
|
||||
|
||||
export interface SdkMigrationResult {
|
||||
state: "not_required" | "cutover" | "deferred";
|
||||
/** `complete`: already on SDK authority; leftovers were repaired. */
|
||||
state: "not_required" | "cutover" | "complete" | "deferred";
|
||||
code?: string;
|
||||
}
|
||||
|
||||
@@ -84,8 +85,7 @@ export async function runSdkMigration(
|
||||
}
|
||||
if (migrated.exitCode !== 0) return deferred(migrated);
|
||||
const state = migrated.document.state;
|
||||
if (state === "cutover" || state === "complete")
|
||||
return { state: "cutover" };
|
||||
if (state === "cutover" || state === "complete") return { state };
|
||||
if (state === "not_required") return { state: "not_required" };
|
||||
return deferred(migrated);
|
||||
});
|
||||
|
||||
Reference in new issue
Block a user