diff --git a/integrations/goose-agent-middleware/LICENSE b/integrations/goose-agent-middleware/LICENSE new file mode 100644 index 00000000..d6456956 --- /dev/null +++ b/integrations/goose-agent-middleware/LICENSE @@ -0,0 +1,202 @@ + + 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 [yyyy] [name of copyright owner] + + 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/goose-agent-middleware/README.md b/integrations/goose-agent-middleware/README.md new file mode 100644 index 00000000..0b8cd217 --- /dev/null +++ b/integrations/goose-agent-middleware/README.md @@ -0,0 +1,113 @@ +# Goose + Agent Middleware + +Runs Agent Middleware (AMW) as a normal Goose **StreamableHttp MCP extension** +that proxies the official `@modelcontextprotocol/server-filesystem` server +(`read_text_file`, `write_file`), plus a Goose **PreToolUse deny hook**. + +Every proxied call goes through permit -> idempotency -> invoke -> receipt. + +## Pieces + +| Piece | What it does | +|---|---| +| `goose-amw-proxy` | StreamableHttp MCP server (protocol 2026-07-28 and older). Verifies the caller, enforces permits, dedupes writes, signs receipts. | +| `goose-amw-hook` | Goose `PreToolUse` hook. A call with no permit is blocked **before it runs**; Goose shows `Tool call denied by policy hook ...: AMW permit denied: `. | +| `plugin/amw-permit/hooks/hooks.json` | Template plugin that wires the hook to `amw__*` tools. | + +The proxy enforces the same permit rules server-side, so a client without the +hook is still refused. + +## Upstream server + +The filesystem server is pinned: `AMW_FS_SERVER_PACKAGE` (default +`@modelcontextprotocol/server-filesystem@2.0.1`) is passed to `npx -y`, so a +new upstream release never runs without an explicit change. + +A failure before the tool call is sent (spawn, session setup, `initialize`) +is safe to retry: AMW restores the idempotency key to its state before the +attempt and issues a `not_dispatched` receipt. Only a failure after dispatch +is recorded as `failed_unknown`. + +## Permit policy + +```json +{"rules": [ + {"id": "write-allowed", "tool": "write_file", "folder": "/path/to/allowed"}, + {"id": "read-sandbox", "tool": "read_text_file", "folder": "/path/to/sandbox"} +]} +``` + +A call is allowed only if a rule names the tool and the resolved `path` is +inside the rule's folder (`..` and symlinks are resolved first). Anything else +is denied with a reason. + +## Idempotency key + +AMW keys a write by **verified caller identity + tool name + canonical +argument hash + `agent-session-id`**. + +Goose (protocol 2026-07-28) sends `_meta["agent-tool-call-request-id"]`, but +that id is the model's tool-call id: when a call fails and the model retries, +Goose sends a **new** id (recorded in `tests/fixtures/g1_gate_calls.jsonl`). +AMW therefore logs it on receipts and in the ledger but never uses it as the +key. A **write without `agent-session-id`** (for example an older protocol with +no `_meta`) is refused with a clear reason and a `refuse` receipt, because it +cannot be deduplicated safely. Reads do not need a session id. + +The caller is verified from the `Authorization: Bearer` header against +`AMW_CALLERS` (a JSON map of `sha256(token)` to caller id). Unverified calls +are refused. + +## Replay window + +`AMW_IDEMPOTENCY_WINDOW_SECONDS` (default **600**, must be > 0). + +* A completed write repeated with the same key **inside** the window is not + run again; AMW returns the stored result and issues a `replayed` receipt. +* The same write made **after** the window runs normally, so deliberate + identical writes later in a session are not swallowed. +* A write whose outcome is unknown (`failed_unknown`: the call was + dispatched and the response was lost, or AMW crashed mid-call) is **never + re-run automatically**, inside or outside the window. Retries are refused + with a reason until an operator resolves the entry. +* A definite upstream error (`failed`, no effect) may be retried. + +Choose a window longer than your client's retry horizon and shorter than the +interval at which you write identical content on purpose. + +## Receipts + +Each decision (allow, deny, replay, refuse) appends an Ed25519-signed receipt +to `$AMW_STATE_DIR/receipts.jsonl`: `tool`, `args_hash`, `decision`, +`outcome`, `issued_at`, plus caller, session, key and the logged request id. +The signature covers the canonical JSON of `body`; the public key is written +to `receipt_public_key.b64`. + +## Goose setup + +```yaml +extensions: + amw: + type: streamable_http + name: amw + uri: http://127.0.0.1:8791/mcp + headers: + Authorization: Bearer ${AMW_TOKEN} + timeout: 120 + enabled: true +``` + +Install `plugin/amw-permit` into `.agents/plugins/` (project) or the Goose +plugins directory, replacing the paths in its `hooks.json`. + +## Tests + +`pytest` needs no external data. `tests/fixtures/g1_gate_calls.jsonl` is the +real call log recorded by AMW during the G1 gate session. +`tests/e2e/fault_proxy.py` is a test-only launcher that injects one fault into +the shipped proxy to record the cut-stream and unknown-outcome sessions; it +is not part of the package. + +An error result from the upstream server for a write is recorded as +`failed_unknown`: the pinned server writes files in place, so the file may be +partly written, and the write is never re-run automatically. diff --git a/integrations/goose-agent-middleware/plugin/amw-permit/hooks/hooks.json b/integrations/goose-agent-middleware/plugin/amw-permit/hooks/hooks.json new file mode 100644 index 00000000..6ee29317 --- /dev/null +++ b/integrations/goose-agent-middleware/plugin/amw-permit/hooks/hooks.json @@ -0,0 +1,17 @@ +{ + "hooks": { + "PreToolUse": [ + { + "matcher": "^amw__", + "hooks": [ + { + "type": "command", + "command": "goose-amw-hook --policy /path/to/policy.json --state-dir /path/to/state", + "on_failure": "block", + "timeout": 10 + } + ] + } + ] + } +} diff --git a/integrations/goose-agent-middleware/pyproject.toml b/integrations/goose-agent-middleware/pyproject.toml new file mode 100644 index 00000000..97f3c65d --- /dev/null +++ b/integrations/goose-agent-middleware/pyproject.toml @@ -0,0 +1,35 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "goose-agent-middleware" +version = "0.1.0" +description = "Agent Middleware as a Goose StreamableHttp MCP extension: permits, PreToolUse deny hook, idempotent writes, signed receipts" +readme = "README.md" +license = "Apache-2.0" +requires-python = ">=3.11" +dependencies = [ + "amw-sdk>=0.4.0", + "mcp>=2.3.0,<3", + "cryptography>=50.0.0", +] + +[project.optional-dependencies] +dev = ["pytest", "ruff"] + +[project.scripts] +goose-amw-proxy = "goose_amw.proxy:main" +goose-amw-hook = "goose_amw.hook:main" + +[tool.hatch.build.targets.wheel] +packages = ["src/goose_amw"] + +[tool.ruff] +line-length = 100 + +[tool.pytest.ini_options] +pythonpath = ["src", "../../sdk/python/src"] + +[tool.uv.sources] +amw-sdk = { path = "../../sdk/python", editable = true } diff --git a/integrations/goose-agent-middleware/src/goose_amw/__init__.py b/integrations/goose-agent-middleware/src/goose_amw/__init__.py new file mode 100644 index 00000000..77f8e9e8 --- /dev/null +++ b/integrations/goose-agent-middleware/src/goose_amw/__init__.py @@ -0,0 +1,3 @@ +"""Agent Middleware for Goose: a governed StreamableHttp MCP proxy.""" + +__all__ = ["canonical", "ledger", "permits", "receipts"] diff --git a/integrations/goose-agent-middleware/src/goose_amw/canonical.py b/integrations/goose-agent-middleware/src/goose_amw/canonical.py new file mode 100644 index 00000000..746d3565 --- /dev/null +++ b/integrations/goose-agent-middleware/src/goose_amw/canonical.py @@ -0,0 +1,36 @@ +"""Canonical JSON and argument hashing.""" + +from __future__ import annotations + +import hashlib +from typing import Any + +from b2a_sdk.receipt_verifier import canonical_json + + +def sha256_hex(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def args_hash(arguments: dict[str, Any] | None) -> str: + return "sha256:" + sha256_hex(canonical_json(arguments or {})) + + +def idempotency_key( + *, caller: str, tool: str, arguments: dict[str, Any] | None, session_id: str +) -> str: + """Fallback key: verified caller + tool + canonical argument hash + agent-session-id. + + ``agent-tool-call-request-id`` is deliberately NOT an input: the G1 gate showed + Goose issues a new id when the model retries a failed call. + """ + material = { + "caller": caller, + "tool": tool, + "args_hash": args_hash(arguments), + "session_id": session_id, + } + return "amw-fb-" + sha256_hex(canonical_json(material)) + + +__all__ = ["args_hash", "canonical_json", "idempotency_key", "sha256_hex"] diff --git a/integrations/goose-agent-middleware/src/goose_amw/hook.py b/integrations/goose-agent-middleware/src/goose_amw/hook.py new file mode 100644 index 00000000..fb523650 --- /dev/null +++ b/integrations/goose-agent-middleware/src/goose_amw/hook.py @@ -0,0 +1,60 @@ +"""Goose PreToolUse hook: deny a tool call that has no AMW permit. + +Goose runs this before the tool call with the hook context JSON on stdin. +Printing ``{"decision":"block","reason":...}`` stops the call before it +reaches the MCP server; Goose shows the reason to the user and the model. + +Usage in ``hooks/hooks.json``:: + + goose-amw-hook --policy /path/policy.json --state-dir /path/state +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +from . import permits +from .canonical import args_hash +from .receipts import ReceiptSigner, load_or_create_key + + +def decide(payload: dict, policy: dict) -> dict: + tool = payload.get("tool_name") or payload.get("matcher_context") or "" + arguments = payload.get("tool_input") or {} + decision = permits.check(policy, tool, arguments) + if decision.allowed: + return {"decision": "allow"} + return {"decision": "block", "reason": f"AMW permit denied: {decision.reason}"} + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--policy", required=True) + ap.add_argument("--state-dir", required=True) + ns = ap.parse_args(argv) + payload = json.load(sys.stdin) + out = decide(payload, permits.load_policy(ns.policy)) + if out["decision"] == "block": + state = Path(ns.state_dir) + signer = ReceiptSigner( + load_or_create_key(state / "receipt_signing_key.pem"), state / "receipts.jsonl" + ) + signer.issue( + tool=payload.get("tool_name") or "", + args_hash=args_hash(payload.get("tool_input") or {}), + decision="deny", + outcome="not_run", + enforced_by="goose_pre_tool_use_hook", + reason=out["reason"], + session_id=payload.get("session_id"), + agent_tool_call_request_id=payload.get("tool_call_id"), + ) + print(json.dumps(out)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/integrations/goose-agent-middleware/src/goose_amw/ledger.py b/integrations/goose-agent-middleware/src/goose_amw/ledger.py new file mode 100644 index 00000000..4cae9a28 --- /dev/null +++ b/integrations/goose-agent-middleware/src/goose_amw/ledger.py @@ -0,0 +1,139 @@ +"""SQLite idempotency ledger with a bounded replay window. + +States: ``in_flight`` -> ``completed`` | ``failed`` | ``failed_unknown``. + +* ``completed``: a retry inside the window replays the stored result; the tool + does not run again. +* ``failed``: the upstream answered with a definite error (no effect); a retry + may run again. +* ``failed_unknown``: the outcome is unclear (transport lost after dispatch). + It is never re-run automatically, inside or outside the window, until an + operator resolves it. +* A ``completed`` entry older than the window no longer dedupes, so an + identical write made on purpose much later runs normally. +""" + +from __future__ import annotations + +import sqlite3 +import threading +import time +from dataclasses import dataclass +from pathlib import Path + +DEFAULT_WINDOW_SECONDS = 600 + + +@dataclass(frozen=True) +class Claim: + """Result of trying to claim a key before running a tool.""" + + action: str # "run" | "replay" | "refuse" + result: str | None = None + status: str | None = None + previous: tuple | None = None # the row as it was before this claim, for release() + + +class Ledger: + def __init__(self, path: str | Path, window_seconds: float = DEFAULT_WINDOW_SECONDS): + if window_seconds <= 0: + raise ValueError("window_seconds must be positive") + self.window_seconds = window_seconds + self._lock = threading.Lock() + self._db = sqlite3.connect(str(path), check_same_thread=False, isolation_level=None) + self._db.execute( + "CREATE TABLE IF NOT EXISTS calls (" + " key TEXT PRIMARY KEY, tool TEXT NOT NULL, args_hash TEXT NOT NULL," + " status TEXT NOT NULL, result TEXT, first_request_id TEXT," + " created_at REAL NOT NULL, updated_at REAL NOT NULL, executions INTEGER NOT NULL)" + ) + + def claim( + self, + key: str, + *, + tool: str, + args_hash: str, + request_id: str | None, + now: float | None = None, + ) -> Claim: + now = time.time() if now is None else now + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + full = self._db.execute("SELECT * FROM calls WHERE key = ?", (key,)).fetchone() + row = None if full is None else (full[3], full[4], full[6]) + if row is not None: + status, result, created_at = row + if status == "failed_unknown": + self._db.execute("COMMIT") + return Claim("refuse", status=status) + if status == "in_flight": + self._db.execute("COMMIT") + return Claim("refuse", status=status) + if status == "completed" and now - created_at < self.window_seconds: + self._db.execute("COMMIT") + return Claim("replay", result=result, status=status) + # failed (no effect) or completed outside the window: run again. + self._db.execute( + "UPDATE calls SET status='in_flight', result=NULL, created_at=?," + " updated_at=?, executions=executions+1 WHERE key=?", + (now, now, key), + ) + else: + self._db.execute( + "INSERT INTO calls VALUES (?, ?, ?, 'in_flight', NULL, ?, ?, ?, 1)", + (key, tool, args_hash, request_id, now, now), + ) + self._db.execute("COMMIT") + return Claim("run", previous=full) + except BaseException: + self._db.execute("ROLLBACK") + raise + + def finish(self, key: str, status: str, result: str | None = None) -> None: + if status not in ("completed", "failed", "failed_unknown"): + raise ValueError(status) + with self._lock: + self._db.execute( + "UPDATE calls SET status=?, result=?, updated_at=? WHERE key=?", + (status, result, time.time(), key), + ) + + def release(self, key: str, claim: Claim) -> None: + """Undo a ``run`` claim whose call was never dispatched upstream. + + The key goes back to exactly its state before the claim (or disappears if it + was new), so a pre-dispatch failure leaves no trace and the call stays retryable. + """ + with self._lock: + self._db.execute("BEGIN IMMEDIATE") + try: + self._db.execute("DELETE FROM calls WHERE key=? AND status='in_flight'", (key,)) + if claim.previous is not None: + marks = ",".join("?" * len(claim.previous)) + self._db.execute(f"INSERT INTO calls VALUES ({marks})", claim.previous) + self._db.execute("COMMIT") + except BaseException: + self._db.execute("ROLLBACK") + raise + + def mark_stale_in_flight_unknown(self) -> int: + """At startup, an ``in_flight`` row means a crash mid-call: outcome unknown.""" + with self._lock: + cur = self._db.execute( + "UPDATE calls SET status='failed_unknown', updated_at=? WHERE status='in_flight'", + (time.time(),), + ) + return cur.rowcount + + def get(self, key: str) -> dict | None: + cur = self._db.execute( + "SELECT key, tool, status, result, executions, first_request_id FROM calls WHERE key=?", + (key,), + ) + row = cur.fetchone() + if row is None: + return None + names = ["key", "tool", "status", "result", "executions", "first_request_id"] + return dict(zip(names, row, strict=True)) diff --git a/integrations/goose-agent-middleware/src/goose_amw/permits.py b/integrations/goose-agent-middleware/src/goose_amw/permits.py new file mode 100644 index 00000000..b9065151 --- /dev/null +++ b/integrations/goose-agent-middleware/src/goose_amw/permits.py @@ -0,0 +1,43 @@ +"""Path-scoped permit rules shared by the proxy and the Goose PreToolUse hook.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class Decision: + allowed: bool + reason: str + + +def load_policy(path: str | Path) -> dict[str, Any]: + return json.loads(Path(path).read_text()) + + +def _bare(tool: str) -> str: + # Goose prefixes extension tools as "__". + return tool.split("__", 1)[1] if "__" in tool else tool + + +def check(policy: dict[str, Any], tool: str, arguments: dict[str, Any] | None) -> Decision: + """Allow only when a rule grants ``tool`` on a path inside its folder.""" + name = _bare(tool) + raw = (arguments or {}).get("path") + if not isinstance(raw, str) or not raw: + return Decision(False, f"no permit: {name} call has no path argument") + if not Path(raw).is_absolute(): + # The upstream server resolves relative (and ~) paths against its own root, + # not our working directory, so only an absolute path can be checked safely. + return Decision(False, f"no permit: {name} path must be absolute, got {raw!r}") + target = Path(raw).resolve() + for rule in policy.get("rules", []): + if rule.get("tool") != name: + continue + folder = Path(rule["folder"]).expanduser().resolve() + if target.is_relative_to(folder): + return Decision(True, f"permit {rule.get('id', name)} allows {name} in {folder}") + return Decision(False, f"no permit for {name} on {target}: outside every permitted folder") diff --git a/integrations/goose-agent-middleware/src/goose_amw/proxy.py b/integrations/goose-agent-middleware/src/goose_amw/proxy.py new file mode 100644 index 00000000..b6e2bf98 --- /dev/null +++ b/integrations/goose-agent-middleware/src/goose_amw/proxy.py @@ -0,0 +1,293 @@ +"""AMW StreamableHttp MCP server for Goose. + +Proxies one read tool and one write tool to the official +``@modelcontextprotocol/server-filesystem`` server. Every call is checked +against the permit policy, keyed for idempotency with the fallback key +(verified caller + tool + canonical argument hash + agent-session-id), and +receives a signed receipt. ``agent-tool-call-request-id`` is logged only. + +Configuration (environment): + +* ``AMW_POLICY`` - permit policy JSON (see README) +* ``AMW_STATE_DIR`` - ledger, receipts and signing key +* ``AMW_CALLERS`` - JSON map of sha256(bearer token) -> caller id +* ``AMW_FS_ROOT`` - root directory handed to the filesystem server +* ``AMW_FS_SERVER_PACKAGE`` - pinned npm spec of the filesystem server, + default ``@modelcontextprotocol/server-filesystem@2.0.1`` +* ``AMW_IDEMPOTENCY_WINDOW_SECONDS`` - replay window, default 600 +* ``AMW_HOST`` / ``AMW_PORT`` - bind address, default 127.0.0.1:8791 +""" + +from __future__ import annotations + +import hmac +import json +import logging +import os +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from mcp.server.mcpserver import Context, MCPServer +from mcp.server.mcpserver.exceptions import ToolError + +from . import permits +from .canonical import args_hash, idempotency_key, sha256_hex +from .ledger import DEFAULT_WINDOW_SECONDS, Ledger +from .receipts import ReceiptSigner, load_or_create_key + +log = logging.getLogger("goose_amw") + +REQUEST_ID_KEY = "agent-tool-call-request-id" +SESSION_ID_KEY = "agent-session-id" + + +class UpstreamError(Exception): + """The upstream answered with a definite error: the call had no effect.""" + + +class UpstreamNotDispatched(UpstreamError): + """Failed before the tool call was sent (spawn, session setup, initialize). + + Nothing reached the upstream tool, so the call is safe to retry and the + idempotency key is left exactly as it was before the attempt. + """ + + +class UpstreamUnknown(Exception): + """The call was dispatched but its outcome is unknown.""" + + +Upstream = Callable[[str, dict[str, Any]], Awaitable[str]] + + +FS_SERVER_PACKAGE = "@modelcontextprotocol/server-filesystem@2.0.1" + + +def filesystem_upstream(fs_root: str, package: str | None = None) -> Upstream: + from mcp.client.stdio import StdioServerParameters + + spec = package or os.environ.get("AMW_FS_SERVER_PACKAGE") or FS_SERVER_PACKAGE + return stdio_upstream(StdioServerParameters(command="npx", args=["-y", spec, fs_root])) + + +def stdio_upstream(params: Any) -> Upstream: + """Call one tool on a stdio MCP server, classifying failures by dispatch point.""" + from mcp.client.session import ClientSession + from mcp.client.stdio import stdio_client + + async def call(tool: str, arguments: dict[str, Any]) -> str: + dispatched = False + res = None + try: + async with stdio_client(params) as (read, write), ClientSession(read, write) as session: + await session.initialize() + dispatched = True + res = await session.call_tool(tool, arguments) + except BaseException as exc: + if res is not None and isinstance(exc, Exception): + # The server already answered; a teardown error does not change that. + log.warning("ignoring stdio teardown error after a result: %r", exc) + elif not dispatched: + raise UpstreamNotDispatched(f"filesystem server unavailable: {exc!r}") from exc + else: + raise UpstreamUnknown(f"{type(exc).__name__}: {exc}") from exc + text = "\n".join(getattr(c, "text", "") for c in res.content) + if res.is_error: + raise UpstreamError(text) + return text + + return call + + +@dataclass +class Config: + policy: dict[str, Any] + state_dir: Path + callers: dict[str, str] + fs_root: str + window_seconds: float = DEFAULT_WINDOW_SECONDS + + @classmethod + def from_env(cls) -> Config: + return cls( + policy=permits.load_policy(os.environ["AMW_POLICY"]), + state_dir=Path(os.environ["AMW_STATE_DIR"]), + callers=json.loads(Path(os.environ["AMW_CALLERS"]).read_text()), + fs_root=os.environ["AMW_FS_ROOT"], + window_seconds=float( + os.environ.get("AMW_IDEMPOTENCY_WINDOW_SECONDS", DEFAULT_WINDOW_SECONDS) + ), + ) + + +def verify_caller(headers: Mapping[str, str] | None, callers: dict[str, str]) -> str: + auth = (headers or {}).get("authorization", "") + if not auth.lower().startswith("bearer "): + raise ToolError("AMW refused: caller not verified (missing bearer token)") + digest = sha256_hex(auth[7:].strip()) + for known, caller in callers.items(): + if hmac.compare_digest(known, digest): + return caller + raise ToolError("AMW refused: caller not verified (unknown token)") + + +def _meta(ctx: Context) -> dict[str, Any]: + m = ctx.request_context.meta + if m is None: + return {} + return dict(m) if isinstance(m, dict) else m.model_dump(by_alias=True, exclude_none=True) + + +class Gateway: + """Permit -> idempotency -> invoke -> receipt for each proxied call.""" + + def __init__( + self, + config: Config, + upstream: Upstream | None = None, + after_commit: Callable[[str, str], Awaitable[None]] | None = None, + ): + config.state_dir.mkdir(parents=True, exist_ok=True) + self.config = config + self.ledger = Ledger(config.state_dir / "ledger.sqlite3", config.window_seconds) + stale = self.ledger.mark_stale_in_flight_unknown() + if stale: + log.warning("marked %d interrupted call(s) failed_unknown", stale) + self.signer = ReceiptSigner( + load_or_create_key(config.state_dir / "receipt_signing_key.pem"), + config.state_dir / "receipts.jsonl", + ) + (config.state_dir / "receipt_public_key.b64").write_text(self.signer.public_key_b64) + self.upstream = upstream or filesystem_upstream(config.fs_root) + # Test-only seam: tests inject faults here. Never set in production. + self.after_commit = after_commit + + async def handle(self, tool: str, arguments: dict[str, Any], ctx: Context, *, write: bool): + meta = _meta(ctx) + request_id = meta.get(REQUEST_ID_KEY) + session_id = meta.get(SESSION_ID_KEY) or None + ahash = args_hash(arguments) + base = {"protocol": ctx.protocol_version, "agent_tool_call_request_id": request_id} + caller = verify_caller(ctx.headers, self.config.callers) + base.update(caller=caller, session_id=session_id) + + decision = permits.check(self.config.policy, tool, arguments) + if not decision.allowed: + self.signer.issue( + tool=tool, + args_hash=ahash, + decision="deny", + outcome="not_run", + reason=decision.reason, + **base, + ) + raise ToolError(f"AMW permit denied: {decision.reason}") + + if not write: + return await self._run(tool, arguments, ahash, base, key=None) + + if not session_id: + reason = ( + f"{tool} needs _meta['{SESSION_ID_KEY}'] to be deduplicated safely; " + "the client sent none, so the write was not run" + ) + self.signer.issue( + tool=tool, + args_hash=ahash, + decision="refuse", + outcome="not_run", + reason=reason, + **base, + ) + raise ToolError(f"AMW refused: {reason}") + + key = idempotency_key(caller=caller, tool=tool, arguments=arguments, session_id=session_id) + base["idempotency_key"] = key + + claim = self.ledger.claim(key, tool=tool, args_hash=ahash, request_id=request_id) + if claim.action == "replay": + self.signer.issue( + tool=tool, args_hash=ahash, decision="allow", outcome="replayed", **base + ) + return claim.result + if claim.action == "refuse": + self.signer.issue( + tool=tool, args_hash=ahash, decision="refuse", outcome=claim.status, **base + ) + raise ToolError( + f"AMW refused to re-run: an earlier identical {tool} call is {claim.status}; " + "its outcome is unknown, so it will not run again automatically" + ) + return await self._run(tool, arguments, ahash, base, key=key, claim=claim) + + async def _run(self, tool, arguments, ahash, base, *, key, claim=None): + try: + result = await self.upstream(tool, arguments) + except UpstreamNotDispatched as exc: + if key: + self.ledger.release(key, claim) + self.signer.issue( + tool=tool, args_hash=ahash, decision="allow", outcome="not_dispatched", **base + ) + raise ToolError(f"upstream unavailable, safe to retry: {exc}") from exc + except UpstreamError as exc: + if key: + # The filesystem server writes in place, so an error result for a + # write may come after a partial write: never claim it had no effect. + self.ledger.finish(key, "failed_unknown") + self.signer.issue( + tool=tool, args_hash=ahash, decision="allow", outcome="failed_unknown", **base + ) + raise ToolError( + f"upstream error; the write may be partly applied and will not re-run: {exc}" + ) from exc + self.signer.issue( + tool=tool, args_hash=ahash, decision="allow", outcome="failed", **base + ) + raise ToolError(f"upstream error: {exc}") from exc + except BaseException: + if key: + self.ledger.finish(key, "failed_unknown") + self.signer.issue( + tool=tool, args_hash=ahash, decision="allow", outcome="failed_unknown", **base + ) + raise + if key: + self.ledger.finish(key, "completed", result) + self.signer.issue(tool=tool, args_hash=ahash, decision="allow", outcome="completed", **base) + if key and self.after_commit is not None: + await self.after_commit(tool, key) + return result + + +def build_server(gateway: Gateway) -> MCPServer: + mcp = MCPServer("amw") + + @mcp.tool() + async def read_text_file(path: str, ctx: Context) -> str: + """Read a UTF-8 text file through AMW (permit-checked, receipted).""" + return await gateway.handle("read_text_file", {"path": path}, ctx, write=False) + + @mcp.tool() + async def write_file(path: str, content: str, ctx: Context) -> str: + """Write a text file through AMW (permit-checked, idempotent, receipted).""" + args = {"path": path, "content": content} + return await gateway.handle("write_file", args, ctx, write=True) + + return mcp + + +def main() -> None: + logging.basicConfig(level=logging.INFO) + server = build_server(Gateway(Config.from_env())) + server.run( + transport="streamable-http", + host=os.environ.get("AMW_HOST", "127.0.0.1"), + port=int(os.environ.get("AMW_PORT", "8791")), + ) + + +if __name__ == "__main__": + main() diff --git a/integrations/goose-agent-middleware/src/goose_amw/receipts.py b/integrations/goose-agent-middleware/src/goose_amw/receipts.py new file mode 100644 index 00000000..96ac674e --- /dev/null +++ b/integrations/goose-agent-middleware/src/goose_amw/receipts.py @@ -0,0 +1,92 @@ +"""Ed25519-signed receipts, one JSON line per tool call decision.""" + +from __future__ import annotations + +import base64 +import json +import os +import time +import uuid +from pathlib import Path +from typing import Any + +from cryptography.exceptions import InvalidSignature +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric.ed25519 import ( + Ed25519PrivateKey, + Ed25519PublicKey, +) + +from .canonical import canonical_json, sha256_hex + + +def load_or_create_key(path: str | Path) -> Ed25519PrivateKey: + p = Path(path) + if p.exists(): + key = serialization.load_pem_private_key(p.read_bytes(), password=None) + assert isinstance(key, Ed25519PrivateKey) + return key + key = Ed25519PrivateKey.generate() + p.parent.mkdir(parents=True, exist_ok=True) + fd = os.open(p, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(fd, "wb") as f: + f.write( + key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + ) + return key + + +def _raw_public(key: Ed25519PublicKey) -> bytes: + return key.public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw) + + +class ReceiptSigner: + def __init__(self, key: Ed25519PrivateKey, log_path: str | Path): + self._key = key + self._log = Path(log_path) + pub = _raw_public(key.public_key()) + self.public_key_b64 = base64.b64encode(pub).decode() + self.kid = "amw-goose-" + sha256_hex(self.public_key_b64)[:16] + + def issue( + self, + *, + tool: str, + args_hash: str, + decision: str, + outcome: str, + **extra: Any, + ) -> dict[str, Any]: + body = { + "receipt_id": str(uuid.uuid4()), + "tool": tool, + "args_hash": args_hash, + "decision": decision, + "outcome": outcome, + "issued_at": time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime()) + "Z", + "alg": "Ed25519", + "kid": self.kid, + **{k: v for k, v in extra.items() if v is not None}, + } + sig = self._key.sign(canonical_json(body).encode()) + receipt = {"body": body, "signature": base64.b64encode(sig).decode()} + self._log.parent.mkdir(parents=True, exist_ok=True) + with self._log.open("a") as f: + f.write(json.dumps(receipt) + "\n") + return receipt + + +def verify(receipt: dict[str, Any], public_key_b64: str) -> bool: + pub = Ed25519PublicKey.from_public_bytes(base64.b64decode(public_key_b64)) + try: + pub.verify( + base64.b64decode(receipt["signature"]), + canonical_json(receipt["body"]).encode(), + ) + except (InvalidSignature, ValueError, KeyError): + return False + return True diff --git a/integrations/goose-agent-middleware/tests/e2e/fault_proxy.py b/integrations/goose-agent-middleware/tests/e2e/fault_proxy.py new file mode 100644 index 00000000..3f53c934 --- /dev/null +++ b/integrations/goose-agent-middleware/tests/e2e/fault_proxy.py @@ -0,0 +1,44 @@ +"""TEST-ONLY launcher: the shipped proxy plus one injected fault. + +Not part of the package. Used to record the real Goose sessions in the PR: + +* ``AMW_TEST_FAULT=cut_after_commit`` - the first write runs and commits, then + its response stream is closed before the result reaches Goose. +* ``AMW_TEST_FAULT=lose_upstream_response`` - the first write is dispatched + to the filesystem server, then its response is dropped (outcome unknown). +""" + +import os + +from goose_amw.proxy import ( + Config, + Gateway, + UpstreamUnknown, + build_server, + filesystem_upstream, +) + +FAULT = os.environ["AMW_TEST_FAULT"] +cfg = Config.from_env() +real = filesystem_upstream(cfg.fs_root) +fired = [] + + +async def upstream(tool, args): + result = await real(tool, args) + if FAULT == "lose_upstream_response" and tool == "write_file" and not fired: + fired.append(tool) + raise UpstreamUnknown("test fault: upstream response dropped after dispatch") + return result + + +async def after_commit(tool, key): + if FAULT == "cut_after_commit" and not fired: + fired.append(key) + raise ConnectionAbortedError("test fault: response stream cut after commit") + + +gw = Gateway(cfg, upstream=upstream, after_commit=after_commit) +build_server(gw).run( + transport="streamable-http", host="127.0.0.1", port=int(os.environ["AMW_PORT"]) +) diff --git a/integrations/goose-agent-middleware/tests/e2e/verify_receipts.py b/integrations/goose-agent-middleware/tests/e2e/verify_receipts.py new file mode 100644 index 00000000..480657ae --- /dev/null +++ b/integrations/goose-agent-middleware/tests/e2e/verify_receipts.py @@ -0,0 +1,28 @@ +"""Verify every receipt in an AMW state dir and print a compact summary.""" + +import json +import sys +from pathlib import Path + +from goose_amw.receipts import verify + +state = Path(sys.argv[1]) +pub = (state / "receipt_public_key.b64").read_text().strip() +bad = 0 +for line in (state / "receipts.jsonl").read_text().splitlines(): + r = json.loads(line) + ok = verify(r, pub) + bad += not ok + b = r["body"] + print( + "OK " if ok else "BAD", + b["issued_at"], + b["tool"], + b["decision"], + b["outcome"], + b["args_hash"][:19], + b.get("enforced_by", "proxy"), + b.get("agent_tool_call_request_id") or "-", + b.get("protocol", "-"), + ) +sys.exit(1 if bad else 0) diff --git a/integrations/goose-agent-middleware/tests/fixtures/g1_gate_calls.jsonl b/integrations/goose-agent-middleware/tests/fixtures/g1_gate_calls.jsonl new file mode 100644 index 00000000..ab3fa361 --- /dev/null +++ b/integrations/goose-agent-middleware/tests/fixtures/g1_gate_calls.jsonl @@ -0,0 +1,3 @@ +{"tool": "write_file", "path": "/Users/sellers/scratch/g1-sandbox/allowed/gate.txt", "protocol": "2026-07-28", "meta": {"progress_token": 1, "agent-session-id": "20261024_4", "agent-working-dir": "/Users/sellers/scratch/g1-sandbox", "agent-tool-call-request-id": "tool_CpDpciFxHvt0zB4Vgyo0wXHe", "io.modelcontextprotocol/protocolVersion": "2026-07-28", "io.modelcontextprotocol/clientInfo": {"name": "goose-cli", "version": "1.56.0"}, "io.modelcontextprotocol/clientCapabilities": {"extensions": {}, "roots": {}, "elicitation": {}}}, "decision": "dropped_mid_stream", "ts": 1792826481.62643} +{"tool": "write_file", "path": "/Users/sellers/scratch/g1-sandbox/allowed/gate.txt", "protocol": "2026-07-28", "meta": {"progress_token": 2, "agent-session-id": "20261024_4", "agent-working-dir": "/Users/sellers/scratch/g1-sandbox", "agent-tool-call-request-id": "tool_Cn0pLB5b7j1weRQrhby6UOhB", "io.modelcontextprotocol/protocolVersion": "2026-07-28", "io.modelcontextprotocol/clientInfo": {"name": "goose-cli", "version": "1.56.0"}, "io.modelcontextprotocol/clientCapabilities": {"extensions": {}, "roots": {}, "elicitation": {}}}, "decision": "allow", "ts": 1792826485.893873} +{"tool": "read_text_file", "path": "/Users/sellers/scratch/g1-sandbox/allowed/gate.txt", "protocol": "2026-07-28", "meta": {"progress_token": 3, "agent-session-id": "20261024_4", "agent-working-dir": "/Users/sellers/scratch/g1-sandbox", "agent-tool-call-request-id": "tool_KAT1peagg5d6MBVBcMhucmZy", "io.modelcontextprotocol/protocolVersion": "2026-07-28", "io.modelcontextprotocol/clientInfo": {"name": "goose-cli", "version": "1.56.0"}, "io.modelcontextprotocol/clientCapabilities": {"extensions": {}, "roots": {}, "elicitation": {}}}, "ts": 1792826493.060245} diff --git a/integrations/goose-agent-middleware/tests/test_core.py b/integrations/goose-agent-middleware/tests/test_core.py new file mode 100644 index 00000000..00eb3cda --- /dev/null +++ b/integrations/goose-agent-middleware/tests/test_core.py @@ -0,0 +1,395 @@ +"""Unit tests. Every value comes from the real G1 gate recording or is made at runtime.""" + +import asyncio +import json +import secrets +import sys +import uuid +from pathlib import Path +from types import SimpleNamespace + +import pytest +from mcp.server.mcpserver.exceptions import ToolError + +from goose_amw import hook, permits +from goose_amw.canonical import args_hash, canonical_json, idempotency_key, sha256_hex +from goose_amw.ledger import DEFAULT_WINDOW_SECONDS, Ledger +from goose_amw.proxy import ( + Config, + Gateway, + UpstreamError, + UpstreamNotDispatched, + UpstreamUnknown, + stdio_upstream, + verify_caller, +) +from goose_amw.receipts import ReceiptSigner, load_or_create_key, verify + +FIXTURE = Path(__file__).parent / "fixtures" / "g1_gate_calls.jsonl" +RECORDS = [json.loads(line) for line in FIXTURE.read_text().splitlines()] +WRITES = [r for r in RECORDS if r["tool"] == "write_file"] +READ = next(r for r in RECORDS if r["tool"] != "write_file") +W = WRITES[0] +SESSION = W["meta"]["agent-session-id"] +ARGS = {"path": W["path"]} # the only argument field the recording holds + + +def rid(record): + return record["meta"]["agent-tool-call-request-id"] + + +def fresh(): + return str(uuid.uuid4()) + + +def test_canonical_hash_ignores_key_order(): + meta = W["meta"] + reordered = dict(reversed(list(meta.items()))) + assert canonical_json(meta) == canonical_json(reordered) + assert args_hash(meta) == args_hash(reordered) + + +def test_key_changes_with_each_input_and_never_uses_request_id(tmp_path): + caller = fresh() + k = idempotency_key(caller=caller, tool=W["tool"], arguments=ARGS, session_id=SESSION) + assert rid(W) not in k + assert k == idempotency_key(caller=caller, tool=W["tool"], arguments=ARGS, session_id=SESSION) + for variant in ( + {"caller": fresh(), "tool": W["tool"], "arguments": ARGS, "session_id": SESSION}, + {"caller": caller, "tool": READ["tool"], "arguments": ARGS, "session_id": SESSION}, + { + "caller": caller, + "tool": W["tool"], + "arguments": {"path": str(tmp_path)}, + "session_id": SESSION, + }, + {"caller": caller, "tool": W["tool"], "arguments": ARGS, "session_id": fresh()}, + ): + assert idempotency_key(**variant) != k + + +def test_recorded_gate_retry_maps_to_same_key_despite_new_request_id(): + """Real G1 gate evidence: Goose sent a new request id on the model's retry.""" + assert len(WRITES) == 2 + assert rid(WRITES[0]) != rid(WRITES[1]) + assert all(w["protocol"] == "2026-07-28" for w in WRITES) + caller = fresh() + keys = { + idempotency_key( + caller=caller, + tool=w["tool"], + arguments={"path": w["path"]}, + session_id=w["meta"]["agent-session-id"], + ) + for w in WRITES + } + assert len(keys) == 1 + + +def _gate_key(): + return idempotency_key(caller=fresh(), tool=W["tool"], arguments=ARGS, session_id=SESSION) + + +def test_ledger_replays_inside_window_and_reruns_after(tmp_path): + led = Ledger(tmp_path / "l.db") + key, ahash, t0 = _gate_key(), args_hash(ARGS), W["ts"] + result = fresh() + assert ( + led.claim(key, tool=W["tool"], args_hash=ahash, request_id=rid(W), now=t0).action == "run" + ) + led.finish(key, "completed", result) + retry = led.claim( + key, tool=W["tool"], args_hash=ahash, request_id=rid(WRITES[1]), now=WRITES[1]["ts"] + ) + assert (retry.action, retry.result) == ("replay", result) + later = t0 + DEFAULT_WINDOW_SECONDS + 1 + assert ( + led.claim(key, tool=W["tool"], args_hash=ahash, request_id=fresh(), now=later).action + == "run" + ) + assert led.get(key)["executions"] == 2 + + +def test_failed_unknown_is_never_rerun(tmp_path): + led = Ledger(tmp_path / "l.db") + key, ahash, t0 = _gate_key(), args_hash(ARGS), W["ts"] + led.claim(key, tool=W["tool"], args_hash=ahash, request_id=rid(W), now=t0) + led.finish(key, "failed_unknown") + for now in (WRITES[1]["ts"], t0 + 10 * DEFAULT_WINDOW_SECONDS): + claim = led.claim(key, tool=W["tool"], args_hash=ahash, request_id=rid(WRITES[1]), now=now) + assert claim.action == "refuse" + assert led.get(key)["executions"] == 1 + + +def test_crash_in_flight_becomes_failed_unknown(tmp_path): + led = Ledger(tmp_path / "l.db") + key = _gate_key() + led.claim(key, tool=W["tool"], args_hash=args_hash(ARGS), request_id=rid(W)) + assert Ledger(tmp_path / "l.db").mark_stale_in_flight_unknown() == 1 + assert led.get(key)["status"] == "failed_unknown" + + +def test_definite_failure_may_rerun(tmp_path): + led = Ledger(tmp_path / "l.db") + key = _gate_key() + led.claim(key, tool=W["tool"], args_hash=args_hash(ARGS), request_id=rid(W)) + led.finish(key, "failed") + claim = led.claim(key, tool=W["tool"], args_hash=args_hash(ARGS), request_id=rid(WRITES[1])) + assert claim.action == "run" + + +def test_window_must_be_positive(tmp_path): + with pytest.raises(ValueError): + Ledger(tmp_path / "l.db", window_seconds=0) + + +def _policy(tmp_path): + allowed = tmp_path / fresh() + allowed.mkdir() + return allowed, {"rules": [{"id": fresh(), "tool": W["tool"], "folder": str(allowed)}]} + + +def test_permits(tmp_path): + allowed, policy = _policy(tmp_path) + inside = str(allowed / fresh()) + prefixed = f"amw__{W['tool']}" + assert permits.check(policy, prefixed, {"path": inside}).allowed + assert not permits.check(policy, prefixed, {"path": f"{allowed}/../{fresh()}"}).allowed + assert not permits.check(policy, prefixed, {"path": str(tmp_path / fresh())}).allowed + assert not permits.check(policy, f"amw__{READ['tool']}", {"path": inside}).allowed + assert not permits.check(policy, W["tool"], {}).allowed + + +def test_hook_blocks_with_reason_and_allows(tmp_path): + allowed, policy = _policy(tmp_path) + name = f"amw__{W['tool']}" + out = hook.decide({"tool_name": name, "tool_input": {"path": str(tmp_path / fresh())}}, policy) + assert out["decision"] == "block" and "no permit" in out["reason"] + ok = {"tool_name": name, "tool_input": {"path": str(allowed / fresh())}} + assert hook.decide(ok, policy) == {"decision": "allow"} + + +def test_receipts_are_signed_and_tamper_evident(tmp_path): + signer = ReceiptSigner(load_or_create_key(tmp_path / "k.pem"), tmp_path / "r.jsonl") + r = signer.issue( + tool=W["tool"], args_hash=args_hash(ARGS), decision="allow", outcome="completed" + ) + assert {"tool", "args_hash", "decision", "issued_at"} <= r["body"].keys() + assert verify(r, signer.public_key_b64) + r["body"]["decision"] = "deny" + assert not verify(r, signer.public_key_b64) + assert (tmp_path / "k.pem").stat().st_mode & 0o077 == 0 + + +def test_write_without_session_id_is_refused_and_not_run(tmp_path): + allowed, policy = _policy(tmp_path) + token = secrets.token_hex(24) + cfg = Config( + policy=policy, + state_dir=tmp_path / "state", + callers={sha256_hex(token): fresh()}, + fs_root=str(tmp_path), + ) + calls = [] + + async def upstream(tool, arguments): + calls.append(tool) + return fresh() + + gw = Gateway(cfg, upstream=upstream) + ctx = SimpleNamespace( + request_context=SimpleNamespace(meta={}), + headers={"authorization": f"Bearer {token}"}, + protocol_version=None, + ) + args = {"path": str(allowed / fresh()), "content": fresh()} + with pytest.raises(ToolError, match="agent-session-id"): + asyncio.run(gw.handle(W["tool"], args, ctx, write=True)) + assert calls == [] + receipt = json.loads((cfg.state_dir / "receipts.jsonl").read_text().splitlines()[-1]) + assert receipt["body"]["decision"] == "refuse" + assert receipt["body"]["outcome"] == "not_run" + assert verify(receipt, gw.signer.public_key_b64) + + +def _gateway(tmp_path, upstream): + allowed, policy = _policy(tmp_path) + token = secrets.token_hex(24) + cfg = Config( + policy=policy, + state_dir=tmp_path / "state", + callers={sha256_hex(token): fresh()}, + fs_root=str(tmp_path), + ) + ctx = SimpleNamespace( + request_context=SimpleNamespace(meta={"agent-session-id": fresh()}), + headers={"authorization": f"Bearer {token}"}, + protocol_version=None, + ) + return ( + Gateway(cfg, upstream=upstream), + ctx, + {"path": str(allowed / fresh()), "content": fresh()}, + ) + + +def test_verify_caller_rejects_missing_and_wrong_tokens(): + token = secrets.token_hex(24) + caller = fresh() + callers = {sha256_hex(token): caller} + assert verify_caller({"authorization": f"Bearer {token}"}, callers) == caller + for headers in ( + None, + {}, + {"authorization": token}, + {"authorization": f"Bearer {secrets.token_hex(24)}"}, + {"authorization": f"Bearer {token[:-1]}"}, + ): + with pytest.raises(ToolError, match="caller not verified"): + verify_caller(headers, callers) + + +def test_unverified_caller_never_reaches_upstream(tmp_path): + calls = [] + + async def upstream(tool, arguments): + calls.append(tool) + return fresh() + + gw, ctx, args = _gateway(tmp_path, upstream) + ctx.headers = {"authorization": f"Bearer {secrets.token_hex(24)}"} + with pytest.raises(ToolError, match="caller not verified"): + asyncio.run(gw.handle(W["tool"], args, ctx, write=True)) + assert calls == [] + + +def test_pre_dispatch_failure_leaves_key_untouched_and_retryable(tmp_path): + attempts = [] + + async def upstream(tool, arguments): + attempts.append(tool) + if len(attempts) == 1: + raise UpstreamNotDispatched(fresh()) + return fresh() + + gw, ctx, args = _gateway(tmp_path, upstream) + with pytest.raises(ToolError, match="safe to retry"): + asyncio.run(gw.handle(W["tool"], args, ctx, write=True)) + assert gw.ledger._db.execute("SELECT COUNT(*) FROM calls").fetchone()[0] == 0 + asyncio.run(gw.handle(W["tool"], args, ctx, write=True)) + rows = gw.ledger._db.execute("SELECT status, executions FROM calls").fetchall() + assert rows == [("completed", 1)] + + +def test_release_restores_the_previous_row(tmp_path): + led = Ledger(tmp_path / "l.db") + key = _gate_key() + led.claim(key, tool=W["tool"], args_hash=args_hash(ARGS), request_id=rid(W)) + led.finish(key, "failed") + before = led.get(key) + claim = led.claim(key, tool=W["tool"], args_hash=args_hash(ARGS), request_id=rid(WRITES[1])) + led.release(key, claim) + assert led.get(key) == before + + +def test_server_that_dies_before_initialize_is_not_dispatched(): + from mcp.client.stdio import StdioServerParameters + + call = stdio_upstream(StdioServerParameters(command=sys.executable, args=["-c", ""])) + with pytest.raises(UpstreamNotDispatched): + asyncio.run(asyncio.wait_for(call(W["tool"], ARGS), 30)) + + +def test_handle_denies_unpermitted_path_without_reaching_upstream(tmp_path): + calls = [] + + async def upstream(tool, arguments): + calls.append(tool) + return fresh() + + gw, ctx, args = _gateway(tmp_path, upstream) + for path in (str(tmp_path / fresh()), f"{Path(args['path']).parent}/../{fresh()}"): + with pytest.raises(ToolError, match="permit denied"): + asyncio.run(gw.handle(W["tool"], {**args, "path": path}, ctx, write=True)) + assert calls == [] + assert gw.ledger._db.execute("SELECT COUNT(*) FROM calls").fetchone()[0] == 0 + lines = (tmp_path / "state" / "receipts.jsonl").read_text().splitlines() + for line in lines: + receipt = json.loads(line) + assert receipt["body"]["decision"] == "deny" + assert receipt["body"]["outcome"] == "not_run" + assert verify(receipt, gw.signer.public_key_b64) + assert len(lines) == 2 + + +def test_write_error_result_is_failed_unknown_and_not_rerun(tmp_path): + calls = [] + + async def upstream(tool, arguments): + calls.append(tool) + raise UpstreamError(fresh()) + + gw, ctx, args = _gateway(tmp_path, upstream) + with pytest.raises(ToolError, match="partly applied"): + asyncio.run(gw.handle(W["tool"], args, ctx, write=True)) + with pytest.raises(ToolError, match="will not run again"): + asyncio.run(gw.handle(W["tool"], args, ctx, write=True)) + assert calls == [W["tool"]] + rows = gw.ledger._db.execute("SELECT status FROM calls").fetchall() + assert rows == [("failed_unknown",)] + + +def test_relative_and_home_paths_are_denied(tmp_path, monkeypatch): + allowed, policy = _policy(tmp_path) + monkeypatch.chdir(allowed) + name = fresh() + for raw in (name, f"./{name}", f"~/{name}"): + decision = permits.check(policy, W["tool"], {"path": raw}) + assert not decision.allowed + assert "must be absolute" in decision.reason + assert permits.check(policy, W["tool"], {"path": str(allowed / name)}).allowed + + +def _fake_stdio(monkeypatch, result, *, fail_on_exit): + import contextlib + + import mcp.client.session + import mcp.client.stdio + + @contextlib.asynccontextmanager + async def client(params): + yield (None, None) + if fail_on_exit: + raise OSError(fresh()) + + class Session: + def __init__(self, read, write): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def initialize(self): + return None + + async def call_tool(self, tool, arguments): + return result + + monkeypatch.setattr(mcp.client.stdio, "stdio_client", client) + monkeypatch.setattr(mcp.client.session, "ClientSession", Session) + + +def test_teardown_error_after_result_keeps_the_result(monkeypatch): + text = fresh() + result = SimpleNamespace(content=[SimpleNamespace(text=text)], is_error=False) + _fake_stdio(monkeypatch, result, fail_on_exit=True) + assert asyncio.run(stdio_upstream(object())(W["tool"], ARGS)) == text + + +def test_teardown_error_without_result_is_unknown(monkeypatch): + _fake_stdio(monkeypatch, None, fail_on_exit=True) + with pytest.raises(UpstreamUnknown): + asyncio.run(stdio_upstream(object())(W["tool"], ARGS))