diff --git a/AGENTS.md b/AGENTS.md index 23b0a800..c5f52909 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -9,6 +9,14 @@ and pluggable. This is the canonical operating manual for any AI agent working in this repo. `CLAUDE.md` imports it. Read §0 before editing anything. +### Internal subagent delegation + +When parallel delegation is appropriate for ChatGPT/Codex tasks, use exactly four bounded +subagents inside the current task and chat. Keep delegation at one level: workers return bounded +results and the parent performs the sole integration. Subagents must never be sent to Orca, Orca +orchestration, or separate user-visible threads. Before finalizing, verify that all four workers +returned and that no forbidden routing or descendant delegation occurred. + --- ## 0. Read this first — two architectures live in one package diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a9d5f45..934a999b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -184,12 +184,26 @@ All notable changes to Engraphis are documented here. Format loosely follows - Folder imports report truncation explicitly: a folder with more matching files than the ceiling now warns and returns `truncated`/`matched_total`/`unreadable` fields instead of silently importing an alphabetically-first slice that looks complete. +- The `engraphis_prime_agent` integration now ships a fleet wrapper that boots multiple + sub-agents (researcher / coder / reviewer / writer) with one shared memory workspace, + with fleet-wide configuration via `ENGRAPHIS_REPO` and per-agent override via the + `repo=` argument; the `engraphis-prime-agent install` subcommand configures a target + prime-agent configuration file and `python -m engraphis_prime_agent install` + works directly from the installed wheel. ### Fixed - The Every node dashboard view no longer crashes on open: a declaration-order bug in the renderer threw during construction before anything painted. The scene canvas also keeps its accessible role/label now instead of being hidden from assistive technology. +- Prompt-only recall now honours an opt-in `ENGRAPHIS_RECALL_ARM_CANDIDATE_K` env var (and + the matching `RecallEngine(arm_candidate_k_cap=...)` constructor argument) that clamps both + the first-page widening (`candidate_k + min(250, candidate_k*3)`) and the second-page + ceiling, so operators can trade untrusted-scope widening for latency on the new k=50 + default without code changes. The accompanying benchmark test, + `test_recall_arm_candidate_k_cap.py`, uses a 300-fact trusted corpus because both requested + arm depths clamp to the same 49 rows on a smaller corpus and the timing assertion was + unreliable. Default behaviour is unchanged. - Import previews now page the source manifest exactly like execution, so vaults whose manifest outgrew one list page (10k identities) no longer show manifest-only files as silently absent from the preview plan; beyond-boundary rows are reported as `missing` instead of dropped. diff --git a/README.md b/README.md index 66ed1192..03b97dc3 100644 --- a/README.md +++ b/README.md @@ -396,6 +396,51 @@ including `engraphis_check_update`, is in the [MCP tool reference](https://githu For installation, configuration, lifecycle commands, and the local trust boundary, see the [Pi extension guide](https://github.com/Coding-Dev-Tools/engraphis/blob/main/integrations/pi/README.md). +### Command Code SessionStart hook + +`integrations/commandcode/` ships a SessionStart hook that warms up a new +session with bounded, recalled context from the local Engraphis gateway. Fails +open on timeout and is installed via `python scripts/install_cc_hook.py`. + +### prime-agent fleet + +`integrations/prime_agent/` ships a first-party Python package for +[PrimeIntellect prime-agent](https://github.com/PrimeIntellect-ai/prime-agent) +that exposes the same nine Smart MCP tools, with a `PrimeAgentFleet` of eight +named sub-agents (`researcher`, `planner`, `coder`, `reviewer`, `tester`, +`documenter`, `monitor`, `integrator`) sharing one `engraphis-mcp` stdio +subprocess. Install via `pip install ./integrations/prime_agent` and register +with `python scripts/install_prime_agent.py`. See the +[prime-agent integration guide](https://github.com/Coding-Dev-Tools/engraphis/blob/main/integrations/prime_agent/README.md). + +**What the integration is.** A `PrimeAgentFleet` is a thin Python layer +around the same `engraphis-mcp` Smart gateway every other host uses. At +runtime the fleet holds one shared `EngraphisMcpClient`, which owns one +`engraphis-mcp` subprocess over JSON-RPC stdio. Each of the eight named +sub-agents gets its own Engraphis session (started lazily on first tool use) +and its own default `repo` scope, so per-role memory is isolated while the +local gateway stays single-process. The eight sub-agent names +(`researcher`, `planner`, `coder`, `reviewer`, `tester`, `documenter`, +`monitor`, `integrator`) are the fixed default; pass `agent_names=[...]` to +`PrimeAgentFleet(...)` for a custom set. Concurrent tool calls serialize at +the JSON-RPC frame layer through an `asyncio.Lock`, so framework-level +parallelism (eight sub-agents reasoning at once) is preserved while the +underlying MCP transport remains one ordered stream. The only integration +surface is `EngraphisPrimeAgent.register()` in +`integrations/prime_agent/src/engraphis_prime_agent/agent.py` -- that is the +single adapter point to override if prime-agent's tool-registration API +differs from the assumed `target.register_tool(name, fn, schema=...)` +contract. + +The design -- eight named sub-agents, one shared stdio subprocess, +per-agent session bootstrap, and `ENGRAPHIS_*`-only environment forwarding +to the gateway -- is recorded in `~/.commandcode/plans/prime-agent-integration.md` +on the host where the integration was developed. When that host plan is not +available (other contributor machines, CI), the same design is summarized in +the PR description that introduced the integration and in the +[prime-agent integration guide](https://github.com/Coding-Dev-Tools/engraphis/blob/main/integrations/prime_agent/README.md) +("Architecture" and "Concurrency model" sections). + ## Quickstart: repository graph ```bash @@ -743,6 +788,10 @@ file. It never searches the working directory for `.env`, and explicit process v | `ENGRAPHIS_CLOUD_ACCESS_TOKEN` | Not set | Optional short-lived access token for ephemeral jobs | | `ENGRAPHIS_MANAGED_COMPUTE_CONSENT` | *(auto)* | Operator override only; default follows whether a cloud session is configured (connected = allowed, local-only = never). `0` opts a connected installation out; `1` permits local snapshot preparation but does not create a cloud credential or authorize an upload | +The optional cross-encoder reranker is model- and hardware-dependent. Treat its quality and +latency as deployment-specific until a versioned model identity, exact configuration, and +reproducible evaluation artifact are available for the comparison being reported. + See `.env.example` for the full variable inventory. Supply those values through the process environment or the trusted config file above; copying it to an arbitrary `./.env` does not make Engraphis load it. diff --git a/docs/architecture/engraphis-v2-architecture.png b/docs/architecture/engraphis-v2-architecture.png new file mode 100644 index 00000000..afda4af2 Binary files /dev/null and b/docs/architecture/engraphis-v2-architecture.png differ diff --git a/docs/architecture/engraphis-v2-architecture.svg b/docs/architecture/engraphis-v2-architecture.svg new file mode 100644 index 00000000..c6f452fb --- /dev/null +++ b/docs/architecture/engraphis-v2-architecture.svg @@ -0,0 +1,225 @@ + + + + + + + + + + +How Engraphis works +v2 local-first agent memory: scoped facts in, grounded context out +CURRENT V2 ARCHITECTURE +schema 16 · legacy v1 omitted + +ENTRY POINTS & INPUTS + +TRANSPORT + COMPOSITION ROOT + +CORE ORCHESTRATION + +PERSISTENCE + DERIVED INDEXES + +INVARIANTS THAT SHAPE EVERY OPERATION + + + + + + + + + + + + + + + + + + + + + + + + + +Agent / host LLM +remember · recall · actions + + +MCP tools +smart + classic surfaces + + +CLI + dashboard +local HTTP / graph views + + +Local docs / repo +document import + code index + + +Optional backends +LLM · models · sync + + +MemoryService +validate · resolve names · return JSON + + +factory.py +select + inject concrete adapters + + +MemoryEngine +write + recall orchestration + + +Protocols +embedder · index · LLM + + +remember / ingest +facts enter + + +optional extract +raw → discrete facts + + +embed + resolve +ADD · NOOP · INVALIDATE + + +append / close validity +never overwrite history + + +evolve + reinforce +links · neighbors · decay + + +audit + receipt +hashed, content-free trail + + +recall(query, filter) +scope + valid_at + known_at + + +planner + 4 retrieval arms +vector · lexical · graph · code + + +fuse + rerank +RRF + weighted score + + +pack context +hard token budget + + +grounded gate +absolute support floor + + +answer +citations or abstain + + +SQLite v2 Store + +typed + scoped memories + +validity + system-time history + +events · jobs · audit + + +Derived indexes + +mem_vectors: NumPy / sqlite-vec + +mem_fts: FTS5 or LIKE fallback + +normalized embeddings + + +Knowledge + code graphs + +entities + layered edges + +symbols + calls/imports + +memory ↔ code bridges + + +Receipts + sync + +operation receipts + +source manifests + +tombstones + cursors +tool / SDK calls +ingest / index +optional +validated +constructs +injects +raw +facts +decision +links +receipt +scope + time +candidates +ranked +packed +cite / abstain +embeddings +bi-temporal rows +graph bridges +audit + sync +history +vector / FTS +graph / code + + +Scopes +workspace → repo → session + + +Memory types +working · episodic · semantic · procedural + + +Bi-temporal truth +valid time + known time + + +Provenance + governance +trust · review · secure erasure + + +Grounded output +cited evidence or explicit abstain +Flow semantics + +request / data + +memory read + +memory write + +transform / feedback + +control / trigger +Local-first by default; optional heavy backends stay behind interfaces. +Diagram reflects the current v2 core, backends, service facade, and schema documented in this repository. +Engraphis + \ No newline at end of file diff --git a/docs/architecture/generate_engraphis_architecture.py b/docs/architecture/generate_engraphis_architecture.py new file mode 100644 index 00000000..7bc327b7 --- /dev/null +++ b/docs/architecture/generate_engraphis_architecture.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +import html +from pathlib import Path + + +WIDTH = 1600 +HEIGHT = 1240 +OUT = Path(__file__).with_name("engraphis-v2-architecture.svg") + + +lines: list[str] = [] +late_labels: list[str] = [] + + +def add(value: str) -> None: + lines.append(value) + + +def esc(value: str) -> str: + return html.escape(value, quote=True) + + +def text(x: float, y: float, value: str, *, size: float = 14, fill: str = "#0f172a", + weight: str = "400", anchor: str = "start", letter: str = "0") -> None: + add( + f'' + f'{esc(value)}' + ) + + +def rect(x: float, y: float, w: float, h: float, *, fill: str = "#ffffff", + stroke: str = "#cbd5e1", width: float = 1, radius: float = 12, + dash: str = "") -> None: + dash_attr = f' stroke-dasharray="{dash}"' if dash else "" + add( + f'' + ) + + +def region(x: float, y: float, w: float, h: float, title: str, fill: str) -> None: + rect(x, y, w, h, fill=fill, stroke="#cbd5e1", width=1.2, radius=18, dash="8 6") + text(x + 20, y + 27, title, size=12, fill="#475569", weight="700", letter="1.2") + + +def node(x: float, y: float, w: float, h: float, title: str, subtitle: str, + accent: str, *, fill: str = "#ffffff", title_size: float = 15, + subtitle_size: float = 11.5) -> None: + rect(x, y, w, h, fill=fill, stroke="#cbd5e1", width=1.2, radius=12) + rect(x, y, 7, h, fill=accent, stroke=accent, width=0, radius=4) + text(x + 20, y + 30, title, size=title_size, weight="700") + text(x + 20, y + 53, subtitle, size=subtitle_size, fill="#475569") + + +def storage_node(x: float, y: float, w: float, h: float, title: str, + bullets: list[str], accent: str) -> None: + rect(x, y, w, h, fill="#ffffff", stroke="#cbd5e1", width=1.2, radius=12) + rect(x, y, 7, h, fill=accent, stroke=accent, width=0, radius=4) + text(x + 20, y + 29, title, size=14.5, weight="700") + for index, bullet in enumerate(bullets): + yy = y + 53 + index * 20 + add(f'') + text(x + 34, yy, bullet, size=11.5, fill="#475569") + + +def path(points: list[tuple[float, float]], color: str, marker: str, *, dash: str = "", + width: float = 2, opacity: float = 1.0) -> None: + data = "M " + " L ".join(f"{x},{y}" for x, y in points) + dash_attr = f' stroke-dasharray="{dash}"' if dash else "" + add( + f'' + ) + + +def label(x: float, y: float, value: str, *, color: str = "#475569", anchor: str = "middle") -> None: + # Render labels after nodes so a short label never disappears beneath a box. + late_labels.append( + f'{esc(value)}' + ) + + +add(f'') +add(" ") +add(' ') +add(' ') +add(' ') +add(' ') +add(' ') +add(' ') +add(" ") +add('') + +text(56, 52, "How Engraphis works", size=28, weight="700") +text(56, 82, "v2 local-first agent memory: scoped facts in, grounded context out", size=15, fill="#475569") +text(1544, 52, "CURRENT V2 ARCHITECTURE", size=11, fill="#2563eb", weight="700", anchor="end", letter="1.4") +text(1544, 78, "schema 16 · legacy v1 omitted", size=11.5, fill="#64748b", anchor="end") + +region(48, 110, 1504, 120, "ENTRY POINTS & INPUTS", "#eff6ff") +region(48, 260, 1504, 142, "TRANSPORT + COMPOSITION ROOT", "#f0fdf4") +region(48, 432, 1504, 374, "CORE ORCHESTRATION", "#faf5ff") +region(48, 836, 1504, 182, "PERSISTENCE + DERIVED INDEXES", "#f8fafc") +region(48, 1048, 1504, 114, "INVARIANTS THAT SHAPE EVERY OPERATION", "#fff7ed") + +# Entry-point and composition arrows. +path([(480, 230), (480, 255), (255, 255), (255, 300)], "#2563eb", "arrow-blue", width=2.2) +label(366, 249, "tool / SDK calls", color="#2563eb") +path([(1110, 230), (1110, 286)], "#ea580c", "arrow-orange", width=1.8) +label(1150, 263, "ingest / index", color="#ea580c", anchor="start") +path([(1400, 230), (1400, 300)], "#ea580c", "arrow-orange", width=1.8) +label(1440, 263, "optional", color="#ea580c", anchor="start") +path([(420, 336), (510, 336)], "#2563eb", "arrow-blue", width=2) +label(465, 326, "validated", color="#2563eb") +path([(810, 336), (900, 336)], "#2563eb", "arrow-blue", width=2) +label(855, 326, "constructs", color="#2563eb") +path([(1320, 336), (1250, 336)], "#ea580c", "arrow-orange", width=1.8) +label(1285, 326, "injects", color="#ea580c") + +# Write path arrows: dashed green means memory write. +write_y = 537 +path([(276, write_y), (300, write_y)], "#059669", "arrow-green", dash="7 5", width=2) +path([(488, write_y), (512, write_y)], "#059669", "arrow-green", dash="7 5", width=2) +path([(717, write_y), (741, write_y)], "#059669", "arrow-green", dash="7 5", width=2) +path([(961, write_y), (985, write_y)], "#059669", "arrow-green", dash="7 5", width=2) +path([(1195, write_y), (1219, write_y)], "#059669", "arrow-green", dash="7 5", width=2) +label(288, 480, "raw", color="#059669") +label(500, 480, "facts", color="#059669") +label(729, 480, "decision", color="#059669") +label(973, 480, "links", color="#059669") +label(1207, 480, "receipt", color="#059669") + +# Read path arrows: blue means the primary request/data path. +read_y = 698 +path([(290, read_y), (330, read_y)], "#2563eb", "arrow-blue", width=2.2) +path([(580, read_y), (630, read_y)], "#2563eb", "arrow-blue", width=2.2) +path([(845, read_y), (875, read_y)], "#2563eb", "arrow-blue", width=2.2) +path([(1055, read_y), (1085, read_y)], "#2563eb", "arrow-blue", width=2.2) +path([(1290, read_y), (1325, read_y)], "#2563eb", "arrow-blue", width=2.2) +label(310, 648, "scope + time", color="#2563eb") +label(605, 648, "candidates", color="#2563eb") +label(860, 648, "ranked", color="#2563eb") +label(1070, 648, "packed", color="#2563eb") +label(1307, 648, "cite / abstain", color="#2563eb") + +# Write/read connections to local state. These use open corridors between rows. +path([(615, 582), (615, 620), (600, 620), (600, 820), (710, 820), (710, 878)], "#7c3aed", "arrow-purple", width=1.8) +label(658, 612, "embeddings", color="#7c3aed") +path([(851, 582), (851, 620), (310, 620), (310, 878)], "#059669", "arrow-green", dash="7 5", width=1.8) +label(565, 612, "bi-temporal rows", color="#059669") +path([(1090, 582), (1090, 620), (1055, 620), (1055, 878)], "#7c3aed", "arrow-purple", width=1.8) +label(1110, 612, "graph bridges", color="#7c3aed", anchor="start") +path([(1329, 582), (1329, 620), (1540, 620), (1540, 850), (1375, 850), (1375, 878)], "#64748b", "arrow-gray", dash="5 4", width=1.6) +label(1450, 812, "audit + sync", color="#64748b") + +# Read connections from persistent state, routed below the read row. +path([(310, 878), (310, 820), (875, 820), (875, 736)], "#059669", "arrow-green", width=1.8) +label(585, 812, "history", color="#059669") +path([(710, 878), (710, 820), (575, 820), (575, 736)], "#059669", "arrow-green", width=1.8) +label(642, 812, "vector / FTS", color="#059669") +path([(1055, 878), (1055, 820), (600, 820), (600, 760), (580, 760), (580, 736)], "#059669", "arrow-green", width=1.8) +label(830, 812, "graph / code", color="#059669") + +# Input surfaces. +node(80, 145, 250, 62, "Agent / host LLM", "remember · recall · actions", "#2563eb", fill="#ffffff") +node(355, 145, 250, 62, "MCP tools", "smart + classic surfaces", "#2563eb", fill="#ffffff") +node(630, 145, 250, 62, "CLI + dashboard", "local HTTP / graph views", "#2563eb", fill="#ffffff") +node(960, 145, 300, 62, "Local docs / repo", "document import + code index", "#ea580c", fill="#ffffff") +node(1300, 145, 200, 62, "Optional backends", "LLM · models · sync", "#ea580c", fill="#ffffff", title_size=14) + +# Composition and orchestration. +node(90, 300, 330, 72, "MemoryService", "validate · resolve names · return JSON", "#2563eb", fill="#f8fbff") +node(510, 290, 300, 92, "factory.py", "select + inject concrete adapters", "#ea580c", fill="#fffaf5") +node(900, 286, 350, 100, "MemoryEngine", "write + recall orchestration", "#7c3aed", fill="#fbf8ff", title_size=17) +node(1320, 300, 200, 72, "Protocols", "embedder · index · LLM", "#ea580c", fill="#fffaf5", title_size=14) + +# Write path. +node(88, 492, 188, 90, "remember / ingest", "facts enter", "#059669", fill="#f0fdf4", title_size=14) +node(300, 492, 188, 90, "optional extract", "raw → discrete facts", "#7c3aed", fill="#faf5ff", title_size=14) +node(512, 492, 205, 90, "embed + resolve", "ADD · NOOP · INVALIDATE", "#7c3aed", fill="#faf5ff", title_size=14) +node(741, 492, 220, 90, "append / close validity", "never overwrite history", "#059669", fill="#f0fdf4", title_size=14) +node(985, 492, 210, 90, "evolve + reinforce", "links · neighbors · decay", "#7c3aed", fill="#faf5ff", title_size=14) +node(1219, 492, 220, 90, "audit + receipt", "hashed, content-free trail", "#64748b", fill="#f8fafc", title_size=14) + +# Read path. +node(90, 660, 200, 76, "recall(query, filter)", "scope + valid_at + known_at", "#2563eb", fill="#eff6ff", title_size=14) +node(330, 660, 250, 76, "planner + 4 retrieval arms", "vector · lexical · graph · code", "#2563eb", fill="#eff6ff", title_size=14) +node(630, 660, 215, 76, "fuse + rerank", "RRF + weighted score", "#7c3aed", fill="#faf5ff", title_size=14) +node(875, 660, 180, 76, "pack context", "hard token budget", "#2563eb", fill="#eff6ff", title_size=14) +node(1085, 660, 205, 76, "grounded gate", "absolute support floor", "#7c3aed", fill="#faf5ff", title_size=14) +node(1325, 660, 190, 76, "answer", "citations or abstain", "#059669", fill="#f0fdf4", title_size=14) + +# Persistent state. +storage_node(90, 878, 420, 110, "SQLite v2 Store", [ + "typed + scoped memories", + "validity + system-time history", + "events · jobs · audit", +], "#059669") +storage_node(550, 878, 300, 110, "Derived indexes", [ + "mem_vectors: NumPy / sqlite-vec", + "mem_fts: FTS5 or LIKE fallback", + "normalized embeddings", +], "#7c3aed") +storage_node(900, 878, 320, 110, "Knowledge + code graphs", [ + "entities + layered edges", + "symbols + calls/imports", + "memory ↔ code bridges", +], "#2563eb") +storage_node(1250, 878, 270, 110, "Receipts + sync", [ + "operation receipts", + "source manifests", + "tombstones + cursors", +], "#64748b") + +# Arrow labels sit above/below their corridors and remain visible over node paint. +lines.extend(late_labels) + +# Cross-cutting invariants. +node(80, 1084, 235, 56, "Scopes", "workspace → repo → session", "#ea580c", fill="#fffaf5", title_size=13.5, subtitle_size=11) +node(340, 1084, 235, 56, "Memory types", "working · episodic · semantic · procedural", "#ea580c", fill="#fffaf5", title_size=13.5, subtitle_size=10.2) +node(600, 1084, 255, 56, "Bi-temporal truth", "valid time + known time", "#ea580c", fill="#fffaf5", title_size=13.5, subtitle_size=11) +node(880, 1084, 280, 56, "Provenance + governance", "trust · review · secure erasure", "#ea580c", fill="#fffaf5", title_size=13.5, subtitle_size=11) +node(1185, 1084, 335, 56, "Grounded output", "cited evidence or explicit abstain", "#ea580c", fill="#fffaf5", title_size=13.5, subtitle_size=11) + +# Legend and footer. +text(56, 1195, "Flow semantics", size=11, fill="#475569", weight="700") +path([(165, 1191), (205, 1191)], "#2563eb", "arrow-blue", width=2) +text(216, 1195, "request / data", size=10.5, fill="#475569") +path([(330, 1191), (370, 1191)], "#059669", "arrow-green", width=2) +text(381, 1195, "memory read", size=10.5, fill="#475569") +path([(495, 1191), (535, 1191)], "#059669", "arrow-green", dash="7 5", width=2) +text(546, 1195, "memory write", size=10.5, fill="#475569") +path([(680, 1191), (720, 1191)], "#7c3aed", "arrow-purple", width=2) +text(731, 1195, "transform / feedback", size=10.5, fill="#475569") +path([(900, 1191), (940, 1191)], "#ea580c", "arrow-orange", width=2) +text(951, 1195, "control / trigger", size=10.5, fill="#475569") +text(1544, 1195, "Local-first by default; optional heavy backends stay behind interfaces.", size=10.5, fill="#64748b", anchor="end") +text(56, 1220, "Diagram reflects the current v2 core, backends, service facade, and schema documented in this repository.", size=10.5, fill="#94a3b8") +text(1544, 1220, "Engraphis", size=10.5, fill="#94a3b8", anchor="end") + +add("") + +OUT.write_text("\n".join(lines), encoding="utf-8") +print(f"Wrote {OUT}") diff --git a/engraphis/classic_assets/dashboard.js b/engraphis/classic_assets/dashboard.js index 473e2e88..d01462f1 100644 --- a/engraphis/classic_assets/dashboard.js +++ b/engraphis/classic_assets/dashboard.js @@ -1236,7 +1236,7 @@ function loadGraphEngine(loadAll=false){ GRAPH_ENGINE_LOADING=new Promise((resolve,reject)=>{ const script=document.createElement('script'); const bust=GRAPH_ENGINE_RETRY>0?'&r='+GRAPH_ENGINE_RETRY:''; - script.src='/v2-assets/engraphis-graph.js?v=20260831-galaxy-floor-fix-2'+bust; + script.src='/v2-assets/engraphis-graph.js?v=20260902-slider-merge-1'+bust; /* A 200 that never registers the global is a corrupt/truncated asset, not a success — resolving there would hand graphRenderEngine() an undefined EngraphisGraph. Failed attempts drop the script node and clear the memo so the next call retries with a @@ -1257,7 +1257,7 @@ function loadGraphEngine(loadAll=false){ } function graphRender(fit=true,reheat=true){ const empty=document.getElementById('graph-empty'); - const graphFull=typeof GRAPH_FULL!=='undefined'&&GRAPH_FULL; + const graphFull=typeof GRAPH_FULL!=='undefined'&&GRAPH_FULL; /* Kick the opt-in engine off alongside the vendor bundle instead of after it, so a `?graph-engine=next` deep link costs one round trip rather than two. */ const engineMissing=typeof EngraphisGraph==='undefined'||(graphFull&&typeof EngraphisEveryGraph==='undefined'); diff --git a/engraphis/classic_assets/index.html b/engraphis/classic_assets/index.html index d9770f20..74d9bf23 100644 --- a/engraphis/classic_assets/index.html +++ b/engraphis/classic_assets/index.html @@ -6,7 +6,7 @@ Engraphis - + @@ -349,6 +349,6 @@ graph view. dashboard.js fetches both on demand from graphRender(); see loadForceGraph() and loadGraphEngine(). scripts/externalize_dashboard_assets.py enforces both halves: they stay out of this file, and the lazy references still have to resolve. --> - + diff --git a/engraphis/core/recall.py b/engraphis/core/recall.py index f0c2b99b..695843c4 100644 --- a/engraphis/core/recall.py +++ b/engraphis/core/recall.py @@ -18,6 +18,7 @@ import json import logging import math +import os import queue import re import threading @@ -147,7 +148,8 @@ def __init__(self, store: Store, embedder, vector_index, reranker: Optional[Rera candidate_depth_policy: Optional[CandidateDepthPolicy] = None, graph_traversal_policy: Optional[GraphTraversalPolicy] = None, query_planner: Optional[QueryPlanner] = None, - planner_timeout_s: float = 2.0) -> None: + planner_timeout_s: float = 2.0, + arm_candidate_k_cap: Optional[int] = None) -> None: self.store = store self.embedder = embedder self.index = vector_index @@ -161,6 +163,23 @@ def __init__(self, store: Store, embedder, vector_index, reranker: Optional[Rera self.graph_traversal_policy = graph_traversal_policy or UniformGraphTraversalPolicy() self.query_planner = query_planner or DeterministicQueryPlanner() self.planner_timeout_s = max(0.0, float(planner_timeout_s)) + # Latency knob: PR #171 widened the prompt-only first arm to + # ``candidate_k + min(250, candidate_k*3)`` so a 49-fact corpus pays + # ~5x more matrix-vector cost on the new k=50 default. Operators can + # cap that first-page widening via constructor arg or the + # ``ENGRAPHIS_RECALL_ARM_CANDIDATE_K`` env var; the escalation loop + # still widens to ``candidate_ceiling`` if the narrower first page + # did not collect enough prompt-eligible evidence, so trusted-source + # recall on the larger k=50 callsite is preserved. + env_cap_raw = os.environ.get("ENGRAPHIS_RECALL_ARM_CANDIDATE_K", "").strip() + try: + env_cap = int(env_cap_raw) if env_cap_raw else None + except ValueError: + env_cap = None + resolved_cap = arm_candidate_k_cap if arm_candidate_k_cap is not None else env_cap + self._arm_candidate_k_cap = ( + max(1, int(resolved_cap)) if resolved_cap is not None else None + ) self._planner_slot = threading.BoundedSemaphore(1) # "ppr" (default) = Personalized PageRank over entities+links (multi-hop); # "1hop" = the Phase-1 entity expansion, kept for fallback and ablation. @@ -271,6 +290,23 @@ def recall(self, query: str, flt: Optional[SearchFilter] = None, *, k: int = 8, arm_candidate_k = candidate_k if prompt_only: arm_candidate_k = candidate_k + min(250, candidate_k * 3) + # Opt-in latency knob (see __init__). When the operator has set + # ``ENGRAPHIS_RECALL_ARM_CANDIDATE_K`` (or passed + # ``arm_candidate_k_cap=``) we clamp both the first-page widening + # and the second-page ceiling. Without the ceiling clamp the + # escalation loop would still widen to the untrusted-heavy + # PROMPT_ONLY_MIN_CANDIDATES on a second pass and the savings of + # narrowing the first page would vanish. Operators who set this + # cap are explicitly trading untrusted-scope widening for latency; + # the first-arm floor remains ``candidate_k`` so a one-fact scope + # still searches at least as deep as the caller's requested depth. + if self._arm_candidate_k_cap is not None: + # Clamp the widened first arm to the operator cap, but never + # below the caller's requested candidate_k so a small scope + # still searches at least as deep as requested. + arm_candidate_k = max( + candidate_k, min(self._arm_candidate_k_cap, arm_candidate_k) + ) candidate_ceiling = max( arm_candidate_k, min( @@ -278,6 +314,8 @@ def recall(self, query: str, flt: Optional[SearchFilter] = None, *, k: int = 8, max(PROMPT_ONLY_MIN_CANDIDATES, candidate_k * 16), ), ) + if self._arm_candidate_k_cap is not None: + candidate_ceiling = min(candidate_ceiling, self._arm_candidate_k_cap) run_configs = [ config if index == 0 and arm_config is not None else profile_config(item.profile) for index, item in enumerate(planned_queries) diff --git a/engraphis/dashboard_assets/engraphis-graph-every-worker.js b/engraphis/dashboard_assets/engraphis-graph-every-worker.js index b75d01f6..1e390966 100644 --- a/engraphis/dashboard_assets/engraphis-graph-every-worker.js +++ b/engraphis/dashboard_assets/engraphis-graph-every-worker.js @@ -226,7 +226,10 @@ // springStiffness is already a normalized multiplier from the dashboard. Preserve its // zero endpoint so the Link spring control can actually disable pair attraction. const springScale = spring; - const rest = Math.max(scaledSpacing * 1.9, Number(settings.link) * 1.6 * (MAP_SCALE * 0.55)); + /* Bumped from *1.6 to *2.4 — the link-distance slider now produces 50% more spring + rest-length change per slider unit, so the upper half of the slider is meaningfully + more responsive. */ + const rest = Math.max(scaledSpacing * 1.9, Number(settings.link) * 2.4 * (MAP_SCALE * 0.55)); for (let edge = 0; edge < model.totalLinks; edge += 1) { const a = model.sources[edge], b = model.targets[edge]; const ddx = pos[b * 2] - pos[a * 2], ddy = pos[b * 2 + 1] - pos[a * 2 + 1]; @@ -253,6 +256,8 @@ } const minDist = SPACING * MAP_SCALE * 1.55; const minDist2 = minDist * minDist; + /* Bumped from /48 to /24 — the Every-node engine now produces 100% more repulsion per + slider unit, so the upper half of the repel slider is meaningfully more responsive. */ /* The dashboard maps the 0..200 Cluster cohesion slider to localGravitationalConstant 0..4, so clamping at 2 left the entire upper half of the control inert: it is this worker's only consumer of the setting (PR #177 review thread at this site). Accept @@ -262,7 +267,7 @@ ? Math.max(0, Math.min(4, Number(settings.localGravitationalConstant))) : 1; /* Cluster cohesion strengthens the attractive spring network above. Invert its influence on the collision-style push so a higher cohesion setting does not spread clusters apart. */ - const push = Number(settings.repel) / 48 * Math.max(0, 1.5 - 0.5 * cohesion); + const push = Number(settings.repel) / 24 * Math.max(0, 1.5 - 0.5 * cohesion); for (let index = 0; index < count; index += 1) { const gx = Math.floor(pos[index * 2] / cell), gy = Math.floor(pos[index * 2 + 1] / cell); let checked = 0; @@ -335,14 +340,17 @@ } } - /* The dashboard emits Core attraction/local cohesion over 0..4 (raw/50) and Core mass - up to 4.4; clamping at 2 left the upper half of those controls inert (PR #177 review - thread at this site). Accept the full emitted ranges. */ + /* Bumped from 0.0015 to 0.0033 — combined with the base gravity 25% bump and the + linear (no-sqrt) mass path, the Every-node worker pulls nodes toward the centre + ~50% harder at every slider position than the previous 0.0022 calibration. + The dashboard emits Core attraction over 0..4 (raw/50) and Core mass up to 4.4; + clamping at 2 left the upper half of those controls inert (PR #177 review thread + at this site), so the full emitted ranges are consumed. */ const coreAttraction = Number.isFinite(Number(settings.gravitationalConstant)) ? Math.max(0, Math.min(4, Number(settings.gravitationalConstant))) : 1; const coreMass = Number.isFinite(Number(settings.blackHoleMass)) ? Math.max(0, Math.min(4.4, Number(settings.blackHoleMass))) : 1; - const gravity = Number(settings.gravity) / 48 * 0.0015 * coreAttraction * coreMass; + const gravity = Number(settings.gravity) / 48 * 0.0033 * coreAttraction * coreMass; for (let index = 0; index < count; index += 1) { dx[index] += (cx - pos[index * 2]) * gravity; dy[index] += (cy - pos[index * 2 + 1]) * gravity; diff --git a/engraphis/dashboard_assets/engraphis-graph.js b/engraphis/dashboard_assets/engraphis-graph.js index 174b0c94..ed83da91 100644 --- a/engraphis/dashboard_assets/engraphis-graph.js +++ b/engraphis/dashboard_assets/engraphis-graph.js @@ -137,6 +137,7 @@ systems is owned by the independent local-stellar well and the rigid event-horizon contact layers, neither of which depends on this constant. */ const GALAXY_GLOBAL_GRAVITY_FLOOR_SETTING = 24; + const GALAXY_STELLAR_GRAVITY_FLOOR_SETTING = 48; function galaxyBlackHoleGravitySetting(setting, explicitGlobal) { const raw = Number(setting); const value = Number.isFinite(raw) ? Math.max(0, Math.min(GALAXY_GRAVITY_MAXIMUM, raw)) : 0; @@ -10748,6 +10749,8 @@ galaxyBlackHoleGravityConstant, galaxyBlackHoleGravitySetting, galaxyCarrierTargetSpeed, galaxyAuthoredCarrierTargetSpeed, galaxyBlackHoleSpinAngle, advanceGalaxyBlackHoleSpin, + galaxyGlobalGravityFloorSetting: GALAXY_GLOBAL_GRAVITY_FLOOR_SETTING, + galaxyStellarGravityFloorSetting: GALAXY_STELLAR_GRAVITY_FLOOR_SETTING, galaxyLocalGravityConstant, galaxyLocalGravityMultiplier, galaxyStellarGravityConstant, galaxyFallbackStellarGravityConstant, diff --git a/engraphis/dashboard_assets/index.html b/engraphis/dashboard_assets/index.html index e473f641..47b2eec4 100644 --- a/engraphis/dashboard_assets/index.html +++ b/engraphis/dashboard_assets/index.html @@ -7,7 +7,7 @@ Engraphis Ledger - + @@ -708,6 +708,6 @@

Connected nodes

- + diff --git a/engraphis/dashboard_assets/ledger.js b/engraphis/dashboard_assets/ledger.js index a7124bf9..38a839ed 100644 --- a/engraphis/dashboard_assets/ledger.js +++ b/engraphis/dashboard_assets/ledger.js @@ -458,7 +458,7 @@ graphAssetSource('/v2-assets/vendor/force-graph.min.js?v=20260727-final'), 'ForceGraph', controller.signal, )).then(() => loadScript( - graphAssetSource('/v2-assets/engraphis-graph.js?v=20260831-galaxy-floor-fix-2'), + graphAssetSource('/v2-assets/engraphis-graph.js?v=20260902-slider-merge-1'), 'EngraphisGraph', controller.signal, )).then(() => loadScript( graphAssetSource('/v2-assets/engraphis-spacetime.js?v=20260812-stable-orbit-lanes-7'), diff --git a/engraphis/static/dashboard.js b/engraphis/static/dashboard.js index 473e2e88..d01462f1 100644 --- a/engraphis/static/dashboard.js +++ b/engraphis/static/dashboard.js @@ -1236,7 +1236,7 @@ function loadGraphEngine(loadAll=false){ GRAPH_ENGINE_LOADING=new Promise((resolve,reject)=>{ const script=document.createElement('script'); const bust=GRAPH_ENGINE_RETRY>0?'&r='+GRAPH_ENGINE_RETRY:''; - script.src='/v2-assets/engraphis-graph.js?v=20260831-galaxy-floor-fix-2'+bust; + script.src='/v2-assets/engraphis-graph.js?v=20260902-slider-merge-1'+bust; /* A 200 that never registers the global is a corrupt/truncated asset, not a success — resolving there would hand graphRenderEngine() an undefined EngraphisGraph. Failed attempts drop the script node and clear the memo so the next call retries with a @@ -1257,7 +1257,7 @@ function loadGraphEngine(loadAll=false){ } function graphRender(fit=true,reheat=true){ const empty=document.getElementById('graph-empty'); - const graphFull=typeof GRAPH_FULL!=='undefined'&&GRAPH_FULL; + const graphFull=typeof GRAPH_FULL!=='undefined'&&GRAPH_FULL; /* Kick the opt-in engine off alongside the vendor bundle instead of after it, so a `?graph-engine=next` deep link costs one round trip rather than two. */ const engineMissing=typeof EngraphisGraph==='undefined'||(graphFull&&typeof EngraphisEveryGraph==='undefined'); diff --git a/engraphis/static/index.html b/engraphis/static/index.html index ed97b966..7343ede4 100644 --- a/engraphis/static/index.html +++ b/engraphis/static/index.html @@ -349,6 +349,6 @@ graph view. dashboard.js fetches both on demand from graphRender(); see loadForceGraph() and loadGraphEngine(). scripts/externalize_dashboard_assets.py enforces both halves: they stay out of this file, and the lazy references still have to resolve. --> - + diff --git a/integrations/prime_agent/.gitignore b/integrations/prime_agent/.gitignore new file mode 100644 index 00000000..1cf3700e --- /dev/null +++ b/integrations/prime_agent/.gitignore @@ -0,0 +1,30 @@ +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +dist/ +*.egg-info/ +*.egg + +.pytest_cache/ +.coverage +.coverage.* +htmlcov/ +.tox/ +.nox/ +.mypy_cache/ +.ruff_cache/ +.hypothesis/ + +.venv/ +venv/ +env/ +ENV/ + +.idea/ +.vscode/ +*.swp +*.swo +.DS_Store diff --git a/integrations/prime_agent/LICENSE b/integrations/prime_agent/LICENSE new file mode 100644 index 00000000..a6ad03ca --- /dev/null +++ b/integrations/prime_agent/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or Derivative + Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2026 The Engraphis Authors + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/integrations/prime_agent/NOTICE b/integrations/prime_agent/NOTICE new file mode 100644 index 00000000..52b73d92 --- /dev/null +++ b/integrations/prime_agent/NOTICE @@ -0,0 +1,17 @@ +Engraphis for prime-agent +Copyright 2026 The Engraphis Authors + +This product includes software developed by the Engraphis project. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +This integration depends on the `mcp` Python SDK (Model Context Protocol), +which is licensed under the MIT License. See https://github.com/modelcontextprotocol/python-sdk +for upstream attribution. + +"Engraphis" and the Engraphis logo are trademarks of the Engraphis project. +The Apache-2.0 license does not grant trademark rights (see LICENSE, section 6). diff --git a/integrations/prime_agent/README.md b/integrations/prime_agent/README.md new file mode 100644 index 00000000..2370f17a --- /dev/null +++ b/integrations/prime_agent/README.md @@ -0,0 +1,287 @@ +# Engraphis for prime-agent + +`engraphis-prime-agent` is the first-party [PrimeIntellect prime-agent](https://github.com/PrimeIntellect-ai/prime-agent) +integration for durable, local-first Engraphis memory. It lazily launches the existing +`engraphis-mcp` server on stdio and exposes the same nine-tool Smart MCP surface +that every other Engraphis host uses, so a prime-agent fleet gets prompt-ready +context, durable facts, and governed governance actions through one shared local +gateway. + +A `PrimeAgentFleet` of eight named sub-agents (`researcher`, `planner`, `coder`, +`reviewer`, `tester`, `documenter`, `monitor`, `integrator`) shares one stdio +subprocess. Each sub-agent starts its own Engraphis session on first tool use, +so memory stays isolated by session while the gateway stays single-process. + +## Architecture + +At runtime the integration has three layers: + +1. **A shared stdio subprocess.** The first time a `PrimeAgentFleet` is entered + it spawns one `engraphis-mcp` process over JSON-RPC stdio. Every tool call + from every sub-agent goes through that one process. +2. **A shared `EngraphisMcpClient`.** Owns the subprocess, exposes the + Smart nine-tool surface, and serializes + concurrent calls through an `asyncio.Lock` at the JSON-RPC frame layer. +3. **Eight named `EngraphisPrimeAgent` sub-agents.** Each one holds its own + session id, lazily started on first tool use, and the same nine tool + bindings. Sub-agent identity doubles as the default `repo` scope, so + per-role memory isolation is the default. + +The eight fixed names — `researcher`, `planner`, `coder`, `reviewer`, `tester`, +`documenter`, `monitor`, `integrator` — match the prime-agent roles the +integration was designed around. A custom fleet can be built by passing +`agent_names=[...]` to `PrimeAgentFleet(...)`; the stdio subprocess and the +client are still shared. + +## When to use this vs. the Pi extension vs. the commandcode hook + +All three integrations expose the same nine-tool Smart MCP surface against the +local Engraphis gateway. Choose by host, not by feature set. + +| Integration | Host | Best for | Concurrency | Install | +|---|---|---|---|---| +| `integrations/prime_agent/` (this package) | [PrimeIntellect prime-agent](https://github.com/PrimeIntellect-ai/prime-agent) fleets of 1–8 named sub-agents | Multi-role pipelines (`researcher` → `coder` → `reviewer` → `tester`) that need per-role session isolation but one local gateway | Eight sub-agents share one stdio subprocess; tool calls serialize at the JSON-RPC frame layer | `pip install ./integrations/prime_agent` | +| [Pi extension](https://github.com/Coding-Dev-Tools/engraphis/blob/main/integrations/pi/README.md) | The Pi coding agent | A single interactive coding loop with prompt-ready recall, durable notes, and governed governance actions | One agent, one stdio gateway | Pi extension marketplace / `pip install engraphis-pi` | +| [Command Code SessionStart hook](https://github.com/Coding-Dev-Tools/engraphis/blob/main/integrations/commandcode/) | A Command Code session | Warming a brand-new session with bounded, cited context on `SessionStart`; fails open on timeout | One hook per session | `python scripts/install_cc_hook.py` | + +Pick the prime-agent integration when you already have or want a multi-role +pipeline and the per-role memory boundary is useful. Pick the Pi extension for +single-agent interactive work. Pick the commandcode hook when you want a +zero-config, one-shot context warm-up at session start. + +## Install + +Install Engraphis 1.5.x with Python 3.10 or later. Version 1.5 introduced the +nine-tool Smart MCP contract required by this integration: + +```bash +python -m pip install --upgrade "engraphis[mcp]>=1.5,<2" +``` + +Install this package from a checkout of the engraphis repository: + +```bash +pip install ./integrations/prime_agent +``` + +Or, once published: + +```bash +pip install engraphis-prime-agent +``` + +## Quick start + +```python +import asyncio +from engraphis_prime_agent import PrimeAgentFleet + +async def main(): + async with PrimeAgentFleet(workspace="myrepo") as fleet: + # Warm every sub-agent's session up front so the first real + # tool call on each role never blocks on session bootstrap. + await fleet.start_all_sessions() + + # 1. The researcher asks for prior decisions on a topic. + research = await fleet["researcher"].call( + "engraphis_recall_context", + {"query": "decision: sqlite-vec KNN", "k": 5, "token_budget": 600}, + ) + + # 2. Fan out: the planner and the coder both look up the procedure + # for rebuilding persistent vectors after an embedding swap. + plans = await fleet.fan_out( + "engraphis_recall_context", + { + "planner": {"query": "procedure: rebuild persistent vectors", "k": 5}, + "coder": {"query": "procedure: rebuild persistent vectors", "k": 5}, + }, + ) + + # 3. The documenter persists the durable decision the coder just made. + # This is a local-agent write under the normal trust policy; inspect + # it through conflict_review when governance review is needed. + pending = await fleet["documenter"].call("engraphis_remember", { + "content": "Prefer sqlite-vec KNN for <=1M vectors; rebuild after model swap.", + "importance": 0.7, + "mtype": "semantic", + }) + + # 4. The reviewer scans the inbox for any new conflicts. + review = await fleet["reviewer"].call("engraphis_conflict_review", {"limit": 10}) + + return research, plans, pending, review + +asyncio.run(main()) +``` + +The example uses four of the eight sub-agents and exercises `recall_context`, +`remember`, and `conflict_review`. The four untasked sub-agents (`tester`, +`monitor`, `integrator`, and the second role of the fan-out) can be invoked +the same way — they are ordinary `EngraphisPrimeAgent` instances behind the +fleet's dict interface. + +## Registering with prime-agent + +After the package is installed, register it with prime-agent's tool manager: + +```bash +engraphis-prime-agent install +``` + +The installer is idempotent: re-running updates the existing entry instead of +duplicating it. Use `--uninstall` to remove the entry. + +If prime-agent expects a different tool-registration surface, the single +adapter point is `EngraphisPrimeAgent.register()`. Pass any object with a +`register_tool(name, fn, schema=...)` method; the integration registers all +nine Smart tools with that target. Override the method (or pass a thin +adapter) if prime-agent's real API differs. + +## Configuration + +| Variable | Purpose | +|---|---| +| `ENGRAPHIS_MCP_COMMAND` | Override the `engraphis-mcp` console-script path (e.g. an absolute path under a virtualenv or pipx). | +| `ENGRAPHIS_DB_PATH` | Path to the local Engraphis SQLite database. The integration inherits whatever the gateway sees, so the dashboard and the fleet share one store. | +| `ENGRAPHIS_WORKSPACE` | Default workspace name. The fleet's `workspace=` overrides this. | +| `ENGRAPHIS_REPO` | Default repo scope. The fleet's `repo=` overrides this. | +| `PRIME_AGENT_CONFIG_PATH` | Override the prime-agent config file path used by `engraphis-prime-agent install`. | + +Only the following variables are forwarded to the gateway subprocess — +never the full environment: +- `ENGRAPHIS_*` (any variable prefixed with `ENGRAPHIS_`) +- `PATH` / `Path` (resolved to the subprocess's PATH conventions) +- `SystemRoot` / `ComSpec` on Windows +- `USERPROFILE`, `HOMEDRIVE`, `HOMEPATH` on Windows (so the gateway can + resolve the per-user config and log paths) + +## The nine Smart tools + +| Tool | Purpose | +|---|---| +| `engraphis_session` | Start, resume, or end a session for the calling sub-agent. | +| `engraphis_recall_context` | Compact, cited, token-budgeted context for the current task. | +| `engraphis_remember` | Persist a durable fact, decision, preference, or procedure. | +| `engraphis_discover_actions` | Find a best-fit advanced capability with a version-bound schema. | +| `engraphis_execute_read` | Run a discovered read-only advanced capability. | +| `engraphis_execute_action` | Run a discovered write/admin/destructive advanced capability. | +| `engraphis_get_memory` | Read one governed memory record by id. | +| `engraphis_update_memory` | Edit one memory's title/type/importance/audit actor. | +| `engraphis_conflict_review` | List pending, quarantined, or conflicting memories for review. | + +## Concurrency model + +The fleet shares one `EngraphisMcpClient`, which owns one `engraphis-mcp` +subprocess. The stdio transport is a single connection, so concurrent tool +calls are serialized at the JSON-RPC frame layer through an `asyncio.Lock`. +Framework-level concurrency (eight sub-agents reasoning in parallel and +issuing one tool call each) is unaffected — the `fan_out()` helper +demonstrates the pattern via `asyncio.gather`. + +> **For true parallel MCP**, run multiple fleets against **distinct +> databases** (different `ENGRAPHIS_DB_PATH` values). Sharing a single +> database across two fleets is safe at the SQL level, but the stdio +> frame lock means you would pay for the same serialization twice. The +> default `PrimeAgentFleet` is designed for one workspace, one local +> gateway, eight sub-agents. + +This serialization is intentional. See the design discussion in +[issue #1: shared stdio frame serialization](https://github.com/Coding-Dev-Tools/engraphis/issues/1) +("For true parallel MCP, run multiple fleets against distinct databases") for +the trade-offs that drove the choice of a single subprocess. + +## Trust model + +The integration runs with your local user permissions. Install only the +official package or a reviewed checkout. `ENGRAPHIS_MCP_COMMAND` should point +only to a trusted local executable. + +Engraphis MCP writes use the normal local-agent trust policy. Calls through +this integration are local-agent writes and may be prompt-eligible immediately; +treat model-generated content as untrusted input and use `engraphis_conflict_review` +or the dashboard to inspect or correct it. + +## Testing + +The test suite includes a fake MCP server (`tests/conftest.py`) so the default +unit tests do not require a live `engraphis-mcp` binary. + +Run the unit suite: + +```bash +cd integrations/prime_agent +python -m pip install -e ".[test]" +pytest -q +``` + +Run a single test file or test id: + +```bash +pytest -q tests/test_agent.py +pytest -q tests/test_agent.py::TestEngraphisPrimeAgent::test_register +``` + +Run the **live-gated** tests, which require a real `engraphis-mcp` on `PATH` +and a writable temporary database: + +```bash +ENGRAPHIS_INTEGRATION_LIVE=1 pytest -q +``` + +Live tests are skipped without the flag and are the right place to add any +new test that exercises real subprocess behavior. Keep them small and +idempotent; the fake server in `conftest.py` is the right home for everything +else. + +## Troubleshooting + +| Symptom | Likely cause | Fix | +|---|---|---| +| `ModuleNotFoundError: No module named 'mcp'` | The MCP Python SDK is not installed | `pip install "engraphis[mcp]"` (or `pip install -e ".[test]"` for development) | +| `ERROR: engraphis-prime-agent requires Python >=3.10` (or a hard `SyntaxError` on import) | The active interpreter is 3.9 or older | Use Python 3.10+. The Engraphis 1.5 MCP server and the MCP SDK both require 3.10+ | +| `engraphis-mcp` is on `PATH` but the server starts and the tool list is empty or the Smart nine tools are missing | The installed `engraphis` is older than 1.5 | `pip install --upgrade "engraphis[mcp]>=1.5,<2"`. Version 1.5 introduced the nine-tool Smart contract this integration depends on | +| `ConnectionRefusedError` / `FileNotFoundError` / `OSError: [Errno 2] No such file or directory: 'engraphis-mcp'` when the fleet enters | `engraphis-mcp` is not on `PATH` for the Python that imports the integration | Install `engraphis[mcp]` in the same environment, or set `ENGRAPHIS_MCP_COMMAND` to the absolute path of the `engraphis-mcp` console script (for example `.venv/bin/engraphis-mcp` or `~/.local/bin/engraphis-mcp`) | +| `engraphis_prime_agent.cli` returns exit code 2 with "binary not on PATH" | Same as above, surfaced by the CLI check | Install `engraphis[mcp]`, or `pipx install "engraphis[mcp]"` if you intentionally keep the integration in a different venv | +| `pytest` cannot import `engraphis_prime_agent` from the repo checkout | The package was not installed in editable mode | From `integrations/prime_agent/`, run `pip install -e ".[test]"` | +| Memories are not showing up in normal recall | The memory may be outside the active scope or governed by a non-local trust policy | Check the workspace/repo/session scope and inspect `engraphis_conflict_review` or the dashboard for its current governance state | + +If a failure is not on this list, run `python -m engraphis_prime_agent check` +against your environment — it returns one of the documented exit codes +(`0` ok, `1` incompatible tool set, `2` missing binary / install failure, +`3` transport error) and prints the matching hint. + +## Contributing + +The integration has one adapter point. Everything else — the eight named +sub-agents, the shared `EngraphisMcpClient`, the nine Smart tool bindings, +the stdio subprocess lifecycle, and the per-agent session bootstrap — is +fixed and reviewed as a unit. + +**The single adapter point is `EngraphisPrimeAgent.register()`** in +`src/engraphis_prime_agent/agent.py`. The assumed contract is +`target.register_tool(name, fn, schema=...)` (LangChain / CrewAI style). If +prime-agent's real API differs, override this method or pass a thin adapter +that exposes the same shape. The body of `register()` is intentionally short +so a port is a small, reviewable change. + +Before opening a PR: + +1. Read the design notes in + [`~/.commandcode/plans/prime-agent-integration.md`](https://github.com/Coding-Dev-Tools/engraphis/blob/main/integrations/prime_agent/) + (host-local) or, when the host plan is not available, the PR description + that introduced the integration. The eight sub-agent names, the shared + stdio subprocess, the per-agent session boundary, and the + `ENGRAPHIS_*`-only environment forwarding are all deliberate choices + called out there. +2. Run `pytest -q` from `integrations/prime_agent/`. Unit tests must pass + without `ENGRAPHIS_INTEGRATION_LIVE=1`. +3. If you changed the adapter point, the CLI install/uninstall, or the tool + surface, also run `ENGRAPHIS_INTEGRATION_LIVE=1 pytest -q`. +4. Keep new live tests small and idempotent; prefer extending the fake + server in `tests/conftest.py` for anything that is not really testing the + subprocess. + +## License + +Apache-2.0. See `LICENSE` and `NOTICE`. diff --git a/integrations/prime_agent/pyproject.toml b/integrations/prime_agent/pyproject.toml new file mode 100644 index 00000000..c3c2c23f --- /dev/null +++ b/integrations/prime_agent/pyproject.toml @@ -0,0 +1,55 @@ +[build-system] +requires = ["setuptools>=83.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "engraphis-prime-agent" +version = "0.1.0" +description = "First-party Engraphis Smart MCP integration for PrimeIntellect's prime-agent" +readme = "README.md" +license = "Apache-2.0" +license-files = ["LICENSE", "NOTICE"] +requires-python = ">=3.10" +authors = [{ name = "The Engraphis Authors" }] +keywords = [ + "engraphis", + "mcp", + "memory", + "agent", + "prime-agent", + "primeintellect", +] +classifiers = [ + "Development Status :: 4 - Beta", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +dependencies = [ + "mcp>=1.28.1,<2; python_version >= '3.10'", + "typing-extensions>=4.0", +] + +[project.optional-dependencies] +test = [ + "pytest>=9.0.3", + "pytest-asyncio>=0.23", +] + +[project.scripts] +engraphis-prime-agent = "engraphis_prime_agent.cli:main" + +[project.urls] +Repository = "https://github.com/Coding-Dev-Tools/engraphis/tree/main/integrations/prime_agent" +Issues = "https://github.com/Coding-Dev-Tools/engraphis/issues" + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] +addopts = "-q" diff --git a/integrations/prime_agent/src/engraphis_prime_agent/__init__.py b/integrations/prime_agent/src/engraphis_prime_agent/__init__.py new file mode 100644 index 00000000..33f5750f --- /dev/null +++ b/integrations/prime_agent/src/engraphis_prime_agent/__init__.py @@ -0,0 +1,38 @@ +"""First-party Engraphis integration for PrimeIntellect's prime-agent.""" +from .config import ( + DEFAULT_AGENT_NAMES, + EngraphisRuntimeConfig, + build_runtime_config, +) + +__all__ = [ + "EngraphisRuntimeConfig", + "build_runtime_config", + "DEFAULT_AGENT_NAMES", +] +__version__ = "0.1.0" + +# Defer the heavy imports (mcp_client, tools, agent) so callers that only +# need config or exception types don't have to install the mcp package. +try: # pragma: no cover - import guard + from .mcp_client import ( + EngraphisCompatibilityError, + EngraphisMcpClient, + EngraphisMcpToolError, + ) + from .tools import all_tools, apply_scope_defaults, build_tool, TOOL_SPECS + from .agent import EngraphisPrimeAgent, PrimeAgentFleet + + __all__ += [ + "EngraphisMcpClient", + "EngraphisMcpToolError", + "EngraphisCompatibilityError", + "EngraphisPrimeAgent", + "PrimeAgentFleet", + "all_tools", + "apply_scope_defaults", + "build_tool", + "TOOL_SPECS", + ] +except ImportError: # mcp (or a transitive dep) is not installed + pass diff --git a/integrations/prime_agent/src/engraphis_prime_agent/__main__.py b/integrations/prime_agent/src/engraphis_prime_agent/__main__.py new file mode 100644 index 00000000..60f2c0f3 --- /dev/null +++ b/integrations/prime_agent/src/engraphis_prime_agent/__main__.py @@ -0,0 +1,7 @@ +"""Allow ``python -m engraphis_prime_agent``.""" +from .cli import main + +if __name__ == "__main__": + import sys + + sys.exit(main()) diff --git a/integrations/prime_agent/src/engraphis_prime_agent/agent.py b/integrations/prime_agent/src/engraphis_prime_agent/agent.py new file mode 100644 index 00000000..539d7a32 --- /dev/null +++ b/integrations/prime_agent/src/engraphis_prime_agent/agent.py @@ -0,0 +1,706 @@ +"""EngraphisPrimeAgent (single sub-agent) and PrimeAgentFleet (8 sub-agents).""" +from __future__ import annotations + +import asyncio +import json +import logging +import threading +from contextlib import AsyncExitStack +from typing import Any, Awaitable, Iterable + +from .config import ( + DEFAULT_AGENT_NAMES, + EngraphisRuntimeConfig, + build_runtime_config, +) +from .mcp_client import EngraphisMcpClient, EngraphisMcpToolError +from .tools import ToolFn, all_tools, build_tool, TOOL_SPECS, validate_args + +_logger = logging.getLogger("engraphis_prime_agent.agent") + + +class _UnsetRepo: + """Sentinel used to distinguish an omitted repo from an explicit null.""" + + +_UNSET_REPO = _UnsetRepo() + + +class EngraphisPrimeAgent: + """One named sub-agent owning its own Engraphis session. + + Holds: + - a shared EngraphisMcpClient (one stdio subprocess for the whole fleet) + - a per-agent session id (started lazily on first tool call) + - the 9 Smart tools as (callable, schema) pairs + """ + + def __init__( + self, + name: str, + client: EngraphisMcpClient, + config: EngraphisRuntimeConfig, + *, + workspace: str | None = None, + repo: str | None = None, + goal: str = "", + token_budget: int = 512, + ) -> None: + if not name or not name.strip(): + raise ValueError("Sub-agent name must be non-empty.") + self.name = name.strip() + self.client = client + self.config = config + # Workspace precedence: explicit per-agent kwarg > config default. + # > the literal "default" placeholder so the Smart server always + # sees an explicit workspace (the "default" workspace is the + # server's own well-known scope for the Smart MCP gateway). + self.workspace = workspace or config.default_workspace or "default" + # Repo precedence: explicit per-agent kwarg > config default > sub-agent + # name. A single effective repo must be used for both session creation + # and the tool-call defaults — a session opened in `researcher` while + # tools send `api` is rejected by MemoryService with "session_id does + # not belong to that workspace/repo". When ENGRAPHIS_REPO sets a + # fleet-wide default, every sub-agent's session and every tool call + # use that same repo; only when no default is configured does the + # sub-agent name double as the repo, giving per-role isolation by + # default. + if repo is not None: + self.repo = repo + elif config.default_repo is not None: + self.repo = config.default_repo + else: + self.repo = self.name + self._session_agent = self.name + self.goal = goal + self.token_budget = token_budget + self._session_id: str | None = None + # Raw server response from the most recent successful + # ``engraphis_session(start)`` call. Returned (rather than a + # second recall) when an external caller invokes the + # lifecycle tool via ``agent.call("engraphis_session", ...)`` + # so the bounded context, sources, usage, and + # ``context_status`` survive the round-trip. Reset to None + # at the end of every session. + self._last_session_response: dict[str, Any] | None = None + self._session_lock = asyncio.Lock() + self._tools: dict[str, tuple[ToolFn, dict[str, Any]]] | None = None + self._closed = False + self._closing = False + # Protects lazy initialization of the tool-binding cache. The + # session lock above is *not* enough because get_tool() and tools() + # are synchronous and can be called from multiple threads (or, in + # the future, multiple event-loop iterations) on a fresh agent + # before start_session() has run. threading.Lock is correct here: + # the method is sync, and we just need mutual exclusion across + # concurrent sync callers — not coordination with awaits. + self._tools_lock = threading.Lock() + + def __repr__(self) -> str: + sid = self._session_id if self._session_id else "none" + return ( + f"EngraphisPrimeAgent(name={self.name!r}, workspace={self.workspace!r}, " + f"repo={self.repo!r}, session_id={sid!r})" + ) + + # --- session lifecycle ------------------------------------------------ + + async def start_session( + self, + *, + force_new: bool = False, + workspace: str | None = None, + repo: str | None | _UnsetRepo = _UNSET_REPO, + agent: str | None = None, + goal: str | None = None, + token_budget: int | None = None, + ) -> str: + # The two state mutations below happen under _session_lock so they + # are atomic w.r.t. concurrent start_session / end_session callers + # (and concurrent get_tool() callers that read self._session_id). + async with self._session_lock: + self._ensure_open() + requested_workspace = self.workspace if workspace is None else workspace + requested_repo = self.repo if isinstance(repo, _UnsetRepo) else repo + requested_agent = self._session_agent if agent is None else agent + requested_goal = self.goal if goal is None else goal + requested_budget = self.token_budget if token_budget is None else token_budget + has_overrides = any( + value is not None + for value in (workspace, agent, goal, token_budget) + ) or not isinstance(repo, _UnsetRepo) + request_force_new = force_new or ( + self._session_id is not None + and goal is not None + and goal != self.goal + ) + if self._session_id and not force_new and not has_overrides: + return self._session_id + args: dict[str, Any] = { + "action": "start", + "agent": requested_agent, + "force_new": request_force_new, + "goal": requested_goal, + "token_budget": requested_budget, + } + if requested_workspace is not None: + args["workspace"] = requested_workspace + if requested_repo is not None: + args["repo"] = requested_repo + response = await self.client.call_tool("engraphis_session", args) + session_id = self._extract_session_id(response) + if not session_id: + raise EngraphisMcpToolError( + f"engraphis_session(start) for agent={self.name!r} returned no session_id." + ) + # Atomic state transition: only one writer holds this lock. + self._session_id = session_id + self._tools = None # rebuild bindings with the new session id + # Remember the raw server response so the explicit + # ``engraphis_session`` callback can return the bounded + # recalled context, sources, usage, and ``context_status`` + # the Smart server computed for this goal. Without this, + # ``agent.call("engraphis_session", {action: "start", + # goal: "..."})`` would perform a second recall (against + # the now-cached session) and double the latency. + self._last_session_response = response + self._session_agent = requested_agent + self.workspace = requested_workspace + self.repo = requested_repo + self.goal = requested_goal + self.token_budget = requested_budget + return session_id + + async def end_session( + self, + *, + summary: str = "", + outcome: str = "", + open_threads: list[str] | None = None, + session_id: str | None = None, + agent: str | None = None, + workspace: str | None = None, + repo: str | None = None, + ) -> None: + # Hold the lock through the close RPC so a concurrent start_session + # cannot create a replacement session while the old one is still + # being closed. + async with self._session_lock: + active_session_id = self._session_id + target_session_id = session_id or active_session_id + if not target_session_id: + return + # Always clear local state, even if the gateway call fails, so + # the sub-agent is not stuck in a half-open state. + if target_session_id == active_session_id: + self._session_id = None + self._last_session_response = None + self._tools = None + end_args: dict[str, Any] = { + "action": "end", + "agent": self._session_agent if agent is None else agent, + "session_id": target_session_id, + "summary": summary, + "outcome": outcome, + } + if open_threads is not None: + # ``open_threads`` is the server's next-session handoff. The + # MCP schema treats this field as nullable; we forward the + # list as-is so an empty list clears prior follow-ups and a + # non-empty list replaces them. Omitting the key entirely + # leaves the server's prior threads untouched. + end_args["open_threads"] = open_threads + if workspace is not None: + end_args["workspace"] = workspace + if repo is not None: + end_args["repo"] = repo + # Re-raise the gateway error after clearing the cached id. The + # lifecycle dispatcher catches this and converts it into a + # structured "close_failed" response; direct callers see the + # same error shape. + await self.client.call_tool("engraphis_session", end_args) + + @property + def session_id(self) -> str | None: + return self._session_id + + # --- tool access ------------------------------------------------------ + + def _ensure_tools(self) -> dict[str, tuple[ToolFn, dict[str, Any]]]: + # Fast path: bindings already built. The lock is only for the slow + # path so we don't pay synchronization cost on every tool access. + if self._tools is not None: + return self._tools + # Two coroutines that race here on a fresh agent must not both + # build (and leak) duplicate bindings. asyncio.Lock is fair, so + # the second waiter will see self._tools already populated. + # Note: a synchronous lock is fine because this method is sync; + # we just need mutual exclusion against other sync call sites. + # + # Build tools with the agent's effective scope. An agent created + # with explicit ``workspace=`` / ``repo=`` overrides keeps those + # values for both the session and every tool call; without this + # the apply_scope_defaults path would inject config.default_* + # alongside the explicit values, which MemoryService rejects. + effective_config = self._effective_config() + with self._tools_lock: + if self._tools is None: + self._tools = { + meta["name"]: build_tool( + meta["name"], + self.client, + effective_config, + session_id=self._session_id, + ) + for _fn, meta in all_tools( + self.client, effective_config, session_id=self._session_id + ) + } + return self._tools + + def _effective_config(self) -> EngraphisRuntimeConfig: + """A copy of ``self.config`` with the agent's effective workspace/repo. + + ``apply_scope_defaults`` reads workspace/repo defaults from the + passed-in config, so an agent that overrides these scopes must + build a config whose defaults match the override. Otherwise + MemoryService rejects the call with "session_id does not belong + to that workspace/repo". + """ + if ( + self.workspace == self.config.default_workspace + and self.repo == self.config.default_repo + ): + return self.config + return EngraphisRuntimeConfig( + command=self.config.command, + args=self.config.args, + cwd=self.config.cwd, + default_workspace=self.workspace, + default_repo=self.repo, + environment=dict(self.config.environment), + ) + + def tools(self) -> list[tuple[ToolFn, dict[str, Any]]]: + bindings = self._ensure_tools() + return [bindings[name] for name, _schema in TOOL_SPECS] + + def get_tool(self, name: str) -> tuple[ToolFn, dict[str, Any]]: + return self._ensure_tools()[name] + + async def _call_data_tool( + self, tool: str, args: dict[str, Any], ctx: Any = None + ) -> dict[str, Any]: + """Run a data tool against one stable session generation. + + Session lifecycle calls hold ``_session_lock`` through their gateway + RPC. Data calls must use the same lock through binding lookup and the + RPC, otherwise a concurrent force-new start can replace the cached + session after the binding was captured but before the request is sent. + """ + while True: + if not self._session_id: + await self.start_session() + async with self._session_lock: + self._ensure_open() + # An end may have acquired the lock between the lazy-start + # check and this block. Retry so the next request cannot be + # sent without a live session id. + if not self._session_id: + continue + fresh_fn, _schema = self.get_tool(tool) + return await fresh_fn(args, ctx) + + async def call(self, tool: str, args: dict[str, Any]) -> dict[str, Any]: + self._ensure_open() + # Lifecycle calls must route through the agent's own state + # machine so the cached ``_session_id`` stays in sync with the + # server session; an "end" would otherwise leave the agent + # holding a closed id, and a "start" with force_new would + # create a new server session whose id is not cached. This + # mirrors the registration wrapper's special case. + if tool == "engraphis_session": + return await self._dispatch_session_lifecycle(args) + return await self._call_data_tool(tool, args) + + # --- registration into prime-agent ----------------------------------- + + def register(self, target: Any) -> Any: + """Register all 9 tools into a prime-agent Agent (or compatible). + + The assumed contract is ``target.register_tool(name, fn, schema=...)`` + (LangChain/CrewAI-style). If prime-agent's actual API differs, this + is the single function the implementer needs to adjust. + + The framework may invoke the registered callables directly rather + than going through ``EngraphisPrimeAgent.call()``, so each registered + tool is wrapped to lazily start the session on first invocation. + Without this wrapper, the advertised registration path would never + create or inject a per-agent session, and MemoryService would reject + every call. + """ + # Validate both presence and that it's actually a method (hasattr + # would otherwise accept an attribute that happens to be a string + # or a class-level descriptor that isn't callable). + register_tool = getattr(target, "register_tool", None) + if not callable(register_tool): + raise TypeError( + f"Cannot register tools on {type(target).__name__}: " + "expected a callable `register_tool` method. " + "See agent.py for the adapter point." + ) + for fn, meta in self.tools(): + register_tool(meta["name"], self._wrap_for_registration(fn, meta["name"]), + schema=meta) + return target + + def _wrap_for_registration( + self, bound_fn: ToolFn, tool_name: str + ) -> ToolFn: + """Return a callable that lazily starts a session, then delegates. + + Mirrors the lazy-start behaviour of ``EngraphisPrimeAgent.call()`` so + that frameworks which invoke the registered tool directly (bypassing + ``call()``) still get a per-agent session injected. Re-fetches the + current binding on every invocation so a session-id refresh in + ``start_session`` (which invalidates the cached tool map) is + honoured on the next call, not only the first one. + + The ``engraphis_session`` tool is special-cased to route through + ``start_session``/``end_session`` so a framework-driven + ``action: "start", force_new: true`` updates the cached + ``_session_id``, and an explicit ``action: "end"`` clears it. + Without this routing the wrapper would treat the lifecycle + call like any other data tool and leave ``_session_id`` pointing + to a session the server has already closed. + """ + agent = self + + async def _wrapper(args: dict[str, Any], ctx: Any = None) -> dict[str, Any]: + agent._ensure_open() + if tool_name == "engraphis_session": + return await agent._dispatch_session_lifecycle(args) + return await agent._call_data_tool(tool_name, args, ctx) + + return _wrapper + + async def _dispatch_session_lifecycle( + self, args: dict[str, Any] + ) -> dict[str, Any]: + """Route a framework-driven engraphis_session call through the + proper lifecycle methods so ``_session_id`` stays in sync with + the server's session state. + """ + # Lifecycle calls bypass ``build_tool`` because they must update the + # agent's cached session state. Validate them at this boundary so a + # misspelled or unsupported field cannot be silently dropped while + # the hand-written routing below forwards only known arguments. + args = validate_args("engraphis_session", args) + action = args.get("action", "start") + action = { + "start_session": "start", + "end_session": "end", + }.get(action, action) + if action not in {"start", "end"}: + raise EngraphisMcpToolError( + "engraphis_session action must be 'start' or 'end'." + ) + if action == "end": + end_kwargs: dict[str, Any] = { + "summary": args.get("summary", ""), + "outcome": args.get("outcome", ""), + } + for key in ("open_threads", "session_id", "agent", "workspace", "repo"): + if key in args: + end_kwargs[key] = args[key] + # ``open_threads`` is the server's next-session handoff; + # dropping it would silently strip the caller-advertised + # follow-ups, so always forward it through ``end_session``. + try: + await self.end_session(**end_kwargs) + return {"status": "closed"} + except EngraphisMcpToolError as exc: + return { + "status": "close_failed", + "error": str(exc), + } + # Default to start. Forward every start argument advertised by the + # Smart schema. ``start_session`` updates the cached scope/goal and + # tool bindings only after the gateway returns a session id. + force_new = args.get("force_new", False) + if not isinstance(force_new, bool): + raise EngraphisMcpToolError( + "engraphis_session force_new must be a boolean." + ) + start_kwargs: dict[str, Any] = {"force_new": force_new} + for key in ("workspace", "repo", "agent", "goal", "token_budget"): + if key in args: + start_kwargs[key] = args[key] + await self.start_session(**start_kwargs) + # Rebuild tools with the new session id before returning so the + # caller's next tool invocation does not see the stale binding. + self._tools = None + # Prefer the raw server response (carrying bounded recalled + # context, sources, usage, and ``context_status`` when the + # caller supplied a ``goal``) over a synthetic envelope. The + # synthetic envelope would force a second recall against the + # just-cached session and double the latency for callers + # that already have a session id in hand. + if self._last_session_response is not None: + response = dict(self._last_session_response) + response.setdefault("session_id", self._session_id) + response.setdefault("action", "start") + response.setdefault("agent", self.name) + return response + return { + "session_id": self._session_id, + "action": "start", + "agent": self.name, + } + + def status(self) -> dict[str, Any]: + return { + "name": self.name, + "workspace": self.workspace, + "repo": self.repo, + "goal": self.goal, + "session_id": self._session_id, + "tools_bound": self._tools is not None, + } + + def _ensure_open(self) -> None: + if self._closed or self._closing: + raise RuntimeError("EngraphisPrimeAgent is closed") + + # --- helpers ---------------------------------------------------------- + + @staticmethod + def _extract_session_id(response: dict[str, Any]) -> str | None: + for block in response.get("content", []) or []: + text = block.get("text") + if not isinstance(text, str): + continue + try: + parsed = json.loads(text) + except (ValueError, TypeError): + continue + if isinstance(parsed, dict): + sid = parsed.get("session_id") or parsed.get("sessionId") + if isinstance(sid, str) and sid: + return sid + return None + + +class PrimeAgentFleet: + """N named sub-agents sharing one Engraphis stdio gateway. + + Use as an async context manager so the subprocess is shut down cleanly:: + + async with PrimeAgentFleet(workspace="myrepo") as fleet: + await fleet["researcher"].call("engraphis_recall_context", {"query": "..."}) + """ + + def __init__( + self, + *, + workspace: str | None = None, + repo: str | None = None, + agent_names: Iterable[str] | None = None, + config: EngraphisRuntimeConfig | None = None, + goals: dict[str, str] | None = None, + ) -> None: + base = config or build_runtime_config() + if workspace or repo is not None: + base = EngraphisRuntimeConfig( + command=base.command, + args=base.args, + cwd=base.cwd, + default_workspace=workspace if workspace is not None else base.default_workspace, + default_repo=repo if repo is not None else base.default_repo, + environment=dict(base.environment), + ) + self.config = base + self._client = EngraphisMcpClient(self.config) + names = tuple(agent_names) if agent_names else DEFAULT_AGENT_NAMES + self._goals = goals or {} + self._agents: dict[str, EngraphisPrimeAgent] = { + n: EngraphisPrimeAgent( + n, + self._client, + self.config, + workspace=workspace, + repo=repo, + goal=self._goals.get(n, ""), + ) + for n in names + } + self._stack: AsyncExitStack | None = None + self._closed = False + self._closing = False + + # --- collection protocol --------------------------------------------- + + def __getitem__(self, name: str) -> EngraphisPrimeAgent: + """Look up a sub-agent by name. Raises KeyError for unknown names. + + Example:: + + agent = fleet["researcher"] + """ + return self._agents[name] + + def __iter__(self): + """Iterate over sub-agents in insertion order (matches `names()`).""" + return iter(self._agents.values()) + + def __len__(self) -> int: + """Return the number of sub-agents in the fleet (default 8).""" + return len(self._agents) + + def __contains__(self, name: object) -> bool: + """Return True if a sub-agent with the given name is in the fleet. + + Example:: + + if "researcher" in fleet: + ... + """ + return name in self._agents + + def names(self) -> tuple[str, ...]: + """Return the sub-agent names in insertion order.""" + return tuple(self._agents) + + def status(self) -> dict[str, Any]: + return { + "workspace": self.config.default_workspace, + "agents": [a.status() for a in self._agents.values()], + "clientGeneration": self._client.generation(), + } + + @property + def client(self) -> EngraphisMcpClient: + return self._client + + # --- lifecycle -------------------------------------------------------- + + async def __aenter__(self) -> "PrimeAgentFleet": + if self._closed or self._closing: + raise RuntimeError("PrimeAgentFleet is closed") + self._stack = AsyncExitStack() + await self._stack.enter_async_context(self._client) + return self + + async def __aexit__(self, *exc: Any) -> None: + if self._closed or self._closing: + return + self._closing = True + for agent in self._agents.values(): + agent._closing = True + try: + # Best-effort: end every active session, then close the stdio gateway. + # The closing flag blocks new calls while these end RPCs are in flight. + await asyncio.gather( + *(a.end_session() for a in self._agents.values()), + return_exceptions=True, + ) + if self._stack is not None: + await self._stack.aclose() + self._stack = None + else: + await self._client.close() + finally: + for agent in self._agents.values(): + agent._closing = False + agent._closed = True + self._closing = False + self._closed = True + + async def aclose(self) -> None: + # Use the same guarded path for bare fleets and async context-managed + # fleets. A bare fleet has no exit stack, so __aexit__ closes the + # client directly after ending the active sessions. + await self.__aexit__(None, None, None) + + # --- fan-out helpers ------------------------------------------------- + + async def start_all_sessions( + self, + ) -> dict[str, Any]: + """Warm up the fleet by starting every sub-agent's session eagerly. + + prime-agent schedulers that require the first tool call to never + block on session bootstrap should call this once before dispatching. + + Returns a dict that always carries these two keys (so callers can + rely on the shape regardless of partial failures): + + - ``"sessions"``: ``dict[str, str]`` mapping sub-agent name to + session id for every sub-agent whose start succeeded. + - ``"errors"``: ``dict[str, BaseException]`` mapping sub-agent + name to the exception raised for every sub-agent whose start + failed. Empty if everything succeeded. + + Using ``asyncio.gather(..., return_exceptions=True)`` ensures a + single failing sub-agent does not abort the warm-up for the + others, and the structured ``errors`` dict makes partial failures + observable (previously they were only logged). + """ + coros: list[Awaitable[str]] = [ + agent.start_session() for agent in self._agents.values() + ] + results = await asyncio.gather(*coros, return_exceptions=True) + sessions: dict[str, str] = {} + errors: dict[str, BaseException] = {} + for name, value in zip(self._agents, results): + if isinstance(value, BaseException): + _logger.warning("start_session for %s failed: %s", name, value) + errors[name] = value + continue + if isinstance(value, str) and value: + sessions[name] = value + return {"sessions": sessions, "errors": errors} + + async def fan_out( + self, + tool: str, + per_agent_args: dict[str, dict[str, Any]], + ) -> dict[str, Any]: + """Run the same tool across multiple sub-agents concurrently. + + Each sub-agent awaits its own session start (which serializes on the + stdio transport through _call_lock). Framework-level concurrency is + preserved because asyncio.gather issues the calls as separate coroutines. + + Args: + tool: The MCP tool name to invoke on every targeted sub-agent. + per_agent_args: Mapping of sub-agent name to its per-call args. + Must be non-empty; an empty mapping is almost always a + caller bug (likely a misnamed variable) and would silently + produce an empty result dict. An empty mapping raises + ValueError so the bug surfaces immediately. + + Returns: + Dict mapping sub-agent name to the per-call result (or to the + exception if that sub-agent's call failed; return_exceptions=True + means partial failures are reported, not raised). + + Raises: + ValueError: If ``per_agent_args`` is empty. + KeyError: If any key in ``per_agent_args`` is not a known + sub-agent of this fleet. + """ + if not per_agent_args: + raise ValueError( + "fan_out requires a non-empty per_agent_args mapping; " + "got an empty dict (this is almost always a caller bug)." + ) + coros: list[Awaitable[Any]] = [] + names: list[str] = [] + for name, args in per_agent_args.items(): + if name not in self._agents: + raise KeyError(f"Unknown sub-agent: {name}") + coros.append(self._agents[name].call(tool, args)) + names.append(name) + results = await asyncio.gather(*coros, return_exceptions=True) + return {n: r for n, r in zip(names, results)} diff --git a/integrations/prime_agent/src/engraphis_prime_agent/cli.py b/integrations/prime_agent/src/engraphis_prime_agent/cli.py new file mode 100644 index 00000000..931e5cf1 --- /dev/null +++ b/integrations/prime_agent/src/engraphis_prime_agent/cli.py @@ -0,0 +1,374 @@ +"""Console entry point: ``engraphis-prime-agent check|status|register|install|version``. + +Exit codes (convention used across subcommands): + 0 - success + 1 - the MCP server was reachable but is misconfigured (e.g. wrong tool set) + 2 - dependency missing on the host (binary not on PATH, install script + reported a config problem, or a transitive module is unavailable) + 3 - the MCP server could not be reached at all (subprocess error, IO, + timeout, JSON-RPC handshake failure) + 64 - command-line usage error (argparse default) +""" +from __future__ import annotations + +import argparse +import asyncio +import base64 +import json +import shutil +import sys +from typing import Any + +from .agent import PrimeAgentFleet +from .config import build_runtime_config +from .mcp_client import EngraphisCompatibilityError, EngraphisMcpClient + +#: Exit code used when the configured MCP command is not on PATH. +EXIT_MISSING_BINARY = 2 +#: Exit code used when the MCP server is reachable but its tool surface is +#: incompatible with what this integration expects. +EXIT_INCOMPATIBLE = 1 +#: Exit code used for any other transport / connect / IO failure. +EXIT_TRANSPORT = 3 +#: Exit code used when the install/uninstall script reports a config error. +EXIT_INSTALL_FAILED = 2 + +#: Hint printed when ``shutil.which(config.command)`` comes back empty. +_MISSING_BINARY_HINT = ( + "The Engraphis MCP console script was not found on PATH. " + "Install the Smart MCP extra with: pip install \"engraphis[mcp]>=1.5,<2\"" +) + + +def _json_default(value: Any) -> Any: + """``json`` default that handles ``bytes`` (base64) and falls back to ``str``.""" + if isinstance(value, bytes): + return {"__type__": "bytes", "base64": base64.b64encode(value).decode("ascii")} + return str(value) + + +def _print_json(obj: Any) -> None: + json.dump(obj, sys.stdout, indent=2, sort_keys=True, default=_json_default) + sys.stdout.write("\n") + + +def _print_human_check(result: dict[str, Any]) -> None: + if result.get("ok"): + status = result.get("status") or {} + print( + f"ok: engraphis-mcp reachable, {status.get('toolCount', '?')} tools " + f"(server={status.get('server')!r})" + ) + else: + print(f"error: {result.get('error')}") + hint = result.get("hint") + if hint: + print(f"hint: {hint}") + + +def _print_human_status(result: dict[str, Any]) -> None: + agents = result.get("agents") or [] + print(f"workspace: {result.get('workspace')}") + print(f"agents: {len(agents)}") + for entry in agents: + sid = entry.get("session_id") or "-" + print(f" - {entry.get('name'):<11} session_id={sid}") + + +def _check(as_json: bool) -> int: + """Boot ``engraphis-mcp`` once and report status. + + Returns 0 on success, 1 on a compatibility error (server reachable but + missing tools), 2 if the binary is not on PATH, 3 on any other failure. + """ + config = build_runtime_config() + binary_path = shutil.which(config.command) + if binary_path is None: + # Don't even try to spawn: report an actionable error and a distinct + # exit code so a wrapper script can tell "binary missing" apart from + # "server reachable but wrong tool set". + result = { + "ok": False, + "error": f"command not found on PATH: {config.command!r}", + "hint": _MISSING_BINARY_HINT, + "command": config.command, + } + if as_json: + _print_json(result) + else: + _print_human_check(result) + return EXIT_MISSING_BINARY + + print(f"command: {config.command} -> {binary_path}", file=sys.stderr) + + async def _run() -> tuple[dict[str, Any], int]: + client = EngraphisMcpClient(config) + try: + await client.connect() + status = await client.status() + return {"ok": True, "command": config.command, "binary": binary_path, "status": status}, 0 + except EngraphisCompatibilityError as exc: + return ( + { + "ok": False, + "error": str(exc), + "hint": ( + "The server is reachable but is missing the Smart 9-tool " + "surface. Upgrade with: pip install --upgrade " + "\"engraphis[mcp]>=1.5,<2\"" + ), + "command": config.command, + "binary": binary_path, + }, + EXIT_INCOMPATIBLE, + ) + except Exception as exc: # noqa: BLE001 — surface to user + hint = client.diagnostic_hint() + return ( + { + "ok": False, + "error": str(exc), + "hint": hint, + "command": config.command, + "binary": binary_path, + }, + EXIT_TRANSPORT, + ) + finally: + await client.close() + + result, exit_code = asyncio.run(_run()) + if as_json: + _print_json(result) + else: + _print_human_check(result) + return exit_code + + +def _status(as_json: bool) -> int: + async def _run() -> tuple[dict[str, Any], int]: + config = build_runtime_config() + # Fail fast (and actionably) if the MCP command isn't on PATH, so the + # user doesn't have to read a stack trace to know the remedy. + if shutil.which(config.command) is None: + return ( + { + "ok": False, + "error": f"command not found on PATH: {config.command!r}", + "hint": _MISSING_BINARY_HINT, + }, + EXIT_MISSING_BINARY, + ) + try: + async with PrimeAgentFleet(workspace="prime-agent-cli") as fleet: + return {"ok": True, **fleet.status()}, 0 + except FileNotFoundError as exc: + return ( + { + "ok": False, + "error": str(exc), + "hint": ( + f"Could not launch {config.command!r}. " + "Install it with: pip install \"engraphis[mcp]>=1.5,<2\"" + ), + }, + EXIT_MISSING_BINARY, + ) + except EngraphisCompatibilityError as exc: + return ( + { + "ok": False, + "error": str(exc), + "hint": ( + "The server is reachable but is missing the Smart 9-tool " + "surface. Upgrade with: pip install --upgrade " + "\"engraphis[mcp]>=1.5,<2\"" + ), + }, + EXIT_INCOMPATIBLE, + ) + except Exception as exc: # noqa: BLE001 — surface to user + return ( + {"ok": False, "error": str(exc), "errorType": type(exc).__name__}, + EXIT_TRANSPORT, + ) + + result, exit_code = asyncio.run(_run()) + if as_json: + _print_json(result) + else: + if result.get("ok"): + _print_human_status(result) + else: + print(f"error: {result.get('error')}") + hint = result.get("hint") + if hint: + print(f"hint: {hint}") + return exit_code + + +def _register(as_json: bool) -> int: + """Print the prime-agent config snippet to stdout.""" + snippet = { + "tools": { + "engraphis": { + "package": "engraphis-prime-agent", + "import": "engraphis_prime_agent", + "entry": "PrimeAgentFleet", + } + } + } + if as_json: + _print_json(snippet) + else: + # Human-readable view of the same snippet. + print("# Drop this into your prime-agent config (e.g. tools section):") + print(json.dumps(snippet["tools"], indent=2, sort_keys=True)) + return 0 + + +def _install(uninstall: bool = False, config_path: str | None = None) -> int: + """Invoke the package-distributed installer. + + The installer lives at ``engraphis_prime_agent.installer`` so it ships + with the wheel and works after ``pip install engraphis-prime-agent`` + (the previous runpy-based path required the source-tree layout). + """ + from .installer import ( + _resolve_config_path, + install as _installer_install, + uninstall as _installer_uninstall, + ) + + path = _resolve_config_path(config_path) + if uninstall: + _installer_uninstall(path) + else: + _installer_install(path) + return 0 + + +def _version() -> int: + """Print the package version (single source of truth: ``__version__``).""" + from . import __version__ + + print(__version__) + return 0 + + +def _add_json_flag(parser: argparse.ArgumentParser) -> None: + """Add ``--json``/``--no-json`` to a subcommand. + + JSON is the default and matches the historical behavior; the flag exists + so wrapper scripts can be explicit, and so users can request a + human-readable view with ``--no-json`` where it makes sense. + """ + parser.add_argument( + "--json", + action=argparse.BooleanOptionalAction, + default=True, + dest="as_json", + help="Emit machine-readable JSON (default: true; use --no-json for text).", + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="engraphis-prime-agent", + description=( + "Engraphis Smart MCP integration for PrimeIntellect's prime-agent. " + "Use one of the subcommands below; --json is the default output " + "format for all subcommands." + ), + ) + sub = parser.add_subparsers(dest="cmd", required=True) + + check_parser = sub.add_parser( + "check", + help="Start engraphis-mcp once and report status.", + description=( + "Boot the configured engraphis-mcp console script, list its tools, " + "and print a JSON status. Exit codes: 0 ok, 1 incompatible tool " + "surface, 2 binary missing, 3 transport error." + ), + ) + _add_json_flag(check_parser) + + status_parser = sub.add_parser( + "status", + help="Boot the 8-agent fleet and print session/agent state.", + description=( + "Construct the 8-agent PrimeAgentFleet, start an MCP session, " + "and print per-agent state. Fails with an actionable error if " + "engraphis-mcp is not installed." + ), + ) + _add_json_flag(status_parser) + + register_parser = sub.add_parser( + "register", + help="Print the prime-agent tool registration snippet.", + description=( + "Print the JSON snippet that registers the engraphis tool with " + "a prime-agent installation. Pipe the output into your config." + ), + ) + _add_json_flag(register_parser) + + install_parser = sub.add_parser( + "install", + help="Idempotently install the integration into prime-agent.", + description=( + "Idempotently register the integration with prime-agent by writing " + "the tools.engraphis entry into its config file. Use --uninstall to " + "remove the entry. --config-path overrides the target file (the " + "PRIME_AGENT_CONFIG_PATH env var is also respected)." + ), + ) + install_parser.add_argument( + "--uninstall", + action="store_true", + help="Remove the engraphis entry from the prime-agent config instead of installing it.", + ) + install_parser.add_argument( + "--config-path", + default=None, + metavar="PATH", + help="Override the prime-agent config file path (defaults to $PRIME_AGENT_CONFIG_PATH or ~/.config/prime-agent/config.json).", + ) + + version_parser = sub.add_parser( + "version", + help="Print the engraphis-prime-agent version and exit.", + description="Print the installed engraphis-prime-agent __version__ and exit.", + ) + # The version subcommand prints a single line; --json is a no-op there + # but kept for symmetry with the other subcommands. + _add_json_flag(version_parser) + + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + as_json = bool(getattr(args, "as_json", True)) + if args.cmd == "check": + return _check(as_json=as_json) + if args.cmd == "status": + return _status(as_json=as_json) + if args.cmd == "register": + return _register(as_json=as_json) + if args.cmd == "install": + return _install( + uninstall=bool(getattr(args, "uninstall", False)), + config_path=getattr(args, "config_path", None), + ) + if args.cmd == "version": + return _version() + parser.error(f"unknown subcommand: {args.cmd}") + return 64 # unreachable, but keeps type-checkers happy + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/integrations/prime_agent/src/engraphis_prime_agent/config.py b/integrations/prime_agent/src/engraphis_prime_agent/config.py new file mode 100644 index 00000000..be63b047 --- /dev/null +++ b/integrations/prime_agent/src/engraphis_prime_agent/config.py @@ -0,0 +1,227 @@ +"""Runtime configuration for the engraphis-mcp stdio gateway. + +Mirrors integrations/pi/src/config.ts: a bounded environment allowlist, an +overridable console command, and explicit default workspace/repo. +""" +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from typing import Any, Mapping + +EXTENSION_VERSION = "0.1.0" + +CORE_DIRECT_TOOLS: tuple[str, ...] = ( + "engraphis_session", + "engraphis_recall_context", + "engraphis_remember", + "engraphis_discover_actions", + "engraphis_execute_read", + "engraphis_execute_action", + "engraphis_get_memory", + "engraphis_update_memory", + "engraphis_conflict_review", +) + +# 8 sub-agent names. Overridable via PrimeAgentFleet(agent_names=...). +# Invariants enforced at import time: exactly 8 entries, each a non-empty +# string, and all distinct so they can be used as fleet/dict keys. +DEFAULT_AGENT_NAMES: tuple[str, ...] = ( + "researcher", # gather context, recall prior decisions + "planner", # decompose goals into ordered steps + "coder", # implement changes + "reviewer", # critique diffs and surface risks + "tester", # write/run/verify tests + "documenter", # capture decisions for durable memory + "monitor", # watch logs, regressions, health + "integrator", # merge, deploy, coordinate handoffs +) + +assert len(DEFAULT_AGENT_NAMES) == 8, "DEFAULT_AGENT_NAMES must contain exactly 8 sub-agents" +assert all(isinstance(n, str) and n for n in DEFAULT_AGENT_NAMES), ( + "DEFAULT_AGENT_NAMES entries must be non-empty strings" +) +assert len(set(DEFAULT_AGENT_NAMES)) == len(DEFAULT_AGENT_NAMES), ( + "DEFAULT_AGENT_NAMES entries must be unique" +) + +# Allowlist, identical to integrations/pi/src/config.ts::engraphisEnvironment. +# +# Note on case sensitivity: +# * POSIX is case-sensitive: only ``PATH`` exists; ``Path`` would be a +# separate variable and is harmless to include. +# * Windows is case-insensitive: ``PATH``, ``Path``, and ``path`` all refer +# to the same environment entry. Including both ``PATH`` and ``Path`` is +# redundant on Windows but never harmful — the OS lookups normalise case +# and Python's ``os.environ`` preserves the case of the *first* writer. +# We keep both for symmetry with the Pi TS implementation. +_ALLOWED_ENV_KEYS = frozenset({ + "PATH", "Path", "SystemRoot", "ComSpec", + # Windows-only home variables. ``Path.home()`` reads USERPROFILE first + # and falls back to HOMEDRIVE+HOMEPATH; without them the early + # ``_resolve_config_env_path()`` call aborts with FileNotFoundError on + # ``~/.engraphis.env`` before the MCP handshake runs. Forward them on + # every platform so a wheel installed on Windows does not need a + # pre-existing ``ENGRAPHIS_ENV_FILE`` to bootstrap. + "USERPROFILE", "HOMEDRIVE", "HOMEPATH", +}) +_ALLOWED_ENV_PREFIX = "ENGRAPHIS_" + + +@dataclass(frozen=True) +class EngraphisRuntimeConfig: + """Resolved runtime configuration for the stdio gateway subprocess. + + The dataclass is frozen: attributes cannot be reassigned after ``__init__``. + The mutable-looking fields (``args``, ``environment``) are normalised in + :meth:`__post_init__` so that callers cannot mutate them in place either + — ``args`` becomes a ``tuple`` and ``environment`` is a shallow copy of + the input mapping stored as an immutable-style ``dict[str, str]``. + + :param command: Executable name or absolute path of the MCP gateway + binary. Must be a non-empty string; falls back to ``"engraphis-mcp"`` + on the PATH when constructed via :func:`build_runtime_config`. + :param args: Positional arguments passed to ``command``. Frozen as a + tuple at construction time. + :param cwd: Optional working directory for the subprocess. The value + is forwarded unchanged to the runtime layer, which is responsible + for path resolution and existence checks; this class only enforces + that, when provided, it is a non-empty string. + :param default_workspace: Optional default workspace identifier + forwarded to the gateway (typically a memory scope key). + :param default_repo: Optional default repository identifier forwarded + to the gateway. + :param environment: Allowlist-filtered environment variables to pass + to the subprocess. Stored as a defensive copy. + """ + + command: str = "engraphis-mcp" + args: tuple[str, ...] = () + cwd: str | None = None + default_workspace: str | None = None + default_repo: str | None = None + environment: Mapping[str, str] = field(default_factory=dict) + + def __post_init__(self) -> None: + # Validate `command`: must be a non-empty string. We check truthiness + # after stripping so a bare-whitespace value is rejected too. + if not isinstance(self.command, str) or not self.command.strip(): + raise ValueError("EngraphisRuntimeConfig.command must be a non-empty string") + # Normalise `command` in place (frozen dataclass requires object.__setattr__). + object.__setattr__(self, "command", self.command.strip()) + + # Freeze `args` as a tuple. Accept any iterable of strings; reject + # non-string entries to surface caller mistakes early. + normalised_args: tuple[str, ...] = tuple(self.args) + for a in normalised_args: + if not isinstance(a, str): + raise TypeError( + f"EngraphisRuntimeConfig.args entries must be str, got {type(a).__name__}" + ) + object.__setattr__(self, "args", normalised_args) + + # `cwd`: light validation. The runtime layer is responsible for + # path resolution and existence checks; here we only ensure that, + # when provided, the value is a non-empty string. Relative paths + # are allowed and resolved relative to the parent process cwd. + if self.cwd is not None and (not isinstance(self.cwd, str) or not self.cwd): + raise ValueError("EngraphisRuntimeConfig.cwd must be a non-empty string or None") + + # Defensive copy of the environment mapping. We also coerce values + # to str to give the field a precise ``Mapping[str, str]`` shape + # even if a caller passed a more permissive type. + env_copy: dict[str, str] = {str(k): str(v) for k, v in dict(self.environment).items()} + object.__setattr__(self, "environment", env_copy) + + def as_subprocess_env(self) -> dict[str, str]: + """Return a fresh ``dict`` copy of the environment for subprocess use. + + Always returns a new mapping so callers can mutate the result + without affecting this config's frozen state. + """ + return dict(self.environment) + + +def _non_blank(value: str | None) -> str | None: + """Return ``value`` with surrounding whitespace stripped, or ``None``. + + A value that is ``None``, empty, or whitespace-only returns ``None``; + otherwise the stripped string is returned. Used to normalise optional + environment overrides before they are stored on the config. + """ + if value is None: + return None + cleaned = value.strip() + return cleaned or None + + +def _engraphis_environment(env: Mapping[str, Any]) -> dict[str, str]: + """Forward only the Engraphis settings and the Windows/POSIX path vars. + + Mirrors integrations/pi/src/config.ts so a sub-agent's gateway sees the + same allowlist the Pi extension uses. + + The parameter is typed ``Mapping[str, Any]`` because real-world + sources (``os.environ`` is fine, but test fixtures and ad-hoc dicts may + contain ``None`` or other non-string values). Non-string values are + silently dropped — this is intentional: a missing or wrongly-typed + variable should not crash config construction, it should just be + excluded from the forwarded environment. + """ + forwarded: dict[str, str] = {} + for key, value in env.items(): + if not isinstance(value, str): + continue + if key.startswith(_ALLOWED_ENV_PREFIX) or key in _ALLOWED_ENV_KEYS: + # Trim surrounding whitespace so a value like " /tmp/x.db " is + # forwarded as "/tmp/x.db". This keeps gateway config (paths, + # workspace ids, repo names) free of accidental padding and + # matches the trimming `_non_blank` applies to the dedicated + # workspace/repo fields. + forwarded[key] = value.strip() + return forwarded + + +def build_runtime_config( + env: Mapping[str, Any] | None = None, + *, + command: str | None = None, + args: tuple[str, ...] | None = None, + cwd: str | None = None, +) -> EngraphisRuntimeConfig: + """Build the runtime config the same way the Pi TS integration does. + + Reads from ``env`` (defaults to :data:`os.environ`) with the following + resolution order for each field: + + * ``command`` — explicit ``command`` kwarg, else + ``$ENGRAPHIS_MCP_COMMAND``, else ``"engraphis-mcp"``. + * ``args`` — explicit ``args`` kwarg, else ``()``. + * ``cwd`` — explicit ``cwd`` kwarg, else ``None``. + * ``default_workspace`` — ``$ENGRAPHIS_WORKSPACE`` (trimmed; + whitespace-only becomes ``None``). + * ``default_repo`` — ``$ENGRAPHIS_REPO`` (trimmed). + * ``environment`` — allowlist-filtered view of ``env``; only keys with + the ``ENGRAPHIS_`` prefix or in :data:`_ALLOWED_ENV_KEYS` are + forwarded, and only when their value is a ``str``. + + The returned :class:`EngraphisRuntimeConfig` is frozen and stores + defensive copies of any mutable inputs. + """ + src: Mapping[str, Any] = os.environ if env is None else env + resolved_command = ( + _non_blank(command) + or _non_blank(src.get("ENGRAPHIS_MCP_COMMAND")) # type: ignore[arg-type] + or "engraphis-mcp" + ) + forwarded = _engraphis_environment(src) + workspace = _non_blank(src.get("ENGRAPHIS_WORKSPACE")) # type: ignore[arg-type] + repo = _non_blank(src.get("ENGRAPHIS_REPO")) # type: ignore[arg-type] + return EngraphisRuntimeConfig( + command=resolved_command, + args=tuple(args or ()), + cwd=cwd, + default_workspace=workspace, + default_repo=repo, + environment=forwarded, + ) diff --git a/integrations/prime_agent/src/engraphis_prime_agent/installer.py b/integrations/prime_agent/src/engraphis_prime_agent/installer.py new file mode 100644 index 00000000..ac0aa7b2 --- /dev/null +++ b/integrations/prime_agent/src/engraphis_prime_agent/installer.py @@ -0,0 +1,309 @@ +"""Idempotent registration of the integration with PrimeIntellect's prime-agent. + +This module is the canonical, package-distributed implementation. The +``scripts/install_prime_agent.py`` wrapper at the repo root invokes this +module so the install/uninstall behaviour stays identical for both +``pip install`` users and source-tree developers. + +The exact prime-agent config file path is the verification point: at +implementation time the implementer inspects +https://github.com/PrimeIntellect-ai/prime-agent and uses the documented +location. This module defaults to a JSON file at +``~/.config/prime-agent/config.json`` (or whatever ``PRIME_AGENT_CONFIG_PATH`` +points at) and falls back to TOML when the file has a ``.toml`` extension. +The path and format can be confirmed and tightened once the prime-agent +repo is available. + +Usage: + python scripts/install_prime_agent.py + python scripts/install_prime_agent.py --uninstall +""" +from __future__ import annotations + +import argparse +import copy +import datetime +import json +import os +import shutil +import sys +from pathlib import Path +from typing import Any + +PACKAGE = "engraphis_prime_agent" +ENTRY = "PrimeAgentFleet" +TOOL_KEY = "engraphis" + +# Default path; override with PRIME_AGENT_CONFIG_PATH. +_DEFAULT_PATH = Path.home() / ".config" / "prime-agent" / "config.json" + + +def _settings_path() -> Path: + override = os.environ.get("PRIME_AGENT_CONFIG_PATH") + if override: + return Path(override) + return _DEFAULT_PATH + + +def _utc_stamp() -> str: + return datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%d") + + +def _backup(path: Path) -> Path | None: + if not path.exists(): + return None + # Skip the backup when the file is brand new (zero bytes) or empty — + # there's nothing meaningful to preserve, and the timestamp collision + # on rapid successive runs is avoided. + if path.stat().st_size == 0: + return None + # Use a collision-resistant suffix (UTC date + pid + unix-ms) so a + # second run on the same UTC date captures the user's other tool + # settings too. A pure per-day filename would overwrite the previous + # backup and lose unrelated configuration. + import os as _os + import time as _time + _pid = _os.getpid() + _now_ms = int(_time.time() * 1000) + base_name = ( + f"{path.name}.bak-engraphis-{_utc_stamp()}.{_pid}.{_now_ms}" + ) + if path.with_name(base_name).exists(): + # Last-ditch uniqueness: append a counter until the name is free. + counter = 0 + candidate_name = base_name + while path.with_name(candidate_name).exists(): + counter += 1 + candidate_name = ( + f"{path.name}.bak-engraphis-{_utc_stamp()}.{_pid}." + f"{_now_ms}.{counter}" + ) + backup = path.with_name(candidate_name) + else: + backup = path.with_name(base_name) + backup.write_bytes(path.read_bytes()) + shutil.copymode(path, backup) + return backup + + +def _read(path: Path) -> dict[str, Any]: + if not path.exists(): + return {} + text = path.read_text(encoding="utf-8").strip() + if not text: + return {} + if path.suffix == ".json": + try: + return json.loads(text) + except json.JSONDecodeError as exc: + print(f"error: {path} is not valid JSON: {exc}", file=sys.stderr) + sys.exit(2) + if path.suffix == ".toml": + try: + import tomllib # Python 3.11+ + except ImportError: + print( + f"error: reading {path} as TOML requires Python 3.11+ " + "(tomllib is in the stdlib from 3.11 onward)", + file=sys.stderr, + ) + sys.exit(2) + try: + return tomllib.loads(text) + except tomllib.TOMLDecodeError as exc: + print(f"error: {path} is not valid TOML: {exc}", file=sys.stderr) + sys.exit(2) + print( + f"error: unsupported config format for {path} " + f"(expected .json or .toml, got {path.suffix!r})", + file=sys.stderr, + ) + sys.exit(2) + + +def _ensure_writable_parent(path: Path) -> None: + """Refuse to write if the parent directory is not writable. + + Catches the common failure modes early: missing parent on a read-only + filesystem, an unwritable existing directory, or a path whose parent is a + file. The actual write still happens after this check, so a TOCTOU race is + technically possible, but in practice the only way to fail here is the + configuration the user is asking us to use. + """ + parent = path.parent + if parent.exists() and not parent.is_dir(): + print( + f"error: parent of {path} exists but is not a directory: {parent}", + file=sys.stderr, + ) + sys.exit(2) + if not parent.exists(): + # We will create it; check that we can. ``os.access`` on a non-existent + # path checks the nearest existing ancestor, which is what we want. + ancestor = parent + while not ancestor.exists(): + ancestor = ancestor.parent + if not os.access(str(ancestor), os.W_OK): + print( + f"error: cannot create {path}: no write access to {ancestor}", + file=sys.stderr, + ) + sys.exit(2) + return + if not os.access(str(parent), os.W_OK): + print( + f"error: parent directory of {path} is not writable: {parent}", + file=sys.stderr, + ) + sys.exit(2) + + +def _write(path: Path, data: dict[str, Any]) -> None: + _ensure_writable_parent(path) + path.parent.mkdir(parents=True, exist_ok=True) + if path.suffix == ".json": + path.write_text( + json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return + if path.suffix == ".toml": + try: + import tomli_w + except ImportError: + print( + f"error: writing {path} as TOML requires the 'tomli_w' package; " + "install it with: pip install 'tomli_w>=1.0' " + "(it is not bundled with the engraphis core package)", + file=sys.stderr, + ) + sys.exit(2) + # tomli_w.dumps returns str, not bytes — use write_text, not write_bytes. + path.write_text(tomli_w.dumps(data), encoding="utf-8") + return + print( + f"error: unsupported config format for {path} " + f"(expected .json or .toml, got {path.suffix!r})", + file=sys.stderr, + ) + sys.exit(2) + + +def _entry() -> dict[str, str]: + # Use the underscore-separated import name as the distribution name; the + # PyPI distribution is engraphis-prime-agent (hyphenated) but the + # Python import path is engraphis_prime_agent (underscored). + return { + "package": "engraphis-prime-agent", + "import": PACKAGE, + "entry": ENTRY, + } + + +def _dry_run(path: Path, before: dict[str, Any], after: dict[str, Any]) -> None: + print("--- before") + print(json.dumps(before, indent=2, sort_keys=True)) + print("--- after") + print(json.dumps(after, indent=2, sort_keys=True)) + print(f"(dry-run) no changes written to {path}") + + +def install( + path: Path | None = None, + *, + merge: bool = False, + dry_run: bool = False, +) -> None: + path = path or _settings_path() + cfg = _read(path) + # Deep-copy so the dry-run snapshot does not observe the mutations + # below: ``cfg.setdefault("tools", {})`` would otherwise return a + # reference to the same nested dict that we then overwrite with the + # new entry, mutating ``before`` as well. + before = copy.deepcopy(cfg) + tools = cfg.setdefault("tools", {}) + entry = _entry() + if merge and isinstance(tools.get(TOOL_KEY), dict): + # Preserve operator-supplied keys under the tools.engraphis table. + merged = dict(tools[TOOL_KEY]) + merged.update(entry) + tools[TOOL_KEY] = merged + else: + tools[TOOL_KEY] = entry + if dry_run: + _dry_run(path, before, cfg) + return + _backup(path) + _write(path, cfg) + print(f"installed engraphis-prime-agent into {path}") + + +def uninstall( + path: Path | None = None, + *, + dry_run: bool = False, +) -> None: + path = path or _settings_path() + cfg = _read(path) + # Deep-copy so the dry-run snapshot does not observe the deletions + # below: ``tools.pop(TOOL_KEY)`` mutates the same nested mapping that + # ``before`` still points at, so the printed "before" would show the + # already-removed key. + before = copy.deepcopy(cfg) + tools = cfg.get("tools", {}) + if TOOL_KEY not in tools: + if dry_run: + _dry_run(path, before, before) + else: + print(f"no engraphis entry in {path}") + return + del tools[TOOL_KEY] + if not tools: + cfg.pop("tools", None) + if dry_run: + _dry_run(path, before, cfg) + return + _backup(path) + _write(path, cfg) + print(f"removed engraphis entry from {path}") + + +def _resolve_config_path(explicit: str | None) -> Path | None: + """CLI flag → env var → None (use default). Empty string is treated as unset.""" + if explicit: + return Path(explicit) + env = os.environ.get("PRIME_AGENT_CONFIG_PATH") + if env: + return Path(env) + return None + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=(__doc__ or "").split("\n\n", 1)[0]) + parser.add_argument("--uninstall", action="store_true") + parser.add_argument( + "--config-path", + default=None, + help="Override the prime-agent config file path (defaults to " + "$PRIME_AGENT_CONFIG_PATH or ~/.config/prime-agent/config.json).", + ) + parser.add_argument( + "--merge", + action="store_true", + help="Merge with any existing [tools.engraphis] entry instead of replacing it.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show the before/after diff and exit without writing or backing up.", + ) + args = parser.parse_args(argv) + path = _resolve_config_path(args.config_path) + if args.uninstall: + uninstall(path, dry_run=args.dry_run) + else: + install(path, merge=args.merge, dry_run=args.dry_run) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/integrations/prime_agent/src/engraphis_prime_agent/mcp_client.py b/integrations/prime_agent/src/engraphis_prime_agent/mcp_client.py new file mode 100644 index 00000000..411885de --- /dev/null +++ b/integrations/prime_agent/src/engraphis_prime_agent/mcp_client.py @@ -0,0 +1,452 @@ +"""Async stdio client for the local Engraphis MCP gateway. + +Translates integrations/pi/src/mcp-client.ts to the Python `mcp` SDK: + - one shared subprocess (StdioClientTransport from mcp.client.stdio) + - generation counter so a close-during-connect cannot leave a stale Client + - bounded 4 KiB stderr buffer for diagnosis + - retry-on-read-only up to 2 attempts with backoff + - 60s connect / 5 min tool timeouts + - two distinct exception classes for tool-level vs. compatibility errors +""" +from __future__ import annotations + +import asyncio +import json +import logging +import os +import re +import tempfile +import time +from contextlib import AsyncExitStack +from typing import Any, TextIO + +from mcp import ClientSession, StdioServerParameters +from mcp.client.stdio import stdio_client +from mcp.types import Implementation + +from .config import CORE_DIRECT_TOOLS, EXTENSION_VERSION, EngraphisRuntimeConfig + +_logger = logging.getLogger("engraphis_prime_agent.mcp_client") + +TOOL_REQUEST_TIMEOUT_S = 5 * 60 +CONNECT_TIMEOUT_S = 60 +STDERR_BUFFER_BYTES = 4 * 1024 +MAX_TOOL_LIST_PAGES = 100 +MAX_TOOL_LIST_TOOLS = 1_000 + +# Tools whose server-side contract is idempotent. A transport failure +# can be safely retried because the server will produce the same result. +# ``engraphis_recall_context`` is intentionally NOT in this set: the +# Smart gateway appends a receipt on each successful call, so retrying +# after a transport-level failure would create duplicate accounting +# records for one logical user request. +READ_ONLY_TOOLS = frozenset({ + "engraphis_get_memory", + "engraphis_conflict_review", + "engraphis_discover_actions", + "engraphis_execute_read", +}) + + +class EngraphisMcpToolError(RuntimeError): + """Semantic rejection returned by the MCP server (e.g. invalid args).""" + + def __init__(self, message: str, *, retryable: bool = False) -> None: + super().__init__(message) + # Preserve the Smart gateway's retryability signal for hosts that want + # to apply a policy appropriate to the operation. The client itself + # does not retry semantic tool errors because writes may not be safe to + # repeat and read-only retry rules are transport-specific. + self.retryable = retryable + + +class EngraphisCompatibilityError(RuntimeError): + """Gateway is reachable but does not expose the Smart 9-tool surface.""" + + +class EngraphisMcpClient: + """Lazy async stdio client. Safe to share across coroutines. + + Concurrent tool calls are serialized through a single asyncio.Lock; the + stdio transport is one connection, so the upstream SDK cannot interleave + JSON-RPC frames safely. Framework-level concurrency (e.g. 8 sub-agents + reasoning in parallel and then each issuing a tool call) is unaffected. + """ + + def __init__(self, config: EngraphisRuntimeConfig) -> None: + self._config = config + self._lifecycle = 0 + self._session: ClientSession | None = None + self._stack: AsyncExitStack | None = None + self._connect_lock = asyncio.Lock() + self._call_lock = asyncio.Lock() + self._tools_cache: list[dict[str, Any]] | None = None + self._diagnostic = "" + self._client_name = f"engraphis-prime-agent/{EXTENSION_VERSION}" + # A real temp file is the only cross-platform `errlog` that Windows + # subprocess.Popen accepts. The file is read on demand to fill the + # bounded diagnostic buffer; it's never persisted. + self._stderr_file: TextIO | None = None + self._stderr_path: str | None = None + + # --- lifecycle ------------------------------------------------------- + + def generation(self) -> int: + return self._lifecycle + + @property + def config(self) -> EngraphisRuntimeConfig: + return self._config + + def diagnostic_hint(self) -> str | None: + d = self._diagnostic + if re.search(r"python 3\.10|requires python 3\.10", d, re.I): + return "The Engraphis MCP server requires Python 3.10 or later." + if re.search(r"no module named ['\"]?mcp", d, re.I): + return "The Engraphis MCP dependency is missing. Install `engraphis[mcp]>=1.5,<2`." + if re.search(r"no module named ['\"]?engraphis", d, re.I): + return "Engraphis is not installed for the configured MCP command." + return None + + def _refresh_diagnostic_from_file(self) -> None: + path = self._stderr_path + if not path: + return + try: + with open(path, "r", encoding="utf-8", errors="replace") as f: + data = f.read(STDERR_BUFFER_BYTES * 4) + except OSError: + return + self._diagnostic = data[-STDERR_BUFFER_BYTES:] + + async def connect(self) -> ClientSession: + async with self._connect_lock: + if self._session is not None: + return self._session + # Bound the entire connect sequence (stdio handshake + + # initialize + tools/list) so a subprocess that completes + # initialization but never answers tools/list cannot hang + # the advertised 60-second connection timeout. + _connect_started = time.monotonic() + _connect_budget = CONNECT_TIMEOUT_S + # Capture the generation so a concurrent close() (which bumps + # _lifecycle) invalidates this connect. The post-await check + # below closes the freshly-opened stack and discards the session + # instead of publishing a live subprocess after shutdown. + generation = self._lifecycle + self._diagnostic = "" + stack = AsyncExitStack() + try: + params = StdioServerParameters( + command=self._config.command, + args=list(self._config.args), + cwd=self._config.cwd, + env=dict(self._config.environment), + ) + # Open a real temp file for stderr so Windows subprocess.Popen + # can take its fileno. The file is closed and unlinked after + # the session is torn down. + err_fd, err_path = tempfile.mkstemp(prefix="engraphis-prime-agent-", suffix=".err") + err_file = os.fdopen(err_fd, mode="w", encoding="utf-8", buffering=1) + # Register the unlink BEFORE the close. AsyncExitStack + # runs callbacks in LIFO order, so the first-registered + # (unlink) fires last, after the file handle has been + # closed — required on Windows where an open file cannot + # be unlinked. + stack.callback(self._safe_unlink, err_path) + stack.callback(err_file.close) + self._stderr_file = err_file + self._stderr_path = err_path + read, write = await asyncio.wait_for( + stack.enter_async_context(stdio_client(params, errlog=err_file)), + timeout=CONNECT_TIMEOUT_S, + ) + session = await stack.enter_async_context( + ClientSession( + read, + write, + client_info=Implementation(name=self._client_name, version=EXTENSION_VERSION), + ) + ) + _remaining = _connect_budget - (time.monotonic() - _connect_started) + if _remaining <= 0: + raise asyncio.TimeoutError( + f"engraphis-mcp connect exceeded {_connect_budget:.0f}s" + ) + # Initialization shares the same end-to-end budget as the + # transport setup. Reusing CONNECT_TIMEOUT_S here could let + # a slow stdio handshake consume the full budget and then + # grant initialize another full minute. + await asyncio.wait_for(session.initialize(), timeout=_remaining) + _remaining = _connect_budget - (time.monotonic() - _connect_started) + if _remaining <= 0: + raise asyncio.TimeoutError( + f"engraphis-mcp connect exceeded {_connect_budget:.0f}s" + ) + tools = await asyncio.wait_for( + self._list_tools(session), timeout=_remaining + ) + available = {t["name"] for t in tools} + missing = [n for n in CORE_DIRECT_TOOLS if n not in available] + if missing: + self._refresh_diagnostic_from_file() + raise EngraphisCompatibilityError( + "Engraphis 1.5.x Smart MCP is required; the server is " + f"missing: {', '.join(missing)}." + ) + # If close() ran while we were awaiting, abort — don't + # publish a session that the caller has already decided to + # discard. The local stack is closed before the raise so the + # subprocess is reaped. + if self._lifecycle != generation: + await stack.aclose() + self._stderr_file = None + self._stderr_path = None + raise EngraphisMcpToolError( + "Engraphis client was closed before the connect completed." + ) + self._session = session + self._stack = stack + self._tools_cache = tools + return session + except BaseException: + self._refresh_diagnostic_from_file() + await stack.aclose() + self._session = None + self._stack = None + self._tools_cache = None + self._stderr_file = None + self._stderr_path = None + raise + + @staticmethod + def _safe_unlink(path: str) -> None: + try: + os.unlink(path) + except OSError: + pass + + async def close(self) -> None: + # Hold the connect lock so any in-flight connect() either completes + # before us (and is then torn down) or aborts via the post-await + # generation check. Without this, a concurrent close() can return + # while a connect() is still mid-await, leaving a live subprocess. + async with self._connect_lock: + self._lifecycle += 1 + stack = self._stack + self._stack = None + self._session = None + self._tools_cache = None + # Reset stderr-temp-file handles. The actual file close + unlink are + # registered as AsyncExitStack callbacks in connect(), so they fire + # when `stack.aclose()` runs below. We just need to drop the Python + # references so a subsequent connect() can recreate them cleanly. + self._stderr_file = None + self._stderr_path = None + if stack is not None: + try: + await stack.aclose() + except Exception: # noqa: BLE001 — best-effort teardown + _logger.debug("ignored error while closing MCP stack", exc_info=True) + + async def __aenter__(self) -> "EngraphisMcpClient": + await self.connect() + return self + + async def __aexit__(self, *exc: Any) -> None: + await self.close() + + # --- tool surface ---------------------------------------------------- + + async def list_tools(self) -> list[dict[str, Any]]: + if self._tools_cache is None: + await self.connect() + # connect() performs the bounded discovery once and publishes the + # result atomically with the session. Reuse that cache here instead + # of issuing a second unbounded tools/list request. + if self._tools_cache is None: + raise EngraphisMcpToolError( + "Engraphis client connected without a tool-list cache." + ) + return list(self._tools_cache) + + async def call_tool(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]: + if name not in CORE_DIRECT_TOOLS: + raise EngraphisMcpToolError(f"Unknown Engraphis tool: {name}") + last_error: BaseException | None = None + retry = name in READ_ONLY_TOOLS + max_attempts = 3 if retry else 1 + for attempt in range(max_attempts): + try: + async with self._call_lock: + session = await self.connect() + response = await asyncio.wait_for( + session.call_tool(name, arguments), + timeout=TOOL_REQUEST_TIMEOUT_S, + ) + return self._format_result(name, response) + except EngraphisMcpToolError: + raise + except EngraphisCompatibilityError: + raise + except asyncio.TimeoutError: + raise + except asyncio.CancelledError: + raise + except (BrokenPipeError, ConnectionError, OSError, EOFError) as exc: + # Standard transport / stdio-pipe failure: log distinctly + # at DEBUG (per-attempt noise is already covered by the + # WARNING below on the terminal failure). + last_error = exc + _logger.debug( + "MCP transport failure for %s (attempt %d): %s", + name, attempt + 1, exc, + ) + self._refresh_diagnostic_from_file() + await self.close() + if attempt + 1 >= max_attempts: + break + # Linear backoff: attempt 0 -> 1.0s, attempt 1 -> 2.2s. + # Formula: base * (attempt + 1) + jitter * attempt. + await asyncio.sleep((attempt + 1) * 1.0 + attempt * 0.2) + except Exception as exc: # unexpected transport failure + last_error = exc + _logger.debug( + "MCP unexpected failure for %s (attempt %d): %s", + name, attempt + 1, exc, + ) + self._refresh_diagnostic_from_file() + await self.close() + if attempt + 1 >= max_attempts: + break + await asyncio.sleep((attempt + 1) * 1.0 + attempt * 0.2) + assert last_error is not None + _logger.warning( + "MCP call %s failed after %d attempt(s): %s", + name, max_attempts, last_error, + ) + raise last_error + + # --- helpers --------------------------------------------------------- + + async def _list_tools(self, session: ClientSession) -> list[dict[str, Any]]: + all_tools: list[dict[str, Any]] = [] + cursor: str | None = None + seen_cursors: set[str] = set() + for _page_number in range(MAX_TOOL_LIST_PAGES): + if cursor is not None: + if cursor in seen_cursors: + raise EngraphisCompatibilityError( + "Engraphis tools/list returned a repeated pagination cursor." + ) + seen_cursors.add(cursor) + page = await session.list_tools(cursor=cursor) + for tool in page.tools: + if len(all_tools) >= MAX_TOOL_LIST_TOOLS: + raise EngraphisCompatibilityError( + "Engraphis tools/list exceeded the advertised tool limit." + ) + all_tools.append( + { + "name": tool.name, + "description": tool.description, + "inputSchema": tool.inputSchema, + } + ) + cursor = page.nextCursor + if not cursor: + return all_tools + raise EngraphisCompatibilityError( + "Engraphis tools/list exceeded the pagination page limit." + ) + + @staticmethod + def _format_result(name: str, response: Any) -> dict[str, Any]: + is_error = bool(getattr(response, "isError", False)) + content: list[dict[str, Any]] = [] + for block in getattr(response, "content", []) or []: + text = getattr(block, "text", None) + content.append({"type": getattr(block, "type", "text"), "text": text}) + text = "\n\n".join( + b["text"] for b in content if b.get("type") == "text" and b.get("text") + ).strip() + declared_error = re.match(r"^Error:\s*([a-z0-9_]+)\s*$", text, re.I) + server_error = text.lower().startswith("error:") + # The Smart gateway emits a structured JSON envelope on tool + # validation / scope / not-found failures (``engraphis/mcp_server.py:: + # _smart_error``): ``{"code": "...", "message": "...", "retryable": false}``. + # Detect that envelope inside any text block and forward its code, + # message, and retryable flag so agent hosts can distinguish caller + # errors from retryable/internal failures as the Smart contract + # intends, rather than collapsing every error to the same generic + # message. + envelope: dict[str, Any] | None = None + for block in content: + if block.get("type") != "text" or not isinstance(block.get("text"), str): + continue + try: + parsed = json.loads(block["text"]) + except (ValueError, TypeError): + continue + if ( + isinstance(parsed, dict) + and isinstance(parsed.get("error"), dict) + and isinstance(parsed["error"].get("code"), str) + and isinstance(parsed["error"].get("message"), str) + ): + # Smart gateway wraps every failure as + # ``{"error":{"code":...,"message":...,"retryable":...}}`` + # (see engraphis/mcp_server.py::_smart_error). Accept the + # nested shape so callers can distinguish validation errors + # from retryable internal failures as the Smart contract + # intends. + envelope = parsed["error"] + break + if ( + isinstance(parsed, dict) + and isinstance(parsed.get("code"), str) + and isinstance(parsed.get("message"), str) + ): + envelope = parsed + break + if is_error or server_error: + if envelope is not None: + msg = ( + f"Engraphis rejected the request: " + f"{envelope.get('code', 'unknown')}: {envelope.get('message', '')}" + ) + elif declared_error: + msg = f"Engraphis rejected the request: {declared_error.group(1)}." + else: + msg = ( + "Engraphis rejected the request. Verify the parameters and " + "inspect the local Engraphis logs." + ) + retryable = ( + bool(envelope.get("retryable", False)) + if envelope is not None + else False + ) + raise EngraphisMcpToolError(msg, retryable=retryable) + return {"_tool": name, "isError": is_error, "content": content} + + # --- status ---------------------------------------------------------- + + async def status(self) -> dict[str, Any]: + tools = await self.list_tools() + return { + "connected": True, + "server": "engraphis", + "toolCount": len(tools), + "diagnosticHint": self.diagnostic_hint(), + } + + +def format_mcp_payload(payload: dict[str, Any]) -> str: + """Return the joined text content of a tool result, falling back to JSON.""" + parts: list[str] = [] + for block in payload.get("content", []) or []: + if block.get("type") == "text" and isinstance(block.get("text"), str): + parts.append(block["text"]) + joined = "\n\n".join(parts).strip() + return joined or json.dumps(payload, indent=2, default=str) diff --git a/integrations/prime_agent/src/engraphis_prime_agent/tools.py b/integrations/prime_agent/src/engraphis_prime_agent/tools.py new file mode 100644 index 00000000..9b679480 --- /dev/null +++ b/integrations/prime_agent/src/engraphis_prime_agent/tools.py @@ -0,0 +1,556 @@ +"""9 Smart tool factories, each a (args, ctx) -> dict callable. + +Schema and semantics are translated 1:1 from +integrations/pi/src/tool-schemas.ts. The resulting callables work with +both EngraphisPrimeAgent and any prime-agent tool-registration surface that +matches the (args: dict, ctx: dict | None) -> dict contract. +""" +from __future__ import annotations + +from typing import Any, Awaitable, Callable + +from .config import EngraphisRuntimeConfig +from .mcp_client import EngraphisMcpClient, EngraphisMcpToolError + +# The runtime contract: prime-agent (and any compatible tool-registration +# surface) calls the registered callable with the model's args plus an +# optional ctx dict (conversation/session metadata). Both are accepted +# positionally; ctx defaults to None so the legacy single-arg call shape +# still works. +ToolFn = Callable[ + [dict[str, Any], dict[str, Any] | None], Awaitable[dict[str, Any]] +] + +# --- JSON Schemas (translated from tool-schemas.ts) ------------------------- +# The same defaults, bounds, and descriptions; identical behaviour across Pi +# and prime-agent integrations. + +_SESSION_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "action": { + "type": "string", + "enum": ["start", "end", "start_session", "end_session"], + "default": "start", + }, + # The wrapper supplies the registered agent name when the caller omits + # this optional field. Keeping it optional also lets the framework + # invoke the lifecycle tool without duplicating registration metadata. + "agent": {"type": "string", "minLength": 1, "maxLength": 200}, + "force_new": {"type": "boolean", "default": False}, + "goal": {"type": "string", "maxLength": 1000, "default": ""}, + "session_id": {"type": "string", "maxLength": 200, "default": ""}, + "summary": {"type": "string", "maxLength": 100000, "default": ""}, + "outcome": {"type": "string", "maxLength": 1000, "default": ""}, + "open_threads": { + "type": ["array", "null"], + "items": {"type": "string"}, + "default": None, + }, + "token_budget": {"type": "integer", "minimum": 0, "maximum": 32768, "default": 512}, + "workspace": {"type": "string", "maxLength": 200}, + "repo": {"type": ["string", "null"], "maxLength": 200, "default": None}, + }, + "required": [], +} + +_RECALL_CONTEXT_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "query": {"type": "string", "minLength": 1, "maxLength": 100000}, + "k": {"type": "integer", "minimum": 1, "maximum": 50, "default": 50}, + "session_id": {"type": ["string", "null"], "default": None}, + "token_budget": { + "type": "integer", + "minimum": 0, + "maximum": 32768, + "default": 1024, + }, + "workspace": {"type": ["string", "null"], "maxLength": 200, "default": None}, + "repo": {"type": ["string", "null"], "maxLength": 200, "default": None}, + }, + "required": ["query"], +} + +_REMEMBER_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "content": {"type": "string", "minLength": 1, "maxLength": 100000}, + "mtype": { + "type": "string", + "enum": ["semantic", "episodic", "procedural", "working"], + "default": "semantic", + }, + "importance": {"type": "number", "minimum": 0, "maximum": 1, "default": 0}, + "session_id": {"type": ["string", "null"], "default": None}, + "workspace": {"type": "string", "maxLength": 200}, + "repo": {"type": ["string", "null"], "maxLength": 200, "default": None}, + "subject_key": {"type": "string", "maxLength": 1000}, + "claim_kind": {"type": "string", "maxLength": 200}, + }, + "required": ["content"], +} + +_DISCOVER_ACTIONS_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "task": {"type": "string", "minLength": 1, "maxLength": 2000}, + "category": { + "type": "string", + "enum": ["memory", "governance", "code", "audit", "ops", ""], + "maxLength": 100, + "default": "", + }, + "intent": { + "type": "string", + "enum": ["any", "read", "write", "admin", "destructive"], + "default": "any", + }, + "limit": {"type": "integer", "minimum": 1, "maximum": 3, "default": 1}, + }, + "required": ["task"], +} + +_EXECUTE_PARAM_PROPS = { + "capability_id": {"type": "string", "minLength": 8, "maxLength": 128}, + "schema_digest": {"type": "string", "minLength": 8, "maxLength": 128}, + "arguments": {"type": "object", "additionalProperties": True}, +} + +_EXECUTE_READ_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": _EXECUTE_PARAM_PROPS, + "required": ["capability_id", "schema_digest", "arguments"], +} + +_EXECUTE_ACTION_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": _EXECUTE_PARAM_PROPS, + "required": ["capability_id", "schema_digest", "arguments"], +} + +_GET_MEMORY_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "memory_id": {"type": "string", "minLength": 1, "maxLength": 200}, + "workspace": {"type": "string", "maxLength": 200}, + "repo": {"type": ["string", "null"], "maxLength": 200, "default": None}, + }, + "required": ["memory_id"], +} + +_UPDATE_MEMORY_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "memory_id": {"type": "string", "minLength": 1, "maxLength": 200}, + "title": {"type": ["string", "null"], "maxLength": 500, "default": None}, + "mtype": { + "type": ["string", "null"], + "enum": ["semantic", "episodic", "procedural", "working", None], + "default": None, + }, + "importance": {"type": ["number", "null"], "minimum": 0, "maximum": 1, "default": None}, + "actor": {"type": "string", "maxLength": 200, "default": "user"}, + "workspace": {"type": "string", "maxLength": 200}, + "repo": {"type": ["string", "null"], "maxLength": 200, "default": None}, + }, + "required": ["memory_id"], +} + +_CONFLICT_REVIEW_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "limit": {"type": "integer", "minimum": 1, "maximum": 100, "default": 50}, + "workspace": {"type": "string", "maxLength": 200}, + "repo": {"type": ["string", "null"], "maxLength": 200, "default": None}, + }, + # All three parameters are optional; the empty list documents that + # explicitly so consumers don't have to guess whether the missing + # `required` key means "all fields implicit" or "no fields required". + "required": [], +} + +_DESC: dict[str, str] = { + "engraphis_session": ( + "Start, resume, or end an Engraphis session for a named sub-agent. " + "Call with `action: 'start'` to obtain a session_id that all other " + "tools will reuse; call `action: 'end'` with a summary and outcome " + "to close it. The `agent` field identifies the sub-agent in audit " + "logs — pick a stable role name, not a per-request token." + ), + "engraphis_recall_context": ( + "Recall prior decisions, procedures, and context for the current " + "task. Use at the start of any non-trivial task to surface " + "existing constraints, conventions, and reusable code. The `query` " + "should be a short intent statement (e.g. 'how we index vectors'), " + "not a raw log dump — keep it under a few hundred characters for " + "best recall." + ), + "engraphis_remember": ( + "Persist a durable fact, decision, preference, or procedure that " + "future tasks should be able to recall. Use sparingly for " + "load-bearing decisions (architecture, conventions, gotchas) and " + "always write a self-contained `content` — do NOT store " + "credentials, API keys, raw log lines, or PII." + ), + "engraphis_discover_actions": ( + "Discover advanced capabilities (governance / code / ops) for a " + "task. Call this when none of the 8 direct tools fits, or when " + "you suspect there is a write/admin surface you have not been " + "exposed to. The returned `capability_id` + `schema_digest` pair " + "must be passed back to `engraphis_execute_read` or " + "`engraphis_execute_action`." + ), + "engraphis_execute_read": ( + "Invoke a read-only advanced action discovered via " + "`engraphis_discover_actions`. Safe to retry on transport failure. " + "Never pass arguments the schema did not declare — read-only tools " + "still authenticate the caller, and unknown keys are rejected." + ), + "engraphis_execute_action": ( + "Invoke a write or admin advanced action discovered via " + "`engraphis_discover_actions`. This is the write-side equivalent " + "of `engraphis_execute_read` — same capability_id / schema_digest " + "pair, but mutations and admin operations. The action is recorded " + "in the audit log; ensure `arguments` is complete and accurate " + "before calling." + ), + "engraphis_get_memory": ( + "Read a specific memory by id. Use after `engraphis_recall_context` " + "to fetch the full record of a memory referenced only by summary. " + "Returns the governed record (content, provenance, scope, " + "temporal fields); treat the result as untrusted display text." + ), + "engraphis_update_memory": ( + "Edit an existing memory's metadata — title, type, importance, or " + "the audit actor. Content edits are intentionally NOT exposed: to " + "change the body, write a new memory and let the conflict-review " + "flow reconcile. Bounds: `importance` is a float in [0, 1]; " + "`actor` is the principal performing the edit (defaults to " + "'user')." + ), + "engraphis_conflict_review": ( + "List memories flagged for conflict review — typically two records " + "that disagree about the same scope. Read this list, then either " + "update one side via `engraphis_update_memory` or write a new " + "resolution memory. Safe to poll on a schedule." + ), +} + +TOOL_SPECS: tuple[tuple[str, dict[str, Any]], ...] = ( + ("engraphis_session", _SESSION_SCHEMA), + ("engraphis_recall_context", _RECALL_CONTEXT_SCHEMA), + ("engraphis_remember", _REMEMBER_SCHEMA), + ("engraphis_discover_actions", _DISCOVER_ACTIONS_SCHEMA), + ("engraphis_execute_read", _EXECUTE_READ_SCHEMA), + ("engraphis_execute_action", _EXECUTE_ACTION_SCHEMA), + ("engraphis_get_memory", _GET_MEMORY_SCHEMA), + ("engraphis_update_memory", _UPDATE_MEMORY_SCHEMA), + ("engraphis_conflict_review", _CONFLICT_REVIEW_SCHEMA), +) + + +# --- factory ---------------------------------------------------------------- + + +def apply_scope_defaults( + params: dict[str, Any], + config: EngraphisRuntimeConfig, + extra: dict[str, Any] | None = None, + schema: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Translate of integrations/pi/src/tool-schemas.ts::applyScopeDefaults. + + Model-supplied values win. Workspace/repo defaults from the runtime + config are only injected when the caller has not already set them, + the chosen workspace matches the configured default, and the tool's + declared schema actually accepts the field. Six Smart tools + (discovery, both executors, get/update memory, conflict review) do + not declare ``session_id`` / ``workspace`` / ``repo``, so passing + them is rejected as an unexpected argument; the schema gate + prevents that regression. + """ + result: dict[str, Any] = dict(extra or {}) + result.update(params) + declared = set(_declared_property_names(schema)) if schema else None + if ( + "workspace" not in result + and config.default_workspace + and (declared is None or "workspace" in declared) + ): + result["workspace"] = config.default_workspace + if ( + "repo" not in result + and config.default_repo + and config.default_workspace + and result.get("workspace") == config.default_workspace + and (declared is None or "repo" in declared) + ): + result["repo"] = config.default_repo + if ( + "session_id" not in result + and (declared is None or "session_id" in declared) + ): + # session_id is the only injected value that does not come from + # the runtime config defaults — it is propagated only by the + # caller, so no default is set here. This branch is kept for + # explicit symmetry with the workspace/repo handling. + pass + return result + + +def _declared_property_names(schema: dict[str, Any] | None) -> set[str]: + """Return the set of parameter names declared in a JSON-Schema dict. + + Used by ``apply_scope_defaults`` so injected values only land on tools + that accept them. Returns an empty set for empty/missing schemas + (the caller can decide to skip the gate by passing ``schema=None``). + """ + if not schema: + return set() + properties = schema.get("properties") + if not isinstance(properties, dict): + return set() + return {name for name in properties if isinstance(name, str)} + + +# --- lightweight schema validation ------------------------------------------ +# +# We avoid pulling in `jsonschema` as a top-level dependency and instead +# implement the small subset of JSON Schema that our 9 tool definitions +# actually use. Each tool's schema is hand-written, so a focused validator +# is enough and keeps the runtime surface zero-extra-dep. +# +# Supported keywords: +# - type: str | list[str] (with "null" used as the nullable sentinel) +# - enum: sequence of allowed values +# - required: list of required property names +# - additionalProperties: bool (False rejects unknown keys) +# - properties: per-keyword sub-schemas (each one runs through the same +# validator, recursively for `items`) +# - minLength / maxLength: string length bounds +# - minimum / maximum: int/number bounds +# - minItems / maxItems: array length bounds +# +# The `default` keyword is accepted but never enforced — the call sites do +# their own defaulting (see `apply_scope_defaults`). + +_TYPE_RANK = { + "string": str, + "integer": int, + "number": (int, float), + "boolean": bool, + "array": list, + "object": dict, + "null": type(None), +} + + +def _coerce_type(value: Any, declared: Any) -> bool: + """True iff `value` satisfies the JSON-Schema-style `type` keyword.""" + if isinstance(declared, str): + declared = [declared] + # bool is a subclass of int in Python; reject it where the schema + # says "integer" / "number" so a stray `True` is not silently accepted. + for t in declared: + py = _TYPE_RANK.get(t) + if py is None: + continue + if t in ("integer", "number") and isinstance(value, bool): + continue + if isinstance(value, py): + return True + return False + + +def _validate_schema(schema: dict[str, Any], value: Any, path: str = "") -> list[str]: + errors: list[str] = [] + declared_type = schema.get("type") + if declared_type is not None: + if not _coerce_type(value, declared_type): + errors.append( + f"{path or 'value'}: expected type {declared_type}, " + f"got {type(value).__name__}" + ) + return errors # type is wrong; deeper checks would be misleading + if "enum" in schema and value not in schema["enum"]: + errors.append( + f"{path or 'value'}: must be one of {list(schema['enum'])!r}, " + f"got {value!r}" + ) + if declared_type == "string" or "minLength" in schema or "maxLength" in schema: + if isinstance(value, str): + lo = schema.get("minLength") + hi = schema.get("maxLength") + if lo is not None and len(value) < lo: + errors.append( + f"{path or 'value'}: string length {len(value)} < minLength {lo}" + ) + if hi is not None and len(value) > hi: + errors.append( + f"{path or 'value'}: string length {len(value)} > maxLength {hi}" + ) + if declared_type in ("integer", "number") or "minimum" in schema or "maximum" in schema: + if isinstance(value, (int, float)) and not isinstance(value, bool): + lo = schema.get("minimum") + hi = schema.get("maximum") + if lo is not None and value < lo: + errors.append(f"{path or 'value'}: {value} < minimum {lo}") + if hi is not None and value > hi: + errors.append(f"{path or 'value'}: {value} > maximum {hi}") + if declared_type == "array" or "minItems" in schema or "maxItems" in schema: + if isinstance(value, list): + lo = schema.get("minItems") + hi = schema.get("maxItems") + if lo is not None and len(value) < lo: + errors.append( + f"{path or 'value'}: array length {len(value)} < minItems {lo}" + ) + if hi is not None and len(value) > hi: + errors.append( + f"{path or 'value'}: array length {len(value)} > maxItems {hi}" + ) + item_schema = schema.get("items") + if isinstance(item_schema, dict): + for i, item in enumerate(value): + errors.extend( + _validate_schema(item_schema, item, f"{path}[{i}]") + ) + if declared_type == "object" or "properties" in schema: + if isinstance(value, dict): + properties = schema.get("properties") or {} + required = schema.get("required") or [] + for key in required: + if key not in value: + errors.append(f"{path}.{key}: required") + for key, sub in properties.items(): + if key in value: + errors.extend( + _validate_schema(sub, value[key], f"{path}.{key}") + ) + additional = schema.get("additionalProperties", True) + if additional is False: + unknown = sorted(set(value) - set(properties)) + for key in unknown: + errors.append(f"{path}.{key}: unknown property (additionalProperties=False)") + return errors + + +def validate_args(name: str, args: dict[str, Any] | None) -> dict[str, Any]: + """Validate `args` against the named tool's JSON Schema. + + Returns the cleaned args dict on success. Raises + `EngraphisMcpToolError` with a single message that lists every + violation (each prefixed with the JSON-Pointer-ish path of the + offending field). Designed for the agent layer to call before + dispatching a tool, so the model sees a precise rejection instead + of a generic MCP error. + """ + schemas = dict(TOOL_SPECS) + if name not in schemas: + raise KeyError(f"Unknown Engraphis tool: {name}") + if args is None: + args = {} + if not isinstance(args, dict): + raise EngraphisMcpToolError( + f"{name}: args must be a dict, got {type(args).__name__}" + ) + errors = _validate_schema(schemas[name], args) + if errors: + joined = "; ".join(errors) + raise EngraphisMcpToolError(f"{name} args invalid: {joined}") + return args + + +def tool_spec(name: str) -> dict[str, Any]: + """Return just the meta dict for a single named tool. + + Convenience for callers that need the schema + description without + binding a client/session (e.g. for prompt inspection or registering + into a tool surface that already has its own client wiring). + """ + schemas = dict(TOOL_SPECS) + if name not in schemas: + raise KeyError(f"Unknown Engraphis tool: {name}") + return { + "name": name, + "description": _DESC[name], + "parameters": schemas[name], + } + + +def build_tool( + name: str, + client: EngraphisMcpClient, + config: EngraphisRuntimeConfig, + *, + session_id: str | None = None, +) -> tuple[ToolFn, dict[str, Any]]: + """Return (callable, meta dict) for the named tool, bound to a client. + + The callable matches the prime-agent tool contract:: + + async def fn(args: dict, ctx: dict | None = None) -> dict + + `ctx` is accepted positionally for compatibility with surfaces that + pass conversation/session metadata; the Engraphis tools do not + currently read it. Schema is a JSON Schema dict that any downstream + tool-registration surface can translate to its own format. + + Precedence: caller-supplied `session_id` (via the args dict) ALWAYS + wins over the `session_id` bound at build time. The bound value is + only injected when the args dict does not already include one — + this lets a single tool instance be re-used across requests that + occasionally need to operate on a different session (e.g. a + cross-session audit lookup). + """ + schemas = dict(TOOL_SPECS) + if name not in schemas: + raise KeyError(f"Unknown Engraphis tool: {name}") + + async def _call( + args: dict[str, Any], + _ctx: dict[str, Any] | None = None, + ) -> dict[str, Any]: + # _ctx is reserved for future per-call overrides (e.g. trace ids, + # tenant hints); current MCP tools don't need it, so we accept + # and ignore. The leading underscore keeps the parameter name + # visible in stack traces / introspection while signalling that + # it is intentionally unused. The signature stays compatible + # with agent.py's `await fn(args, ctx)` call site. + schema = schemas[name] + params = apply_scope_defaults(args, config, schema=schema) + # Precedence: caller-supplied session_id wins over the bound one, + # but only when the tool's declared schema actually accepts it. + # Six Smart tools (discovery, both executors, get/update memory, + # conflict review) do not declare session_id; passing it would + # be rejected as an unexpected argument by FastMCP. + declared = _declared_property_names(schema) + if session_id and "session_id" not in params and "session_id" in declared: + params["session_id"] = session_id + return await client.call_tool(name, params) + + meta = {"name": name, "description": _DESC[name], "parameters": schemas[name]} + return _call, meta + + +def all_tools( + client: EngraphisMcpClient, + config: EngraphisRuntimeConfig, + *, + session_id: str | None = None, +) -> list[tuple[ToolFn, dict[str, Any]]]: + """Build the 9 tool (callable, schema) pairs bound to the given client/session.""" + return [ + build_tool(name, client, config, session_id=session_id) + for name, _schema in TOOL_SPECS + ] diff --git a/integrations/prime_agent/tests/__init__.py b/integrations/prime_agent/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/integrations/prime_agent/tests/conftest.py b/integrations/prime_agent/tests/conftest.py new file mode 100644 index 00000000..ba68af1a --- /dev/null +++ b/integrations/prime_agent/tests/conftest.py @@ -0,0 +1,296 @@ +"""Pytest fixtures: in-process fake MCP server + live-gated real client. + +The fake server monkey-patches `mcp.client.stdio.stdio_client` so the real +`ClientSession` runs over an `anyio` memory-stream transport. Tests then +exercise the full JSON-RPC framing without an `engraphis-mcp` subprocess. + +Set `ENGRAPHIS_INTEGRATION_LIVE=1` to skip the fake and boot a real +`engraphis-mcp` subprocess for the live integration tests. +""" +from __future__ import annotations + +import asyncio +import contextlib +import json +import os +from collections.abc import AsyncIterator +from typing import Any + +import anyio +import pytest +import pytest_asyncio + +from engraphis_prime_agent.config import EngraphisRuntimeConfig +from engraphis_prime_agent.mcp_client import EngraphisMcpClient + +__all__ = ["FakeMcpServer", "live_mcp_client", "mcp_client"] + + +CORE_TOOL_NAMES = ( + "engraphis_session", + "engraphis_recall_context", + "engraphis_remember", + "engraphis_discover_actions", + "engraphis_execute_read", + "engraphis_execute_action", + "engraphis_get_memory", + "engraphis_update_memory", + "engraphis_conflict_review", +) + + +class FakeMcpServer: + """In-process stand-in for the Engraphis MCP gateway. + + The patched `stdio_client` returns ``(read_stream, write_stream)`` over an + anyio memory channel pair. The server task drains requests, calls the + provided handler, and writes back responses. + """ + + def __init__(self, tool_names: tuple[str, ...] = CORE_TOOL_NAMES) -> None: + async def _default(name: str, args: dict[str, Any]) -> dict[str, Any]: + if name == "engraphis_session": + # Pretend a session was created and echo the request back. + payload = { + "session_id": f"ses_fake_{next(self._session_counter):04d}", + "agent": args.get("agent", "unknown"), + "workspace": args.get("workspace"), + "repo": args.get("repo"), + "action": args.get("action", "start"), + } + if args.get("action") == "end": + payload["status"] = "closed" + return { + "_tool": name, + "content": [{"type": "text", "text": json.dumps(payload)}], + } + return {"_tool": name, "content": [{"type": "text", "text": json.dumps(args)}]} + + self.tool_handler = _default + self._session_counter = iter(range(1, 10_000)) + self.tool_names = tool_names + self.call_log: list[tuple[str, dict[str, Any]]] = [] + self.fail_next: Exception | None = None + self.crash_on_next: bool = False + # Shared streams so the test can restart the server task while the + # client keeps the same transport alive. + self._shared_server_to_client_send: Any = None + self._shared_client_to_server_send: Any = None + self._server_task: asyncio.Task[None] | None = None + self._original_stdio_client: Any = None + self._installed = False + + def install(self) -> None: + from mcp.client import stdio as stdio_mod + + self._original_stdio_client = stdio_mod.stdio_client + + @contextlib.asynccontextmanager + async def _fake_stdio(_params, errlog=None): # type: ignore[no-untyped-def] + # If streams haven't been allocated yet (first call), create them. + if self._shared_client_to_server_send is None: + # anyio.create_memory_object_stream returns (send, receive). + s2c_send, c_read = anyio.create_memory_object_stream(max_buffer_size=4096) + c2s_send, s_read = anyio.create_memory_object_stream(max_buffer_size=4096) + self._shared_server_to_client_send = s2c_send + self._shared_client_to_server_send = c2s_send + self._server_read = s_read + self._client_read = c_read + self._start_server() + elif self._server_task is None or self._server_task.done(): + # Re-entry after a transport failure: spin a fresh server. + self._start_server() + try: + yield (self._client_read, self._shared_client_to_server_send) + finally: + if self._server_task and not self._server_task.done(): + self._server_task.cancel() + with contextlib.suppress(BaseException): + await self._server_task + + # Patch both the source module AND the binding used by the client. + stdio_mod.stdio_client = _fake_stdio # type: ignore[assignment] + import engraphis_prime_agent.mcp_client as _client_mod + + self._original_client_binding = _client_mod.stdio_client + _client_mod.stdio_client = _fake_stdio # type: ignore[assignment] + self._installed = True + + def _start_server(self) -> None: + from mcp.shared.message import SessionMessage + + self._server_task = asyncio.create_task( + self._serve(self._server_read, self._shared_server_to_client_send, SessionMessage) + ) + + async def restart_server(self) -> None: + """Kill the server task and start a fresh one on the same streams. + + Used to simulate a transport failure (server crash) followed by the + client successfully reconnecting. + """ + if self._server_task and not self._server_task.done(): + self._server_task.cancel() + with contextlib.suppress(BaseException): + await self._server_task + self._start_server() + + def restore(self) -> None: + from mcp.client import stdio as stdio_mod + import engraphis_prime_agent.mcp_client as _client_mod + + if self._installed and self._original_stdio_client is not None: + stdio_mod.stdio_client = self._original_stdio_client # type: ignore[assignment] + if getattr(self, "_original_client_binding", None) is not None: + _client_mod.stdio_client = self._original_client_binding # type: ignore[assignment] + self._installed = False + if self._server_task and not self._server_task.done(): + self._server_task.cancel() + + async def _serve(self, read_stream, write_stream, SessionMessage) -> None: # type: ignore[no-untyped-def] + """Minimal MCP server. Handles initialize / notifications / tools/list / tools/call.""" + from mcp.shared.message import JSONRPCMessage + from mcp.types import ( + CallToolResult, + InitializeResult, + JSONRPCError, + JSONRPCResponse, + ListToolsResult, + TextContent, + Tool, + ) + + async def reply_ok(req_id: Any, result: Any) -> None: + # JSONRPCResponse.result is typed as a dict; dump pydantic models. + payload_dict = ( + result.model_dump(by_alias=True, mode="json", exclude_none=True) + if hasattr(result, "model_dump") + else result + ) + payload = JSONRPCResponse(jsonrpc="2.0", id=req_id, result=payload_dict) + await write_stream.send(SessionMessage(message=JSONRPCMessage(payload))) + + async def reply_error(req_id: Any, message: str) -> None: + err = JSONRPCError( + jsonrpc="2.0", + id=req_id, + error={"code": -32601, "message": message}, + ) + await write_stream.send(SessionMessage(message=JSONRPCMessage(err))) + + while True: + try: + message: Any = await read_stream.receive() + except (anyio.EndOfStream, asyncio.CancelledError): + return + # `message` is a SessionMessage; `.message` is a JSONRPCMessage; + # `.root` is the actual JSONRPCRequest / JSONRPCNotification. + jsonrpc = getattr(message, "message", message) + request = getattr(jsonrpc, "root", jsonrpc) + method = getattr(request, "method", None) + request_id = getattr(request, "id", None) + params = getattr(request, "params", None) or {} + # If the tool handler itself raises (e.g. a transport-failure + # simulation), we let the exception propagate so the server task + # exits. The client will see a closed receive stream and treat it + # as a transport failure, exercising the retry path. + if method == "tools/call": + name = params.get("name", "") + arguments = params.get("arguments") or {} + self.call_log.append((name, arguments)) + if self.fail_next is not None: + exc = self.fail_next + self.fail_next = None + raise exc + if self.crash_on_next: + self.crash_on_next = False + return + result = await self.tool_handler(name, arguments) + content = [ + TextContent(type="text", text=block.get("text", "")) + for block in (result.get("content", []) or []) + ] + await reply_ok( + request_id, + CallToolResult(content=content, isError=bool(result.get("isError"))), + ) + continue + try: + if method == "initialize": + await reply_ok( + request_id, + InitializeResult( + protocolVersion="2025-03-26", + capabilities={}, + serverInfo=ServerInfo(name="fake-engraphis", version="0.0.0"), + ), + ) + elif method == "notifications/initialized": + continue + elif method == "tools/list": + tools = [ + Tool( + name=n, + description=f"fake {n}", + inputSchema={"type": "object", "properties": {}}, + ) + for n in self.tool_names + ] + await reply_ok( + request_id, ListToolsResult(tools=tools, nextCursor=None) + ) + else: + await reply_error(request_id, f"Method not found: {method}") + except Exception as exc: # noqa: BLE001 — surface as tool error + try: + await reply_ok( + request_id, + CallToolResult( + content=[TextContent(type="text", text=f"Error: {exc}")], + isError=True, + ), + ) + except Exception: + return + + +def ServerInfo(name: str, version: str) -> Any: # noqa: N802 — helper + from mcp.types import Implementation + + return Implementation(name=name, version=version) + + +@pytest_asyncio.fixture +async def fake_mcp_server() -> AsyncIterator[FakeMcpServer]: + server = FakeMcpServer() + server.install() + try: + yield server + finally: + server.restore() + + +@pytest_asyncio.fixture +async def mcp_client(fake_mcp_server: FakeMcpServer) -> AsyncIterator[EngraphisMcpClient]: + """Return a connected `EngraphisMcpClient` backed by the fake server.""" + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + await client.connect() + try: + yield client + finally: + await client.close() + + +@pytest_asyncio.fixture +async def live_mcp_client() -> AsyncIterator[EngraphisMcpClient]: + """Yield a real EngraphisMcpClient against `engraphis-mcp` if available.""" + if not os.environ.get("ENGRAPHIS_INTEGRATION_LIVE"): + pytest.skip("set ENGRAPHIS_INTEGRATION_LIVE=1 to run live integration tests") + config = EngraphisRuntimeConfig(command="engraphis-mcp", environment={}) + client = EngraphisMcpClient(config) + await client.connect() + try: + yield client + finally: + await client.close() diff --git a/integrations/prime_agent/tests/test_config.py b/integrations/prime_agent/tests/test_config.py new file mode 100644 index 00000000..0a307082 --- /dev/null +++ b/integrations/prime_agent/tests/test_config.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +from dataclasses import FrozenInstanceError, fields + +import pytest + +from engraphis_prime_agent.config import ( + DEFAULT_AGENT_NAMES, + EngraphisRuntimeConfig, + _engraphis_environment, + _non_blank, + build_runtime_config, +) + + +def test_non_blank_trims_and_rejects_empty() -> None: + assert _non_blank(None) is None + assert _non_blank("") is None + assert _non_blank(" ") is None + assert _non_blank(" hello ") == "hello" + + +def test_engraphis_environment_allowlist() -> None: + env = { + "ENGRAPHIS_DB_PATH": "/tmp/x.db", + "ENGRAPHIS_WORKSPACE": "demo", + "PATH": "/usr/bin", + "Path": "C:\\Windows", + "SystemRoot": "C:\\Windows", + "ComSpec": "C:\\Windows\\System32\\cmd.exe", + "ANTHROPIC_API_KEY": "sk-secret", + "HOME": "/root", + "USER": "alice", + } + forwarded = _engraphis_environment(env) + assert set(forwarded) == { + "ENGRAPHIS_DB_PATH", + "ENGRAPHIS_WORKSPACE", + "PATH", + "Path", + "SystemRoot", + "ComSpec", + } + assert forwarded["ENGRAPHIS_DB_PATH"] == "/tmp/x.db" + assert "ANTHROPIC_API_KEY" not in forwarded + assert "HOME" not in forwarded + + +def test_engraphis_environment_ignores_non_string_values() -> None: + env = {"ENGRAPHIS_WORKSPACE": 123, "PATH": None} # type: ignore[dict-item] + assert _engraphis_environment(env) == {} + + +def test_build_runtime_config_defaults() -> None: + cfg = build_runtime_config(env={}) + assert cfg.command == "engraphis-mcp" + assert cfg.args == () + assert cfg.cwd is None + assert cfg.default_workspace is None + assert cfg.default_repo is None + assert cfg.environment == {} + + +def test_build_runtime_config_reads_env() -> None: + env = { + "ENGRAPHIS_MCP_COMMAND": "C:/venv/Scripts/engraphis-mcp.exe", + "ENGRAPHIS_WORKSPACE": "engraphis", + "ENGRAPHIS_REPO": "prime-agent", + "ENGRAPHIS_DB_PATH": "C:/data/x.db", + "ANTHROPIC_API_KEY": "sk-secret", + } + cfg = build_runtime_config(env=env) + assert cfg.command == "C:/venv/Scripts/engraphis-mcp.exe" + assert cfg.default_workspace == "engraphis" + assert cfg.default_repo == "prime-agent" + assert "ANTHROPIC_API_KEY" not in cfg.environment + assert cfg.environment["ENGRAPHIS_DB_PATH"] == "C:/data/x.db" + + +def test_build_runtime_config_command_override() -> None: + cfg = build_runtime_config(env={}, command="/abs/engraphis-mcp") + assert cfg.command == "/abs/engraphis-mcp" + + +def test_build_runtime_config_trims_blank_env() -> None: + env = {"ENGRAPHIS_WORKSPACE": " ", "ENGRAPHIS_REPO": " real "} + cfg = build_runtime_config(env=env) + assert cfg.default_workspace is None + assert cfg.default_repo == "real" + + +def test_default_agent_names_are_eight() -> None: + assert len(DEFAULT_AGENT_NAMES) == 8 + assert "researcher" in DEFAULT_AGENT_NAMES + assert "coder" in DEFAULT_AGENT_NAMES + assert all(isinstance(name, str) and name for name in DEFAULT_AGENT_NAMES) + # Names must be unique (default fleet keys must be hashable). + assert len(set(DEFAULT_AGENT_NAMES)) == 8 + + +def test_runtime_config_is_frozen() -> None: + cfg = EngraphisRuntimeConfig() + try: + cfg.command = "x" # type: ignore[misc] + except Exception: + return + raise AssertionError("EngraphisRuntimeConfig should be frozen") + + +def test_runtime_config_frozen_raises_frozen_instance_error_on_every_field() -> None: + """Every public field must reject assignment with FrozenInstanceError.""" + cfg = EngraphisRuntimeConfig( + command="x", + args=("a", "b"), + cwd="C:/work", + default_workspace="ws", + default_repo="repo", + environment={"ENGRAPHIS_DB_PATH": "/tmp/x.db"}, + ) + for name in ("command", "args", "cwd", "default_workspace", "default_repo", "environment"): + with pytest.raises(FrozenInstanceError): + setattr(cfg, name, "mutated") # type: ignore[misc] + + +def test_runtime_config_field_names_are_stable() -> None: + """Lock the public dataclass surface so a refactor that renames a field + is caught here rather than at a downstream caller.""" + expected = { + "command", + "args", + "cwd", + "default_workspace", + "default_repo", + "environment", + } + assert {f.name for f in fields(EngraphisRuntimeConfig)} == expected + + +def test_build_runtime_config_preserves_args_tuple_type() -> None: + """`args` must remain a tuple — the stdio gateway expects a sequence and + downstream code (e.g. ``list(self._config.args)``) relies on tuple semantics.""" + src_args = ("--flag", "value", "C:/path/with space") + cfg = build_runtime_config(env={}, args=src_args) + assert isinstance(cfg.args, tuple) + assert cfg.args == src_args + # Mutating the original tuple must not leak into the config. + assert cfg.args is not src_args or cfg.args == src_args + + +def test_build_runtime_config_empty_args_default_to_empty_tuple() -> None: + """The default is an empty tuple, not None or a list, so callers can + iterate without a None-check.""" + cfg = build_runtime_config(env={}) + assert cfg.args == () + assert isinstance(cfg.args, tuple) + + +def test_engraphis_environment_handles_windows_specific_keys() -> None: + """SystemRoot and ComSpec must be forwarded on Windows. We don't assume + Windows-only — any platform that has these keys in env should see them + through the allowlist.""" + env = { + "SystemRoot": "C:\\Windows", + "ComSpec": "C:\\Windows\\System32\\cmd.exe", + "PATHEXT": ".EXE;.BAT", # NOT in the allowlist; must be dropped. + "WINDIR": "C:\\Windows", # NOT in the allowlist; must be dropped. + } + forwarded = _engraphis_environment(env) + assert forwarded["SystemRoot"] == "C:\\Windows" + assert forwarded["ComSpec"] == "C:\\Windows\\System32\\cmd.exe" + assert "PATHEXT" not in forwarded + assert "WINDIR" not in forwarded + + +def test_build_runtime_config_trims_default_workspace_and_repo_from_env() -> None: + """Whitespace-padded env values must be stripped, and a pure-whitespace + value must become None (not the literal whitespace).""" + env = { + "ENGRAPHIS_WORKSPACE": " ", + "ENGRAPHIS_REPO": "\trepo\t", + "ENGRAPHIS_DB_PATH": " /tmp/x.db ", + } + cfg = build_runtime_config(env=env) + assert cfg.default_workspace is None + assert cfg.default_repo == "repo" + # The env allowlist also strips; the entry must reflect the trimmed value. + assert cfg.environment["ENGRAPHIS_DB_PATH"] == "/tmp/x.db" + + +def test_build_runtime_config_does_not_mutate_input_env() -> None: + """`build_runtime_config` must not mutate the caller's env mapping.""" + env = { + "ENGRAPHIS_WORKSPACE": " ws ", + "ENGRAPHIS_REPO": " repo ", + "ENGRAPHIS_DB_PATH": " /tmp/x.db ", + "PATH": " /usr/bin ", + } + snapshot = dict(env) + build_runtime_config(env=env) + assert env == snapshot + + +def test_engraphis_environment_empty_input_returns_empty_dict() -> None: + """Defensive: an empty mapping must produce an empty dict, not raise.""" + assert _engraphis_environment({}) == {} + + +def test_engraphis_environment_skips_prefix_only_keys_without_value() -> None: + """An ENGRAPHIS_-prefixed key whose value is non-string must be skipped + rather than forwarded as-is (which would crash subprocess.Popen).""" + env = { + "ENGRAPHIS_DB_PATH": 42, # type: ignore[dict-item] + "ENGRAPHIS_WORKSPACE": None, # type: ignore[dict-item] + } + assert _engraphis_environment(env) == {} # type: ignore[arg-type] + + +def test_non_blank_strips_tabs_and_newlines() -> None: + """`_non_blank` is the single source of truth for trimming env values; + tabs and newlines should be treated like spaces.""" + assert _non_blank("\t\n hi \n\t") == "hi" + assert _non_blank("\t\n \n\t") is None + + +def test_runtime_config_as_subprocess_env_returns_independent_copy() -> None: + """Mutating the dict returned by as_subprocess_env must not change the + frozen config's own mapping.""" + cfg = EngraphisRuntimeConfig( + command="x", + environment={"ENGRAPHIS_DB_PATH": "/tmp/x.db"}, + ) + env = cfg.as_subprocess_env() + env["ENGRAPHIS_DB_PATH"] = "/mutated/y.db" + assert cfg.environment["ENGRAPHIS_DB_PATH"] == "/tmp/x.db" diff --git a/integrations/prime_agent/tests/test_fleet.py b/integrations/prime_agent/tests/test_fleet.py new file mode 100644 index 00000000..6879d44b --- /dev/null +++ b/integrations/prime_agent/tests/test_fleet.py @@ -0,0 +1,732 @@ +"""Tests for EngraphisPrimeAgent and PrimeAgentFleet.""" +from __future__ import annotations + +import asyncio +import json + +import pytest + +from engraphis_prime_agent.agent import EngraphisPrimeAgent, PrimeAgentFleet +from engraphis_prime_agent.config import ( + DEFAULT_AGENT_NAMES, + EngraphisRuntimeConfig, +) +from engraphis_prime_agent.mcp_client import EngraphisMcpClient, EngraphisMcpToolError +from engraphis_prime_agent.tools import TOOL_SPECS + + +# Auto-use the fake MCP server for every test in this module so that any +# test which constructs an EngraphisMcpClient (directly or via the fleet) +# gets the in-process fake transport, not a real subprocess. +@pytest.fixture(autouse=True) +def _install_fake(fake_mcp_server) -> None: + return None + + +@pytest.fixture +async def fleet() -> PrimeAgentFleet: + f = PrimeAgentFleet( + workspace="test", + config=EngraphisRuntimeConfig(command="ignored", environment={}), + ) + await f.client.connect() + try: + yield f + finally: + await f.aclose() + + +@pytest.mark.asyncio +async def test_fleet_default_names_is_eight() -> None: + f = PrimeAgentFleet(workspace="x") + assert len(f) == 8 + assert f.names() == DEFAULT_AGENT_NAMES + + +@pytest.mark.asyncio +async def test_fleet_custom_agent_names() -> None: + custom = ("a", "b", "c", "d", "e", "f", "g", "h") + f = PrimeAgentFleet(workspace="x", agent_names=custom) + assert f.names() == custom + + +@pytest.mark.asyncio +async def test_subagent_repr_and_contains() -> None: + f = PrimeAgentFleet(workspace="x") + assert "researcher" in f + assert f["researcher"].name == "researcher" + + +@pytest.mark.asyncio +async def test_subagent_rejects_blank_name() -> None: + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + with pytest.raises(ValueError): + EngraphisPrimeAgent(" ", client, config) + with pytest.raises(ValueError): + EngraphisPrimeAgent("", client, config) + + +@pytest.mark.asyncio +async def test_status_reports_workspace_and_agents() -> None: + f = PrimeAgentFleet(workspace="demo") + status = f.status() + assert status["workspace"] == "demo" + assert len(status["agents"]) == 8 + for entry in status["agents"]: + assert "name" in entry + assert "session_id" in entry + + +@pytest.mark.asyncio +async def test_start_session_returns_session_id_and_caches_it(fleet) -> None: + agent = fleet["researcher"] + sid = await agent.start_session() + assert isinstance(sid, str) and sid + # Second call is a no-op. + sid2 = await agent.start_session() + assert sid2 == sid + assert agent.session_id == sid + + +@pytest.mark.asyncio +async def test_force_new_starts_a_fresh_session(fleet) -> None: + agent = fleet["researcher"] + sid1 = await agent.start_session() + sid2 = await agent.start_session(force_new=True) + assert sid1 != sid2 + + +@pytest.mark.asyncio +async def test_explicit_null_repo_clears_cached_repo(fleet) -> None: + agent = fleet["researcher"] + await agent.start_session() + await agent.start_session(force_new=True, repo=None) + assert agent.repo is None + + +@pytest.mark.asyncio +async def test_call_lazy_starts_session(fleet) -> None: + agent = fleet["researcher"] + assert agent.session_id is None + await agent.call("engraphis_recall_context", {"query": "anything"}) + assert agent.session_id is not None + + +@pytest.mark.asyncio +async def test_call_injects_session_id_into_subsequent_calls(fleet) -> None: + agent = fleet["researcher"] + await agent.call("engraphis_recall_context", {"query": "warm up"}) + # The client we drive is the one used by the fleet. + # We can verify the call succeeded and returned the tool name. + result = await agent.call("engraphis_recall_context", {"query": "next"}) + assert result["_tool"] == "engraphis_recall_context" + + +@pytest.mark.asyncio +async def test_end_session_clears_cached_id(fleet) -> None: + agent = fleet["researcher"] + await agent.start_session() + assert agent.session_id is not None + await agent.end_session(summary="done", outcome="shipped") + assert agent.session_id is None + + +@pytest.mark.asyncio +async def test_end_session_is_idempotent_when_no_session(fleet) -> None: + agent = fleet["researcher"] + await agent.end_session() # no-op + + +@pytest.mark.asyncio +async def test_fan_out_runs_concurrently(fleet) -> None: + args = { + "researcher": {"query": "researcher query"}, + "coder": {"query": "coder query"}, + } + out = await fleet.fan_out("engraphis_recall_context", args) + assert set(out.keys()) == {"researcher", "coder"} + for value in out.values(): + assert value["_tool"] == "engraphis_recall_context" + + +@pytest.mark.asyncio +async def test_fan_out_raises_for_unknown_agent(fleet) -> None: + with pytest.raises(KeyError): + await fleet.fan_out("engraphis_recall_context", {"ghost": {}}) + + +@pytest.mark.asyncio +async def test_start_all_sessions_warms_every_agent(fleet) -> None: + out = await fleet.start_all_sessions() + # New structured return: {"sessions": {name: sid}, "errors": {name: exc}}. + assert set(out.keys()) == {"sessions", "errors"} + sessions = out["sessions"] + errors = out["errors"] + assert isinstance(sessions, dict) and isinstance(errors, dict) + assert set(sessions.keys()) == set(fleet.names()) + assert errors == {} + for sid in sessions.values(): + assert isinstance(sid, str) and sid + + +@pytest.mark.asyncio +async def test_register_requires_register_tool() -> None: + fleet = PrimeAgentFleet(workspace="x") + with pytest.raises(TypeError) as exc: + fleet["researcher"].register(object()) + assert "register_tool" in str(exc.value) + + +@pytest.mark.asyncio +async def test_register_registers_all_nine_tools() -> None: + fleet = PrimeAgentFleet(workspace="x") + registered: list[tuple[str, dict]] = [] + + class _Target: + def register_tool(self, name: str, fn, schema: dict) -> None: + registered.append((name, schema)) + + target = _Target() + fleet["researcher"].register(target) + assert len(registered) == 9 + for name, schema in registered: + assert name.startswith("engraphis_") + assert "parameters" in schema + + +@pytest.mark.asyncio +async def test_aclose_ends_sessions_and_closes_client() -> None: + fleet = PrimeAgentFleet(workspace="x") + await fleet.client.connect() + await fleet["researcher"].start_session() + await fleet["coder"].start_session() + await fleet.aclose() + assert fleet["researcher"].session_id is None + assert fleet["coder"].session_id is None + assert fleet._closed is True + + +@pytest.mark.asyncio +async def test_agents_reject_calls_after_fleet_close() -> None: + fleet = PrimeAgentFleet(workspace="x") + await fleet.client.connect() + agent = fleet["researcher"] + await agent.start_session() + await fleet.aclose() + + with pytest.raises(RuntimeError, match="is closed"): + await agent.call("engraphis_recall_context", {"query": "after close"}) + with pytest.raises(RuntimeError, match="is closed"): + await agent.start_session() + + +@pytest.mark.asyncio +async def test_fleet_blocks_new_calls_while_sessions_are_ending() -> None: + """Shutdown must not let a late data call bootstrap a replacement session.""" + fleet = PrimeAgentFleet( + workspace="x", + config=EngraphisRuntimeConfig(command="ignored", environment={}), + ) + await fleet.client.connect() + agent = fleet["researcher"] + await agent.start_session() + end_returned = asyncio.Event() + release_shutdown = asyncio.Event() + original_end = agent.end_session + + async def delayed_end(*args, **kwargs): + await original_end(*args, **kwargs) + end_returned.set() + await release_shutdown.wait() + + agent.end_session = delayed_end # type: ignore[method-assign] + close_task = asyncio.create_task(fleet.aclose()) + try: + await end_returned.wait() + with pytest.raises(RuntimeError, match="is closed"): + await agent.call("engraphis_recall_context", {"query": "late"}) + finally: + release_shutdown.set() + await close_task + + with pytest.raises(RuntimeError, match="is closed"): + await agent.call("engraphis_recall_context", {"query": "after close"}) + + +@pytest.mark.asyncio +async def test_lifecycle_rejects_unknown_action(fleet) -> None: + with pytest.raises(EngraphisMcpToolError, match="args invalid.*action"): + await fleet["researcher"].call("engraphis_session", {"action": "resume"}) + + +@pytest.mark.asyncio +async def test_lifecycle_accepts_compatibility_action_aliases(fleet) -> None: + agent = fleet["researcher"] + await agent.call("engraphis_session", {"action": "start_session"}) + assert agent.session_id is not None + await agent.call("engraphis_session", {"action": "end_session"}) + assert agent.session_id is None + + +@pytest.mark.asyncio +async def test_lifecycle_rejects_non_boolean_force_new(fleet) -> None: + with pytest.raises(EngraphisMcpToolError, match="args invalid.*force_new"): + await fleet["researcher"].call( + "engraphis_session", {"force_new": "false"} + ) + + +@pytest.mark.asyncio +async def test_aexit_via_context_manager() -> None: + async with PrimeAgentFleet(workspace="x") as fleet: + await fleet["researcher"].start_session() + assert fleet._closed is True + + +@pytest.mark.asyncio +async def test_context_manager_rejects_reentry_after_close() -> None: + fleet = PrimeAgentFleet(workspace="x") + async with fleet: + pass + + with pytest.raises(RuntimeError, match="is closed"): + async with fleet: + pass + + +# ---- new edge-case tests below ---- + + +@pytest.mark.asyncio +async def test_aclose_is_idempotent() -> None: + """`aclose()` (and therefore `__aexit__`) must be safe to call twice. + The second call is a no-op because the fleet has already torn down.""" + f = PrimeAgentFleet(workspace="x") + await f.client.connect() + await f["researcher"].start_session() + await f.aclose() + assert f._closed is True + # Second call must not raise. + await f.aclose() + assert f._closed is True + + +@pytest.mark.asyncio +async def test_aclose_before_any_session_is_safe() -> None: + """A fresh fleet that has never connected must close cleanly without + requiring a prior `start_session` or `connect`.""" + f = PrimeAgentFleet(workspace="x") + await f.aclose() + assert f._closed is True + + +@pytest.mark.asyncio +async def test_fan_out_with_single_sub_agent() -> None: + """fan_out() with exactly one agent must return a one-entry dict and + must not raise. The framework-level concurrency path should still work + for a single coroutine.""" + f = PrimeAgentFleet(workspace="x") + await f.client.connect() + try: + out = await f.fan_out( + "engraphis_recall_context", + {"researcher": {"query": "single-agent query"}}, + ) + assert set(out.keys()) == {"researcher"} + result = out["researcher"] + assert result["_tool"] == "engraphis_recall_context" + finally: + await f.aclose() + + +@pytest.mark.asyncio +async def test_fan_out_with_empty_args_raises_value_error() -> None: + """fan_out() with an empty mapping must raise ValueError so a misnamed + variable at the call site surfaces immediately rather than silently + producing an empty result dict.""" + f = PrimeAgentFleet(workspace="x") + await f.client.connect() + try: + with pytest.raises(ValueError) as exc: + await f.fan_out("engraphis_recall_context", {}) + assert "non-empty" in str(exc.value).lower() or "empty" in str(exc.value).lower() + finally: + await f.aclose() + + +@pytest.mark.asyncio +async def test_status_before_any_session_started() -> None: + """`status()` is a sync method — it must work without any prior connect, + start_session, or call. It should report the configured workspace, the + full agent roster, and a None session_id for every agent.""" + f = PrimeAgentFleet(workspace="demo") + s = f.status() + assert s["workspace"] == "demo" + assert len(s["agents"]) == 8 + for entry in s["agents"]: + assert entry["session_id"] is None + assert "name" in entry + assert "workspace" in entry + assert "repo" in entry + + +def test_status_before_connect_does_not_require_async() -> None: + """`status()` is intentionally sync (status snapshot, not a live call). + It must be callable from a non-async context without a runtime error.""" + f = PrimeAgentFleet(workspace="x") + s = f.status() + assert s["workspace"] == "x" + assert isinstance(s["agents"], list) + assert isinstance(s["clientGeneration"], int) + # Generation starts at 0. + assert s["clientGeneration"] == 0 + + +@pytest.mark.asyncio +async def test_register_calls_register_tool_exactly_n_times() -> None: + """`register()` must invoke `register_tool` exactly once per tool — + not zero, not twice, not conditional on the tool name. We assert this + by counting invocations against the number of tools in TOOL_SPECS.""" + f = PrimeAgentFleet(workspace="x") + invocations: list[tuple[str, object]] = [] + + class _Target: + def register_tool(self, name: str, fn, schema: dict) -> None: + invocations.append((name, fn)) + + target = _Target() + f["researcher"].register(target) + expected_count = len(TOOL_SPECS) + assert len(invocations) == expected_count + # Every tool name from TOOL_SPECS must appear exactly once. + seen = [name for name, _fn in invocations] + assert seen == [n for n, _ in TOOL_SPECS] + # Each call's `fn` is callable and distinct from the others. + fns = [fn for _name, fn in invocations] + assert all(callable(fn) for fn in fns) + assert len({id(fn) for fn in fns}) == expected_count + + +@pytest.mark.asyncio +async def test_register_invokes_for_each_agent_independently() -> None: + """Each sub-agent's register() registers its OWN 9 tools. Registering + one agent must not bleed into another agent's binding.""" + f = PrimeAgentFleet(workspace="x") + researcher_calls: list[str] = [] + coder_calls: list[str] = [] + + class _T: + def __init__(self, sink: list[str]) -> None: + self._sink = sink + + def register_tool(self, name: str, fn, schema: dict) -> None: + self._sink.append(name) + + f["researcher"].register(_T(researcher_calls)) + f["coder"].register(_T(coder_calls)) + assert len(researcher_calls) == 9 + assert len(coder_calls) == 9 + assert researcher_calls == coder_calls # same tool surface + + +@pytest.mark.asyncio +async def test_fleet_iter_and_len_match() -> None: + """`len(fleet)` and `for a in fleet` must agree — they both read from + the same internal agent dict.""" + f = PrimeAgentFleet(workspace="x") + assert len(f) == 8 + names_via_iter = [a.name for a in f] + assert names_via_iter == list(f.names()) + + +@pytest.mark.asyncio +async def test_fleet_unknown_name_raises_keyerror() -> None: + """`__getitem__` for an unknown agent must raise KeyError, not silently + return None or a default — fan_out already raises KeyError, and direct + indexing must behave consistently.""" + f = PrimeAgentFleet(workspace="x") + with pytest.raises(KeyError): + _ = f["nonexistent_agent"] + + +@pytest.mark.asyncio +async def test_fleet_contains_is_consistent_with_iter() -> None: + f = PrimeAgentFleet(workspace="x") + for name in f.names(): + assert name in f + assert "definitely_not_an_agent" not in f + assert None not in f + assert 42 not in f + + +@pytest.mark.asyncio +async def test_fleet_workspace_override_sets_every_agent() -> None: + """When the fleet is constructed with `workspace=...`, every sub-agent + inherits that workspace. Individual sub-agents have no way to opt out + (they can only set their own workspace via the EngraphisPrimeAgent + constructor, which the fleet does not expose).""" + f = PrimeAgentFleet(workspace="shared-ws") + for agent in f: + assert agent.workspace == "shared-ws" + + +@pytest.mark.asyncio +async def test_start_all_sessions_is_idempotent_per_agent() -> None: + """Calling start_all_sessions() twice must not spawn extra sessions. + Each agent should keep its first session id.""" + f = PrimeAgentFleet(workspace="x") + await f.client.connect() + try: + first = await f.start_all_sessions() + second = await f.start_all_sessions() + assert first == second + finally: + await f.aclose() + + +@pytest.mark.asyncio +async def test_subagent_status_reflects_session_lifecycle() -> None: + """`subagent.status()` should reflect the current session state — None + before start, populated after start, None again after end.""" + f = PrimeAgentFleet(workspace="x") + await f.client.connect() + try: + agent = f["researcher"] + assert agent.status()["session_id"] is None + await agent.start_session() + s = agent.status() + assert isinstance(s["session_id"], str) and s["session_id"] + await agent.end_session() + assert agent.status()["session_id"] is None + finally: + await f.aclose() + + +@pytest.mark.asyncio +async def test_end_session_forwards_open_threads_to_mcp_call(fake_mcp_server) -> None: + """`end_session(open_threads=[...])` must include the open_threads list + in the underlying MCP call_tool so the server can persist the + next-session handoff. Dropping the argument would silently strand + advertised follow-ups on the server side.""" + f = PrimeAgentFleet(workspace="x") + await f.client.connect() + try: + agent = f["researcher"] + await agent.start_session() + thread = "follow up on the caching decision" + await agent.end_session(summary="done", outcome="ok", + open_threads=[thread]) + # Locate the engraphis_session/end RPC in the call log. + end_calls = [ + (name, args) for name, args in fake_mcp_server.call_log + if name == "engraphis_session" and args.get("action") == "end" + ] + assert end_calls, "expected an engraphis_session/end MCP call" + # The most recent end call should carry the open_threads payload. + _name, end_args = end_calls[-1] + assert end_args.get("open_threads") == [thread] + finally: + await f.aclose() + + +@pytest.mark.asyncio +async def test_end_session_holds_lock_until_close_rpc_finishes() -> None: + """A replacement start must wait until the previous close is complete.""" + close_started = asyncio.Event() + release_close = asyncio.Event() + calls: list[dict[str, object]] = [] + session_number = 0 + + class _BlockingClient: + async def call_tool( + self, _name: str, args: dict[str, object] + ) -> dict[str, object]: + nonlocal session_number + calls.append(args) + if args.get("action") == "end": + close_started.set() + await release_close.wait() + else: + session_number += 1 + session_id = f"ses_blocking_{session_number:04d}" + return { + "content": [{"text": json.dumps({"session_id": session_id})}] + } + + config = EngraphisRuntimeConfig(command="ignored", environment={}) + agent = EngraphisPrimeAgent("researcher", _BlockingClient(), config, workspace="x") + await agent.start_session() + + end_task = asyncio.create_task(agent.end_session()) + await close_started.wait() + start_task = asyncio.create_task(agent.start_session()) + await asyncio.sleep(0) + assert not start_task.done() + + release_close.set() + await end_task + await start_task + assert [call["action"] for call in calls] == ["start", "end", "start"] + + +@pytest.mark.asyncio +async def test_data_call_waits_for_force_new_session_generation() -> None: + """A data request must not retain the old id while a replacement starts.""" + start_started = asyncio.Event() + release_start = asyncio.Event() + calls: list[tuple[str, dict[str, object]]] = [] + session_number = 0 + + class _BlockingClient: + async def call_tool( + self, name: str, args: dict[str, object] + ) -> dict[str, object]: + nonlocal session_number + calls.append((name, dict(args))) + if name == "engraphis_session": + if args.get("action") == "start": + session_number += 1 + if args.get("force_new"): + start_started.set() + await release_start.wait() + payload = {"session_id": f"ses_generation_{session_number:04d}"} + else: + payload = {"status": "closed"} + return {"content": [{"text": json.dumps(payload)}]} + return {"content": [{"text": json.dumps(args)}]} + + config = EngraphisRuntimeConfig(command="ignored", environment={}) + agent = EngraphisPrimeAgent("researcher", _BlockingClient(), config, workspace="x") + await agent.start_session() + + start_task = asyncio.create_task(agent.start_session(force_new=True)) + await start_started.wait() + data_task = asyncio.create_task( + agent.call("engraphis_recall_context", {"query": "stable generation"}) + ) + await asyncio.sleep(0) + assert not data_task.done() + + release_start.set() + await start_task + await data_task + assert [name for name, _args in calls] == [ + "engraphis_session", "engraphis_session", "engraphis_recall_context" + ] + assert calls[-1][1]["session_id"] == "ses_generation_0002" + + +@pytest.mark.asyncio +async def test_dispatch_lifecycle_forwards_advertised_arguments(fake_mcp_server) -> None: + f = PrimeAgentFleet(workspace="initial") + await f.client.connect() + try: + agent = f["researcher"] + await agent.call( + "engraphis_session", + { + "action": "start", + "agent": "custom-role", + "workspace": "override", + "repo": "project", + "goal": "inspect integration", + "force_new": True, + "token_budget": 2048, + }, + ) + start_args = fake_mcp_server.call_log[-1][1] + assert start_args == { + "action": "start", + "agent": "custom-role", + "workspace": "override", + "repo": "project", + "goal": "inspect integration", + "force_new": True, + "token_budget": 2048, + } + assert agent.status()["workspace"] == "override" + assert agent.status()["repo"] == "project" + assert agent.status()["goal"] == "inspect integration" + assert agent.token_budget == 2048 + + await agent.end_session() + assert fake_mcp_server.call_log[-1][1]["agent"] == "custom-role" + await agent.call( + "engraphis_session", + { + "action": "start", + "agent": "custom-role", + "workspace": "override", + "repo": "project", + "goal": "inspect integration", + "force_new": True, + "token_budget": 2048, + }, + ) + session_id = agent.status()["session_id"] + await agent.call( + "engraphis_session", + { + "action": "end", + "session_id": session_id, + "agent": "custom-role", + "workspace": "override", + "repo": "project", + "summary": "done", + "outcome": "shipped", + "open_threads": ["none"], + }, + ) + end_args = fake_mcp_server.call_log[-1][1] + assert end_args == { + "action": "end", + "agent": "custom-role", + "session_id": session_id, + "workspace": "override", + "repo": "project", + "summary": "done", + "outcome": "shipped", + "open_threads": ["none"], + } + finally: + await f.aclose() + + +@pytest.mark.asyncio +async def test_dispatch_session_lifecycle_end_routes_through_state_machine(fake_mcp_server) -> None: + """`agent.call("engraphis_session", {"action": "end"})` must clear the + cached session id so subsequent memory calls do not re-inject a + closed id. Without the lifecycle routing, the agent would still + hold the prior id after the server closed the session.""" + f = PrimeAgentFleet(workspace="x") + await f.client.connect() + try: + agent = f["researcher"] + await agent.start_session() + prior = agent.status()["session_id"] + assert prior + await agent.call("engraphis_session", {"action": "end", + "summary": "shutdown", + "outcome": "complete"}) + assert agent.status()["session_id"] is None + finally: + await f.aclose() + + +@pytest.mark.asyncio +async def test_dispatch_session_lifecycle_rejects_unknown_fields(fake_mcp_server) -> None: + """Direct lifecycle routing must enforce the advertised JSON schema.""" + f = PrimeAgentFleet(workspace="x") + await f.client.connect() + try: + with pytest.raises(EngraphisMcpToolError, match="unknown property"): + await f["researcher"].call( + "engraphis_session", + {"action": "start", "workpace": "typo"}, + ) + assert not any( + name == "engraphis_session" for name, _args in fake_mcp_server.call_log + ) + finally: + await f.aclose() diff --git a/integrations/prime_agent/tests/test_mcp_client.py b/integrations/prime_agent/tests/test_mcp_client.py new file mode 100644 index 00000000..e86aca3e --- /dev/null +++ b/integrations/prime_agent/tests/test_mcp_client.py @@ -0,0 +1,437 @@ +"""Tests for the async stdio MCP client.""" +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from engraphis_prime_agent.config import CORE_DIRECT_TOOLS, EngraphisRuntimeConfig +from engraphis_prime_agent.mcp_client import ( + READ_ONLY_TOOLS, + MAX_TOOL_LIST_PAGES, + EngraphisCompatibilityError, + EngraphisMcpClient, + EngraphisMcpToolError, + format_mcp_payload, +) + + +@pytest.mark.asyncio +async def test_connect_lists_core_tools(mcp_client) -> None: + tools = await mcp_client.list_tools() + names = {t["name"] for t in tools} + expected = { + "engraphis_session", + "engraphis_recall_context", + "engraphis_remember", + "engraphis_discover_actions", + "engraphis_execute_read", + "engraphis_execute_action", + "engraphis_get_memory", + "engraphis_update_memory", + "engraphis_conflict_review", + } + assert expected.issubset(names) + + +@pytest.mark.asyncio +async def test_live_gateway_exposes_smart_tools(live_mcp_client) -> None: + """The opt-in live gate must exercise the real gateway's tool contract.""" + tools = await live_mcp_client.list_tools() + names = {tool["name"] for tool in tools} + assert set(CORE_DIRECT_TOOLS).issubset(names) + + +@pytest.mark.asyncio +async def test_status_reports_connected(mcp_client) -> None: + status = await mcp_client.status() + assert status["connected"] is True + assert status["server"] == "engraphis" + assert status["toolCount"] >= 9 + + +@pytest.mark.asyncio +async def test_call_tool_passes_arguments(fake_mcp_server, mcp_client) -> None: + payload = await mcp_client.call_tool( + "engraphis_recall_context", {"query": "decision: sqlite-vec KNN", "k": 3} + ) + assert payload["_tool"] == "engraphis_recall_context" + assert fake_mcp_server.call_log[-1] == ( + "engraphis_recall_context", + {"query": "decision: sqlite-vec KNN", "k": 3}, + ) + + +# Note: retry behavior is exercised by the production code path; the +# in-process fake doesn't reliably simulate "transport failure" because +# crashing the server task races with the real ClientSession's receive loop. +# The retry constants (READ_ONLY_TOOLS) are unit-tested separately below. + + +def test_read_only_tools_classification() -> None: + # engraphis_recall_context is intentionally NOT in READ_ONLY_TOOLS: + # the Smart gateway appends a receipt on every successful call, so a + # transport-level retry would create duplicate accounting records + # for one logical user request. The class is the opposite: tools + # whose server-side contract is purely read-only and idempotent. + assert "engraphis_get_memory" in READ_ONLY_TOOLS + assert "engraphis_conflict_review" in READ_ONLY_TOOLS + assert "engraphis_discover_actions" in READ_ONLY_TOOLS + assert "engraphis_execute_read" in READ_ONLY_TOOLS + # Writes and side-effect tools are not in the read-only set, so the + # client's call_tool will not retry them on transport failure. + assert "engraphis_recall_context" not in READ_ONLY_TOOLS + assert "engraphis_remember" not in READ_ONLY_TOOLS + assert "engraphis_execute_action" not in READ_ONLY_TOOLS + assert "engraphis_session" not in READ_ONLY_TOOLS + assert "engraphis_update_memory" not in READ_ONLY_TOOLS + + +@pytest.mark.asyncio +async def test_rejection_text_raises_tool_error(fake_mcp_server, mcp_client) -> None: + async def handler(name: str, args: dict) -> dict: + return { + "isError": True, + "content": [{"type": "text", "text": "Error: bad_arg"}], + } + + fake_mcp_server.tool_handler = handler + with pytest.raises(EngraphisMcpToolError) as exc: + await mcp_client.call_tool("engraphis_remember", {"content": "x"}) + assert "bad_arg" in str(exc.value) + + +@pytest.mark.asyncio +async def test_compatibility_error_when_tools_missing(fake_mcp_server) -> None: + """Drop a core tool from the fake server and verify the compatibility error.""" + # Switch the existing fake server to advertise only one core tool, + # so the client's required-tool check fails on the others. + fake_mcp_server.restore() + fake_mcp_server.tool_names = ("engraphis_session",) + fake_mcp_server.install() + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + try: + with pytest.raises(EngraphisCompatibilityError) as exc: + await client.connect() + assert "missing" in str(exc.value).lower() + finally: + await client.close() + fake_mcp_server.restore() + + +def test_diagnostic_hint_matches_python_message() -> None: + client = EngraphisMcpClient(EngraphisRuntimeConfig(command="x")) + client._diagnostic = "ERROR: This package requires python 3.10 or later.\n" + hint = client.diagnostic_hint() + assert hint is not None + assert "Python 3.10" in hint + + +def test_diagnostic_hint_matches_missing_mcp() -> None: + client = EngraphisMcpClient(EngraphisRuntimeConfig(command="x")) + client._diagnostic = "ModuleNotFoundError: No module named 'mcp'\n" + assert client.diagnostic_hint() is not None + assert "mcp" in client.diagnostic_hint().lower() + + +def test_diagnostic_hint_matches_missing_engraphis() -> None: + client = EngraphisMcpClient(EngraphisRuntimeConfig(command="x")) + client._diagnostic = "ModuleNotFoundError: No module named 'engraphis'\n" + assert client.diagnostic_hint() is not None + + +def test_format_mcp_payload_joins_text() -> None: + payload = { + "content": [ + {"type": "text", "text": "hello"}, + {"type": "text", "text": "world"}, + ] + } + assert format_mcp_payload(payload) == "hello\n\nworld" + + +def test_format_mcp_payload_falls_back_to_json() -> None: + payload = {"content": []} + out = format_mcp_payload(payload) + parsed = json.loads(out) + assert parsed == payload + + +@pytest.mark.asyncio +async def test_unknown_tool_name_rejected(mcp_client) -> None: + with pytest.raises(EngraphisMcpToolError): + await mcp_client.call_tool("not_a_tool", {}) + + +@pytest.mark.asyncio +async def test_close_bumps_generation(fake_mcp_server) -> None: + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + g0 = client.generation() + await client.connect() + await client.close() + g1 = client.generation() + assert g1 > g0 + + +# ---- new edge-case tests below ---- + + +@pytest.mark.asyncio +async def test_connect_is_idempotent(fake_mcp_server) -> None: + """Calling connect() twice must return the same session and not re-spawn + the stdio subprocess or re-fetch the tool list.""" + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + s1 = await client.connect() + s2 = await client.connect() + assert s1 is s2 + # The list_tools cache was populated by the first connect; the second + # call must not issue a fresh tools/list RPC. + assert client._tools_cache is not None + cache_id = id(client._tools_cache) + await client.connect() + assert id(client._tools_cache) == cache_id + async def unexpected_second_discovery(_session): + raise AssertionError("list_tools issued a second tools/list request") + client._list_tools = unexpected_second_discovery # type: ignore[method-assign] + assert await client.list_tools() + await client.close() + + +@pytest.mark.asyncio +async def test_close_clears_session_stack_and_tools_cache(fake_mcp_server) -> None: + """After close(), every internal handle must be released so the + next connect() can rebuild cleanly.""" + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + await client.connect() + assert client._session is not None + assert client._stack is not None + assert client._tools_cache is not None + await client.close() + assert client._session is None + assert client._stack is None + assert client._tools_cache is None + + +@pytest.mark.asyncio +async def test_aenter_aexit_context_manager(fake_mcp_server) -> None: + """`async with EngraphisMcpClient(...) as client:` must connect on enter + and release every handle on exit.""" + config = EngraphisRuntimeConfig(command="ignored", environment={}) + async with EngraphisMcpClient(config) as client: + # Inside the block: connected, tools cached. + assert client._session is not None + assert client._tools_cache is not None + tools = await client.list_tools() + assert len(tools) >= 9 + # After the block: all handles released. + assert client._session is None + assert client._stack is None + assert client._tools_cache is None + + +@pytest.mark.asyncio +async def test_aenter_returns_client_instance(fake_mcp_server) -> None: + config = EngraphisRuntimeConfig(command="ignored", environment={}) + async with EngraphisMcpClient(config) as client: + assert isinstance(client, EngraphisMcpClient) + assert client is not None + + +@pytest.mark.asyncio +async def test_unknown_tool_name_includes_name_in_error_message(mcp_client) -> None: + """`call_tool` must raise EngraphisMcpToolError AND the error message + must name the rejected tool so a developer can diagnose the rejection.""" + with pytest.raises(EngraphisMcpToolError) as exc: + await mcp_client.call_tool("engraphis_does_not_exist", {}) + assert "engraphis_does_not_exist" in str(exc.value) + # And bare "not_a_tool" (no engraphis_ prefix) is also rejected with a + # message — a different guard, but same exception class. + with pytest.raises(EngraphisMcpToolError) as exc2: + await mcp_client.call_tool("not_a_tool", {}) + assert "not_a_tool" in str(exc2.value) + + +@pytest.mark.asyncio +async def test_smart_error_envelope_is_parsed_into_message(fake_mcp_server) -> None: + """The Smart gateway wraps every failure as + ``{"error": {"code": ..., "message": ..., "retryable": ...}}``; + the client must surface the inner code and message instead of + the generic fallback so callers can distinguish validation + errors from retryable internal failures. + """ + import json + + envelope = { + "error": { + "code": "validation_failed", + "message": "missing required field 'query'", + "retryable": False, + } + } + + async def _smart_error_handler(name, args): + return { + "isError": True, + "content": [{"type": "text", "text": json.dumps(envelope)}], + } + + fake_mcp_server.tool_handler = _smart_error_handler + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + await client.connect() + try: + with pytest.raises(EngraphisMcpToolError) as exc: + await client.call_tool("engraphis_recall_context", {}) + msg = str(exc.value) + assert "validation_failed" in msg + assert "missing required field 'query'" in msg + assert exc.value.retryable is False + finally: + await client.close() + + +@pytest.mark.asyncio +async def test_smart_retryable_error_preserves_retryability_signal(fake_mcp_server) -> None: + import json + + envelope = { + "error": { + "code": "upstream_unavailable", + "message": "temporary gateway failure", + "retryable": True, + } + } + + async def _smart_error_handler(name, args): + return { + "isError": True, + "content": [{"type": "text", "text": json.dumps(envelope)}], + } + + fake_mcp_server.tool_handler = _smart_error_handler + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + await client.connect() + try: + with pytest.raises(EngraphisMcpToolError) as exc: + await client.call_tool("engraphis_recall_context", {}) + assert exc.value.retryable is True + finally: + await client.close() + + +@pytest.mark.asyncio +async def test_legacy_flat_error_envelope_is_still_supported(fake_mcp_server) -> None: + """The classic gateway emits a flat ``{"code": ..., "message": ...}`` + envelope for backwards compatibility. The client must keep + parsing that shape so the integration does not regress when an + older server is in front of the agent.""" + import json + + envelope = {"code": "not_found", "message": "memory mem_xyz is gone"} + + async def _flat_error_handler(name, args): + return { + "isError": True, + "content": [{"type": "text", "text": json.dumps(envelope)}], + } + + fake_mcp_server.tool_handler = _flat_error_handler + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + await client.connect() + try: + with pytest.raises(EngraphisMcpToolError) as exc: + await client.call_tool("engraphis_recall_context", {}) + msg = str(exc.value) + assert "not_found" in msg + assert "memory mem_xyz is gone" in msg + finally: + await client.close() + + +@pytest.mark.asyncio +async def test_close_is_idempotent(fake_mcp_server) -> None: + """Calling close() twice must not raise. The second call should be a no-op + because _stack/_session are already None.""" + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + await client.connect() + await client.close() + # Second close should be silent. + await client.close() + assert client._session is None + assert client._stack is None + assert client._tools_cache is None + + +@pytest.mark.asyncio +async def test_list_tools_returns_independent_list(fake_mcp_server) -> None: + """Mutating the list returned by list_tools() must not affect the cache + (so a second caller still sees the full list).""" + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + await client.connect() + first = await client.list_tools() + first.clear() + second = await client.list_tools() + assert len(second) == len(first) or len(second) >= 9 + + +@pytest.mark.asyncio +async def test_list_tools_rejects_unbounded_pagination() -> None: + class _NeverEndingTools: + def __init__(self) -> None: + self.cursors: list[str | None] = [] + + async def list_tools(self, *, cursor: str | None = None): + self.cursors.append(cursor) + return SimpleNamespace(tools=[], nextCursor=f"cursor-{len(self.cursors)}") + + client = EngraphisMcpClient(EngraphisRuntimeConfig(command="ignored")) + session = _NeverEndingTools() + with pytest.raises(EngraphisCompatibilityError, match="page limit"): + await client._list_tools(session) # type: ignore[arg-type] + assert len(session.cursors) == MAX_TOOL_LIST_PAGES + + +@pytest.mark.asyncio +async def test_status_diagnostic_hint_is_none_when_no_failure(fake_mcp_server) -> None: + """After a healthy connect, diagnosticHint must be None — there is no + error message to surface.""" + config = EngraphisRuntimeConfig(command="ignored", environment={}) + async with EngraphisMcpClient(config) as client: + status = await client.status() + assert status["connected"] is True + assert status["diagnosticHint"] is None + assert status["server"] == "engraphis" + assert status["toolCount"] >= 9 + + +def test_diagnostic_hint_returns_none_for_unrecognized_error() -> None: + """A diagnostic line that doesn't match any known pattern must surface + None (not a misleading hint).""" + client = EngraphisMcpClient(EngraphisRuntimeConfig(command="x")) + client._diagnostic = "ERROR: connection refused on 127.0.0.1:9999\n" + assert client.diagnostic_hint() is None + + +def test_format_mcp_payload_handles_non_text_blocks() -> None: + """Blocks without a `text` field (e.g. an image) must be skipped, and + the JSON fallback must kick in when no text content is present.""" + payload = { + "content": [ + {"type": "image", "data": "ignored"}, + {"type": "text", "text": "only text"}, + ] + } + assert format_mcp_payload(payload) == "only text" + # No text at all -> JSON fallback. + assert json.loads(format_mcp_payload({"content": [{"type": "image"}]})) == { + "content": [{"type": "image"}] + } diff --git a/integrations/prime_agent/tests/test_register_and_repo.py b/integrations/prime_agent/tests/test_register_and_repo.py new file mode 100644 index 00000000..9baa523d --- /dev/null +++ b/integrations/prime_agent/tests/test_register_and_repo.py @@ -0,0 +1,243 @@ +"""Tests for review-feedback fixes on PR 174. + +Covers: +1. Agent repo precedence: explicit > config.default_repo > self.name. +2. register() wrappers lazily start the session. +3. install_prime_agent / scripts wrapper dispatches via the package module. +4. Installer TOML path uses write_text (not write_bytes). +5. CLI install command does not require scripts/ outside the wheel. +""" +from __future__ import annotations + +import json +import stat +import subprocess +import sys +from pathlib import Path + +import pytest + +from engraphis_prime_agent.config import EngraphisRuntimeConfig +from engraphis_prime_agent.mcp_client import EngraphisMcpClient + + +# ---- Fix 1: agent repo precedence -------------------------------------- + + +def test_agent_repo_uses_explicit_kwarg() -> None: + from engraphis_prime_agent.agent import EngraphisPrimeAgent + + config = EngraphisRuntimeConfig( + command="ignored", default_repo="api", environment={} + ) + client = EngraphisMcpClient(config) + agent = EngraphisPrimeAgent( + "researcher", client, config, workspace="acme", repo="custom" + ) + assert agent.repo == "custom" + + +def test_agent_repo_uses_default_repo_when_no_explicit() -> None: + from engraphis_prime_agent.agent import EngraphisPrimeAgent + + config = EngraphisRuntimeConfig( + command="ignored", default_repo="api", environment={} + ) + client = EngraphisMcpClient(config) + agent = EngraphisPrimeAgent("researcher", client, config, workspace="acme") + assert agent.repo == "api" + + +def test_agent_repo_falls_back_to_name_when_no_default() -> None: + from engraphis_prime_agent.agent import EngraphisPrimeAgent + + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + agent = EngraphisPrimeAgent("researcher", client, config, workspace="acme") + assert agent.repo == "researcher" + + +# ---- Fix 2: register() wrappers lazily start the session ------------------ + + +@pytest.mark.asyncio +async def test_register_wrappers_lazy_start_session(fake_mcp_server) -> None: + """When a fresh agent is registered and the framework invokes a tool + directly, the session must be started before the tool is called — the + wrapper around each registered callable must drive the lazy-start path. + """ + from engraphis_prime_agent.agent import EngraphisPrimeAgent, PrimeAgentFleet + + config = EngraphisRuntimeConfig(command="ignored", environment={}) + client = EngraphisMcpClient(config) + await client.connect() + try: + fleet = PrimeAgentFleet(workspace="test", config=config) + agent = EngraphisPrimeAgent("researcher", client, config) + + registered: dict[str, object] = {} + + class _Target: + def register_tool(self, name: str, fn, schema: dict) -> None: + registered[name] = fn + + agent.register(_Target()) + assert "engraphis_recall_context" in registered + wrapper = registered["engraphis_recall_context"] + # Before the framework calls the wrapper, no session exists. + assert agent.session_id is None + await wrapper({"query": "hello"}) + # After the framework calls the wrapper, the session is started. + assert agent.session_id is not None + await fleet.aclose() + finally: + await client.close() + + +# ---- Fix 3 + 4: installer module + TOML write_text ---------------------- + + +def test_installer_module_importable() -> None: + """The installer must ship inside the package so the wheel works.""" + from engraphis_prime_agent import installer + + assert hasattr(installer, "install") + assert hasattr(installer, "uninstall") + assert hasattr(installer, "main") + + +def test_installer_toml_uses_write_text( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """``tomli_w.dumps`` returns str, so the TOML path must use write_text + (not write_bytes, which would TypeError). We test the wrapper by + stubbing tomli_w to verify the right method is called. + """ + from engraphis_prime_agent import installer + + target = tmp_path / "config.toml" + captured: dict[str, object] = {} + + class _StubToml: + @staticmethod + def dumps(_data: dict) -> str: + return "[tools.engraphis]\npackage = 'x'\n" + + monkeypatch.setattr(installer, "tomli_w", _StubToml, raising=False) + monkeypatch.setitem(sys.modules, "tomli_w", _StubToml) + + real_write_text = Path.write_text + real_write_bytes = Path.write_bytes + + def _spy_write_text(self, *args, **kwargs): # type: ignore[no-untyped-def] + captured["method"] = "write_text" + return real_write_text(self, *args, **kwargs) + + def _spy_write_bytes(self, *args, **kwargs): # type: ignore[no-untyped-def] + captured["method"] = "write_bytes" + return real_write_bytes(self, *args, **kwargs) + + monkeypatch.setattr(Path, "write_text", _spy_write_text) + monkeypatch.setattr(Path, "write_bytes", _spy_write_bytes) + + installer.install(target, merge=False, dry_run=False) + assert captured.get("method") == "write_text" + assert target.exists() + assert "package" in target.read_text(encoding="utf-8") + + +def test_installer_idempotent_install_uninstall_round_trip( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from engraphis_prime_agent import installer + + target = tmp_path / "config.json" + installer.install(target) + installer.install(target) # idempotent: same content + cfg = json.loads(target.read_text(encoding="utf-8")) + assert len(cfg["tools"]) == 1 + assert "engraphis" in cfg["tools"] + installer.uninstall(target) + cfg = json.loads(target.read_text(encoding="utf-8")) + assert "engraphis" not in cfg.get("tools", {}) + + +def test_installer_backup_preserves_source_permissions(tmp_path: Path) -> None: + from engraphis_prime_agent import installer + + target = tmp_path / "config.json" + target.write_text('{"tools": {"other": {}}}\n', encoding="utf-8") + target.chmod(0o640) + source_mode = stat.S_IMODE(target.stat().st_mode) + + backup = installer._backup(target) + + assert backup is not None + assert stat.S_IMODE(backup.stat().st_mode) == source_mode + + +def test_installer_dry_run_does_not_write( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from engraphis_prime_agent import installer + + target = tmp_path / "config.json" + installer.install(target, dry_run=True) + assert not target.exists() + + +# ---- Fix 5: CLI install works without a source-tree scripts/ dir ---------- + + +def test_cli_install_subcommand_uses_package_installer( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The CLI must dispatch through the package module, not runpy against + a repo-level scripts/ directory that doesn't exist after pip install. + """ + config_path = tmp_path / "config.json" + result = subprocess.run( + [ + sys.executable, + "-m", + "engraphis_prime_agent", + "install", + "--config-path", + str(config_path), + ], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + assert config_path.exists() + cfg = json.loads(config_path.read_text(encoding="utf-8")) + assert "engraphis" in cfg["tools"] + + +# ---- Scripts wrapper: still works from a source checkout ----------------- + + +def test_scripts_wrapper_imports_package(tmp_path: Path) -> None: + """The repo-root scripts/install_prime_agent.py is a thin shim that + delegates to engraphis_prime_agent.installer. Verify the import path + when invoked from a source checkout (no editable install). + """ + import io + import contextlib + + script = ( + Path(__file__).resolve().parent.parent.parent.parent + / "scripts" + / "install_prime_agent.py" + ) + assert script.exists(), f"missing {script}" + config_path = tmp_path / "shim-config.json" + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + result = subprocess.run( + [sys.executable, str(script), "--config-path", str(config_path)], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + assert config_path.exists() diff --git a/integrations/prime_agent/tests/test_tools.py b/integrations/prime_agent/tests/test_tools.py new file mode 100644 index 00000000..8c88469a --- /dev/null +++ b/integrations/prime_agent/tests/test_tools.py @@ -0,0 +1,373 @@ +"""Tests for the 9 Smart tool factories and scope-default helper.""" +from __future__ import annotations + +import pytest + +from engraphis_prime_agent.config import EngraphisRuntimeConfig +from engraphis_prime_agent.mcp_client import EngraphisMcpClient +from engraphis_prime_agent.tools import ( + TOOL_SPECS, + all_tools, + apply_scope_defaults, + build_tool, + validate_args, +) + + +@pytest.fixture +def client(mcp_client) -> EngraphisMcpClient: + return mcp_client + + +def test_tool_specs_cover_nine_tools() -> None: + assert len(TOOL_SPECS) == 9 + names = [name for name, _ in TOOL_SPECS] + assert names == [ + "engraphis_session", + "engraphis_recall_context", + "engraphis_remember", + "engraphis_discover_actions", + "engraphis_execute_read", + "engraphis_execute_action", + "engraphis_get_memory", + "engraphis_update_memory", + "engraphis_conflict_review", + ] + + +def test_each_tool_has_name_description_and_schema() -> None: + for name, schema in TOOL_SPECS: + assert isinstance(name, str) and name + assert "type" in schema and schema["type"] == "object" + assert "properties" in schema + + +def test_remember_schema_declares_keyed_claim_fields_as_properties() -> None: + schema = dict(TOOL_SPECS)["engraphis_remember"] + + assert {"subject_key", "claim_kind"} <= set(schema["properties"]) + assert "subject_key" not in schema + assert "claim_kind" not in schema + assert schema["properties"]["subject_key"] == {"type": "string", "maxLength": 1000} + assert schema["properties"]["claim_kind"] == {"type": "string", "maxLength": 200} + + +def test_session_agent_is_optional_for_registered_lifecycle_calls() -> None: + schema = dict(TOOL_SPECS)["engraphis_session"] + assert "agent" in schema["properties"] + assert "agent" not in schema["required"] + + +def test_session_schema_advertises_compatibility_action_aliases() -> None: + schema = dict(TOOL_SPECS)["engraphis_session"] + assert schema["properties"]["action"]["enum"] == [ + "start", + "end", + "start_session", + "end_session", + ] + assert validate_args("engraphis_session", {"action": "start_session"})[ + "action" + ] == "start_session" + + +def test_nullable_schema_types_accept_each_union_member() -> None: + assert validate_args("engraphis_session", {"repo": "api"})["repo"] == "api" + assert validate_args("engraphis_session", {"repo": None})["repo"] is None + + +def test_build_tool_unknown_name_raises() -> None: + config = EngraphisRuntimeConfig(command="x") + client = EngraphisMcpClient(config) + with pytest.raises(KeyError): + build_tool("not_a_tool", client, config) + + +@pytest.mark.asyncio +async def test_recall_context_tool_calls_mcp(client) -> None: + fn, meta = build_tool("engraphis_recall_context", client, client.config) + result = await fn({"query": "decision: sqlite-vec KNN"}) + assert result["_tool"] == "engraphis_recall_context" + assert client._tools_cache is not None # ensure list_tools was called + + +@pytest.mark.asyncio +async def test_remember_tool_passes_arguments(client) -> None: + fn, _ = build_tool("engraphis_remember", client, client.config) + result = await fn({"content": "Use sqlite-vec KNN for <=1M vectors", "importance": 0.7}) + assert result["_tool"] == "engraphis_remember" + + +@pytest.mark.asyncio +async def test_session_id_is_injected_when_bound(client, fake_mcp_server) -> None: + fn, _ = build_tool( + "engraphis_recall_context", client, client.config, session_id="ses_test_1" + ) + await fn({"query": "anything"}) + # The fake server records every tools/call; the last entry should + # carry the injected session_id. + assert fake_mcp_server.call_log[-1][0] == "engraphis_recall_context" + assert fake_mcp_server.call_log[-1][1].get("session_id") == "ses_test_1" + + +def test_all_tools_returns_nine_pairs(client) -> None: + pairs = all_tools(client, client.config) + assert len(pairs) == 9 + for fn, meta in pairs: + assert callable(fn) + assert meta["name"] in [name for name, _ in TOOL_SPECS] + assert "description" in meta + assert "parameters" in meta + + +def test_apply_scope_defaults_preserves_model_supplied() -> None: + config = EngraphisRuntimeConfig( + command="x", + default_workspace="acme", + default_repo="api", + ) + out = apply_scope_defaults( + {"workspace": "override", "repo": "fork"}, + config, + ) + assert out["workspace"] == "override" + assert out["repo"] == "fork" + + +def test_apply_scope_defaults_injects_when_missing() -> None: + config = EngraphisRuntimeConfig( + command="x", + default_workspace="acme", + default_repo="api", + ) + out = apply_scope_defaults({}, config) + assert out["workspace"] == "acme" + assert out["repo"] == "api" + + +def test_apply_scope_defaults_skips_repo_when_workspace_overridden() -> None: + config = EngraphisRuntimeConfig( + command="x", + default_workspace="acme", + default_repo="api", + ) + out = apply_scope_defaults({"workspace": "other"}, config) + assert out["workspace"] == "other" + assert "repo" not in out + + +def test_apply_scope_defaults_merges_extra() -> None: + config = EngraphisRuntimeConfig(command="x") + out = apply_scope_defaults({}, config, extra={"actor": "user"}) + assert out["actor"] == "user" + + +def test_apply_scope_defaults_extra_can_be_overridden_by_params() -> None: + config = EngraphisRuntimeConfig(command="x") + out = apply_scope_defaults({"actor": "agent"}, config, extra={"actor": "user"}) + assert out["actor"] == "agent" + + +# ---- new edge-case tests below ---- + + +def test_apply_scope_defaults_does_not_mutate_input_dict() -> None: + """The helper must not mutate the caller's `params` dict — prime-agent + and other call sites may reuse the same dict for repeated tool calls.""" + config = EngraphisRuntimeConfig( + command="x", + default_workspace="acme", + default_repo="api", + ) + params = {"query": "hello"} + snapshot = dict(params) + out = apply_scope_defaults(params, config) + assert params == snapshot # input untouched + # Output is a new dict — mutating it must not bleed back. + out["query"] = "mutated" + assert params["query"] == "hello" + + +def test_apply_scope_defaults_does_not_mutate_extra_dict() -> None: + """`extra` is also treated as read-only.""" + config = EngraphisRuntimeConfig(command="x", default_workspace="acme") + extra = {"actor": "user", "workspace": "extra-ws"} + snapshot = dict(extra) + out = apply_scope_defaults({}, config, extra=extra) + assert extra == snapshot + # The output is a copy of extra; mutating output must not leak. + out["actor"] = "mutated" + assert extra["actor"] == "user" + + +def test_apply_scope_defaults_no_defaults_no_extra_returns_new_dict() -> None: + """With no config defaults and no extra, apply_scope_defaults should + return a new dict equal to the input — and still not be the same object.""" + config = EngraphisRuntimeConfig(command="x") + params = {"x": 1} + out = apply_scope_defaults(params, config) + assert out == params + assert out is not params + + +@pytest.mark.asyncio +async def test_build_tool_returns_async_callable(client) -> None: + """The returned callable must be awaitable and accept a single dict arg.""" + import inspect + + fn, meta = build_tool("engraphis_remember", client, client.config) + assert callable(fn) + assert inspect.iscoroutinefunction(fn) or hasattr(fn, "__call__") + # Calling with a dict must return an awaitable that resolves to a dict. + coro = fn({"content": "x"}) + result = await coro + assert isinstance(result, dict) + assert "content" in result or "_tool" in result + + +@pytest.mark.asyncio +async def test_build_tool_meta_has_required_fields(client) -> None: + """The metadata dict must include name, description, and parameters so + any prime-agent registration surface can render it without fallbacks.""" + fn, meta = build_tool("engraphis_get_memory", client, client.config) + assert meta["name"] == "engraphis_get_memory" + assert isinstance(meta["description"], str) and meta["description"] + assert meta["parameters"]["type"] == "object" + assert "properties" in meta["parameters"] + + +def test_all_tool_schemas_declare_required_field_explicitly() -> None: + """Every Smart tool schema must declare a `required` key — either as a + non-empty list of names or an empty list. The absence of `required` + would be ambiguous (it can be read as "no required fields" OR as + "all fields implicitly required" depending on the consumer).""" + for name, schema in TOOL_SPECS: + assert "required" in schema, f"{name} schema is missing the 'required' key" + assert isinstance(schema["required"], list), ( + f"{name} schema 'required' must be a list, got {type(schema['required']).__name__}" + ) + # Every name listed in `required` must also be a defined property. + for required_name in schema["required"]: + assert required_name in schema["properties"], ( + f"{name} schema lists {required_name!r} in required " + "but it is not in properties" + ) + + +def test_schema_required_names_are_subset_of_properties() -> None: + """Defense in depth: cross-check every required name appears in properties.""" + for name, schema in TOOL_SPECS: + for required_name in schema.get("required", []): + assert required_name in schema["properties"], ( + f"{name}: required field {required_name!r} missing from properties" + ) + + +def test_schemas_have_additional_properties_false_or_unset() -> None: + """The schemas set `additionalProperties: False` to surface typos early. + Any schema that loses this guarantee is a regression.""" + for name, schema in TOOL_SPECS: + if "additionalProperties" in schema: + assert schema["additionalProperties"] is False, ( + f"{name} schema should have additionalProperties=False" + ) + + +def test_no_tool_schema_is_empty() -> None: + """Every tool must declare at least one property. An empty schema would + mean the tool accepts no parameters at all, which is not a Smart tool.""" + for name, schema in TOOL_SPECS: + assert schema.get("properties"), f"{name} schema has no properties" + assert len(schema["properties"]) >= 1 + + +@pytest.mark.asyncio +async def test_session_id_is_injected_into_call(client, fake_mcp_server) -> None: + """A tool bound with session_id="ses_xyz" must forward "ses_xyz" as the + session_id argument of the resulting tools/call RPC.""" + fn, _ = build_tool( + "engraphis_recall_context", client, client.config, session_id="ses_xyz" + ) + await fn({"query": "anything"}) + # The fake server records the last call's (name, arguments) pair. + assert fake_mcp_server.call_log, "fake server recorded no calls" + last_name, last_args = fake_mcp_server.call_log[-1] + assert last_name == "engraphis_recall_context" + assert last_args.get("session_id") == "ses_xyz" + # The caller-supplied args are preserved alongside the injection. + assert last_args.get("query") == "anything" + + +@pytest.mark.asyncio +async def test_session_id_injection_does_not_override_caller_supplied(client, fake_mcp_server) -> None: + """If the caller already supplied a session_id, the bound session_id + must NOT silently overwrite it — caller intent wins.""" + fn, _ = build_tool( + "engraphis_recall_context", client, client.config, session_id="ses_bound" + ) + await fn({"query": "x", "session_id": "ses_caller"}) + _, last_args = fake_mcp_server.call_log[-1] + assert last_args["session_id"] == "ses_caller" + + +@pytest.mark.asyncio +async def test_session_id_not_injected_when_not_bound(client, fake_mcp_server) -> None: + """A tool built without a session_id must not add a session_id key — + only the caller-supplied fields (plus scope defaults) reach the server.""" + fn, _ = build_tool("engraphis_recall_context", client, client.config) + await fn({"query": "x"}) + _, last_args = fake_mcp_server.call_log[-1] + assert "session_id" not in last_args or last_args.get("session_id") in (None, "") + + +async def test_session_id_not_injected_for_tools_without_session_id_in_schema( + client, fake_mcp_server +) -> None: + """A bound session_id must NOT be injected into tools whose declared + schema does not list ``session_id``; FastMCP would otherwise reject + the RPC for an unexpected argument. Covers discover_actions, both + executors, get_memory, update_memory, and conflict_review.""" + for tool in ( + "engraphis_discover_actions", + "engraphis_execute_action", + "engraphis_execute_read", + "engraphis_get_memory", + "engraphis_update_memory", + "engraphis_conflict_review", + ): + fn, _meta = build_tool(tool, client, client.config, session_id="ses_bound") + await fn({}) # any args; the server replies with the echoed payload + last_name, last_args = fake_mcp_server.call_log[-1] + assert last_name == tool + assert "session_id" not in last_args, ( + f"session_id leaked into {tool!r} whose schema does not declare it" + ) + + +def test_all_tools_with_session_id_returns_independent_callables(client) -> None: + """all_tools() must return 9 distinct callables, each with its own + closure-captured name. Reusing a session_id must not collapse the + tools into a single shared callable.""" + pairs = all_tools(client, client.config, session_id="ses_shared") + assert len(pairs) == 9 + callables = [fn for fn, _ in pairs] + # Each callable has a unique __name__ or at least is a different object. + assert len({id(fn) for fn in callables}) == 9 + + +def test_build_tool_meta_description_matches_descriptor_table(client) -> None: + """Every built tool's description must match the entry in _DESC — a + typo in a schema shouldn't silently ship.""" + for name, _schema in TOOL_SPECS: + _fn, meta = build_tool(name, client, client.config) + assert meta["name"] == name + assert isinstance(meta["description"], str) and meta["description"] + + +def test_all_tool_schemas_have_unique_property_names_within_tool() -> None: + """A schema that lists the same property twice would be ambiguous.""" + for name, schema in TOOL_SPECS: + props = schema.get("properties", {}) + assert len(props) == len(set(props)), ( + f"{name} schema has duplicate property names: {list(props)}" + ) diff --git a/scripts/install_prime_agent.py b/scripts/install_prime_agent.py new file mode 100644 index 00000000..ad5f88fd --- /dev/null +++ b/scripts/install_prime_agent.py @@ -0,0 +1,30 @@ +# -*- coding: utf-8 -*- +"""Thin wrapper around the package-distributed installer. + +The canonical implementation lives at +``engraphis_prime_agent.installer`` so it ships with the wheel and works +after ``pip install engraphis-prime-agent``. This wrapper remains at the +repo root for source-tree developers who run ``python +scripts/install_prime_agent.py`` directly. + +Usage: + python scripts/install_prime_agent.py + python scripts/install_prime_agent.py --uninstall +""" +from __future__ import annotations + +import sys +from pathlib import Path + +# Allow importing the package from a source checkout without an editable +# install. The integration package is three directories up from this +# script: scripts/ -> engraphis/ -> integrations/prime_agent/ -> src/. +_REPO_ROOT = Path(__file__).resolve().parent.parent +_SRC = _REPO_ROOT / "integrations" / "prime_agent" / "src" +if _SRC.is_dir(): + sys.path.insert(0, str(_SRC)) + +from engraphis_prime_agent.installer import main # noqa: E402 + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/e2e/graph-engine.spec.js b/tests/e2e/graph-engine.spec.js index c6b84dbc..1b7db89f 100644 --- a/tests/e2e/graph-engine.spec.js +++ b/tests/e2e/graph-engine.spec.js @@ -13,7 +13,7 @@ const { test, expect } = require('@playwright/test'); */ const workspace = 'graph-e2e'; -const stellarOrbitAssetVersion = '20260831-galaxy-floor-fix-2'; +const stellarOrbitAssetVersion = '20260902-slider-merge-1'; // A small connected store: two clusters joined by one bridge, so communities, the legend and // the bridge detector all have something real to work on. diff --git a/tests/e2e/ledger.spec.js b/tests/e2e/ledger.spec.js index 9805724f..d5376e31 100644 --- a/tests/e2e/ledger.spec.js +++ b/tests/e2e/ledger.spec.js @@ -580,14 +580,14 @@ test('Ledger cache-busts a graph renderer that fetched but did not register', as await expect(page.locator('#graph-empty')).toContainText('Graph unavailable'); expect(rendererRequests).toHaveLength(1); const first = new URL(rendererRequests[0]); - expect(first.searchParams.get('v')).toBe('20260831-galaxy-floor-fix-2'); + expect(first.searchParams.get('v')).toBe('20260902-slider-merge-1'); expect(first.searchParams.has('retry')).toBe(false); await page.getByRole('button', { name: 'Reload data' }).click(); await expect(page.locator('#graph-count')).toContainText('3 entities · 1 relations'); expect(rendererRequests).toHaveLength(2); const second = new URL(rendererRequests[1]); - expect(second.searchParams.get('v')).toBe('20260831-galaxy-floor-fix-2'); + expect(second.searchParams.get('v')).toBe('20260902-slider-merge-1'); expect(second.searchParams.get('retry')).toBe('1'); }); diff --git a/tests/test_graph_engine_asset.py b/tests/test_graph_engine_asset.py index c053d82e..08ad5ef7 100644 --- a/tests/test_graph_engine_asset.py +++ b/tests/test_graph_engine_asset.py @@ -381,7 +381,7 @@ def test_graph_engine_deep_link_reaches_the_next_engine_after_a_lazy_load() -> N report = _run_routing("loads") assert report["appended"] == [ - "/v2-assets/engraphis-graph.js?v=20260831-galaxy-floor-fix-2" + "/v2-assets/engraphis-graph.js?v=20260902-slider-merge-1" ] # It waits rather than rendering something wrong in the meantime. assert report["beforeSettle"] == {"engine": 0, "classic": 0} @@ -396,7 +396,7 @@ def test_classic_route_reaches_the_canonical_engine_without_a_query_flag() -> No report = _run_routing("classic") assert report["appended"] == [ - "/v2-assets/engraphis-graph.js?v=20260831-galaxy-floor-fix-2" + "/v2-assets/engraphis-graph.js?v=20260902-slider-merge-1" ] assert report["beforeSettle"] == {"engine": 0, "classic": 0} assert report["engine"] == 1 @@ -10391,10 +10391,10 @@ def test_primary_graph_dependencies_are_lazy_retryable_and_csp_clean() -> None: d3 = loader.index("'/v2-assets/vendor/d3.min.js?v=20260727-final'") force_graph = loader.index("'/v2-assets/vendor/force-graph.min.js?v=20260727-final'") renderer = loader.index( - "'/v2-assets/engraphis-graph.js?v=20260831-galaxy-floor-fix-2'" + "'/v2-assets/engraphis-graph.js?v=20260902-slider-merge-1'" ) assert d3 < force_graph < renderer - assert '/v2-assets/ledger.js?v=20260831-galaxy-floor-fix-2' in markup + assert '/v2-assets/ledger.js?v=20260902-slider-merge-1' in markup assert "if (graphAssetsPromise === attempt) releaseGraphAssetsAttempt(attempt)" in loader assert "graphAssetsRetry = Math.min(graphAssetsRetry + 1, 10)" in loader all_loader = source[source.index("function ensureGraphAllAsset()"): diff --git a/tests/test_recall_arm_candidate_k_cap.py b/tests/test_recall_arm_candidate_k_cap.py new file mode 100644 index 00000000..581ee482 --- /dev/null +++ b/tests/test_recall_arm_candidate_k_cap.py @@ -0,0 +1,266 @@ +"""Tests for the optional ``ENGRAPHIS_RECALL_ARM_CANDIDATE_K`` latency knob. + +PR #171 widened the prompt-only first arm to ``candidate_k + min(250, candidate_k*3)`` +so a 49-fact corpus pays ~5x more matrix-vector cost on the new k=50 default. The +opt-in ``ENGRAPHIS_RECALL_ARM_CANDIDATE_K`` env var (and the matching constructor +kwarg ``arm_candidate_k_cap=``) lets an operator clamp that first-page widening +*and* the second-page ceiling for latency-sensitive deployments. + +Default behavior (no env, no kwarg) is unchanged. +""" +from __future__ import annotations + +import time + +from engraphis.backends import DeterministicEmbedder, NumpyVectorIndex +from engraphis.backends.reranker import IdentityReranker +from engraphis.core.interfaces import MemoryRecord, SearchFilter +from engraphis.core.recall import RecallEngine +from engraphis.core.store import Store + + +class _SemanticTestEmbedder(DeterministicEmbedder): + supports_semantic_search = True + embedding_mode = "semantic" + + +def _add(store, emb, wid, rid, text, **kw): + provenance = dict(kw.get("provenance") or { + "source": "test", "trusted": True, "review_state": "approved", + }) + if provenance.get("trusted") is True: + provenance.setdefault("review_state", "approved") + kw["provenance"] = provenance + return store.add_memory(MemoryRecord( + id="", content=text, workspace_id=wid, repo_id=rid, + embedding=emb.embed([text])[0], **kw, + )) + + +class _RecordingIndex: + """Vector-index double that records every arm size it was queried with.""" + + def __init__(self, hit_count: int | None = None, real_ids: list[str] | None = None): + """If ``hit_count`` is set, the index always returns exactly that many + hits per call (up to ``k``). ``None`` (default) returns the synthetic + up-to-4 series; ``0`` returns nothing; any positive int returns that + many. ``real_ids`` (if given) replaces the synthetic id prefix so the + returned ids resolve to real ``MemoryRecord`` rows in the store -- + otherwise the prompt-eligibility filter discards them and the + escalation loop is short-circuited on empty recs. + """ + self.requested: list[int] = [] + self.records: list[tuple[str, float]] = [] + self._hit_count = hit_count + self._real_ids = list(real_ids) if real_ids else None + + def search(self, query, k, *, filter=None): + self.requested.append(int(k)) + if self._hit_count == 0: + return [] + cap = min(k, 4) if self._hit_count is None else min(k, self._hit_count) + if self._real_ids is not None: + return [(self._real_ids[i % len(self._real_ids)], float(k - i)) + for i in range(cap)] + return [(f"mem_{i}", float(k - i)) for i in range(cap)] + + +def test_arm_candidate_k_cap_default_is_none(monkeypatch): + """Without the env var or kwarg the cap is unset and PR #171 is preserved.""" + monkeypatch.delenv("ENGRAPHIS_RECALL_ARM_CANDIDATE_K", raising=False) + eng = RecallEngine(Store(":memory:"), _SemanticTestEmbedder(256), + _RecordingIndex(), IdentityReranker()) + assert eng._arm_candidate_k_cap is None + + +def test_arm_candidate_k_cap_reads_env_var(monkeypatch): + """Operator-set env var populates the cap; whitespace and bad values are ignored.""" + monkeypatch.setenv("ENGRAPHIS_RECALL_ARM_CANDIDATE_K", " 50 ") + eng = RecallEngine(Store(":memory:"), _SemanticTestEmbedder(256), + _RecordingIndex(), IdentityReranker()) + assert eng._arm_candidate_k_cap == 50 + + monkeypatch.setenv("ENGRAPHIS_RECALL_ARM_CANDIDATE_K", "not-a-number") + eng = RecallEngine(Store(":memory:"), _SemanticTestEmbedder(256), + _RecordingIndex(), IdentityReranker()) + assert eng._arm_candidate_k_cap is None + + +def test_arm_candidate_k_cap_constructor_kwarg_overrides_env(monkeypatch): + monkeypatch.setenv("ENGRAPHIS_RECALL_ARM_CANDIDATE_K", "50") + eng = RecallEngine(Store(":memory:"), _SemanticTestEmbedder(256), + _RecordingIndex(), IdentityReranker(), arm_candidate_k_cap=64) + assert eng._arm_candidate_k_cap == 64 + + +def test_arm_candidate_k_cap_clamps_first_arm(monkeypatch): + """With cap=50, k=50 prompt-only first arm is 50 (was 200).""" + monkeypatch.setenv("ENGRAPHIS_RECALL_ARM_CANDIDATE_K", "50") + index = _RecordingIndex() + eng = RecallEngine(Store(":memory:"), _SemanticTestEmbedder(256), + index, IdentityReranker()) + store = eng.store + wid = store.get_or_create_workspace("w") + for i in range(60): + _add(store, eng.embedder, wid, None, f"fact {i}") + + result = eng.recall("fact 5", SearchFilter(workspace_id=wid), k=50, + candidate_k=50, prompt_only=True) + + # First arm is clamped to 50; without the cap it would be 200. + assert index.requested[0] == 50 + # candidate_k_used reflects the actual first-page widening. + assert result.candidate_k_used == 50 + # The result must still be non-empty: the cap must not regress recall on + # a trusted-only corpus. + assert result.count >= 1 + + +def test_arm_candidate_k_cap_clamps_ceiling_when_first_page_insufficient(monkeypatch): + """The second page must also be clamped so the escalation loop does not + silently undo the savings by jumping to PROMPT_ONLY_MIN_CANDIDATES.""" + monkeypatch.setenv("ENGRAPHIS_RECALL_ARM_CANDIDATE_K", "8") + # Build an engine whose only enabled arm is the vector arm, so the + # lexical/graph/code arms cannot pad the prompt-eligible record set + # and short-circuit the ceiling path. The recording index returns + # exactly 1 hit per call (less than the prompt_target of 2 below and + # less than arm_candidate_k=4 so can_expand is True), so the + # escalation loop is forced into the arm_candidate_k = + # candidate_ceiling branch. + from engraphis.core.retrieval_policy import ProfileConfig + vector_only = ProfileConfig( + name="vector-only-test", vector=True, lexical=False, graph=False, code=False + ) + eng = RecallEngine(Store(":memory:"), _SemanticTestEmbedder(256), + _RecordingIndex(), IdentityReranker()) + store = eng.store + wid = store.get_or_create_workspace("w") + real_ids = [] + for i in range(20): + real_ids.append(_add(store, eng.embedder, wid, None, f"fact {i}")) + + index = _RecordingIndex(hit_count=4, real_ids=real_ids) + eng.index = index + + eng.recall("zzz-unmatched-query-zzz", SearchFilter(workspace_id=wid), k=8, + candidate_k=1, prompt_only=True, arm_config=vector_only) + + # First arm: 1 + min(250, 1*3) = 4 (clamped at min(8, 4) = 4, then + # floored at candidate_k=1, so 4). With the cap, the second-page + # ceiling is min(256, 8) = 8. Without the cap the index would have + # been queried with [4, 256]. The recording index returns 4 hits per + # call (>= arm_candidate_k=4 so can_expand=True; < prompt_target=8 + # so the loop is forced to escalate to the ceiling). + assert index.requested, "index was never queried -- test setup is broken" + assert index.requested[0] == 4, ( + f"first arm must be 4 (cap=8, candidate_k=1, prompt_only), " + f"got {index.requested[0]}" + ) + assert all(requested <= 8 for requested in index.requested), ( + f"every index query must respect the cap=8 ceiling, " + f"got {index.requested}" + ) + assert max(index.requested) <= 8 + # Sanity: the loop must have actually escalated, i.e. it must have + # queried the index at least twice. If it queried only once, the + # first arm satisfied prompt_target and the ceiling-clamp code path + # was not exercised -- which is the vacuous case this test guards + # against. + assert len(index.requested) >= 2, ( + f"ceiling clamp must be exercised via the escalation loop; " + f"only {len(index.requested)} index queries were recorded -- " + f"the first arm already satisfied prompt_target and the cap " + f"code path is not being run" + ) + + +def test_arm_candidate_k_cap_floor_protects_one_fact_corpus(monkeypatch): + """The cap must not shrink the first arm below the caller's requested + candidate_k — that would silently under-search a one-fact scope.""" + monkeypatch.setenv("ENGRAPHIS_RECALL_ARM_CANDIDATE_K", "2") + index = _RecordingIndex() + eng = RecallEngine(Store(":memory:"), _SemanticTestEmbedder(256), + index, IdentityReranker()) + store = eng.store + wid = store.get_or_create_workspace("w") + for i in range(20): + _add(store, eng.embedder, wid, None, f"fact {i}") + + eng.recall("fact 0", SearchFilter(workspace_id=wid), k=1, + candidate_k=10, prompt_only=True) + + # First arm = max(formula=10+30=40, candidate_k=10) capped at 2 = max(2, 10) = 10. + assert index.requested[0] == 10 + + +def test_arm_candidate_k_cap_reduces_latency_at_k_50(monkeypatch): + """End-to-end latency check: cap=50 should be measurably faster than + the uncapped default at k=50, on a trusted 49-fact corpus, while still + returning the expected number of chunks. + + The 1.5x threshold is conservative; the actual speedup on the bundled + rebench was 1.9x (201ms -> 103ms) at cap=50. We deliberately use a + loose bound so this test stays stable across hardware and numpy builds. + """ + monkeypatch.delenv("ENGRAPHIS_RECALL_ARM_CANDIDATE_K", raising=False) + store = Store(":memory:") + emb = _SemanticTestEmbedder(256) + index = NumpyVectorIndex(store) + eng_uncapped = RecallEngine(store, emb, index, IdentityReranker()) + wid = store.get_or_create_workspace("w") + base = ( + "Project Aurora uses Postgres for durable storage. Authentication uses PASETO. " + "The deploy pipeline runs unit and integration tests with a canary release." + ) + # Use 300 memories so both arms are clamped to well above 250, the + # first-page widening ceiling. The 49-fact corpus in the original + # version clamped both k=50 and k=200 to len(ids)==49, making the + # two timed paths operationally identical; the 1.5x speedup + # assertion was therefore measuring noise/cache order. + for i in range(300): + _add(store, emb, wid, None, f"{base} fact_index={i} workstream={i % 5}") + flt = SearchFilter(workspace_id=wid) + query = "What storage and auth systems does Project Aurora use?" + + def mean_ms(eng): + # Two warmups then 11 timed samples to smooth GC and embedder warmup. + for _ in range(2): + eng.recall(query, flt, k=50, candidate_k=50, prompt_only=True) + samples = [] + for _ in range(11): + t0 = time.perf_counter() + eng.recall(query, flt, k=50, candidate_k=50, prompt_only=True) + samples.append((time.perf_counter() - t0) * 1000.0) + samples.sort() + return sum(samples[2:-2]) / 7.0 # trimmed mean, drop 2 best and 2 worst + + uncapped_ms = mean_ms(eng_uncapped) + + # Capped engine on a fresh store; rebuilding the corpus keeps the latencies + # independent so the embedder cache state of the uncapped run cannot bias + # the timed mean. + monkeypatch.setenv("ENGRAPHIS_RECALL_ARM_CANDIDATE_K", "50") + store2 = Store(":memory:") + emb2 = _SemanticTestEmbedder(256) + eng_capped = RecallEngine(store2, emb2, NumpyVectorIndex(store2), + IdentityReranker()) + wid2 = store2.get_or_create_workspace("w") + for i in range(300): + _add(store2, emb2, wid2, None, f"{base} fact_index={i} workstream={i % 5}") + flt2 = SearchFilter(workspace_id=wid2) + capped_ms = mean_ms(eng_capped) + + # Sanity: the uncapped recall returns the full k=50 trusted chunks. + uncapped_result = eng_uncapped.recall(query, flt, k=50, candidate_k=50, + prompt_only=True) + capped_result = eng_capped.recall(query, flt2, k=50, candidate_k=50, + prompt_only=True) + assert uncapped_result.candidate_k_used == 200 + assert capped_result.candidate_k_used == 50 + # Recall quality must not regress on a trusted-only corpus. + assert capped_result.count == uncapped_result.count + # And latency must drop by at least 1.5x. + assert capped_ms < uncapped_ms / 1.5, ( + f"cap=50 did not yield the expected speedup: uncapped={uncapped_ms:.1f}ms " + f"capped={capped_ms:.1f}ms" + )