Coverage for astrocyte/harness/memories.py: 96%

1"""``astrocyte memory`` — see and remove what automatic memory captured.

3Capture is automatic, so inspection has to be easy: without it the only way

4to check what an agent will be reminded of, or to take back something said

5in passing, is to open the database by hand. Every command defaults to the

6bank of the project you are standing in — the same bank the hooks and the MCP

7server use — so ``astrocyte memory`` answers "what does my agent remember

8about this repo?".

10Search runs the agent daemon's recall path (no LLM query expansion), so it

11costs nothing and returns what a prompt would be matched against.

13``forget`` erases. The core forget is a soft delete (history stays visible to

14``as_of`` recall), which would leave a removed key's text in the database

15file; where the store can purge, the forgotten rows are then erased from disk.

16"""

18from __future__ import annotations

20import hashlib

21import json

22import os

23import re

24import sys

25from argparse import Namespace

26from datetime import datetime, timezone

27from pathlib import Path

28from typing import Any

30from .paths import config_path

31from .project import project_bank, project_root

33LIST_DEFAULT = 20

34SEARCH_DEFAULT = 8

35MIN_PREFIX = 4 # shortest id prefix `forget` accepts

38def open_local(cfg_path: Path, *, noisy_bank_detection: bool = True) -> tuple[Any, Any]:

39 """``(pipeline, brain)`` for the local config, wired the way the agent

40 daemon needs it: no LLM call on recall, none in the background.

42 ``noisy_bank_detection=False`` is for the user's own commands: noisy-bank

43 detection guards against an *agent* writing junk or looping, and a bulk

44 ``astrocyte memory import`` of short doc sections is neither — flagging

45 it would only print a warning at the person who asked for the import.

46 """

47 from astrocyte import Astrocyte

48 from astrocyte.config import load_config

49 from astrocyte.wiring import build_pipeline

51 config = load_config(str(cfg_path))

52 if not noisy_bank_detection:

53 without_noisy_bank_detection(config)

54 # Recall-time query expansion calls the LLM (5–9 s via a CLI provider).

55 # Observation consolidation issues an LLM call per retained memory in the

56 # background — on a CLI provider, the user's subscription.

57 pipeline = build_pipeline(config, enable_multi_query_expansion=False, enable_observation_consolidation=False)

58 brain = Astrocyte(config)

59 brain.set_pipeline(pipeline)

60 return pipeline, brain

63def without_noisy_bank_detection(config: Any) -> Any:

64 """``config`` with noisy-bank detection off everywhere (in place)."""

65 config.signal_quality.noisy_bank.enabled = False

66 for bank in (config.banks or {}).values():

67 if bank.signal_quality is not None:

68 bank.signal_quality.noisy_bank.enabled = False

69 return config

72def _bank(args: Namespace) -> str:

73 return args.bank or project_bank(str(Path(args.project or os.getcwd()).expanduser().resolve()))

76def _when(dt: datetime | None) -> str:

77 return dt.date().isoformat() if dt else " "

80def _line(text: str) -> str:

81 from .agentd import render_memory

83 return render_memory(text, None).removeprefix("- ")

86def _source(metadata: Any) -> str:

87 return str((metadata or {}).get("source") or "")

90async def _all_items(store: Any, bank: str) -> list[Any]:

91 items, offset = [], 0

92 while True:

93 page = await store.list_vectors(bank, offset=offset, limit=500)

94 items += page

95 if len(page) < 500:

96 return items

97 offset += 500

100# ── commands ─────────────────────────────────────────────────────────────

101

102

103async def _list(args: Namespace, pipeline: Any, _brain: Any) -> int:

104 store, bank = pipeline.vector_store, _bank(args)

105 recent = getattr(store, "list_recent_vectors", None)

106 items = await recent(bank, limit=args.limit) if recent else (await store.list_vectors(bank, limit=args.limit))

107 from . import team

108 from .agentd import who_saved

109

110 teammates = team.pulled_ids(bank)

111 saved_by = {i.id: who_saved(i.metadata) for i in items if i.id in teammates}

112 if args.json:

113 print(

114 json.dumps(

115 [

116 {

117 "id": i.id,

118 "text": i.text,

119 "source": _source(i.metadata),

120 "saved_by": saved_by.get(i.id),

121 "when": (i.occurred_at or i.retained_at).isoformat()

122 if (i.occurred_at or i.retained_at)

123 else None,

124 }

125 for i in items

126 ],

127 indent=2,

128 )

129 )

130 return 0

131 if not items:

132 print(f"No memories in {bank}.")

133 return 0

134 count = f"{len(items)} memor{'y' if len(items) == 1 else 'ies'}"

135 print(f"{bank} — {len(items)} most recent" if len(items) == args.limit else f"{bank} — {count}")

136 for i in items:

137 who = f"({saved_by[i.id]}) " if saved_by.get(i.id) else ""

138 print(f" {i.id[:8]} {_when(i.occurred_at or i.retained_at)} {_source(i.metadata):<11} {who}{_line(i.text)}")

139 print("\nRemove one with: astrocyte memory forget <id>")

140 return 0

141

142

143async def _search(args: Namespace, _pipeline: Any, brain: Any) -> int:

144 bank = _bank(args)

145 result = await brain.recall(args.query, bank_id=bank, max_results=args.limit)

146 hits = [h for h in result.hits if h.text]

147 if args.json:

148 print(json.dumps([{"id": h.memory_id, "score": round(h.score, 4), "text": h.text} for h in hits], indent=2))

149 return 0

150 if not hits:

151 print(f"Nothing in {bank} matches.")

152 return 0

153 for h in hits:

154 print(f" {(h.memory_id or '')[:8]} {_when(h.occurred_at or h.retained_at)} {_line(h.text)}")

155 return 0

156

157

158async def _forget(args: Namespace, pipeline: Any, brain: Any) -> int:

159 bank = _bank(args)

160 items = await _all_items(pipeline.vector_store, bank)

161 if args.all:

162 if not items:

163 print(f"No memories in {bank}.")

164 return 0

165 if not args.yes:

166 print(

167 f"This permanently removes all {len(items)} memories in {bank}.\nRe-run with --yes to confirm.",

168 file=sys.stderr,

169 )

170 return 1

171 joined = _joined(args, bank)

172 if joined is not None:

173 # This machine's copy only; the team's stays (and isn't pulled back).

174 erased = await joined[0].forget_shared(joined[1], bank, [i.id for i in items], pipeline, brain,

175 team_wide=False)

176 print(f"Removed {erased} memories from {bank} on this machine. The team's copies stay on the "

177 "gateway; erase one for everyone with: astrocyte memory forget <id> --team")

178 return 0

179 result = await brain.forget(bank, scope="all")

180 print(f"Removed {result.deleted_count} memories from {bank}.{await _erase(pipeline, bank, None)}")

181 return 0

182 if not args.ids:

183 print("Name the memories to remove (ids from `astrocyte memory`), or use --all.", file=sys.stderr)

184 return 2

185 if any(len(p) < MIN_PREFIX for p in args.ids):

186 print(f" ✗ give at least {MIN_PREFIX} characters of each id", file=sys.stderr)

187 return 2

188 chosen = _resolve_ids(args, items)

189 if chosen is None:

190 return 1

191 joined = _joined(args, bank)

192 if joined is not None:

193 team, entry = joined

194 shared = [i for i in chosen if team.SyncState.load(bank).shared(i)]

195 team_wide, local = getattr(args, "team", False), getattr(args, "local", False)

196 if shared and not (team_wide or local):

197 who = {i.id: i for i in items}

198 print("Shared with the team: " + ", ".join(f"{i[:8]}" + _by(who[i], team) for i in shared) + ".\n"

199 "Erase for everyone with --team (needs forget permission on the gateway), "

200 "or on this machine only with --local.", file=sys.stderr)

201 return 1

202 if shared or local:

203 try:

204 erased = await team.forget_shared(entry, bank, chosen, pipeline, brain, team_wide=bool(team_wide))

205 except team.TeamError as e:

206 print(f"astrocyte memory forget: {e}", file=sys.stderr)

207 return 1

208 where = "for everyone on the team" if team_wide and shared else "on this machine"

209 noun = "memory" if erased == 1 else "memories"

210 print(f"Removed {erased} {noun} from {bank} {where}. Erased from disk.")

211 return 0

212 result = await brain.forget(bank, memory_ids=chosen)

213 noun = "memory" if result.deleted_count == 1 else "memories"

214 print(f"Removed {result.deleted_count} {noun} from {bank}.{await _erase(pipeline, bank, chosen)}")

215 return 0

216

217

218def _resolve_ids(args: Namespace, items: list[Any]) -> list[str] | None:

219 """Unique id prefixes to ids, as `forget` takes them; None after printing why not."""

220 chosen: list[str] = []

221 for prefix in args.ids:

222 if len(prefix) < MIN_PREFIX:

223 print(f" ✗ {prefix}: give at least {MIN_PREFIX} characters of the id", file=sys.stderr)

224 return None

225 matches = [i.id for i in items if i.id.startswith(prefix)]

226 if len(matches) != 1:

227 why = "no memory in this project has that id" if not matches else f"matches {len(matches)} memories"

228 print(f" ✗ {prefix}: {why}", file=sys.stderr)

229 return None

230 chosen.append(matches[0])

231 return chosen

232

233

234async def _share(args: Namespace, pipeline: Any, _brain: Any) -> int:

235 from . import team

236

237 bank = _bank(args)

238 chosen = _resolve_ids(args, await _all_items(pipeline.vector_store, bank))

239 if chosen is None:

240 return 1

241 try:

242 ids = await team.share(bank, chosen, pipeline.vector_store)

243 except team.TeamError as e:

244 print(f"astrocyte memory share: {e}", file=sys.stderr)

245 return 1

246 if not ids:

247 print("Nothing to share: those are teammates' memories.")

248 return 0

249 noun = "memory" if len(ids) == 1 else "memories (with the chunks retained together)"

250 when = ("at the next sync (`astrocyte team sync` sends it now)" if _joined(args, bank)

251 else "once this project joins a team (`astrocyte team join`)")

252 print(f"Shared {len(ids)} {noun}: goes to the team {when}.")

253 return 0

254

255

256async def _unshare(args: Namespace, pipeline: Any, _brain: Any) -> int:

257 from . import team

258

259 bank = _bank(args)

260 chosen = _resolve_ids(args, await _all_items(pipeline.vector_store, bank))

261 if chosen is None:

262 return 1

263 joined = _joined(args, bank)

264 try:

265 ids, removed = await team.unshare(joined[1] if joined else None, bank, chosen, pipeline.vector_store)

266 except team.TeamError as e:

267 print(f"astrocyte memory unshare: {e}", file=sys.stderr)

268 return 1

269 gone = f" Taken off the team's gateway: {removed}; teammates' copies go at their next sync." if removed else ""

270 print(f"Kept {len(ids)} {'memory' if len(ids) == 1 else 'memories'} on this machine only.{gone}")

271 return 0

272

273

274def _joined(args: Namespace, bank: str) -> tuple[Any, dict[str, Any]] | None:

275 """``(team module, membership)`` when this bank is shared with a team."""

276 from . import team

277

278 cfg = getattr(args, "config", None)

279 entry = team.load_memberships(Path(cfg).expanduser() if cfg else config_path()).get(bank)

280 return (team, entry) if entry is not None else None

281

282

283def _by(item: Any, team: Any) -> str:

284 from .agentd import who_saved

285

286 who = who_saved(item.metadata) if item.id in team.pulled_ids(item.bank_id) else None

287 return f" (saved by {who})" if who else ""

288

289

290async def _erase(pipeline: Any, bank: str, ids: list[str] | None) -> str:

291 purge = getattr(pipeline.vector_store, "purge", None)

292 if purge is None:

293 return f"\n{type(pipeline.vector_store).__name__} keeps forgotten memories for history; they no longer recall."

294 await purge(bank, ids)

295 return " Erased from disk."

296

297

298async def _banks(args: Namespace, pipeline: Any, _brain: Any) -> int:

299 lister = getattr(pipeline.vector_store, "list_banks", None)

300 if lister is None:

301 print(f"{type(pipeline.vector_store).__name__} can't enumerate banks; name one with --bank.", file=sys.stderr)

302 return 2

303 banks = await lister()

304 here = _bank(Namespace(bank=None, project=args.project))

305 if args.json:

306 print(

307 json.dumps(

308 [{"bank": b, "memories": n, "newest": t.isoformat() if t else None} for b, n, t in banks], indent=2

309 )

310 )

311 return 0

312 if not banks:

313 print("No memories yet.")

314 return 0

315 for b, n, t in banks:

316 print(f" {'→' if b == here else ' '} {b:<48} {n:>6} newest {_when(t)}")

317 return 0

318

319

320# ── import / export ──────────────────────────────────────────────────────

321

322IMPORT_SUFFIXES = (".md", ".mdx", ".markdown", ".txt")

323IMPORT_SKIP_DIRS = frozenset({"node_modules", "venv", ".venv", "dist", "build", "__pycache__", "site-packages"})

324IMPORT_MAX_FILE_BYTES = 512 * 1024

325SECTION_MAX_CHARS = 3_000

326_HEADING = re.compile(r"^#{1,3}\s+\S")

327_FENCE = re.compile(r"^\s*(```|~~~)")

328

329

330def import_files(paths: list[str]) -> list[Path]:

331 """Files named, plus the Markdown/text files under directories named —

332 skipping hidden, vendored and build directories and very large files."""

333 found: list[Path] = []

334 for raw in paths:

335 p = Path(raw).expanduser()

336 if p.is_file():

337 found.append(p)

338 elif p.is_dir():

339 for f in sorted(p.rglob("*")):

340 rel = f.relative_to(p).parts

341 if any(part.startswith(".") or part in IMPORT_SKIP_DIRS for part in rel[:-1]):

342 continue

343 if f.is_file() and f.suffix.lower() in IMPORT_SUFFIXES and f.stat().st_size <= IMPORT_MAX_FILE_BYTES:

344 found.append(f)

345 else:

346 raise FileNotFoundError(raw)

347 return list(dict.fromkeys(found))

348

349

350def sections(text: str) -> list[str]:

351 """Split Markdown at #–### headings (not inside code fences), each section

352 keeping its heading; long sections are split again at blank lines. One

353 focused memory per topic recalls better than one memory per file."""

354 parts: list[list[str]] = [[]]

355 fenced = False

356 for line in text.splitlines():

357 if _FENCE.match(line):

358 fenced = not fenced

359 if not fenced and _HEADING.match(line) and any(x.strip() for x in parts[-1]):

360 parts.append([])

361 parts[-1].append(line)

362 out: list[str] = []

363 for part in parts:

364 body = "\n".join(part).strip()

365 while len(body) > SECTION_MAX_CHARS:

366 cut = body.rfind("\n\n", 0, SECTION_MAX_CHARS)

367 cut = cut if cut > SECTION_MAX_CHARS // 3 else SECTION_MAX_CHARS

368 out.append(body[:cut].strip())

369 body = body[cut:].strip()

370 if body:

371 out.append(body)

372 return out

373

374

375def _digest(text: str) -> str:

376 return hashlib.sha256(text.encode()).hexdigest()[:16]

377

378

379def _label(path: Path, args: Namespace) -> str:

380 """How an imported file is identified: relative to the project root, so

381 the same file in another clone of the repo is recognised."""

382 root = project_root(str(Path(args.project or os.getcwd()).expanduser().resolve()))

383 try:

384 return str(path.resolve().relative_to(root))

385 except ValueError:

386 return str(path.resolve())

387

388

389async def _import(args: Namespace, pipeline: Any, brain: Any) -> int:

390 """Seed a project's memory from documents, as a sync: re-importing a file

391 adds its new sections and removes the ones no longer in it."""

392 bank = _bank(args)

393 if len(args.paths) == 1 and args.paths[0].endswith(".ama.jsonl"):

394 return await _import_archive(args, brain, bank)

395 try:

396 files = import_files(args.paths)

397 except FileNotFoundError as e:

398 print(f" ✗ no such file or directory: {e}", file=sys.stderr)

399 return 2

400 if not files:

401 print("Nothing to import (Markdown and text files: " + ", ".join(IMPORT_SUFFIXES) + ").")

402 return 0

403 held: dict[str, dict[str, list[str]]] = {} # import path → section hash → memory ids

404 for item in await _all_items(pipeline.vector_store, bank):

405 md = item.metadata or {}

406 if md.get("import_path") and md.get("import_hash"):

407 held.setdefault(str(md["import_path"]), {}).setdefault(str(md["import_hash"]), []).append(item.id)

408 added = kept = removed = duplicates = 0

409 for f in files:

410 label = _label(f, args)

411 try:

412 text = f.read_text(encoding="utf-8")

413 except (OSError, UnicodeDecodeError) as e:

414 print(f" ✗ {label}: {e}", file=sys.stderr)

415 continue

416 wanted = {_digest(s): s for s in sections(text)}

417 have = held.get(label, {})

418 new = [(h, s) for h, s in wanted.items() if h not in have]

419 stale = [mid for h, ids in have.items() if h not in wanted for mid in ids]

420 kept += len(wanted) - len(new)

421 if args.dry_run:

422 print(f" {label}: {len(new)} to add, {len(stale)} to remove, {len(wanted) - len(new)} unchanged")

423 added, removed = added + len(new), removed + len(stale)

424 continue

425 when = datetime.fromtimestamp(f.stat().st_mtime, tz=timezone.utc)

426 stored = dup = 0

427 for digest, section in new:

428 # The pipeline's dedup checks the store too, so a section already

429 # said elsewhere — in this run or an earlier one — is skipped.

430 result = await brain.retain(

431 section,

432 bank_id=bank,

433 occurred_at=when,

434 source="import",

435 metadata={"source": "import", "import_path": label, "import_hash": digest},

436 )

437 if result.stored:

438 stored += 1

439 elif getattr(result, "deduplicated", False):

440 dup += 1

441 if stale:

442 removed += (await brain.forget(bank, memory_ids=stale)).deleted_count

443 await _erase(pipeline, bank, stale)

444 added, duplicates = added + stored, duplicates + dup

445 dup_note = f", {dup} duplicates skipped" if dup else ""

446 print(f" ✓ {label}: {stored} added, {len(stale)} removed, {len(wanted) - len(new)} unchanged{dup_note}")

447 verb = "Would add" if args.dry_run else "Added"

448 dup_note = f", skipped {duplicates} near-duplicates of existing memories" if duplicates else ""

449 print(f"\n{verb} {added}, removed {removed}, kept {kept}{dup_note} in {bank}.")

450 return 0

451

452

453async def _import_archive(args: Namespace, brain: Any, bank: str) -> int:

454 path = Path(args.paths[0]).expanduser().resolve()

455 if args.dry_run:

456 print(f" would import {path} into {bank}")

457 return 0

458 result = await brain.import_bank(bank, str(path), allowed_roots=[str(path.parent)])

459 print(f"Imported {result.imported} memories into {bank} ({result.skipped} skipped).")

460 return 0 if not getattr(result, "errors", None) else 1

461

462

463async def _export(args: Namespace, pipeline: Any, _brain: Any) -> int:

464 """Write the bank as an Astrocyte Memory Archive (AMA JSONL), read straight

465 from the store so nothing is missed. Owner-only: it holds conversations."""

466 from astrocyte.portability import AMA_VERSION

467

468 bank = _bank(args)

469 items = await _all_items(pipeline.vector_store, bank)

470 dest = Path(args.file).expanduser()

471 dest.parent.mkdir(parents=True, exist_ok=True)

472 fd = os.open(dest, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)

473 with os.fdopen(fd, "w", encoding="utf-8") as fh:

474 fh.write(

475 json.dumps(

476 {

477 "_ama_version": AMA_VERSION,

478 "bank_id": bank,

479 "exported_at": datetime.now(timezone.utc).isoformat(),

480 "provider": "astrocyte-cli",

481 "memory_count": len(items),

482 }

483 )

484 + "\n"

485 )

486 for item in items:

487 record: dict[str, Any] = {"id": item.id, "text": item.text, "bank_id": bank}

488 if item.fact_type:

489 record["fact_type"] = item.fact_type

490 if item.tags:

491 record["tags"] = item.tags

492 if item.metadata:

493 record["metadata"] = item.metadata

494 if item.metadata.get("source"):

495 record["source"] = item.metadata["source"]

496 if item.occurred_at or item.retained_at:

497 record["occurred_at"] = (item.occurred_at or item.retained_at).isoformat()

498 fh.write(json.dumps(record, default=str) + "\n")

499 print(f"Exported {len(items)} memories from {bank} to {dest}.")

500 return 0

501

502

503_COMMANDS = {"list": _list, "search": _search, "forget": _forget, "banks": _banks, "import": _import, "export": _export,

504 "share": _share, "unshare": _unshare}

505

506

507def run(args: Namespace) -> int:

508 cfg = Path(args.config).expanduser() if args.config else config_path()

509 if not cfg.is_file():

510 print(f"No Astrocyte config at {cfg}. Run: astrocyte setup", file=sys.stderr)

511 return 2

512 pipeline, brain = open_local(cfg, noisy_bank_detection=False)

513

514 async def go() -> int:

515 try:

516 return await _COMMANDS[args.memory_command or "list"](args, pipeline, brain)

517 finally:

518 close = getattr(pipeline.vector_store, "close", None)

519 if close is not None:

520 await close()

521

522 import asyncio # here, not at import: `astrocyte hook` builds this parser too

523

524 return asyncio.run(go())

525

526

527def register(sub) -> None:

528 memory = sub.add_parser(

529 "memory",

530 help="See, search and remove what your agents remember (default: this project's recent memories)",

531 description="Every command works on the memory bank of the current project — the one automatic "

532 "memory and the MCP server use — unless --bank or --project says otherwise.",

533 )

534

535 def scope(p, *, limit: int | None = None) -> None:

536 p.add_argument("--project", help="project directory (default: the current directory)")

537 p.add_argument("--bank", help="a bank id instead of a project (see: astrocyte memory banks)")

538 p.add_argument("--config", help="config path (default: ~/.config/astrocyte/astrocyte.yaml)")

539 p.add_argument("--json", action="store_true", help="machine-readable output")

540 if limit:

541 p.add_argument("-n", "--limit", type=int, default=limit, help=f"how many to show (default {limit})")

542

543 scope(memory, limit=LIST_DEFAULT)

544 memory.set_defaults(func=run, memory_command=None)

545 cmds = memory.add_subparsers(dest="memory_command")

546 scope(cmds.add_parser("list", help="most recent memories"), limit=LIST_DEFAULT)

547 search = cmds.add_parser("search", help="memories matching a query, as an agent prompt would see them")

548 search.add_argument("query")

549 scope(search, limit=SEARCH_DEFAULT)

550 forget = cmds.add_parser("forget", help="remove memories by id (or prefix), or all of this project's")

551 forget.add_argument("ids", nargs="*", help="ids or unique prefixes, from `astrocyte memory`")

552 forget.add_argument("--all", action="store_true", help="every memory in the bank")

553 forget.add_argument("--yes", action="store_true", help="confirm --all")

554 scope_ = forget.add_mutually_exclusive_group()

555 scope_.add_argument("--team", action="store_true",

556 help="a memory shared with the team: erase it for everyone (needs forget permission)")

557 scope_.add_argument("--local", action="store_true",

558 help="a memory shared with the team: erase it on this machine only")

559 scope(forget)

560 scope(cmds.add_parser("banks", help="every bank with memories, this project's marked →"))

561 share = cmds.add_parser("share", help="share memories with the team that would stay here "

562 "(a captured turn, or one tagged private)")

563 share.add_argument("ids", nargs="+", help="ids or unique prefixes, from `astrocyte memory`")

564 scope(share)

565 unshare = cmds.add_parser("unshare", help="keep memories on this machine only (taken off the team's "

566 "gateway if they are already there)")

567 unshare.add_argument("ids", nargs="+", help="ids or unique prefixes, from `astrocyte memory`")

568 scope(unshare)

569 imp = cmds.add_parser(

570 "import",

571 help="seed memory from Markdown/text files or directories (or an .ama.jsonl archive)",

572 description="Each file is split at its headings into one memory per section. Re-importing a file "

573 "syncs it: new sections are added, sections no longer in the file are removed. Credentials are "

574 "redacted as with any memory. Typical first step: astrocyte memory import CLAUDE.md AGENTS.md docs/",

575 )

576 imp.add_argument("paths", nargs="+", help="files, directories, or one .ama.jsonl archive")

577 imp.add_argument("--dry-run", action="store_true", help="show what would change, change nothing")

578 scope(imp)

579 exp = cmds.add_parser("export", help="write the bank to an .ama.jsonl archive (portable; re-import elsewhere)")

580 exp.add_argument("file", help="destination, e.g. project.ama.jsonl")

581 scope(exp)