Coverage for astrocyte/portability.py: 87%
191 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""Memory portability — AMA (Astrocyte Memory Archive) export and import.
3AMA is a newline-delimited JSON (JSONL) format. Line 1 is the header,
4subsequent lines are individual memories. Streamable, self-describing,
5and FFI-safe (plain JSON, no Python-specific types).
7See docs/_design/memory-portability.md for the full specification.
8"""
10from __future__ import annotations
12import json
13import logging
14import os
15from dataclasses import dataclass
16from datetime import datetime, timezone
17from pathlib import Path
18from typing import Awaitable, Callable, Literal
20from astrocyte._log_safety import safe as _safe_log
21from astrocyte.types import Metadata, RecallRequest, RecallResult, RetainRequest, VectorItem
23logger = logging.getLogger("astrocyte.portability")
25# ---------------------------------------------------------------------------
26# Path containment (CWE-022)
27# ---------------------------------------------------------------------------
28#
29# Path(path).resolve() canonicalises but does NOT contain — a caller can
30# pass /etc/passwd and resolve() returns it unchanged. ``_safe_resolve``
31# validates that the resolved path stays within an explicit allow-list.
32#
33# Allow-list resolution order:
34# 1. ``allowed_roots`` kwarg passed to the public function
35# 2. ``ASTROCYTE_PORTABILITY_ROOTS`` env var (os.pathsep-joined)
36#
37# When neither (1) nor (2) is configured, ``_safe_resolve`` REFUSES to
38# return a path unless the caller has explicitly opted into uncontained
39# mode via ``allow_uncontained=True``. This eliminates the silent
40# "no containment" gap that CodeQL CWE-022 (py/path-injection) flags
41# and forces every caller to make a conscious security decision.
42#
43# Recommended usage:
44# * Server / gateway code: set ``ASTROCYTE_PORTABILITY_ROOTS`` and
45# leave ``allow_uncontained=False``. Untrusted HTTP input cannot
46# escape the configured roots.
47# * Library / CLI / unit tests with caller-controlled paths: pass
48# ``allowed_roots=[<known dir>]`` explicitly.
49# * Trusted internal call sites that genuinely need any path: pass
50# ``allow_uncontained=True`` to make the decision audit-able.
52_PORTABILITY_ROOTS_ENV = "ASTROCYTE_PORTABILITY_ROOTS"
54# Null byte and ASCII control characters never have a legitimate place in
55# a filesystem path. Reject them up front; resolve() does NOT strip them.
56_ILLEGAL_PATH_CHAR_ORDS = frozenset(range(0x00, 0x20)) | {0x7F}
59def _portability_roots() -> list[Path]:
60 """Read containment roots from the environment."""
61 raw = os.environ.get(_PORTABILITY_ROOTS_ENV, "")
62 return [Path(p).expanduser().resolve() for p in raw.split(os.pathsep) if p]
65def _resolve_contained(
66 path: str | Path,
67 *,
68 allowed_roots: list[str | Path] | None = None,
69 allow_uncontained: bool = False,
70) -> Path:
71 """Resolve ``path`` and verify it stays within an allowed root.
73 See module docstring for the allow-list resolution order and the
74 ``allow_uncontained`` opt-in semantics.
76 Raises:
77 ValueError: If the path contains illegal control characters,
78 escapes every allowed root, or no containment is configured
79 and the caller did not pass ``allow_uncontained=True``.
80 """
81 path_str = os.fspath(path)
82 if any(ord(c) in _ILLEGAL_PATH_CHAR_ORDS for c in path_str):
83 raise ValueError(f"Portability path contains illegal control character: {path_str!r}")
84 # ``path_str`` is user-controlled, but the taint is neutralized by
85 # the allow-list containment check below — ``resolved`` is matched
86 # against ``allowed_roots`` (or the env-configured
87 # ``_portability_roots()``) before the function returns. Callers
88 # cannot opt out without passing ``allow_uncontained=True``
89 # explicitly. CodeQL's taint tracker doesn't see through allow-list
90 # logic, so this whole file is in CodeQL's ``paths-ignore`` (see
91 # ``.github/codeql/codeql-config.yml``). Threat model is locked by
92 # ``tests/test_portability.py::TestPathContainment``.
93 resolved = Path(path_str).expanduser().resolve()
94 roots: list[Path]
95 if allowed_roots:
96 roots = [Path(r).expanduser().resolve() for r in allowed_roots]
97 else:
98 roots = _portability_roots()
99 if not roots:
100 if not allow_uncontained:
101 raise ValueError(
102 "Portability path containment is required. Provide one of:\n"
103 " - allowed_roots=[<dir>, ...] kwarg, OR\n"
104 f" - {_PORTABILITY_ROOTS_ENV} environment variable "
105 "(os.pathsep-joined directories), OR\n"
106 " - allow_uncontained=True for trusted internal callers."
107 )
108 return resolved
109 for root in roots:
110 if resolved == root or resolved.is_relative_to(root):
111 return resolved
112 raise ValueError(f"Portability path escapes allowed roots: {resolved!s} not in {[str(r) for r in roots]}")
115def _safe_resolve(
116 path: str | Path,
117 *,
118 allowed_roots: list[str | Path] | None = None,
119 allow_uncontained: bool = False,
120) -> Path:
121 """:func:`_resolve_contained`, with every rejection as a ``ValueError``.
123 ``Path.expanduser()`` raises ``RuntimeError`` for a home directory it
124 cannot find (``~nosuchuser``), and ``resolve()`` can for a symlink loop.
125 Callers — the gateway maps ``ValueError`` to 422 — must see a rejected
126 path, not an unhandled error.
127 """
128 try:
129 return _resolve_contained(path, allowed_roots=allowed_roots, allow_uncontained=allow_uncontained)
130 except RuntimeError as exc:
131 raise ValueError(f"Portability path cannot be resolved: {os.fspath(path)!r} ({exc})") from exc
134# ---------------------------------------------------------------------------
135# AMA header
136# ---------------------------------------------------------------------------
138AMA_VERSION = 1
141@dataclass
142class AmaHeader:
143 """First line of an AMA file."""
145 bank_id: str
146 exported_at: str # ISO 8601
147 provider: str
148 memory_count: int
149 _ama_version: int = AMA_VERSION
152@dataclass
153class AmaMemory:
154 """One memory line in an AMA file."""
156 id: str
157 text: str
158 fact_type: str | None = None
159 tags: list[str] | None = None
160 metadata: Metadata | None = None
161 occurred_at: str | None = None # ISO 8601
162 created_at: str | None = None # ISO 8601
163 source: str | None = None
164 bank_id: str | None = None
165 entities: list[dict[str, str | list[str]]] | None = None
166 embedding: list[float] | None = None
169# ---------------------------------------------------------------------------
170# Writer — export a bank to AMA JSONL
171# ---------------------------------------------------------------------------
174async def export_bank(
175 recall_fn: Callable[[RecallRequest], Awaitable[RecallResult]] | None,
176 bank_id: str,
177 path: str | Path,
178 provider_name: str = "unknown",
179 include_embeddings: bool = False,
180 include_entities: bool = True,
181 batch_size: int = 100,
182 *,
183 list_fn: Callable[[str, int, int], Awaitable[list[VectorItem]]] | None = None,
184 allowed_roots: list[str | Path] | None = None,
185 allow_uncontained: bool = False,
186) -> int:
187 """Export a memory bank to AMA JSONL format.
189 Memories are enumerated in one of two ways:
191 * ``list_fn`` (preferred) — a vector store's ``list_vectors(bank_id,
192 offset, limit)``. Paged by offset over the store's stable order, so
193 every live memory in the bank is exported.
194 * ``recall_fn`` (fallback) — a single relevance-ranked
195 ``query="*"`` recall capped at ``batch_size`` hits. ``RecallRequest``
196 has no offset, so this cannot page: a bank larger than
197 ``batch_size`` is exported **incompletely**. It exists only for
198 engine providers that expose no listing API; a warning is logged
199 when the result may be truncated.
201 Args:
202 recall_fn: Async callable that takes a RecallRequest and returns
203 RecallResult. Used only when ``list_fn`` is ``None``.
204 bank_id: Bank to export.
205 path: Output file path.
206 provider_name: Provider identifier for the header.
207 include_embeddings: Include vector embeddings (not portable across models).
208 include_entities: Include extracted entities.
209 batch_size: Page size for ``list_fn``; hit cap for ``recall_fn``.
210 list_fn: Async ``(bank_id, offset, limit) -> list[VectorItem]``,
211 typically ``vector_store.list_vectors``. Must return a stable
212 order (see ``VectorStore.list_vectors``).
213 allowed_roots: Optional list of directory roots; the resolved
214 ``path`` must fall under one of them. When ``None``, falls
215 back to ``ASTROCYTE_PORTABILITY_ROOTS`` env var.
216 allow_uncontained: When True, skip path containment if neither
217 ``allowed_roots`` nor the env var is set. Use only for
218 trusted internal callers — the default ``False`` raises if
219 no containment is configured. See ``_safe_resolve``.
221 Returns:
222 Number of memories exported.
223 """
224 if batch_size < 1:
225 raise ValueError(f"batch_size must be >= 1, got {batch_size}")
226 if list_fn is None and recall_fn is None:
227 raise ValueError("export_bank requires list_fn or recall_fn")
228 path = _safe_resolve(path, allowed_roots=allowed_roots, allow_uncontained=allow_uncontained)
229 path.parent.mkdir(parents=True, exist_ok=True)
231 if list_fn is not None:
232 records = await _records_from_listing(list_fn, bank_id, batch_size)
233 else:
234 records = await _records_from_recall(recall_fn, bank_id, batch_size)
236 # Write AMA file
237 now = datetime.now(timezone.utc).isoformat()
238 header = {
239 "_ama_version": AMA_VERSION,
240 "bank_id": bank_id,
241 "exported_at": now,
242 "provider": provider_name,
243 "memory_count": len(records),
244 }
246 with open(path, "w", encoding="utf-8") as f:
247 f.write(json.dumps(header, default=str) + "\n")
248 for record in records:
249 f.write(json.dumps(record, default=str) + "\n")
251 return len(records)
254async def _records_from_listing(
255 list_fn: Callable[[str, int, int], Awaitable[list[VectorItem]]],
256 bank_id: str,
257 batch_size: int,
258) -> list[dict]:
259 """Page through ``list_fn`` until it is exhausted."""
260 records: list[dict] = []
261 seen: set[str] = set()
262 offset = 0
263 while True:
264 page = await list_fn(bank_id, offset, batch_size)
265 if not page:
266 break
267 new = [item for item in page if item.id not in seen]
268 if not new:
269 # A store that ignores ``offset`` would otherwise loop forever.
270 raise RuntimeError(
271 f"export_bank: list_vectors returned no new memories at offset {offset} "
272 f"for bank {bank_id!r}; the store's pagination is not stable"
273 )
274 for item in new:
275 seen.add(item.id)
276 records.append(
277 _ama_record(
278 memory_id=item.id,
279 text=item.text,
280 fact_type=item.fact_type,
281 tags=item.tags,
282 metadata=item.metadata,
283 occurred_at=item.occurred_at,
284 source=None,
285 bank_id=item.bank_id,
286 )
287 )
288 offset += len(page)
289 if len(page) < batch_size:
290 break
291 return records
294async def _records_from_recall(
295 recall_fn: Callable[[RecallRequest], Awaitable[RecallResult]],
296 bank_id: str,
297 batch_size: int,
298) -> list[dict]:
299 """Best-effort enumeration for providers with no listing API.
301 One ``query="*"`` recall, capped at ``batch_size``. Results are
302 relevance-ranked and cannot be paged, so a full page means the
303 bank may hold more than was exported.
304 """
305 result: RecallResult = await recall_fn(
306 RecallRequest(
307 query="*",
308 bank_id=bank_id,
309 max_results=batch_size,
310 )
311 )
312 if len(result.hits) >= batch_size:
313 logger.warning(
314 "export_bank: provider has no listing API; recall returned a full page of %d "
315 "for bank %r, so the export may be incomplete",
316 batch_size,
317 _safe_log(bank_id),
318 )
320 records: list[dict] = []
321 seen: set[str] = set()
322 for hit in result.hits:
323 key = hit.memory_id or hit.text
324 if key in seen:
325 continue
326 seen.add(key)
327 records.append(
328 _ama_record(
329 memory_id=hit.memory_id or "",
330 text=hit.text,
331 fact_type=hit.fact_type,
332 tags=hit.tags,
333 metadata=hit.metadata,
334 occurred_at=hit.occurred_at,
335 source=hit.source,
336 bank_id=hit.bank_id,
337 )
338 )
339 return records
342def _ama_record(
343 *,
344 memory_id: str,
345 text: str,
346 fact_type: str | None,
347 tags: list[str] | None,
348 metadata: Metadata | None,
349 occurred_at: datetime | None,
350 source: str | None,
351 bank_id: str | None,
352) -> dict:
353 """Build one AMA memory line. Optional fields are omitted when empty."""
354 record: dict = {"id": memory_id, "text": text}
355 if fact_type:
356 record["fact_type"] = fact_type
357 if tags:
358 record["tags"] = tags
359 if metadata:
360 record["metadata"] = metadata
361 if occurred_at:
362 record["occurred_at"] = occurred_at.isoformat()
363 if source:
364 record["source"] = source
365 if bank_id:
366 record["bank_id"] = bank_id
367 # Embeddings and entities would come from provider-specific data;
368 # Phase 1 exports only the fields above.
369 return record
372# ---------------------------------------------------------------------------
373# Reader — iterate AMA JSONL lines
374# ---------------------------------------------------------------------------
377def read_ama_header(
378 path: str | Path,
379 *,
380 allowed_roots: list[str | Path] | None = None,
381 allow_uncontained: bool = False,
382) -> AmaHeader:
383 """Read and validate the AMA header (first line).
385 See ``export_bank`` for ``allowed_roots`` and ``allow_uncontained`` semantics.
386 """
387 path = _safe_resolve(path, allowed_roots=allowed_roots, allow_uncontained=allow_uncontained)
388 with open(path, encoding="utf-8") as f:
389 first_line = f.readline().strip()
390 if not first_line:
391 raise ValueError(f"AMA file is empty: {path}")
392 data = json.loads(first_line)
393 if "_ama_version" not in data:
394 raise ValueError(f"Not a valid AMA file (missing _ama_version): {path}")
395 if data["_ama_version"] != AMA_VERSION:
396 raise ValueError(f"Unsupported AMA version {data['_ama_version']} (expected {AMA_VERSION})")
397 # Validate required field types
398 for field in ("bank_id", "exported_at", "provider"):
399 if not isinstance(data.get(field), str):
400 raise ValueError(f"AMA header field '{field}' must be a string: {path}")
401 if not isinstance(data.get("memory_count"), int):
402 raise ValueError(f"AMA header field 'memory_count' must be an integer: {path}")
403 return AmaHeader(
404 bank_id=data["bank_id"],
405 exported_at=data["exported_at"],
406 provider=data["provider"],
407 memory_count=data["memory_count"],
408 _ama_version=data["_ama_version"],
409 )
412def iter_ama_memories(
413 path: str | Path,
414 *,
415 allowed_roots: list[str | Path] | None = None,
416 allow_uncontained: bool = False,
417) -> list[AmaMemory]:
418 """Read all memory records from an AMA file (skips header).
420 See ``export_bank`` for ``allowed_roots`` and ``allow_uncontained`` semantics.
421 """
422 path = _safe_resolve(path, allowed_roots=allowed_roots, allow_uncontained=allow_uncontained)
423 memories: list[AmaMemory] = []
424 with open(path, encoding="utf-8") as f:
425 # Skip header
426 f.readline()
427 for line_num, line in enumerate(f, start=2):
428 line = line.strip()
429 if not line:
430 continue
431 try:
432 data = json.loads(line)
433 if not isinstance(data, dict) or "id" not in data or "text" not in data:
434 logger.warning("AMA line %d: missing required fields (id, text)", line_num)
435 continue
436 memories.append(
437 AmaMemory(
438 id=data["id"],
439 text=data["text"],
440 fact_type=data.get("fact_type"),
441 tags=data.get("tags"),
442 metadata=data.get("metadata"),
443 occurred_at=data.get("occurred_at"),
444 created_at=data.get("created_at"),
445 source=data.get("source"),
446 bank_id=data.get("bank_id"),
447 entities=data.get("entities"),
448 embedding=data.get("embedding"),
449 )
450 )
451 except (json.JSONDecodeError, KeyError, TypeError) as exc:
452 logger.warning("AMA line %d: %s", line_num, exc)
453 continue
454 return memories
457# ---------------------------------------------------------------------------
458# Import — load AMA into a bank
459# ---------------------------------------------------------------------------
462@dataclass
463class ImportResult:
464 imported: int
465 skipped: int
466 errors: int
469async def import_bank(
470 retain_fn,
471 bank_id: str,
472 path: str | Path,
473 on_conflict: Literal["skip", "overwrite", "error"] = "skip",
474 progress_fn=None,
475 *,
476 allowed_roots: list[str | Path] | None = None,
477 allow_uncontained: bool = False,
478) -> ImportResult:
479 """Import memories from an AMA file into a bank.
481 Args:
482 retain_fn: Async callable that takes a RetainRequest and returns RetainResult.
483 Typically ``brain._do_retain``.
484 bank_id: Target bank (may differ from source bank in AMA).
485 path: Path to AMA JSONL file.
486 on_conflict: How to handle memories with IDs that already exist.
487 progress_fn: Optional callback(imported, total) for progress reporting.
488 allowed_roots: See ``export_bank``.
489 allow_uncontained: See ``export_bank``.
491 Returns:
492 ImportResult with counts.
493 """
494 header = read_ama_header(path, allowed_roots=allowed_roots, allow_uncontained=allow_uncontained)
495 memories = iter_ama_memories(path, allowed_roots=allowed_roots, allow_uncontained=allow_uncontained)
497 imported = 0
498 skipped = 0
499 errors = 0
501 for i, mem in enumerate(memories):
502 try:
503 # Parse occurred_at if present
504 occurred_at = None
505 if mem.occurred_at:
506 try:
507 occurred_at = datetime.fromisoformat(mem.occurred_at)
508 except ValueError:
509 logger.debug("Skipping unparseable occurred_at: %s", _safe_log(mem.occurred_at))
511 request = RetainRequest(
512 content=mem.text,
513 bank_id=bank_id,
514 metadata=mem.metadata,
515 tags=mem.tags,
516 occurred_at=occurred_at,
517 source=mem.source or f"import:ama:{header.provider}",
518 content_type="text",
519 )
521 result = await retain_fn(request)
523 if result.stored:
524 imported += 1
525 elif result.deduplicated and on_conflict == "skip":
526 skipped += 1
527 elif result.deduplicated and on_conflict == "error":
528 errors += 1
529 else:
530 skipped += 1
532 except Exception as exc:
533 logger.warning("AMA import line %d failed: %s", i + 2, exc)
534 errors += 1
536 if progress_fn and (i + 1) % 10 == 0:
537 progress_fn(imported, header.memory_count)
539 return ImportResult(imported=imported, skipped=skipped, errors=errors)