Coverage for astrocyte/mcp.py: 52%
349 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""Astrocyte MCP server — exposes memory tools via Model Context Protocol.
3Usage:
4 astrocyte-mcp --config astrocyte.yaml
5 astrocyte-mcp --config astrocyte.yaml --transport sse --port 8080
7See docs/_design/mcp-server.md for the full specification.
9Note: Access control in MCP deployments typically uses one principal from
10``mcp.principal`` (defaults to ``agent:mcp``), or pass
11``astrocyte_context=...`` to :func:`create_mcp_server` when your host maps
12callers to identity. Per-caller identity still requires transport-level auth
13at the MCP host.
14"""
16from __future__ import annotations
18import argparse
19import json
20import logging
21import os
22import sys
23from datetime import datetime
24from pathlib import Path
25from typing import Any
27from fastmcp import FastMCP
29from astrocyte._astrocyte import Astrocyte
30from astrocyte._mcp_identity import JwtIdentityMiddleware, build_jwt_middleware
31from astrocyte.config import AstrocyteConfig, load_config
32from astrocyte.types import AstrocyteContext
34logger = logging.getLogger("astrocyte.mcp")
36# ---------------------------------------------------------------------------
37# Server factory
38# ---------------------------------------------------------------------------
41def server_instructions(mcp_cfg) -> str:
42 """What the agent is told this server offers — only tools it registers.
44 It used to advertise memory_reflect and memory_forget unconditionally, so
45 with the defaults (forget hidden) an agent asked to forget something went
46 looking for a tool that does not exist.
47 """
48 uses = ["memory_retain to store information", "memory_recall to search memories"]
49 if mcp_cfg.expose_reflect:
50 uses.append("memory_reflect to synthesize answers")
51 if mcp_cfg.expose_forget:
52 uses.append("memory_forget to remove memories")
53 text = "Astrocyte memory server. Use " + ", ".join(uses[:-1]) + ", and " + uses[-1] + "."
54 if not mcp_cfg.expose_forget:
55 text += (
56 " This server cannot delete memories. If the user asks to remove one, find it with "
57 "memory_recall and give them its memory_id; on a local install they remove it with "
58 "`astrocyte memory forget <memory_id> --bank <bank_id>`, which erases it from disk."
59 )
60 return text
63def create_mcp_server(
64 brain: Astrocyte,
65 config: AstrocyteConfig,
66 *,
67 astrocyte_context: AstrocyteContext | None = None,
68 jwt_middleware: JwtIdentityMiddleware | None = None,
69 default_bank_fallback: str | None = None,
70) -> FastMCP:
71 """Create a FastMCP server wired to an Astrocyte instance.
73 Identity resolution order (per call):
75 1. If ``astrocyte_context`` was passed at server-creation time, it is
76 used as-is (legacy embedding pattern). This takes precedence over
77 all header-based resolution so tests and embedded deployments that
78 pre-bind identity are unaffected.
79 2. Else if ``jwt_middleware`` is set OR
80 ``config.identity.jwt_middleware.enabled`` is true, per-call
81 identity is resolved from the inbound request's Authorization
82 header (identity spec §3 Gap 1 wiring). See
83 ``docs/_plugins/jwt-identity-middleware.md`` for operator setup.
84 3. Else the server falls back to the static ``mcp.principal`` label
85 (default ``agent:mcp``) — the pre-middleware behavior.
87 Args:
88 jwt_middleware: Optional pre-built middleware. When None but the
89 config enables JWT middleware, the server constructs one via
90 :func:`build_jwt_middleware`. Pass explicitly to inject a
91 test decoder or bypass PyJWT.
92 """
94 mcp_cfg = config.mcp
95 mcp = FastMCP(name="astrocyte-memory", instructions=server_instructions(mcp_cfg))
97 # Static context — still built for the pre-middleware path (legacy
98 # deployments that don't enable JWT middleware). When middleware is
99 # enabled, this is only used if header resolution is somehow bypassed.
100 static_ctx = astrocyte_context or AstrocyteContext(
101 principal=mcp_cfg.principal or "agent:mcp",
102 )
104 # Build or accept JWT middleware. Explicit > config > none.
105 if jwt_middleware is None and config.identity.jwt_middleware.enabled:
106 jwt_middleware = build_jwt_middleware(config.identity.jwt_middleware)
108 def _resolve_context() -> AstrocyteContext:
109 """Resolve the AstrocyteContext for the current request.
111 Called on every tool invocation so the middleware sees each call's
112 headers. Static ``astrocyte_context`` (when set at server creation)
113 always wins — that preserves backward compatibility for embedded
114 hosts that pre-bind identity.
115 """
116 # An explicit pre-bound context always wins (embedded / test use).
117 if astrocyte_context is not None:
118 return astrocyte_context
119 # Middleware path — pull headers, resolve identity, build ctx.
120 if jwt_middleware is not None:
121 try:
122 from fastmcp.server.dependencies import get_http_headers
124 headers = get_http_headers(include={"authorization"})
125 except Exception:
126 # stdio transport or pre-request context — no headers.
127 # Middleware policy decides how to treat the missing
128 # header (fail-closed vs anonymous).
129 headers = {}
130 return jwt_middleware.resolve(headers)
131 # No middleware, no pre-bound ctx — fall through to static.
132 return static_ctx
134 # Default bank. ``default_bank_fallback`` (the project bank, for a local
135 # stdio server) applies only when the config names none: without it a
136 # memory_retain that omits bank_id fails, and the agent's memories would
137 # never meet the ones its hooks capture.
138 default_bank = mcp_cfg.default_bank_id or default_bank_fallback
140 _MAX_TAGS = 20
141 _MAX_TAG_LENGTH = 255
143 def _resolve_bank(bank_id: str | None) -> str:
144 if bank_id:
145 return bank_id
146 if default_bank:
147 return default_bank
148 raise ValueError("bank_id is required (no default_bank_id configured)")
150 def _validate_tags(tags: list[str] | None) -> list[str] | None:
151 if not tags:
152 return tags
153 if len(tags) > _MAX_TAGS:
154 raise ValueError(f"Too many tags ({len(tags)}), max {_MAX_TAGS}")
155 for tag in tags:
156 if not isinstance(tag, str) or len(tag) > _MAX_TAG_LENGTH:
157 raise ValueError(f"Invalid tag (must be string, max {_MAX_TAG_LENGTH} chars)")
158 return tags
160 # ── memory_retain ──────────────────────────────────────────────
162 @mcp.tool()
163 async def memory_retain(
164 content: str,
165 bank_id: str | None = None,
166 tags: list[str] | None = None,
167 metadata: dict[str, str] | None = None,
168 occurred_at: str | None = None,
169 ) -> str:
170 """Store content into memory.
172 Args:
173 content: The text to memorize.
174 bank_id: Memory bank to store in. Uses default if omitted.
175 tags: Optional tags for filtering.
176 metadata: Optional key-value metadata.
177 occurred_at: ISO 8601 timestamp of when the event happened.
178 """
179 try:
180 bid = _resolve_bank(bank_id)
181 tags = _validate_tags(tags)
182 kwargs: dict[str, Any] = {}
183 if occurred_at:
184 try:
185 kwargs["occurred_at"] = datetime.fromisoformat(occurred_at)
186 except ValueError:
187 return json.dumps({"stored": False, "error": "Invalid occurred_at format; expected ISO 8601"})
188 result = await brain.retain(
189 content,
190 bank_id=bid,
191 tags=tags,
192 metadata=metadata,
193 context=_resolve_context(),
194 **kwargs,
195 )
196 if result.stored:
197 return json.dumps({"stored": True, "memory_id": result.memory_id})
198 return json.dumps({"stored": False, "error": "Failed to store memory"})
199 except Exception as exc:
200 logger.exception("memory_retain failed")
201 return json.dumps({"stored": False, "error": type(exc).__name__})
203 # ── memory_recall ──────────────────────────────────────────────
205 @mcp.tool()
206 async def memory_recall(
207 query: str,
208 bank_id: str | None = None,
209 banks: list[str] | None = None,
210 strategy: str | None = None,
211 max_results: int = 10,
212 tags: list[str] | None = None,
213 session_id: str | None = None,
214 ) -> str:
215 """Retrieve relevant memories for a query — RAW retrieval, no synthesis.
217 Returns a ranked list of memory ``hits``. The caller is expected to
218 compose its own answer from these hits (this is the "retrieval
219 primitive" — you, the calling agent, are the answerer).
221 **Use `memory_recall` when**: you want to read the underlying
222 memories and synthesize them yourself, OR you need to show citations
223 / quotes verbatim, OR you need fine control over which memories
224 end up in your final response.
226 **Use `memory_reflect` instead when**: you want a finished answer
227 you can pass straight through to the user. Reflect does its own
228 synthesis via an agentic loop (multi-step retrieval + reasoning)
229 and returns one composed answer with cited sources. Cheaper for
230 you than calling `memory_recall` + a follow-up LLM round-trip.
232 Args:
233 query: Natural language query.
234 bank_id: Single memory bank to search. Uses default if omitted.
235 banks: Multiple banks to search (overrides bank_id).
236 strategy: Multi-bank strategy: "cascade", "parallel", or "first_match".
237 max_results: Maximum number of results.
238 tags: Optional tag filter.
239 session_id: M31 Fix 2 — opaque session identifier. When the
240 calling application knows which conversation session the
241 query targets, pass it to scope retrieval to that
242 session's facts/sections. ``None`` (default) preserves
243 cross-session retrieval.
244 """
245 try:
246 max_results = max(1, min(max_results, mcp_cfg.max_results_limit))
247 tags = _validate_tags(tags)
249 if banks:
250 result = await brain.recall(
251 query,
252 banks=banks,
253 strategy=strategy,
254 max_results=max_results,
255 tags=tags,
256 context=_resolve_context(),
257 session_id=session_id, # M31 Fix 2
258 )
259 else:
260 bid = _resolve_bank(bank_id)
261 result = await brain.recall(
262 query,
263 bank_id=bid,
264 max_results=max_results,
265 tags=tags,
266 context=_resolve_context(),
267 session_id=session_id, # M31 Fix 2
268 )
270 hits = [
271 {
272 "text": h.text,
273 "score": round(h.score, 4),
274 "fact_type": h.fact_type,
275 "bank_id": h.bank_id,
276 "memory_id": h.memory_id,
277 }
278 for h in result.hits
279 ]
280 return json.dumps(
281 {
282 "hits": hits,
283 "total_available": result.total_available,
284 "truncated": result.truncated,
285 }
286 )
287 except Exception as exc:
288 logger.exception("memory_recall failed")
289 return json.dumps({"hits": [], "total_available": 0, "error": type(exc).__name__})
291 # ── memory_reflect ─────────────────────────────────────────────
293 if mcp_cfg.expose_reflect:
295 @mcp.tool()
296 async def memory_reflect(
297 query: str,
298 bank_id: str | None = None,
299 banks: list[str] | None = None,
300 strategy: str | None = None,
301 max_tokens: int | None = None,
302 include_sources: bool = True,
303 ) -> str:
304 """Get a finished, cited answer to a question from memory.
306 Returns ``{"answer": <synthesized text>, "sources": [...]}``.
307 The ``answer`` IS the response — Astrocyte's internal agentic
308 loop ran multi-step retrieval + reasoning and composed it.
310 **How to use the `answer` field:**
311 - Quote it verbatim, OR paraphrase it into your own voice, OR
312 return it directly to the user.
313 - **Do NOT pass it back through another LLM for re-synthesis.**
314 That double-synthesis pattern loses citation fidelity, doubles
315 cost, and produces strictly worse answers than either pure
316 path. If you want raw memories to compose yourself, use
317 `memory_recall` instead.
318 - `sources` carries the cited memory chunks for traceability /
319 UI rendering — show them as citations next to the answer,
320 don't re-feed them into a synth prompt.
322 **Use `memory_reflect` when**: the user asked a question and you
323 want a complete, citation-grounded answer back. Especially good
324 for multi-hop questions ("how did X relate to Y?"), synthesis
325 questions ("summarize what we discussed about Z"), and any
326 question where the answer requires reasoning across several
327 memories.
329 **Use `memory_recall` instead when**: you want raw memories
330 (e.g. to show as a search result list) or you want to do your
331 own synthesis with strong opinions on prompt / format.
333 Args:
334 query: The question to answer from memory.
335 bank_id: Single memory bank. Uses default if omitted.
336 banks: Multiple banks (overrides bank_id).
337 strategy: Multi-bank strategy: "cascade", "parallel", or "first_match".
338 max_tokens: Maximum tokens for the synthesized answer.
339 include_sources: Whether to include source memories
340 (default True — needed for showing citations).
341 """
342 try:
343 if banks:
344 result = await brain.reflect(
345 query,
346 banks=banks,
347 strategy=strategy,
348 max_tokens=max_tokens,
349 context=_resolve_context(),
350 include_sources=include_sources,
351 )
352 else:
353 bid = _resolve_bank(bank_id)
354 result = await brain.reflect(
355 query,
356 bank_id=bid,
357 max_tokens=max_tokens,
358 context=_resolve_context(),
359 include_sources=include_sources,
360 )
362 out: dict[str, Any] = {"answer": result.answer}
363 if include_sources and result.sources:
364 out["sources"] = [
365 {"text": s.text, "score": round(s.score, 4), "bank_id": s.bank_id} for s in result.sources
366 ]
367 return json.dumps(out)
368 except Exception as exc:
369 logger.exception("memory_reflect failed")
370 return json.dumps({"answer": "", "error": type(exc).__name__})
372 # ── memory_forget ──────────────────────────────────────────────
374 if mcp_cfg.expose_forget:
376 @mcp.tool()
377 async def memory_forget(
378 bank_id: str | None = None,
379 memory_ids: list[str] | None = None,
380 tags: list[str] | None = None,
381 ) -> str:
382 """Remove memories from a bank.
384 Args:
385 bank_id: Memory bank. Uses default if omitted.
386 memory_ids: Specific memory IDs to delete.
387 tags: Delete memories matching these tags.
388 """
389 try:
390 bid = _resolve_bank(bank_id)
391 tags = _validate_tags(tags)
392 result = await brain.forget(
393 bid,
394 memory_ids=memory_ids,
395 tags=tags,
396 context=_resolve_context(),
397 )
398 return json.dumps(
399 {
400 "deleted_count": result.deleted_count,
401 "archived_count": result.archived_count,
402 }
403 )
404 except Exception as exc:
405 logger.exception("memory_forget failed")
406 return json.dumps({"deleted_count": 0, "error": type(exc).__name__})
408 # ── memory_history ─────────────────────────────────────────────
410 @mcp.tool()
411 async def memory_history(
412 query: str,
413 bank_id: str | None = None,
414 as_of: str | None = None,
415 max_results: int = 10,
416 tags: list[str] | None = None,
417 ) -> str:
418 """Reconstruct what the agent knew at a past point in time (M9 time travel).
420 Args:
421 query: Recall query to run against the historical snapshot.
422 bank_id: Bank to query. Uses default if omitted.
423 as_of: ISO 8601 UTC datetime — memories retained after this moment
424 are hidden (e.g. "2025-01-01T00:00:00Z").
425 max_results: Maximum hits to return (default 10).
426 tags: Optional tag filter applied on top of the time filter.
427 """
428 try:
429 bid = _resolve_bank(bank_id)
430 tags = _validate_tags(tags)
431 if not as_of:
432 return json.dumps({"hits": [], "error": "as_of (ISO 8601) is required"})
433 try:
434 as_of_dt = datetime.fromisoformat(as_of)
435 except ValueError:
436 return json.dumps({"hits": [], "error": "Invalid as_of format; expected ISO 8601"})
437 result = await brain.history(
438 query,
439 bid,
440 as_of_dt,
441 max_results=max(1, min(max_results, mcp_cfg.max_results_limit)),
442 tags=tags,
443 context=_resolve_context(),
444 )
445 hits = [
446 {
447 "text": h.text,
448 "score": round(h.score, 4),
449 "fact_type": h.fact_type,
450 "bank_id": h.bank_id,
451 "memory_id": h.memory_id,
452 }
453 for h in result.hits
454 ]
455 return json.dumps(
456 {
457 "hits": hits,
458 "total_available": result.total_available,
459 "truncated": result.truncated,
460 "as_of": result.as_of.isoformat(),
461 "bank_id": result.bank_id,
462 }
463 )
464 except Exception as exc:
465 logger.exception("memory_history failed")
466 return json.dumps({"hits": [], "error": type(exc).__name__})
468 # ── memory_audit ────────────────────────────────────────────────
470 @mcp.tool()
471 async def memory_audit(
472 scope: str,
473 bank_id: str | None = None,
474 max_memories: int = 50,
475 tags: list[str] | None = None,
476 ) -> str:
477 """Identify knowledge gaps for a topic in a memory bank (M10).
479 Use this to discover what the agent does not know about a subject
480 before relying on recall for important decisions.
482 Args:
483 scope: Natural-language topic to audit (e.g. "Alice's employment history").
484 bank_id: Bank to audit. Uses default if omitted.
485 max_memories: Memories to retrieve and scan (default 50).
486 tags: Optional tag filter to narrow the retrieved memories.
487 """
488 try:
489 bid = _resolve_bank(bank_id)
490 tags = _validate_tags(tags)
491 result = await brain.audit(
492 scope,
493 bid,
494 max_memories=max(1, min(max_memories, 200)),
495 tags=tags,
496 )
497 gaps = [{"topic": g.topic, "severity": g.severity, "reason": g.reason} for g in result.gaps]
498 return json.dumps(
499 {
500 "scope": result.scope,
501 "bank_id": result.bank_id,
502 "coverage_score": round(result.coverage_score, 3),
503 "memories_scanned": result.memories_scanned,
504 "gaps": gaps,
505 }
506 )
507 except Exception as exc:
508 logger.exception("memory_audit failed")
509 return json.dumps({"gaps": [], "coverage_score": 0.0, "error": type(exc).__name__})
511 # ── memory_compile ─────────────────────────────────────────────
513 @mcp.tool()
514 async def memory_compile(
515 bank_id: str | None = None,
516 scope: str | None = None,
517 ) -> str:
518 """Compile raw memories into structured wiki pages (M8).
520 Synthesises a wiki page for each detected topic scope using the LLM.
521 Call this periodically to distil accumulated memories into a curated
522 knowledge base that recall can surface ahead of raw fragments.
524 Args:
525 bank_id: Bank to compile. Uses default if omitted.
526 scope: Compile only memories tagged with this scope string.
527 Omit to trigger full scope discovery (tag grouping +
528 embedding cluster labelling across the whole bank).
529 """
530 try:
531 bid = _resolve_bank(bank_id)
532 result = await brain.compile(bid, scope=scope)
533 out: dict[str, Any] = {
534 "bank_id": result.bank_id,
535 "pages_created": result.pages_created,
536 "pages_updated": result.pages_updated,
537 "scopes_compiled": result.scopes_compiled,
538 "noise_memories": result.noise_memories,
539 "tokens_used": result.tokens_used,
540 "elapsed_ms": result.elapsed_ms,
541 }
542 if result.error:
543 out["error"] = result.error
544 return json.dumps(out)
545 except Exception as exc:
546 logger.exception("memory_compile failed")
547 return json.dumps({"pages_created": 0, "pages_updated": 0, "error": type(exc).__name__})
549 # ── memory_graph_search ────────────────────────────────────────
551 @mcp.tool()
552 async def memory_graph_search(
553 query: str,
554 bank_id: str | None = None,
555 limit: int = 10,
556 ) -> str:
557 """Search the knowledge graph for entities matching a name.
559 Returns matching entities with their IDs. Use the IDs with
560 memory_graph_neighbors to traverse connected memories.
562 Args:
563 query: Entity name or partial name to search for.
564 bank_id: Bank whose graph to search. Uses default if omitted.
565 limit: Maximum number of entities to return.
566 """
567 try:
568 bid = _resolve_bank(bank_id)
569 entities = await brain.graph_search(query, bid, limit=limit)
570 return json.dumps(
571 {
572 "entities": [
573 {
574 "id": e.id,
575 "name": e.name,
576 "entity_type": e.entity_type,
577 "aliases": e.aliases or [],
578 }
579 for e in entities
580 ]
581 }
582 )
583 except Exception as exc:
584 logger.exception("memory_graph_search failed")
585 return json.dumps({"entities": [], "error": type(exc).__name__})
587 # ── memory_graph_neighbors ─────────────────────────────────────
589 @mcp.tool()
590 async def memory_graph_neighbors(
591 entity_ids: list[str],
592 bank_id: str | None = None,
593 max_depth: int = 2,
594 limit: int = 20,
595 ) -> str:
596 """Traverse the knowledge graph from seed entity IDs.
598 Walks the entity graph up to max_depth hops from each seed entity
599 and returns memories attached to discovered entities, scored by
600 proximity. Use memory_graph_search first to resolve entity IDs.
602 Args:
603 entity_ids: Seed entity IDs to start traversal from.
604 bank_id: Bank whose graph to traverse. Uses default if omitted.
605 max_depth: Maximum traversal depth (default 2).
606 limit: Maximum number of memory hits to return.
607 """
608 try:
609 bid = _resolve_bank(bank_id)
610 hits = await brain.graph_neighbors(entity_ids, bid, max_depth=max_depth, limit=limit)
611 return json.dumps(
612 {
613 "hits": [
614 {
615 "memory_id": h.memory_id,
616 "text": h.text,
617 "connected_entities": h.connected_entities,
618 "depth": h.depth,
619 "score": round(h.score, 4),
620 }
621 for h in hits
622 ]
623 }
624 )
625 except Exception as exc:
626 logger.exception("memory_graph_neighbors failed")
627 return json.dumps({"hits": [], "error": type(exc).__name__})
629 # ── admin tools (expose_admin=true) ────────────────────────────
631 if mcp_cfg.expose_admin:
633 @mcp.tool()
634 async def memory_lifecycle(
635 bank_id: str | None = None,
636 ) -> str:
637 """Run TTL lifecycle sweep on a bank — archives and deletes expired memories.
639 Args:
640 bank_id: Bank to sweep. Uses default if omitted.
641 """
642 try:
643 bid = _resolve_bank(bank_id)
644 result = await brain.run_lifecycle(bid)
645 return json.dumps(
646 {
647 "bank_id": bid,
648 "archived_count": result.archived_count,
649 "deleted_count": result.deleted_count,
650 "skipped_count": result.skipped_count,
651 }
652 )
653 except Exception as exc:
654 logger.exception("memory_lifecycle failed")
655 return json.dumps({"archived_count": 0, "deleted_count": 0, "error": type(exc).__name__})
657 @mcp.tool()
658 async def memory_bank_health(
659 bank_id: str | None = None,
660 ) -> str:
661 """Get health score and issues for a memory bank.
663 Args:
664 bank_id: Bank to assess. Uses default if omitted. Pass "__all__"
665 to get health for every bank that has recorded operations.
666 """
667 try:
668 if bank_id == "__all__":
669 results = await brain.all_bank_health()
670 return json.dumps(
671 {
672 "banks": [
673 {
674 "bank_id": r.bank_id,
675 "score": round(r.score, 3),
676 "status": r.status,
677 "issues": [
678 {"severity": i.severity, "code": i.code, "message": i.message} for i in r.issues
679 ],
680 }
681 for r in results
682 ]
683 }
684 )
685 bid = _resolve_bank(bank_id)
686 result = await brain.bank_health(bid)
687 return json.dumps(
688 {
689 "bank_id": result.bank_id,
690 "score": round(result.score, 3),
691 "status": result.status,
692 "issues": [
693 {"severity": i.severity, "code": i.code, "message": i.message} for i in result.issues
694 ],
695 "metrics": result.metrics,
696 }
697 )
698 except Exception as exc:
699 logger.exception("memory_bank_health failed")
700 return json.dumps({"score": 0.0, "error": type(exc).__name__})
702 @mcp.tool()
703 async def memory_hold_set(
704 hold_id: str,
705 reason: str,
706 bank_id: str | None = None,
707 set_by: str = "agent:mcp",
708 ) -> str:
709 """Place a bank under legal hold — blocks memory_forget until released.
711 Args:
712 hold_id: Unique identifier for this hold (used to release it later).
713 reason: Human-readable reason for the hold.
714 bank_id: Bank to place on hold. Uses default if omitted.
715 set_by: Actor label, default "agent:mcp".
716 """
717 try:
718 bid = _resolve_bank(bank_id)
719 hold = await brain.place_legal_hold(bid, hold_id, reason, set_by=set_by)
720 return json.dumps(
721 {
722 "hold_id": hold.hold_id,
723 "bank_id": hold.bank_id,
724 "reason": hold.reason,
725 "set_by": hold.set_by,
726 "set_at": hold.set_at.isoformat(),
727 }
728 )
729 except Exception as exc:
730 logger.exception("memory_hold_set failed")
731 return json.dumps({"error": type(exc).__name__})
733 @mcp.tool()
734 async def memory_hold_release(
735 hold_id: str,
736 bank_id: str | None = None,
737 ) -> str:
738 """Release a legal hold from a bank.
740 Args:
741 hold_id: The hold ID to release.
742 bank_id: Bank to release the hold from. Uses default if omitted.
743 """
744 try:
745 bid = _resolve_bank(bank_id)
746 released = await brain.lift_legal_hold(bid, hold_id)
747 return json.dumps({"bank_id": bid, "hold_id": hold_id, "released": released})
748 except Exception as exc:
749 logger.exception("memory_hold_release failed")
750 return json.dumps({"released": False, "error": type(exc).__name__})
752 # ── memory_banks ───────────────────────────────────────────────
754 @mcp.tool()
755 async def memory_banks() -> str:
756 """List available memory banks."""
757 # For now, return banks from config. In future, query provider.
758 bank_ids: list[str] = []
759 if config.banks:
760 bank_ids = list(config.banks.keys())
761 if default_bank and default_bank not in bank_ids:
762 bank_ids.insert(0, default_bank)
763 return json.dumps({"banks": bank_ids, "default": default_bank})
765 # ── M21 — Mental-model CRUD ─────────────────────────────────────
766 #
767 # Hindsight parity: the 4 mental-model tools mirror Hindsight's
768 # ``_register_create_mental_model`` / ``_update_mental_model`` /
769 # ``_delete_mental_model`` MCP surface, with our addition of
770 # ``memory_list_mental_models`` for inspection. Requires
771 # :meth:`Astrocyte.set_mental_model_store` to have been called
772 # (gateway hosts wire this from the configured Postgres store).
774 @mcp.tool()
775 async def memory_list_mental_models(
776 bank_id: str | None = None,
777 scope: str | None = None,
778 ) -> str:
779 """List mental models in a bank.
781 Mental models are curated, refreshable summaries — e.g. "user
782 prefers async updates", "Project X status: blocked on review".
783 Use this to discover what authoritative summaries exist before
784 composing an answer.
786 Args:
787 bank_id: Bank to list. Uses default if omitted.
788 scope: Optional scope filter (e.g. ``"person:alice"``).
789 """
790 try:
791 bid = _resolve_bank(bank_id)
792 models = await brain.list_mental_models(
793 bank_id=bid,
794 scope=scope,
795 context=_resolve_context(),
796 )
797 out = [
798 {
799 "model_id": m.model_id,
800 "title": m.title,
801 "scope": m.scope,
802 "kind": m.kind,
803 "revision": m.revision,
804 "refreshed_at": m.refreshed_at.isoformat() if m.refreshed_at else None,
805 "source_count": len(m.source_ids),
806 }
807 for m in models
808 ]
809 return json.dumps({"models": out})
810 except Exception as exc:
811 logger.exception("memory_list_mental_models failed")
812 return json.dumps({"models": [], "error": type(exc).__name__})
814 @mcp.tool()
815 async def memory_create_mental_model(
816 model_id: str,
817 title: str,
818 content: str | None = None,
819 sections: list[dict] | None = None,
820 scope: str = "bank",
821 bank_id: str | None = None,
822 ) -> str:
823 """Author a new mental model.
825 Provide ONE of:
827 - ``content`` — raw markdown string; the structured doc is
828 parsed lazily on first refresh (legacy path).
829 - ``sections`` — list of section dicts matching the
830 ``StructuredDocument`` JSON shape; the structured doc is
831 populated immediately so future ``memory_update_mental_model``
832 calls can target sections by id (recommended path).
834 Use ``memory_create_directive`` instead for hard-rule
835 directives (those are mental models with ``kind="directive"``).
836 """
837 try:
838 bid = _resolve_bank(bank_id)
839 model = await brain.create_mental_model(
840 model_id=model_id,
841 title=title,
842 content=content,
843 sections=sections,
844 scope=scope,
845 bank_id=bid,
846 context=_resolve_context(),
847 )
848 return json.dumps(
849 {
850 "model_id": model.model_id,
851 "revision": model.revision,
852 "kind": model.kind,
853 "title": model.title,
854 }
855 )
856 except Exception as exc:
857 logger.exception("memory_create_mental_model failed")
858 return json.dumps({"model_id": None, "error": type(exc).__name__})
860 @mcp.tool()
861 async def memory_update_mental_model(
862 model_id: str,
863 operations: list[dict],
864 bank_id: str | None = None,
865 ) -> str:
866 """Apply structured delta operations to an existing mental model.
868 Operations let you modify a document without re-emitting
869 unchanged content — invaluable for refresh flows where the
870 LLM would otherwise drift on sections it didn't intend to
871 touch. Each operation is a dict with an ``op`` key; see
872 :mod:`astrocyte.pipeline.delta_ops` for the full schema:
874 - ``append_block`` / ``insert_block`` / ``replace_block`` /
875 ``remove_block`` — block-level edits within a section.
876 - ``add_section`` / ``remove_section`` / ``replace_section_blocks``
877 / ``rename_section`` — section-level edits.
879 Conservative-failure contract: invalid ops (unknown section
880 ids, out-of-range indices, malformed payloads) drop with a
881 logged reason in the response's ``skipped`` list; the document
882 never gets worse than its input.
884 Returns ``{"changed": bool, "revision": int, "applied": [...],
885 "skipped": [...]}`` so the caller can audit which ops landed.
886 """
887 try:
888 bid = _resolve_bank(bank_id)
889 result = await brain.update_mental_model(
890 model_id=model_id,
891 operations=operations,
892 bank_id=bid,
893 context=_resolve_context(),
894 )
895 if result is None:
896 return json.dumps({"changed": False, "error": "model not found"})
897 model, summary = result
898 return json.dumps(
899 {
900 "changed": summary.get("changed", False),
901 "revision": model.revision,
902 "applied": summary.get("applied", []),
903 "skipped": summary.get("skipped", []),
904 }
905 )
906 except Exception as exc:
907 logger.exception("memory_update_mental_model failed")
908 return json.dumps({"changed": False, "error": type(exc).__name__})
910 @mcp.tool()
911 async def memory_delete_mental_model(
912 model_id: str,
913 bank_id: str | None = None,
914 ) -> str:
915 """Soft-delete a mental model. Returns whether it existed."""
916 try:
917 bid = _resolve_bank(bank_id)
918 ok = await brain.delete_mental_model(
919 model_id,
920 bank_id=bid,
921 context=_resolve_context(),
922 )
923 return json.dumps({"deleted": ok})
924 except Exception as exc:
925 logger.exception("memory_delete_mental_model failed")
926 return json.dumps({"deleted": False, "error": type(exc).__name__})
928 @mcp.tool()
929 async def memory_refresh_mental_model(
930 model_id: str,
931 new_source_ids: list[str],
932 bank_id: str | None = None,
933 ) -> str:
934 """Re-derive a mental model from its sources after new memories have been retained.
936 Hindsight parity for ``mental_model.refresh()``. When new
937 memories that pertain to an existing model land via retain,
938 call this to fold them into the model's ``source_ids`` and
939 bump the revision. Use ``memory_update_mental_model`` for
940 targeted edits via delta ops; use this when the model should
941 be re-derived from an expanded provenance set.
943 Returns ``{"model_id", "revision", "source_count"}`` on
944 success or ``{"error": "model not found"}`` if the model
945 doesn't exist.
946 """
947 try:
948 bid = _resolve_bank(bank_id)
949 model = await brain.refresh_mental_model(
950 model_id,
951 new_source_ids,
952 bank_id=bid,
953 context=_resolve_context(),
954 )
955 if model is None:
956 return json.dumps({"model_id": model_id, "error": "model not found"})
957 return json.dumps(
958 {
959 "model_id": model.model_id,
960 "revision": model.revision,
961 "source_count": len(model.source_ids),
962 }
963 )
964 except Exception as exc:
965 logger.exception("memory_refresh_mental_model failed")
966 return json.dumps({"model_id": model_id, "error": type(exc).__name__})
968 @mcp.tool()
969 async def memory_create_directive(
970 rule_text: str,
971 directive_id: str | None = None,
972 scope: str = "bank",
973 bank_id: str | None = None,
974 ) -> str:
975 """Author a user-curated hard-rule directive.
977 Directives are mental models with ``kind="directive"``. They
978 participate in the same ``search_mental_models`` retrieval as
979 general / preference models, but the discriminator signals to
980 the answerer that the rule should be applied as a preference
981 override rather than as one input among many.
983 Architecturally replaces the M18a-2 ``directive_compile``
984 auto-extraction path that was deprecated in M19 (auto-compressed
985 directives replicated a −30pp SSP regression because they
986 overrode original preference nuance). User-authored directives
987 avoid that failure mode by definition — the user gets exactly
988 the rule they intended.
989 """
990 try:
991 bid = _resolve_bank(bank_id)
992 model = await brain.create_directive(
993 rule_text=rule_text,
994 directive_id=directive_id,
995 scope=scope,
996 bank_id=bid,
997 context=_resolve_context(),
998 )
999 return json.dumps(
1000 {
1001 "model_id": model.model_id,
1002 "kind": model.kind,
1003 "revision": model.revision,
1004 }
1005 )
1006 except Exception as exc:
1007 logger.exception("memory_create_directive failed")
1008 return json.dumps({"model_id": None, "error": type(exc).__name__})
1010 # ── M21 — Observation CRUD ─────────────────────────────────────
1012 @mcp.tool()
1013 async def memory_list_observations(
1014 bank_id: str | None = None,
1015 scope: str | None = None,
1016 trend: str | None = None,
1017 limit: int = 100,
1018 ) -> str:
1019 """List observations in a bank, optionally filtered by trend.
1021 Observations are auto-consolidated distillations of raw
1022 memories (e.g. "user prefers iced coffee in summer"). Each
1023 carries a computed ``trend`` — one of ``new`` /
1024 ``strengthening`` / ``stable`` / ``weakening`` / ``stale``
1025 derived from the evidence timestamps. Use ``trend`` to filter
1026 for fresh insights (``new``, ``strengthening``) or to audit
1027 stale patterns that may no longer apply.
1028 """
1029 try:
1030 bid = _resolve_bank(bank_id)
1031 rows = await brain.list_observations(
1032 bank_id=bid,
1033 scope=scope,
1034 trend=trend,
1035 limit=limit,
1036 context=_resolve_context(),
1037 )
1038 return json.dumps({"observations": rows})
1039 except Exception as exc:
1040 logger.exception("memory_list_observations failed")
1041 return json.dumps({"observations": [], "error": type(exc).__name__})
1043 @mcp.tool()
1044 async def memory_get_observation(
1045 observation_id: str,
1046 bank_id: str | None = None,
1047 ) -> str:
1048 """Fetch one observation by id. Returns null if not found."""
1049 try:
1050 bid = _resolve_bank(bank_id)
1051 row = await brain.get_observation(
1052 observation_id,
1053 bank_id=bid,
1054 context=_resolve_context(),
1055 )
1056 return json.dumps({"observation": row})
1057 except Exception as exc:
1058 logger.exception("memory_get_observation failed")
1059 return json.dumps({"observation": None, "error": type(exc).__name__})
1061 @mcp.tool()
1062 async def memory_create_observation(
1063 text: str,
1064 source_ids: list[str] | None = None,
1065 scope: str | None = None,
1066 confidence: float = 0.9,
1067 bank_id: str | None = None,
1068 ) -> str:
1069 """Hand-author an observation (bypasses the autonomous consolidator).
1071 Use this when an agent or user wants to assert a fact directly
1072 without going through retain → consolidate. The observation
1073 becomes immediately visible to the recall pipeline's
1074 ``observation`` strategy and to ``search_observations`` in
1075 the agentic reflect loop.
1076 """
1077 try:
1078 bid = _resolve_bank(bank_id)
1079 row = await brain.create_observation(
1080 text=text,
1081 source_ids=source_ids,
1082 scope=scope,
1083 confidence=confidence,
1084 bank_id=bid,
1085 context=_resolve_context(),
1086 )
1087 return json.dumps({"observation": row})
1088 except Exception as exc:
1089 logger.exception("memory_create_observation failed")
1090 return json.dumps({"observation": None, "error": type(exc).__name__})
1092 @mcp.tool()
1093 async def memory_delete_observation(
1094 observation_id: str,
1095 bank_id: str | None = None,
1096 ) -> str:
1097 """Delete one observation from the dedicated observations bank."""
1098 try:
1099 bid = _resolve_bank(bank_id)
1100 ok = await brain.delete_observation(
1101 observation_id,
1102 bank_id=bid,
1103 context=_resolve_context(),
1104 )
1105 return json.dumps({"deleted": ok})
1106 except Exception as exc:
1107 logger.exception("memory_delete_observation failed")
1108 return json.dumps({"deleted": False, "error": type(exc).__name__})
1110 # ── memory_health ──────────────────────────────────────────────
1112 @mcp.tool()
1113 async def memory_health() -> str:
1114 """Check memory system health."""
1115 status = await brain.health()
1116 return json.dumps(
1117 {
1118 "healthy": status.healthy,
1119 "message": status.message,
1120 "latency_ms": round(status.latency_ms, 2) if status.latency_ms else None,
1121 }
1122 )
1124 return mcp
1127# ---------------------------------------------------------------------------
1128# CLI entry point
1129# ---------------------------------------------------------------------------
1132def main() -> None:
1133 """CLI entry point for astrocyte-mcp."""
1134 parser = argparse.ArgumentParser(
1135 prog="astrocyte-mcp",
1136 description="Astrocyte MCP server — memory tools for AI agents",
1137 )
1138 parser.add_argument(
1139 "--config",
1140 default=None,
1141 help="Path to astrocyte.yaml (default: $ASTROCYTE_CONFIG, else ~/.config/astrocyte/astrocyte.yaml)",
1142 )
1143 parser.add_argument(
1144 "--transport",
1145 choices=["stdio", "sse"],
1146 default="stdio",
1147 help="MCP transport (default: stdio)",
1148 )
1149 parser.add_argument(
1150 "--port",
1151 type=int,
1152 default=8080,
1153 help="Port for SSE transport (default: 8080)",
1154 )
1155 args = parser.parse_args()
1157 from astrocyte.harness.paths import config_path
1158 from astrocyte.wiring import build_astrocyte
1160 path = Path(args.config).expanduser() if args.config else config_path()
1161 if not path.is_file():
1162 # MCP clients surface stderr when a server fails to start; make it
1163 # actionable rather than a traceback.
1164 print(
1165 f"astrocyte-mcp: no config at {path}.\n"
1166 " Run `astrocyte setup` to create a local memory store and wire your agents,\n"
1167 " or pass --config /path/to/astrocyte.yaml.",
1168 file=sys.stderr,
1169 )
1170 sys.exit(2)
1172 # Pipeline, stores and access grants, wired the same way as every other
1173 # deployment (a bare Astrocyte(config) has no pipeline and stores nothing).
1174 config = load_config(str(path))
1175 brain = build_astrocyte(config)
1177 # A stdio server is one process per agent session, launched in the
1178 # session's directory, so that directory identifies the project. A shared
1179 # SSE server's cwd identifies nothing, so it keeps requiring a bank.
1180 fallback = None
1181 if args.transport == "stdio":
1182 from astrocyte.harness.project import project_bank
1184 fallback = project_bank(os.getcwd())
1185 mcp = create_mcp_server(brain, config, default_bank_fallback=fallback)
1187 # Run
1188 if args.transport == "stdio":
1189 mcp.run(transport="stdio")
1190 else:
1191 mcp.run(transport="sse", port=args.port)
1194if __name__ == "__main__":
1195 main()