From 0c72225311179e1ee02fa0b6534dbeea48d3500e Mon Sep 17 00:00:00 2001 From: KunjShah95 Date: Sun, 26 Jul 2026 22:58:12 +0530 Subject: [PATCH 1/3] fix: corrupt JSON store files must not crash server startup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The small file-backed stores (Inbox, inbox routing, mention threads, channel subscriptions, unrouted/dead-letter, unattended flags, per-persona/per-session connection maps, and user risk overrides) each loaded their JSON file with an unguarded json.loads inside __init__ — and every one is constructed eagerly in SessionManager.__init__. So a single truncated or malformed file turned into an exception that aborted server startup entirely, recoverable only by finding and deleting the file by hand. Several of the same stores also wrote their files non-atomically (write_text in place), which is exactly what produces the corrupt file after a crash/kill/disk-full mid-write. This generalizes the resilient pattern already used inline by ChannelBuffer and WorkspaceTrustStore into two shared helpers (coworker/jsonstore.py): - read_json(): missing/corrupt file -> caller default, moving the bad file aside to .corrupt so it is neither lost nor re-hit on the next boot. - write_json_atomic(): temp sibling + os.replace, atomic on POSIX and Windows. Every affected store now uses them. Adds tests covering the helpers and each store's construct-from-corrupt-file path plus a save round-trip. Related to the WakeStore-specific report in #158 (that fix, #159, covers only selfwake.py); this addresses the same failure mode across the rest of the persistence layer. Co-Authored-By: Claude Opus 4.8 --- coworker/connections.py | 41 ++++----- coworker/inbox.py | 19 ++--- coworker/inbox_routing.py | 36 ++++---- coworker/jsonstore.py | 79 ++++++++++++++++++ coworker/mentions.py | 16 ++-- coworker/overrides.py | 25 ++---- coworker/subscriptions.py | 15 ++-- coworker/unattended.py | 12 +-- coworker/unrouted.py | 16 ++-- tests/test_json_store_resilience.py | 124 ++++++++++++++++++++++++++++ 10 files changed, 266 insertions(+), 117 deletions(-) create mode 100644 coworker/jsonstore.py create mode 100644 tests/test_json_store_resilience.py diff --git a/coworker/connections.py b/coworker/connections.py index 9088b53f..1a81bc21 100644 --- a/coworker/connections.py +++ b/coworker/connections.py @@ -22,11 +22,12 @@ from __future__ import annotations -import json import threading from pathlib import Path from typing import Optional +from .jsonstore import read_json, write_json_atomic + class PersonaConnectionStore: """``{persona_id: {connector: bool}}`` — the per-persona default on/off for each connector.""" @@ -38,21 +39,14 @@ def __init__(self, path: Optional[str | Path] = None) -> None: self._load() def _load(self) -> None: - if self.path and self.path.is_file(): - data = json.loads(self.path.read_text(encoding="utf-8")) - self._rows = { - pid: {str(c): bool(v) for c, v in (row or {}).items()} - for pid, row in data.get("personas", {}).items() - } + data = read_json(self.path, {}) or {} + self._rows = { + pid: {str(c): bool(v) for c, v in (row or {}).items()} + for pid, row in data.get("personas", {}).items() + } def _save(self) -> None: - if not self.path: - return - self.path.parent.mkdir(parents=True, exist_ok=True) - self.path.write_text( - json.dumps({"personas": self._rows}, indent=2), - encoding="utf-8", - ) + write_json_atomic(self.path, {"personas": self._rows}) # -- queries ---------------------------------------------------------------- def get(self, persona_id: str) -> dict[str, bool]: @@ -111,21 +105,14 @@ def __init__(self, path: Optional[str | Path] = None) -> None: self._load() def _load(self) -> None: - if self.path and self.path.is_file(): - data = json.loads(self.path.read_text(encoding="utf-8")) - self._rows = { - sid: {str(c): bool(v) for c, v in (row or {}).items()} - for sid, row in data.get("sessions", {}).items() - } + data = read_json(self.path, {}) or {} + self._rows = { + sid: {str(c): bool(v) for c, v in (row or {}).items()} + for sid, row in data.get("sessions", {}).items() + } def _save(self) -> None: - if not self.path: - return - self.path.parent.mkdir(parents=True, exist_ok=True) - self.path.write_text( - json.dumps({"sessions": self._rows}, indent=2), - encoding="utf-8", - ) + write_json_atomic(self.path, {"sessions": self._rows}) # -- queries ---------------------------------------------------------------- def get(self, session_id: str) -> dict[str, bool]: diff --git a/coworker/inbox.py b/coworker/inbox.py index 924e707a..fc96e0f4 100644 --- a/coworker/inbox.py +++ b/coworker/inbox.py @@ -22,6 +22,8 @@ from pathlib import Path from typing import Any, Optional +from .jsonstore import read_json, write_json_atomic + KIND_APPROVAL = "approval" KIND_QUESTION = "question" KIND_NOTIFICATION = "notification" @@ -97,19 +99,14 @@ def __init__(self, path: Optional[str | Path] = None) -> None: # -- persistence ------------------------------------------------------------ def _load(self) -> None: - if self.path and self.path.is_file(): - data = json.loads(self.path.read_text(encoding="utf-8")) - for raw in data.get("items", []): - item = InboxItem(**raw) - self._items[item.id] = item + data = read_json(self.path, {}) or {} + for raw in data.get("items", []): + item = InboxItem(**raw) + self._items[item.id] = item def _save(self) -> None: - if not self.path: - return - self.path.parent.mkdir(parents=True, exist_ok=True) - self.path.write_text( - json.dumps({"items": [asdict(i) for i in self._items.values()]}, indent=2), - encoding="utf-8", + write_json_atomic( + self.path, {"items": [asdict(i) for i in self._items.values()]} ) # -- adding ----------------------------------------------------------------- diff --git a/coworker/inbox_routing.py b/coworker/inbox_routing.py index ef115e95..104c178f 100644 --- a/coworker/inbox_routing.py +++ b/coworker/inbox_routing.py @@ -11,13 +11,14 @@ from __future__ import annotations -import json import re import threading from dataclasses import asdict, dataclass from pathlib import Path from typing import Callable, Optional +from .jsonstore import read_json, write_json_atomic + DEFAULT_INBOX = "default" # Embeds the item id in a delivered message. Emitted as [ow:…] since the bot's rebrand # to OpenWorker (2026-07-22); the legacy [ocw:…] spelling stays parseable so replies to @@ -44,28 +45,21 @@ def __init__(self, path: Optional[str | Path] = None) -> None: self._load() def _load(self) -> None: - if self.path and self.path.is_file(): - data = json.loads(self.path.read_text(encoding="utf-8")) - for raw in data.get("bindings", []): - b = InboxBinding(**raw) - self._bindings[b.name] = b - self._persona_default = dict(data.get("persona_default", {})) - self._session_override = dict(data.get("session_override", {})) + data = read_json(self.path, {}) or {} + for raw in data.get("bindings", []): + b = InboxBinding(**raw) + self._bindings[b.name] = b + self._persona_default = dict(data.get("persona_default", {})) + self._session_override = dict(data.get("session_override", {})) def _save(self) -> None: - if not self.path: - return - self.path.parent.mkdir(parents=True, exist_ok=True) - self.path.write_text( - json.dumps( - { - "bindings": [asdict(b) for b in self._bindings.values()], - "persona_default": self._persona_default, - "session_override": self._session_override, - }, - indent=2, - ), - encoding="utf-8", + write_json_atomic( + self.path, + { + "bindings": [asdict(b) for b in self._bindings.values()], + "persona_default": self._persona_default, + "session_override": self._session_override, + }, ) # -- config ----------------------------------------------------------------- diff --git a/coworker/jsonstore.py b/coworker/jsonstore.py new file mode 100644 index 00000000..a65dcc44 --- /dev/null +++ b/coworker/jsonstore.py @@ -0,0 +1,79 @@ +"""Durable JSON persistence helpers shared by the file-backed stores. + +The tiny JSON-backed stores (Inbox, subscriptions, mention threads, connection +overrides, …) are all constructed in ``SessionManager.__init__``, and each one loads its +file eagerly there. That makes their read/write robustness a *startup* concern, and two +failure modes bite: + +- **A corrupt/truncated file must never block startup.** An unguarded ``json.loads`` in a + store's ``_load`` turns one unreadable file into a server that won't boot at all — + recoverable only by finding and deleting the file by hand. ``read_json`` returns the + caller's default instead, moving the bad file aside (``.corrupt``) so it is neither + lost nor silently re-hit on the next load. +- **A save must be atomic.** Writing in place leaves a half-written file if the process dies + mid-write — which is exactly what produces the corrupt file above. ``write_json_atomic`` + writes a sibling temp file and ``os.replace``s it into place (atomic on POSIX and Windows). + +This just generalizes the pattern already applied inline to ``ChannelBuffer`` and +``WorkspaceTrustStore``, so every store gets the same treatment instead of a few. +""" + +from __future__ import annotations + +import json +import logging +import os +from pathlib import Path +from typing import Any + +logger = logging.getLogger("coworker.jsonstore") + + +def read_json(path: str | Path | None, default: Any = None) -> Any: + """Load JSON from ``path``, tolerating a missing or unreadable file. + + Returns ``default`` when the path is ``None``, absent, or unreadable/corrupt. A corrupt + file is moved aside to ``.corrupt`` (best effort) so the bad data is preserved for + inspection and the next ``write_json_atomic`` starts from a clean slate rather than the + store re-hitting the same parse error on every boot. + """ + if path is None: + return default + p = Path(path) + if not p.is_file(): + return default + try: + return json.loads(p.read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + logger.warning("ignoring unreadable JSON at %s (%s); moving it aside", p, exc) + _quarantine(p) + return default + + +def _quarantine(path: Path) -> None: + try: + path.replace(path.with_name(path.name + ".corrupt")) + except OSError: + pass + + +def write_json_atomic( + path: str | Path | None, data: Any, *, indent: int | None = 2 +) -> None: + """Serialize ``data`` and write it to ``path`` atomically, creating parent dirs. + + A no-op when ``path`` is ``None`` (the in-memory-only store case). Writes to a sibling + temp file first, then ``os.replace``s it over the target, so a crash mid-write can never + leave a truncated file that would brick the next startup. + """ + if path is None: + return + p = Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + tmp = p.with_name(f".{p.name}.{os.getpid()}.tmp") + try: + tmp.write_text(json.dumps(data, indent=indent), encoding="utf-8") + os.replace(tmp, p) + except BaseException: + tmp.unlink(missing_ok=True) + raise diff --git a/coworker/mentions.py b/coworker/mentions.py index cd4a8717..bbce22ef 100644 --- a/coworker/mentions.py +++ b/coworker/mentions.py @@ -15,12 +15,13 @@ from __future__ import annotations -import json import threading from dataclasses import asdict, dataclass from pathlib import Path from typing import Optional +from .jsonstore import read_json, write_json_atomic + @dataclass class MentionThread: @@ -37,18 +38,11 @@ def __init__(self, path: Optional[str | Path] = None) -> None: self._load() def _load(self) -> None: - if self.path and self.path.is_file(): - data = json.loads(self.path.read_text(encoding="utf-8")) - self._threads = [MentionThread(**raw) for raw in data.get("threads", [])] + data = read_json(self.path, {}) or {} + self._threads = [MentionThread(**raw) for raw in data.get("threads", [])] def _save(self) -> None: - if not self.path: - return - self.path.parent.mkdir(parents=True, exist_ok=True) - self.path.write_text( - json.dumps({"threads": [asdict(t) for t in self._threads]}, indent=2), - encoding="utf-8", - ) + write_json_atomic(self.path, {"threads": [asdict(t) for t in self._threads]}) # -- mutations -------------------------------------------------------------- def set(self, thread_target: str, session_id: str, channel: str) -> MentionThread: diff --git a/coworker/overrides.py b/coworker/overrides.py index 0abb2cf8..8925c99b 100644 --- a/coworker/overrides.py +++ b/coworker/overrides.py @@ -11,12 +11,12 @@ from __future__ import annotations -import json from dataclasses import dataclass from fnmatch import fnmatchcase from pathlib import Path from typing import Callable, Optional +from .jsonstore import read_json, write_json_atomic from .risk import RiskClass @@ -39,32 +39,19 @@ def __init__(self, path: Optional[str | Path] = None) -> None: self._rules: list[_Rule] = self._load() def _load(self) -> list[_Rule]: - if not (self.path and self.path.is_file()): - return [] - data = json.loads(self.path.read_text(encoding="utf-8")) + data = read_json(self.path, {}) or {} rules = [] for r in data.get("rules", []): try: rules.append(_Rule(str(r["pattern"]), RiskClass(str(r["risk"])))) - except (KeyError, ValueError): + except (KeyError, ValueError, TypeError): continue # skip malformed rules rather than failing the whole store return rules def save(self) -> None: - if not self.path: - return - self.path.parent.mkdir(parents=True, exist_ok=True) - self.path.write_text( - json.dumps( - { - "rules": [ - {"pattern": r.pattern, "risk": r.risk.value} - for r in self._rules - ] - }, - indent=2, - ), - encoding="utf-8", + write_json_atomic( + self.path, + {"rules": [{"pattern": r.pattern, "risk": r.risk.value} for r in self._rules]}, ) def set_rule(self, pattern: str, risk: RiskClass | str) -> None: diff --git a/coworker/subscriptions.py b/coworker/subscriptions.py index f8b5d33a..3ec5961d 100644 --- a/coworker/subscriptions.py +++ b/coworker/subscriptions.py @@ -25,6 +25,8 @@ from pathlib import Path from typing import Optional +from .jsonstore import read_json, write_json_atomic + @dataclass class Subscription: @@ -42,17 +44,12 @@ def __init__(self, path: Optional[str | Path] = None) -> None: self._load() def _load(self) -> None: - if self.path and self.path.is_file(): - data = json.loads(self.path.read_text(encoding="utf-8")) - self._subs = [Subscription(**raw) for raw in data.get("subscriptions", [])] + data = read_json(self.path, {}) or {} + self._subs = [Subscription(**raw) for raw in data.get("subscriptions", [])] def _save(self) -> None: - if not self.path: - return - self.path.parent.mkdir(parents=True, exist_ok=True) - self.path.write_text( - json.dumps({"subscriptions": [asdict(s) for s in self._subs]}, indent=2), - encoding="utf-8", + write_json_atomic( + self.path, {"subscriptions": [asdict(s) for s in self._subs]} ) # -- mutations -------------------------------------------------------------- diff --git a/coworker/unattended.py b/coworker/unattended.py index 6b60873f..8224e9d0 100644 --- a/coworker/unattended.py +++ b/coworker/unattended.py @@ -8,25 +8,21 @@ from __future__ import annotations -import json import threading from pathlib import Path from typing import Optional +from .jsonstore import read_json, write_json_atomic + class UnattendedRegistry: def __init__(self, path: Optional[str | Path] = None) -> None: self.path = Path(path) if path else None self._lock = threading.Lock() - self._flags: dict[str, bool] = {} - if self.path and self.path.is_file(): - self._flags = dict(json.loads(self.path.read_text(encoding="utf-8"))) + self._flags: dict[str, bool] = dict(read_json(self.path, {}) or {}) def _save(self) -> None: - if not self.path: - return - self.path.parent.mkdir(parents=True, exist_ok=True) - self.path.write_text(json.dumps(self._flags, indent=2), encoding="utf-8") + write_json_atomic(self.path, self._flags) def is_unattended(self, session_id: str) -> bool: return bool(self._flags.get(session_id, False)) diff --git a/coworker/unrouted.py b/coworker/unrouted.py index e17328ac..38103916 100644 --- a/coworker/unrouted.py +++ b/coworker/unrouted.py @@ -11,13 +11,14 @@ from __future__ import annotations -import json import threading import time from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Optional +from .jsonstore import read_json, write_json_atomic + @dataclass class UnroutedItem: @@ -37,18 +38,11 @@ def __init__(self, path: Optional[str | Path] = None, *, cap: int = 200) -> None self._load() def _load(self) -> None: - if self.path and self.path.is_file(): - data = json.loads(self.path.read_text(encoding="utf-8")) - self._items = [UnroutedItem(**raw) for raw in data.get("items", [])] + data = read_json(self.path, {}) or {} + self._items = [UnroutedItem(**raw) for raw in data.get("items", [])] def _save(self) -> None: - if not self.path: - return - self.path.parent.mkdir(parents=True, exist_ok=True) - self.path.write_text( - json.dumps({"items": [asdict(i) for i in self._items]}, indent=2), - encoding="utf-8", - ) + write_json_atomic(self.path, {"items": [asdict(i) for i in self._items]}) def record(self, source: str, sender: str, text: str, reason: str) -> UnroutedItem: item = UnroutedItem( diff --git a/tests/test_json_store_resilience.py b/tests/test_json_store_resilience.py new file mode 100644 index 00000000..1e36e3e6 --- /dev/null +++ b/tests/test_json_store_resilience.py @@ -0,0 +1,124 @@ +"""A corrupt on-disk JSON file must never crash a store's construction (startup). + +Every file-backed store is built eagerly in ``SessionManager.__init__`` and loads its file +there, so an unreadable file used to raise straight out of the constructor and take the whole +server down. These tests pin the resilient contract for the shared helpers and for each store: +a garbage file loads as empty (and is quarantined), and a normal save round-trips atomically. +""" + +from __future__ import annotations + +from coworker.connections import PersonaConnectionStore, SessionConnectionStore +from coworker.inbox import InboxStore +from coworker.inbox_routing import DEFAULT_INBOX, InboxRouting +from coworker.jsonstore import read_json, write_json_atomic +from coworker.mentions import MentionSessionStore +from coworker.overrides import RiskOverrideStore +from coworker.risk import RiskClass +from coworker.subscriptions import SubscriptionStore +from coworker.unattended import UnattendedRegistry +from coworker.unrouted import UnroutedStore + +CORRUPT = '{"items": [ {"broken": ' # truncated mid-object + + +# -- jsonstore helpers ---------------------------------------------------------- +def test_read_json_missing_returns_default(tmp_path): + assert read_json(tmp_path / "nope.json", {"a": 1}) == {"a": 1} + assert read_json(None, []) == [] + + +def test_read_json_valid_round_trips(tmp_path): + p = tmp_path / "ok.json" + write_json_atomic(p, {"x": [1, 2, 3]}) + assert read_json(p) == {"x": [1, 2, 3]} + + +def test_read_json_corrupt_returns_default_and_quarantines(tmp_path): + p = tmp_path / "bad.json" + p.write_text(CORRUPT, encoding="utf-8") + assert read_json(p, {"fallback": True}) == {"fallback": True} + # The bad file is moved aside, not silently clobbered or left to re-crash next boot. + assert not p.exists() + assert (tmp_path / "bad.json.corrupt").read_text(encoding="utf-8") == CORRUPT + + +def test_write_json_atomic_leaves_no_temp_file(tmp_path): + p = tmp_path / "out.json" + write_json_atomic(p, {"k": "v"}) + assert read_json(p) == {"k": "v"} + # No stray sibling temp files left behind. + assert [f.name for f in tmp_path.iterdir()] == ["out.json"] + + +def test_write_json_atomic_none_path_is_noop(): + write_json_atomic(None, {"anything": 1}) # must not raise + + +# -- every store survives a corrupt file at construction ------------------------ +def _write_corrupt(path): + path.write_text(CORRUPT, encoding="utf-8") + return path + + +def test_inbox_store_survives_corrupt_file(tmp_path): + p = _write_corrupt(tmp_path / "inbox.json") + store = InboxStore(p) # must not raise + assert store.list() == [] + assert (tmp_path / "inbox.json.corrupt").exists() + + +def test_inbox_routing_survives_corrupt_file(tmp_path): + p = _write_corrupt(tmp_path / "inbox_routing.json") + routing = InboxRouting(p) + assert routing.route_for("s1") == DEFAULT_INBOX + + +def test_mentions_store_survives_corrupt_file(tmp_path): + p = _write_corrupt(tmp_path / "mentions.json") + assert MentionSessionStore(p).all() == [] + + +def test_subscription_store_survives_corrupt_file(tmp_path): + p = _write_corrupt(tmp_path / "subs.json") + assert SubscriptionStore(p).all() == [] + + +def test_unrouted_store_survives_corrupt_file(tmp_path): + p = _write_corrupt(tmp_path / "unrouted.json") + assert UnroutedStore(p).list() == [] + + +def test_unattended_registry_survives_corrupt_file(tmp_path): + p = _write_corrupt(tmp_path / "unattended.json") + assert UnattendedRegistry(p).sessions() == [] + + +def test_connection_stores_survive_corrupt_file(tmp_path): + persona = PersonaConnectionStore(_write_corrupt(tmp_path / "persona.json")) + session = SessionConnectionStore(_write_corrupt(tmp_path / "session.json")) + assert persona.get("p1") == {} + assert session.get("s1") == {} + + +def test_override_store_survives_corrupt_file(tmp_path): + p = _write_corrupt(tmp_path / "overrides.json") + assert RiskOverrideStore(p).resolve("some_tool") is None + + +# -- round-trip after recovery: a corrupt file heals on the next save ---------- +def test_store_save_after_corrupt_load_round_trips(tmp_path): + p = _write_corrupt(tmp_path / "subs.json") + store = SubscriptionStore(p) # corrupt file quarantined, in-memory empty + store.subscribe("session-1", "slack:C123") + # A fresh store reads back exactly what was saved — the file is valid again. + assert [(s.session_id, s.channel) for s in SubscriptionStore(p).all()] == [ + ("session-1", "slack:C123") + ] + + +def test_override_store_save_round_trips(tmp_path): + p = tmp_path / "overrides.json" + store = RiskOverrideStore(p) + store.set_rule("mcp__notion__*", RiskClass.READ) + assert RiskOverrideStore(p).resolve("mcp__notion__search") == RiskClass.READ From ab3b447361e2d1c366ac5d04fc4adb3e703abd91 Mon Sep 17 00:00:00 2001 From: KunjShah95 Date: Mon, 27 Jul 2026 11:06:53 +0530 Subject: [PATCH 2/3] Implement fan-out subagent guard middleware Add GuardMiddleware layer in the engine's authorization flow to enforce safety policies for tool calls, limit concurrent subagents, and log all decisions for audit. - GuardRuleSet: load YAML rules, evaluate tool calls with command match patterns (contains/prefix/regex), argument matching, and fan-out limits - GuardLogger: dedicated audit log with timestamped ALLOW/DENY entries - GuardConfigLoader: load/cache YAML config with change detection for hot reload - GuardMiddleware: layered authorization combining PermissionEngine, GuardRuleSet, and GuardLogger; fail-safe on exceptions - Engine integration: inject guard_middleware parameter into TurnEngine, route tool authorization through middleware, track start/end for fan-out counting - Agent wiring: build GuardMiddleware for code agents with config path configurable via COWORKER_GUARD_CONFIG env var - Default config/guard.yaml with safety rules (block recursive deletes, limit explorer fanout to 3 concurrent) - 31 unit tests for all guard components (all passing) --- .gitignore | 4 + config/guard.yaml | 73 +++++ coworker/CLAUDE.md | 209 ++++++++++++++ coworker/agent.py | 21 ++ coworker/engine.py | 22 +- coworker/guard/__init__.py | 15 + coworker/guard/config_loader.py | 71 +++++ coworker/guard/logger.py | 93 ++++++ coworker/guard/middleware.py | 164 +++++++++++ coworker/guard/ruleset.py | 193 +++++++++++++ tests/test_guard.py | 494 ++++++++++++++++++++++++++++++++ 11 files changed, 1356 insertions(+), 3 deletions(-) create mode 100644 config/guard.yaml create mode 100644 coworker/CLAUDE.md create mode 100644 coworker/guard/__init__.py create mode 100644 coworker/guard/config_loader.py create mode 100644 coworker/guard/logger.py create mode 100644 coworker/guard/middleware.py create mode 100644 coworker/guard/ruleset.py create mode 100644 tests/test_guard.py diff --git a/.gitignore b/.gitignore index df96f790..daf55dd6 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,7 @@ __pycache__/ build/ dist/ .coverage +docs/superpowers/specs/ + +# Guard middleware audit log (generated at runtime) +guard.log \ No newline at end of file diff --git a/config/guard.yaml b/config/guard.yaml new file mode 100644 index 00000000..0933e5e1 --- /dev/null +++ b/config/guard.yaml @@ -0,0 +1,73 @@ +# Guard rules — safety policies for tool calls. +# +# Rules are evaluated in order. The FIRST matching rule wins. +# When no rule matches, the base PermissionEngine decision applies. +# +# Each rule can match on: +# - tool name (required) +# - command patterns: contains / prefix / regex +# - argument key-value pairs +# - fan-out concurrency limits (for subagent/spawn tools) +# +# Actions: allow | deny | limit + +rules: + # --- Exec (shell) safety --- + + - name: block-recursive-delete + tool: run_shell + match: + command: + contains: "rm -rf " + action: deny + reason: "Block recursive delete commands" + + - name: block-force-delete + tool: run_shell + match: + command: + contains: "rm -f" + action: deny + reason: "Block force delete commands" + + - name: block-destructive-disk-wipe + tool: run_shell + match: + command: + contains: "dd if=" + action: deny + reason: "Block raw disk write commands" + + - name: block-chmod-recursive + tool: run_shell + match: + command: + contains: "chmod -R" + action: deny + reason: "Block recursive permission changes" + + - name: block-chown-recursive + tool: run_shell + match: + command: + contains: "chown -R" + action: deny + reason: "Block recursive owner changes" + + - name: block-mkfs + tool: run_shell + match: + command: + contains: "mkfs" + action: deny + reason: "Block filesystem creation commands" + + # --- Subagent fan-out limits --- + # Limit concurrent explorer subagents to prevent resource exhaustion. + + - name: limit-explorer-fanout + tool: explore + fanout: + max_concurrent: 3 + action: deny + reason: "Limit concurrent explorer subagents" diff --git a/coworker/CLAUDE.md b/coworker/CLAUDE.md new file mode 100644 index 00000000..42805bcd --- /dev/null +++ b/coworker/CLAUDE.md @@ -0,0 +1,209 @@ +# coworker/ Claude Code Guide + +## Overview +This directory contains the Python backend for OpenWorker - the core agent engine that handles: +- Agent reasoning and decision making +- Tool execution and permission checking +- Memory and conversation storage +- Integration with external services via connectors +- Model provider management and routing +- Skill system and automation framework +- WebSocket server for GUI communication + +## Key Subsystems + +### Agent System (`agents/`) +- `base.py`: Base agent class defining the interface +- `chat.py`: Conversational agent for general assistance +- `code.py`: Coding assistant with development tool expertise +- `cowork.py`: General-purpose coworker agent +- `myhelper.py`: Personal assistant for individual tasks +- `registry.py`: Agent discovery and instantiation + +### Core Engine (`engine.py`) +- Main turn processing loop +- Agent interaction and tool execution coordination +- Permission checking and user approval workflows +- Action planning and execution tracking + +### Memory System (`memory/`) +- `sqlite_store.py`: Persistent conversation and context storage +- `base.py`: Memory interface definition +- `tools.py`: Memory-related agent tools + +### Connectors System (`connectors/`) +- **Adapters**: Real-time messaging platforms (Slack, Telegram) +- **Senders**: Stateless message sending to various services +- **Gateway**: Inbound webhook handling and routing +- **Relay Clients**: Cloud relay connections for managed services +- **Tool Definitions**: MCP-compatible tool specifications +- **Adapters**: Third-party service integrations (GitHub, Jira, etc.) + +### Model Providers (`providers/`) +- **Base Interface**: Standardized provider contract +- **Specific Providers**: OpenAI, Anthropic, Google Gemini, Ollama, etc. +- **Router**: Model routing, caching, and capabilities detection +- **Registry**: Provider discovery and configuration +- **Matrix**: Curated model recommendations and labels + +### Tools System (`tools/`) +- Individual tool implementations (file operations, web search, etc.) +- Tool permissions and risk assessment +- Tool composition and chaining capabilities +- Custom tool development framework + +### Skills System (`skills/`) +- Progressive capability disclosure mechanism +- Skill categorization and dependency tracking +- Skill activation and deactivation workflows +- Built-in skills for common agent capabilities + +### Automation System (`automation/`) +- Scheduled task management (`scheduler.py`) +- Persistent storage for automation rules (`store.py`) +- Automation-specific tools (`tools.py`) +- Data models for scheduled work (`models.py`) + +### Guardrails (`guard/`) +- `ruleset.py`: GuardRuleSet — load YAML rules, evaluate tool calls with command match patterns (contains/prefix/regex), argument matching, and fan-out concurrency limits +- `logger.py`: GuardLogger — dedicated audit log with timestamped ALLOW/DENY entries per tool call +- `config_loader.py`: GuardConfigLoader — load/cache YAML config with file-change detection for hot reload +- `middleware.py`: GuardMiddleware — layered authorization combining PermissionEngine, GuardRuleSet, and GuardLogger; fail-safe on exceptions (deny + needs_user) + +### Security and Permissions +- `permissions.py`: Main permission evaluation engine +- `risk.py`: Risk classification system for tools +- `secrets.py`: Secure credential storage and management +- `config.py`: Allowed command lists and security policies + +## Development Guidelines + +### Adding New Agents +1. Inherit from `agents.base.BaseAgent` +2. Implement required methods: `initialize()`, `step()`, `shutdown()` +3. Register in `agents.registry.AGENT_REGISTRY` +4. Follow existing patterns for prompt construction and tool usage +5. Add unit tests in `tests/test_*agent*.py` + +### Adding New Connectors +1. Choose the appropriate integration pattern: + - Real-time: Implement `BasePlatformAdapter` (like Slack/Telegram) + - Stateless: Extend `connectors.senders` + - Polling: Implement periodic checkers in `connectors.integration_tools` +2. Follow security patterns for credential handling +3. Add to `connectors.make_adapter()` factory function +4. Create corresponding tests in `tests/test_*connectors*.py` + +### Adding New Model Providers +1. Implement `ProviderClient` interface from `providers.base` +2. Add descriptor to `providers.registry.DESCRIPTORS` +3. Handle API key resolution following existing patterns +4. Implement both `complete()` and `stream()` methods +5. Add model capability detection via `providers.capabilities` +6. Add to `providers.matrix.MATRIX` if it meets curated criteria +7. Add comprehensive tests in `tests/test_*provider*.py` + +### Adding New Tools +1. Implement function following tool signature conventions +2. Add appropriate risk classification in `risk.py` +3. Register in the appropriate tool category +4. Add schema documentation for agent consumption +5. Write tests in `tests/test_tools_*.py` +6. Consider permission implications and test accordingly + +### Modifying Permissions System +1. Understand the `RiskClass` system in `risk.py` +2. Follow the principle of least privilege +3. Test both allowed and denied scenarios +4. Consider impact on different modes (PLAN, INTERACTIVE, AUTO) +5. Update corresponding tests in `tests/test_permissions_*.py` + +### Working with the Agent Loop +1. Understand the flow in `engine.py`: perceive → think → act → learn +2. Respect the permission system - never bypass user approval for consequential actions +3. Handle `needs_user` decisions properly by pausing for input +4. Manage state properly across turns using the memory system +5. Follow the tool result format for consistent agent consumption + +## Code Quality Standards + +### Type Hints +- Use Python 3.8+ type hints consistently +- Import typing constructs only when needed +- Use `from __future__ import annotations` for forward references +- Be explicit about Optional types and default values + +### Error Handling +- Catch specific exceptions rather than bare `except:` +- Provide meaningful error messages to users +- Log unexpected errors for debugging +- Fail gracefully and maintain system stability +- Use custom exception types when appropriate + +### Logging +- Use the standard `logging` module +- Follow existing logger naming conventions (`logging.getLogger("coworker.subsystem")`) +- Log at appropriate levels (DEBUG, INFO, WARNING, ERROR) +- Avoid logging sensitive information (tokens, passwords, etc.) + +### Performance Considerations +- Avoid blocking operations in async contexts +- Use efficient data structures and algorithms +- Consider caching for expensive operations +- Profile before optimizing - measure first +- Be mindful of memory usage in long-running processes + +### Testing Practices +- Write tests before implementing new features (TDD when possible) +- Achieve meaningful test coverage (>80% for new code) +- Test both happy paths and error conditions +- Use fixtures to reduce test setup duplication +- Mock external dependencies appropriately +- Test edge cases and boundary conditions + +## Common Patterns to Follow + +### Dependency Injection +- Prefer constructor injection over global state +- Make dependencies explicit in function signatures +- Use interfaces/abstract classes for pluggable components +- Keep classes focused and loosely coupled + +### Configuration Handling +- Use dataclasses for configuration objects +- Provide sensible defaults +- Validate configuration at startup +- Allow override via environment variables or user settings + +### Resource Management +- Properly initialize and clean up resources +- Use context managers (`with` statements) when applicable +- Handle async resource cleanup in `__aexit__` or similar +- Prevent resource leaks in long-running processes + +### Event Handling +- Follow observer/pub-sub patterns for loose coupling +- Use asyncio primitives for event coordination +- Handle event ordering and race conditions appropriately +- Provide clear unsubscription/disposal mechanisms + +## Debugging and Troubleshooting + +### Logging Strategies +- Enable debug logging for specific subsystems when needed +- Trace request/response flows through the system +- Monitor performance metrics and resource usage +- Correlate logs across components using request IDs + +### Common Issues +- Circular imports: Refactor to break dependencies +- Async/await confusion: Ensure proper coroutine handling +- Permission denied errors: Review risk classifications and allowlists +- Model provider failures: Check API keys and network connectivity +- Memory leaks: Monitor object retention and cleanup + +### Diagnostic Tools +- Use Python's built-in `debugger` or IDE debugging tools +- Profile performance with `cProfile` or similar +- Inspect state through available introspection methods +- Check logs in the application data directory \ No newline at end of file diff --git a/coworker/agent.py b/coworker/agent.py index d24f5c59..edd3d6cd 100644 --- a/coworker/agent.py +++ b/coworker/agent.py @@ -6,6 +6,7 @@ from __future__ import annotations +import os from pathlib import Path from typing import Any, Optional @@ -38,6 +39,8 @@ from .tools.subagent import explorer_tools from .web import make_web_fetch_tool, make_web_search_tool from .workspace_trust import WorkspaceTrustStore +from .guard.middleware import GuardMiddleware +from .guard.ruleset import GuardRuleSet from .tools.shell import LocalExecutor from .tools.todo import TodoList @@ -306,6 +309,23 @@ def context_provider() -> str: parts.append(ctx) return "\n\n".join(parts) + # Build guard middleware for agent families that fan out subagents. + guard_middleware: Optional[GuardMiddleware] = None + if agent.family == "code" and ws is not None: + guard_config = Path( + os.environ.get("COWORKER_GUARD_CONFIG", str(state_dir() / "guard.yaml")) + ) + guard_log = state_dir() / "guard.log" + # Load rules from default guard config; empty file = no-op. + ruleset = GuardRuleSet.load_rules(guard_config) + guard_middleware = GuardMiddleware( + permissions=permissions, + config_path=guard_config, + log_path=guard_log, + ruleset=ruleset, + agent_id=agent.name, + ) + engine = TurnEngine( provider=provider, registry=registry, @@ -325,6 +345,7 @@ def context_provider() -> str: directory_requester=directory_requester, plan_approver=plan_approver, question_asker=question_asker, + guard_middleware=guard_middleware, ) engine.executor = executor # type: ignore[attr-defined] engine.todo = todo # type: ignore[attr-defined] diff --git a/coworker/engine.py b/coworker/engine.py index 78a88316..4d429edc 100644 --- a/coworker/engine.py +++ b/coworker/engine.py @@ -73,6 +73,8 @@ def __init__( question_asker: Optional[ Callable[[dict[str, Any]], "Awaitable[dict[str, Any]]"] ] = None, + # Guard middleware for fan-out subagent safety policies. + guard_middleware: Optional[Any] = None, # Called (thread-safe, best-effort) when the user stops the turn — e.g. the # executor's kill for a running shell command. interrupt_hooks: Optional[list[Callable[[], None]]] = None, @@ -114,6 +116,7 @@ def __init__( # tool_call.id → the standing rule that auto-allowed it ("tool → target"), so the # TOOL_FINISHED event can carry the note to the tool card (§25). self._standing_notes: dict[str, str] = {} + self.guard_middleware = guard_middleware self._interrupt_hooks: list[Callable[[], None]] = list(interrupt_hooks or []) # -- external controls ------------------------------------------------------ @@ -484,21 +487,29 @@ async def _handle_tool_calls( if concurrent: for tool_call in concurrent: + if self.guard_middleware is not None: + self.guard_middleware.track_start(tool_call.name) yield Event(EventType.TOOL_STARTED, {"name": tool_call.name}) self._audit(tool_call, stage="started") outcomes = await asyncio.gather( *[asyncio.to_thread(self._execute_sync, tc) for tc in concurrent] ) for tool_call, (result, status) in zip(concurrent, outcomes): + if self.guard_middleware is not None: + self.guard_middleware.track_end(tool_call.name) yield self._record_result(tool_call, result, status) for tool_call in serial: if self._cancel.is_set(): yield self._interrupted_tool(tool_call) continue + if self.guard_middleware is not None: + self.guard_middleware.track_start(tool_call.name) yield Event(EventType.TOOL_STARTED, {"name": tool_call.name}) self._audit(tool_call, stage="started") result, status = await asyncio.to_thread(self._execute_sync, tool_call) + if self.guard_middleware is not None: + self.guard_middleware.track_end(tool_call.name) yield self._record_result(tool_call, result, status) def _interrupted_tool(self, tool_call: ToolCall) -> Event: @@ -532,9 +543,14 @@ async def _authorize(self, tool_call: ToolCall) -> "AsyncIterator[Event | bool]" spec = self.registry.get(tool_call.name) metadata = spec.metadata if spec else None - decision = self.permissions.evaluate( - tool_call.name, tool_call.arguments, metadata - ) + if self.guard_middleware is not None: + decision = self.guard_middleware.evaluate( + tool_call.name, tool_call.arguments, metadata + ) + else: + decision = self.permissions.evaluate( + tool_call.name, tool_call.arguments, metadata + ) allowed = decision.allowed reason = decision.reason diff --git a/coworker/guard/__init__.py b/coworker/guard/__init__.py new file mode 100644 index 00000000..43f84fbc --- /dev/null +++ b/coworker/guard/__init__.py @@ -0,0 +1,15 @@ +# Guard package for fan-out subagent guardrails + +from .ruleset import GuardRuleSet, GuardRule, GuardDecision +from .logger import GuardLogger +from .config_loader import GuardConfigLoader +from .middleware import GuardMiddleware + +__all__ = [ + "GuardRuleSet", + "GuardRule", + "GuardDecision", + "GuardLogger", + "GuardConfigLoader", + "GuardMiddleware", +] diff --git a/coworker/guard/config_loader.py b/coworker/guard/config_loader.py new file mode 100644 index 00000000..df8fc864 --- /dev/null +++ b/coworker/guard/config_loader.py @@ -0,0 +1,71 @@ +"""GuardConfigLoader — load and cache guard YAML configuration. + +Supports change detection via file mtime so the rule set can be hot-reloaded +without restarting the engine (optional for MVP). +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, Optional + + +class GuardConfigLoader: + """Reads guard YAML config from disk, caches it, and detects file changes. + + Usage:: + + loader = GuardConfigLoader("config/guard.yaml") + cfg = loader.load_config() # first load + cfg = loader.load_config() # cached, unless mtime changed + """ + + def __init__(self, config_path: str | Path): + self.config_path = Path(config_path) + self._cached: Optional[dict[str, Any]] = None + self._mtime: float = 0.0 + + # ------------------------------------------------------------------ + # Loading + # ------------------------------------------------------------------ + + def load_config(self) -> dict[str, Any]: + """Load and cache the YAML config. + + Returns the cached dict when the file hasn't changed since the last + load. Returns an empty dict when the file is missing. + """ + if not self._has_changed(): + return self._cached or {} + + try: + import yaml + + with open(self.config_path, "r") as f: + self._cached = yaml.safe_load(f) or {} + self._mtime = os.path.getmtime(self.config_path) + except (FileNotFoundError, PermissionError): + self._cached = {} + self._mtime = 0.0 + except yaml.YAMLError as exc: + self._cached = {} + self._mtime = 0.0 + + return self._cached or {} + + # ------------------------------------------------------------------ + # Change detection + # ------------------------------------------------------------------ + + def has_changed(self) -> bool: + """True when the file has been modified since the last :meth:`load_config`.""" + try: + current = os.path.getmtime(self.config_path) + return current != self._mtime + except (FileNotFoundError, PermissionError): + return self._cached is not None + + def _has_changed(self) -> bool: + """Internal version (used by :meth:`load_config`).""" + return self.has_changed() diff --git a/coworker/guard/logger.py b/coworker/guard/logger.py new file mode 100644 index 00000000..8ab41728 --- /dev/null +++ b/coworker/guard/logger.py @@ -0,0 +1,93 @@ +"""GuardLogger — dedicated logger for guard decisions with audit trail. + +Writes structured, timestamped entries to a configurable log file so every +allow/deny decision can be traced post-hoc. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any, Optional + +_DEFAULT_LOG_PATH = "guard.log" + + +class GuardLogger: + """Logs guard decisions to a dedicated file. + + Each entry carries: timestamp, level, status (ALLOW/DENY), agent id, + tool name, matching rule, and reason. + """ + + def __init__( + self, + log_path: str | Path = _DEFAULT_LOG_PATH, + level: int = logging.INFO, + ): + self._logger = logging.getLogger("coworker.guard") + self._logger.setLevel(level) + self._logger.handlers.clear() # don't duplicate on re-init + + path = Path(log_path) + path.parent.mkdir(parents=True, exist_ok=True) + + handler = logging.FileHandler(str(path), encoding="utf-8") + handler.setLevel(level) + formatter = logging.Formatter( + "%(asctime)s [%(levelname)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + handler.setFormatter(formatter) + self._logger.addHandler(handler) + + # Also propagate warnings+ to the root logger so errors appear in the + # main log during development. Keep INFO silent at root. + self._logger.propagate = False + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def log_decision( + self, + *, + agent_id: str = "", + tool_name: str, + arguments: dict[str, Any], + allowed: bool, + reason: str = "", + rule: str = "", + needs_user: bool = False, + ) -> None: + """Persist a guard decision record to the log file. + + Args: + agent_id: Short identifier for the running agent. + tool_name: The tool that was called (e.g. ``run_shell``, ``explore``). + arguments: The tool's argument dict. + allowed: Whether the call was allowed. + reason: Human-readable reason for the decision. + rule: Name of the rule that matched, if any. + needs_user: Whether the decision defers to the user. + """ + status = "ALLOW" if allowed else "DENY" + parts = [ + status, + f"agent={agent_id}", + f"tool={tool_name}", + ] + if rule: + parts.append(f"rule={rule}") + if reason: + parts.append(f"reason={reason}") + if needs_user: + parts.append("needs_user=1") + + self._logger.info(" | ".join(parts)) + + def set_level(self, level: int) -> None: + """Change the log level at runtime (e.g. for debugging).""" + self._logger.setLevel(level) + for handler in self._logger.handlers: + handler.setLevel(level) diff --git a/coworker/guard/middleware.py b/coworker/guard/middleware.py new file mode 100644 index 00000000..90fe8621 --- /dev/null +++ b/coworker/guard/middleware.py @@ -0,0 +1,164 @@ +"""GuardMiddleware — combine GuardRuleSet, GuardLogger, and PermissionEngine. + +The middleware wraps the existing ``PermissionEngine`` (from +``coworker.permissions``) and applies guard-specific rules on top — deny +patterns, allow patterns, and fan-out concurrency limits for subagent tools. +All decisions are logged to the guard log for audit. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Optional + +from ..permissions import PermissionEngine, Decision +from .ruleset import GuardRuleSet, GuardDecision +from .logger import GuardLogger +from .config_loader import GuardConfigLoader + + +class GuardMiddleware: + """Layered authorization: base permissions → guard rules → audit log. + + Usage:: + + middleware = GuardMiddleware( + permissions, + config_path="config/guard.yaml", + log_path="guard.log", + agent_id="code", + ) + decision = middleware.evaluate("run_shell", {"command": "rm -rf /"}, metadata) + """ + + def __init__( + self, + permissions: PermissionEngine, + *, + config_path: str | Path = "", + log_path: str | Path = "guard.log", + ruleset: Optional[GuardRuleSet] = None, + logger: Optional[GuardLogger] = None, + agent_id: str = "", + ): + self.permissions = permissions + self.agent_id = agent_id + + self._ruleset = ruleset or GuardRuleSet() + self._logger = logger or GuardLogger(log_path=log_path) + self._config_loader: Optional[GuardConfigLoader] = ( + GuardConfigLoader(config_path) if config_path else None + ) + + # ------------------------------------------------------------------ + # Evaluation + # ------------------------------------------------------------------ + + def evaluate( + self, + tool_name: str, + arguments: dict[str, Any], + metadata: Any = None, + ) -> Decision: + """Evaluate a tool call against base permissions, then guard rules. + + Returns a ``Decision`` compatible with the engine's ``_authorize()`` + flow. On exception, fails *safe* (deny + needs_user) so no dangerous + call slips through due to a bug in the guard layer. + """ + try: + # 1. Hot-reload rules if config changed (optional). + self._maybe_reload() + + # 2. Base permission decision (read-only mode, path scoping, etc.). + base_decision = self.permissions.evaluate( + tool_name, arguments, metadata + ) + + if not base_decision.allowed: + # Blocked by permissions — log and return unchanged. + self._logger.log_decision( + agent_id=self.agent_id, + tool_name=tool_name, + arguments=arguments, + allowed=False, + reason=base_decision.reason, + rule="", + ) + return base_decision + + # 3. Guard rule check (deny patterns, fan-out limits). + guard_decision = self._ruleset.evaluate(tool_name, arguments) + + if not guard_decision.allowed: + # Blocked by guard rule. + self._logger.log_decision( + agent_id=self.agent_id, + tool_name=tool_name, + arguments=arguments, + allowed=False, + reason=guard_decision.reason, + rule=guard_decision.rule, + ) + return Decision( + allowed=False, + reason=f"guard: {guard_decision.reason}", + needs_user=guard_decision.needs_user, + rule=guard_decision.rule, + ) + + # 4. Allowed by both layers. + self._logger.log_decision( + agent_id=self.agent_id, + tool_name=tool_name, + arguments=arguments, + allowed=True, + reason=base_decision.reason, + rule=guard_decision.rule, + ) + return base_decision + + except Exception as exc: + # Fail safe: block the call and flag it for user attention. + try: + self._logger.log_decision( + agent_id=self.agent_id, + tool_name=tool_name, + arguments=arguments, + allowed=False, + reason=f"guard middleware error: {exc}", + rule="", + needs_user=True, + ) + except Exception: + pass # best-effort: logging must not mask the safety decision + return Decision( + allowed=False, + reason=f"guard middleware error: {exc}", + needs_user=True, + ) + + # ------------------------------------------------------------------ + # Fan-out tracking + # ------------------------------------------------------------------ + + def track_start(self, tool_name: str) -> None: + """Notify the rule set that *tool_name* has started executing.""" + self._ruleset.track_start(tool_name) + + def track_end(self, tool_name: str) -> None: + """Notify the rule set that *tool_name* has finished executing.""" + self._ruleset.track_end(tool_name) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _maybe_reload(self) -> None: + """If a config file is configured and has changed, reload rules.""" + if self._config_loader is None: + return + if not self._config_loader.has_changed(): + return + # Reload the ruleset from the config path. + self._ruleset = GuardRuleSet.load_rules(self._config_loader.config_path) diff --git a/coworker/guard/ruleset.py b/coworker/guard/ruleset.py new file mode 100644 index 00000000..c8bdfc82 --- /dev/null +++ b/coworker/guard/ruleset.py @@ -0,0 +1,193 @@ +"""GuardRuleSet — load and evaluate guard rules from YAML configuration. + +Supports allow/deny rules with command match patterns (contains, prefix, regex), +argument key-value matches, and fan-out concurrency limits for subagent tools. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Optional + +import yaml + + +@dataclass +class GuardRule: + """A single guard rule evaluated against tool calls.""" + + name: str + tool: str + action: str # "allow", "deny", or "limit" + reason: str = "" + + # Command match patterns (any supplied must all match for the rule to fire) + match_command_contains: Optional[str] = None + match_command_prefix: Optional[str] = None + match_command_regex: Optional[str] = None + + # Key-value arg matcher: tool argument name → expected value (str compare) + match_args: dict[str, str] = field(default_factory=dict) + + # Fan-out concurrency limit (relevant for subagent/spawn tools) + max_concurrent: Optional[int] = None + + +@dataclass +class GuardDecision: + """Outcome of evaluating a tool call against the rule set.""" + + allowed: bool + reason: str = "" + rule: str = "" # name of the matching rule, if any + needs_user: bool = False # True → surface should prompt the user for approval + + +class GuardRuleSet: + """Loads guard rules from YAML and evaluates tool calls against them. + + Thread-safe for fan-out counting if used from a single event-loop thread; + the engine serializes tool-authorization calls, so races on counters are + not expected in practice. + """ + + def __init__(self, rules: Optional[list[GuardRule]] = None): + self._rules: list[GuardRule] = list(rules or []) + # Tool name → number of currently-active (started, not yet finished) calls. + # Used to enforce fan-out concurrency limits. + self._fanout_counters: dict[str, int] = {} + + # ------------------------------------------------------------------ + # Loading + # ------------------------------------------------------------------ + + @classmethod + def load_rules(cls, path: str | Path) -> GuardRuleSet: + """Load rules from a YAML file. + + Expected schema:: + + rules: + - name: block-recursive-delete + tool: run_shell + match: + command: + contains: "rm -rf /" + action: deny + reason: "Block recursive delete outside workspace" + + Returns an empty set when the file is missing or empty. + """ + p = Path(path) + if not p.is_file(): + return cls() + + with open(p, "r") as f: + raw = yaml.safe_load(f) + + rules: list[GuardRule] = [] + for entry in (raw or {}).get("rules", []): + match = entry.get("match") or {} + command = match.get("command") or {} if isinstance(match, dict) else {} + fanout = entry.get("fanout") or {} if isinstance(entry, dict) else {} + args_match = match.get("args") or {} if isinstance(match, dict) else {} + + rules.append( + GuardRule( + name=str(entry.get("name", "unnamed")), + tool=str(entry.get("tool", "")), + action=str(entry.get("action", "deny")), + reason=str(entry.get("reason", "")), + match_command_contains=command.get("contains"), + match_command_prefix=command.get("prefix"), + match_command_regex=command.get("regex"), + match_args={ + str(k): str(v) for k, v in args_match.items() + }, + max_concurrent=fanout.get("max_concurrent"), + ) + ) + + return cls(rules) + + # ------------------------------------------------------------------ + # Evaluation + # ------------------------------------------------------------------ + + def evaluate( + self, + tool_name: str, + arguments: dict[str, Any], + ) -> GuardDecision: + """Check *tool_name* with *arguments* against every registered rule. + + Returns the first rule's decision that matches the call. When no rule + matches, returns ``GuardDecision(allowed=True)`` (no opinion). + """ + command = str(arguments.get("command", "")) + + for rule in self._rules: + # Tool name filter (empty = match any) + if rule.tool and rule.tool != tool_name: + continue + + # --- command match patterns --- + if rule.match_command_contains and rule.match_command_contains not in command: + continue + if rule.match_command_prefix and not command.startswith(rule.match_command_prefix): + continue + if rule.match_command_regex and not re.search(rule.match_command_regex, command): + continue + + # --- argument key-value match --- + if rule.match_args: + if not all( + str(arguments.get(k)) == str(v) + for k, v in rule.match_args.items() + ): + continue + + # --- fan-out concurrency limit --- + if rule.max_concurrent is not None: + current = self._fanout_counters.get(tool_name, 0) + if current >= rule.max_concurrent: + # Over the limit — apply the rule's action (deny/limit). + return GuardDecision( + allowed=False, + reason=f"{rule.reason} (max {rule.max_concurrent} concurrent)", + rule=rule.name, + ) + # Within limits — this rule doesn't apply; continue to next. + # A rule with max_concurrent but no other match patterns + # is purely a fan-out throttle, not a blanket allow/deny. + continue + + # --- action (only reached when no fan-out limit is set) --- + if rule.action == "allow": + return GuardDecision( + allowed=True, reason=rule.reason, rule=rule.name + ) + # "deny" or "limit" → block + return GuardDecision( + allowed=False, + reason=rule.reason or f"blocked by rule: {rule.name}", + rule=rule.name, + ) + + return GuardDecision(allowed=True) + + # ------------------------------------------------------------------ + # Fan-out tracking + # ------------------------------------------------------------------ + + def track_start(self, tool_name: str) -> None: + """Increment the active-count for *tool_name* (tool is starting).""" + self._fanout_counters[tool_name] = self._fanout_counters.get(tool_name, 0) + 1 + + def track_end(self, tool_name: str) -> None: + """Decrement the active-count for *tool_name* (tool finished).""" + current = self._fanout_counters.get(tool_name, 0) + if current > 0: + self._fanout_counters[tool_name] = current - 1 diff --git a/tests/test_guard.py b/tests/test_guard.py new file mode 100644 index 00000000..a5f4f12f --- /dev/null +++ b/tests/test_guard.py @@ -0,0 +1,494 @@ +"""Tests for the guard middleware — GuardRuleSet, GuardLogger, GuardConfigLoader, +and GuardMiddleware integration. + +Follows the same patterns as test_permissions_risk.py: parametrize, tmp_path for +file operations, SimpleNamespace for metadata. +""" + +from __future__ import annotations + +import logging +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest +import yaml + +from coworker.guard.config_loader import GuardConfigLoader +from coworker.guard.logger import GuardLogger +from coworker.guard.middleware import GuardMiddleware +from coworker.guard.ruleset import GuardDecision, GuardRule, GuardRuleSet +from coworker.permissions import Mode, PermissionEngine + + +# ============================================================================= +# GuardRuleSet +# ============================================================================= + + +class TestGuardRuleSet: + def test_empty_ruleset_allows_all(self): + rset = GuardRuleSet() + d = rset.evaluate("any_tool", {}) + assert d.allowed + assert not d.reason + + def test_deny_rule_blocks_by_tool_name(self): + rset = GuardRuleSet( + [GuardRule(name="block-test", tool="dangerous_tool", action="deny")] + ) + d = rset.evaluate("dangerous_tool", {}) + assert not d.allowed + assert d.rule == "block-test" + + def test_deny_rule_does_not_affect_other_tools(self): + rset = GuardRuleSet( + [GuardRule(name="block-test", tool="dangerous_tool", action="deny")] + ) + d = rset.evaluate("safe_tool", {}) + assert d.allowed + + def test_allow_rule_allows_tool(self): + rset = GuardRuleSet( + [GuardRule(name="allow-safe", tool="safe_tool", action="allow")] + ) + d = rset.evaluate("safe_tool", {}) + assert d.allowed + assert d.rule == "allow-safe" + + def test_command_contains_match(self): + rset = GuardRuleSet( + [ + GuardRule( + name="block-rmrf", + tool="run_shell", + action="deny", + match_command_contains="rm -rf ", + ) + ] + ) + assert not rset.evaluate("run_shell", {"command": "rm -rf /"}).allowed + assert not rset.evaluate("run_shell", {"command": "rm -rf ~"}).allowed + # Not blocked: different command + assert rset.evaluate("run_shell", {"command": "ls -la"}).allowed + + def test_command_prefix_match(self): + rset = GuardRuleSet( + [ + GuardRule( + name="allow-git-status", + tool="run_shell", + action="allow", + match_command_prefix="git status", + ) + ] + ) + assert rset.evaluate("run_shell", {"command": "git status"}).allowed + assert rset.evaluate("run_shell", {"command": "git status -s"}).allowed + # Not matched: different prefix + assert rset.evaluate("run_shell", {"command": "git push"}).allowed + # Not matched: not the right tool + assert rset.evaluate("read_file", {"path": "x"}).allowed + + def test_command_regex_match(self): + rset = GuardRuleSet( + [ + GuardRule( + name="allow-git-read-only", + tool="run_shell", + action="allow", + match_command_regex=r"^git (status|diff|log|branch|show)", + ) + ] + ) + assert rset.evaluate("run_shell", {"command": "git status"}).allowed + assert rset.evaluate("run_shell", {"command": "git diff HEAD"}).allowed + assert rset.evaluate("run_shell", {"command": "git log --oneline"}).allowed + # Not matched: write git command + assert rset.evaluate("run_shell", {"command": "git push"}).allowed + # Not matched: doesn't start with git + + def test_arg_key_value_match(self): + rset = GuardRuleSet( + [ + GuardRule( + name="block-specific-tool-arg", + tool="custom_tool", + action="deny", + match_args={"mode": "dangerous"}, + ) + ] + ) + assert not rset.evaluate("custom_tool", {"mode": "dangerous"}).allowed + assert rset.evaluate("custom_tool", {"mode": "safe"}).allowed + assert rset.evaluate("custom_tool", {}).allowed + + def test_all_match_conditions_must_align(self): + rset = GuardRuleSet( + [ + GuardRule( + name="pick-only", + tool="custom_tool", + action="deny", + match_command_contains="danger", + match_args={"mode": "unsafe"}, + ) + ] + ) + # Both match → deny + assert not rset.evaluate("custom_tool", {"command": "danger", "mode": "unsafe"}).allowed + # Only one matches → no match + assert rset.evaluate("custom_tool", {"command": "danger", "mode": "safe"}).allowed + + def test_first_matching_rule_wins(self): + rset = GuardRuleSet( + [ + GuardRule(name="first", tool="test_tool", action="deny"), + GuardRule(name="second", tool="test_tool", action="allow"), + ] + ) + d = rset.evaluate("test_tool", {}) + assert not d.allowed + assert d.rule == "first" + + def test_fanout_limit_blocks_after_max(self): + rset = GuardRuleSet( + [ + GuardRule( + name="limit-explore", + tool="explore", + action="deny", + max_concurrent=2, + ) + ] + ) + # First two start + rset.track_start("explore") + rset.track_start("explore") + # Third is blocked + d = rset.evaluate("explore", {}) + assert not d.allowed + assert "concurrent" in d.reason.lower() + + def test_fanout_limit_allows_after_end(self): + rset = GuardRuleSet( + [ + GuardRule( + name="limit-explore", + tool="explore", + action="deny", + max_concurrent=2, + ) + ] + ) + rset.track_start("explore") + rset.track_start("explore") + rset.track_end("explore") + # One slot freed + d = rset.evaluate("explore", {}) + assert d.allowed + + def test_fanout_track_end_does_not_go_negative(self): + rset = GuardRuleSet( + [ + GuardRule( + name="limit-explore", + tool="explore", + action="deny", + max_concurrent=2, + ) + ] + ) + rset.track_end("explore") # no-op, not negative + d = rset.evaluate("explore", {}) + assert d.allowed + + def test_load_rules_from_yaml(self, tmp_path): + yaml_path = tmp_path / "guard.yaml" + yaml_path.write_text( + yaml.dump( + { + "rules": [ + { + "name": "block-dd", + "tool": "run_shell", + "match": {"command": {"contains": "dd if="}}, + "action": "deny", + "reason": "Block dd", + }, + { + "name": "allow-git-log", + "tool": "run_shell", + "match": {"command": {"prefix": "git log"}}, + "action": "allow", + "reason": "Allow git log", + }, + ] + } + ) + ) + rset = GuardRuleSet.load_rules(yaml_path) + assert len(rset._rules) == 2 + assert not rset.evaluate("run_shell", {"command": "dd if=/dev/sda"}).allowed + assert rset.evaluate("run_shell", {"command": "git log --oneline"}).allowed + + def test_load_rules_missing_file(self, tmp_path): + rset = GuardRuleSet.load_rules(tmp_path / "nonexistent.yaml") + assert len(rset._rules) == 0 + assert rset.evaluate("anything", {}).allowed + + +# ============================================================================= +# GuardLogger +# ============================================================================= + + +class TestGuardLogger: + def test_log_decision_writes_to_file(self, tmp_path): + log_path = tmp_path / "guard.log" + logger = GuardLogger(log_path=str(log_path)) + logger.log_decision( + agent_id="test-agent", + tool_name="run_shell", + arguments={"command": "rm -rf /"}, + allowed=False, + reason="blocked by guard rule", + rule="block-dangerous", + ) + # Force flush + for handler in logger._logger.handlers: + handler.flush() + content = log_path.read_text(encoding="utf-8") + assert "DENY" in content + assert "test-agent" in content + assert "run_shell" in content + assert "block-dangerous" in content + + def test_log_decision_allow(self, tmp_path): + log_path = tmp_path / "guard.log" + logger = GuardLogger(log_path=str(log_path)) + logger.log_decision( + agent_id="agent-1", + tool_name="read_file", + arguments={"path": "/tmp/x"}, + allowed=True, + ) + for handler in logger._logger.handlers: + handler.flush() + content = log_path.read_text(encoding="utf-8") + assert "ALLOW" in content + assert "agent-1" in content + assert "read_file" in content + + def test_logger_custom_level(self, tmp_path): + log_path = tmp_path / "guard.log" + logger = GuardLogger(log_path=str(log_path), level=logging.WARNING) + assert logger._logger.level == logging.WARNING + # Logging at INFO level should not appear + logger.log_decision( + agent_id="test", tool_name="x", arguments={}, allowed=True + ) + for handler in logger._logger.handlers: + handler.flush() + content = log_path.read_text(encoding="utf-8") + assert content == "" # WARNING level, so INFO is ignored + + +# ============================================================================= +# GuardConfigLoader +# ============================================================================= + + +class TestGuardConfigLoader: + def test_load_config_returns_dict(self, tmp_path): + yaml_path = tmp_path / "guard.yaml" + yaml_path.write_text( + yaml.dump({"rules": [{"name": "test", "tool": "x", "action": "deny"}]}) + ) + loader = GuardConfigLoader(yaml_path) + cfg = loader.load_config() + assert "rules" in cfg + assert len(cfg["rules"]) == 1 + + def test_load_config_missing_file_returns_empty(self, tmp_path): + loader = GuardConfigLoader(tmp_path / "nonexistent.yaml") + cfg = loader.load_config() + assert cfg == {} + + def test_has_changed_false_on_first_load(self, tmp_path): + yaml_path = tmp_path / "guard.yaml" + yaml_path.write_text(yaml.dump({"rules": []})) + loader = GuardConfigLoader(yaml_path) + loader.load_config() + # No change + assert not loader.has_changed() + + def test_has_changed_true_after_file_modification(self, tmp_path): + yaml_path = tmp_path / "guard.yaml" + yaml_path.write_text(yaml.dump({"rules": []})) + loader = GuardConfigLoader(yaml_path) + loader.load_config() + # Modify the file + yaml_path.write_text(yaml.dump({"rules": [{"name": "new", "tool": "x", "action": "deny"}]})) + assert loader.has_changed() + + def test_cached_config_returned_without_change(self, tmp_path): + yaml_path = tmp_path / "guard.yaml" + yaml_path.write_text(yaml.dump({"rules": [{"name": "original"}]})) + loader = GuardConfigLoader(yaml_path) + first = loader.load_config() + # Modify file behind our back + yaml_path.write_text(yaml.dump({"rules": [{"name": "modified"}]})) + # Without calling load_config again, has_changed is True + assert loader.has_changed() + # But a second load will pick it up + second = loader.load_config() + assert second["rules"][0]["name"] == "modified" + + +# ============================================================================= +# GuardMiddleware +# ============================================================================= + + +class TestGuardMiddleware: + def test_middleware_allows_low_risk_tool(self, tmp_path): + permissions = PermissionEngine(workspace_root=tmp_path) + middleware = GuardMiddleware(permissions) + d = middleware.evaluate("read_file", {"path": "x"}) + assert d.allowed + assert not d.needs_user + + def test_middleware_blocks_with_guard_rule(self, tmp_path): + permissions = PermissionEngine(workspace_root=tmp_path, mode=Mode.AUTO) + ruleset = GuardRuleSet( + [GuardRule(name="block-all", tool="run_shell", action="deny")] + ) + middleware = GuardMiddleware(permissions, ruleset=ruleset) + d = middleware.evaluate("run_shell", {"command": "anything"}) + assert not d.allowed + assert "guard:" in d.reason + + def test_middleware_uses_permission_engine_decision(self, tmp_path): + permissions = PermissionEngine( + workspace_root=tmp_path, mode=Mode.PLAN # read-only mode + ) + middleware = GuardMiddleware(permissions) + d = middleware.evaluate("run_shell", {"command": "rm -rf /"}) + assert not d.allowed + # Should be blocked by plan mode (not by guard rule) + assert "read-only" in d.reason + + def test_middleware_guard_rule_wins_over_permissions(self, tmp_path): + permissions = PermissionEngine( + workspace_root=tmp_path, mode=Mode.AUTO + ) + ruleset = GuardRuleSet( + [ + GuardRule( + name="block-dangerous-command", + tool="run_shell", + action="deny", + match_command_contains="rm -rf", + ) + ] + ) + middleware = GuardMiddleware(permissions, ruleset=ruleset) + d = middleware.evaluate("run_shell", {"command": "rm -rf /tmp"}) + assert not d.allowed + assert "guard:" in d.reason + + def test_middleware_logs_decision(self, tmp_path): + permissions = PermissionEngine(workspace_root=tmp_path) + log_path = tmp_path / "guard.log" + middleware = GuardMiddleware( + permissions, log_path=str(log_path), agent_id="test-agent" + ) + middleware.evaluate("run_shell", {"command": "echo hi"}) + for handler in middleware._logger._logger.handlers: + handler.flush() + content = log_path.read_text(encoding="utf-8") + # Should have a log entry + assert "test-agent" in content + + def test_middleware_exception_fails_safe(self, tmp_path): + """On exception, middleware returns deny + needs_user (fail safe).""" + permissions = PermissionEngine(workspace_root=tmp_path) + + class _BrokenRuleset: + def evaluate(self, tool_name, arguments): + raise RuntimeError("something went wrong") + + def track_start(self, tool_name): + pass + + def track_end(self, tool_name): + pass + + def track_start(self, tool_name): + pass + + def track_end(self, tool_name): + pass + + middleware = GuardMiddleware(permissions, ruleset=_BrokenRuleset()) # type: ignore + d = middleware.evaluate("read_file", {"path": "x"}) + assert not d.allowed + assert d.needs_user + assert "error" in d.reason + + def test_middleware_fanout_tracking(self, tmp_path): + permissions = PermissionEngine(workspace_root=tmp_path, mode=Mode.AUTO) + ruleset = GuardRuleSet( + [ + GuardRule( + name="limit-explore", + tool="explore", + action="deny", + max_concurrent=1, + ) + ] + ) + middleware = GuardMiddleware(permissions, ruleset=ruleset) + # One active + middleware.track_start("explore") + # Second should be blocked + d = middleware.evaluate("explore", {"task": "research"}) + assert not d.allowed + # After end, allowed again + middleware.track_end("explore") + d = middleware.evaluate("explore", {"task": "research"}) + assert d.allowed + + def test_middleware_config_reload(self, tmp_path): + permissions = PermissionEngine(workspace_root=tmp_path, mode=Mode.AUTO) + yaml_path = tmp_path / "guard.yaml" + # Initial config: no rules + yaml_path.write_text(yaml.dump({"rules": []})) + middleware = GuardMiddleware( + permissions, + config_path=str(yaml_path), + agent_id="test", + ) + # Initially everything is allowed + assert middleware.evaluate("run_shell", {"command": "anything"}).allowed + # Update config with a deny rule + yaml_path.write_text( + yaml.dump( + { + "rules": [ + { + "name": "block-all", + "tool": "run_shell", + "action": "deny", + "reason": "blocked", + } + ] + } + ) + ) + # Next evaluation should pick up the change + d = middleware.evaluate("run_shell", {"command": "anything"}) + assert not d.allowed + assert "guard:" in d.reason From 472e07d42e6a3c30694c8f7a029a1d6367f0d214 Mon Sep 17 00:00:00 2001 From: KunjShah95 Date: Mon, 27 Jul 2026 11:33:15 +0530 Subject: [PATCH 3/3] Make fanout limit enforceable via eager track_start + add integration tests - track_start() now returns bool: checks the fanout limit BEFORE incrementing the counter, returns False if max_concurrent would be exceeded - GuardMiddleware.track_start() returns bool, logs denied decisions for audit - Engine concurrent/serial paths check track_start return value and deny tools that exceed the fanout limit (TOOL_FINISHED with status="denied") - Added is_fanout_block field to GuardDecision for robust fanout detection - New integration test: 3 concurrent tools, max_concurrent=2, threading.Barrier proves first two run concurrently while third is denied by fanout limit - New integration test: broken ruleset raises during evaluate(), middleware fails safe (needs_user=True), engine emits PERMISSION_REQUIRED - Fixed _BrokenRuleset in test_guard.py (duplicate methods, missing return) - Added fail-safe integration test for middleware exception path --- coworker/engine.py | 65 +++++- coworker/guard/middleware.py | 23 +- coworker/guard/ruleset.py | 18 +- tests/test_engine.py | 403 +++++++++++++++++++++++++++++++++-- tests/test_guard.py | 11 +- 5 files changed, 484 insertions(+), 36 deletions(-) diff --git a/coworker/engine.py b/coworker/engine.py index 4d429edc..aeed8838 100644 --- a/coworker/engine.py +++ b/coworker/engine.py @@ -486,25 +486,72 @@ async def _handle_tool_calls( serial = [tc for tc in cleared if tc not in concurrent] if concurrent: + to_run: list[ToolCall] = [] for tool_call in concurrent: if self.guard_middleware is not None: - self.guard_middleware.track_start(tool_call.name) + if not self.guard_middleware.track_start(tool_call.name): + # Fan-out limit exceeded — deny immediately. + self.messages.append( + _tool_error_message( + tool_call, + "guard: fanout limit exceeded (max concurrent)", + ) + ) + yield Event( + EventType.TOOL_FINISHED, + { + "name": tool_call.name, + "status": "denied", + "reason": "guard: fanout limit exceeded", + }, + ) + self._audit( + tool_call, + stage="finished", + status="denied", + reason="fanout limit exceeded", + ) + continue + to_run.append(tool_call) yield Event(EventType.TOOL_STARTED, {"name": tool_call.name}) self._audit(tool_call, stage="started") - outcomes = await asyncio.gather( - *[asyncio.to_thread(self._execute_sync, tc) for tc in concurrent] - ) - for tool_call, (result, status) in zip(concurrent, outcomes): - if self.guard_middleware is not None: - self.guard_middleware.track_end(tool_call.name) - yield self._record_result(tool_call, result, status) + if to_run: + outcomes = await asyncio.gather( + *[asyncio.to_thread(self._execute_sync, tc) for tc in to_run] + ) + for tool_call, (result, status) in zip(to_run, outcomes): + if self.guard_middleware is not None: + self.guard_middleware.track_end(tool_call.name) + yield self._record_result(tool_call, result, status) for tool_call in serial: if self._cancel.is_set(): yield self._interrupted_tool(tool_call) continue if self.guard_middleware is not None: - self.guard_middleware.track_start(tool_call.name) + if not self.guard_middleware.track_start(tool_call.name): + # Fan-out limit exceeded — deny immediately. + self.messages.append( + _tool_error_message( + tool_call, + "guard: fanout limit exceeded (max concurrent)", + ) + ) + yield Event( + EventType.TOOL_FINISHED, + { + "name": tool_call.name, + "status": "denied", + "reason": "guard: fanout limit exceeded", + }, + ) + self._audit( + tool_call, + stage="finished", + status="denied", + reason="fanout limit exceeded", + ) + continue yield Event(EventType.TOOL_STARTED, {"name": tool_call.name}) self._audit(tool_call, stage="started") result, status = await asyncio.to_thread(self._execute_sync, tool_call) diff --git a/coworker/guard/middleware.py b/coworker/guard/middleware.py index 90fe8621..12de7bd8 100644 --- a/coworker/guard/middleware.py +++ b/coworker/guard/middleware.py @@ -142,9 +142,26 @@ def evaluate( # Fan-out tracking # ------------------------------------------------------------------ - def track_start(self, tool_name: str) -> None: - """Notify the rule set that *tool_name* has started executing.""" - self._ruleset.track_start(tool_name) + def track_start(self, tool_name: str) -> bool: + """Try to start tracking *tool_name*. + + Returns ``True`` if the tool was allowed to start (counter incremented), + or ``False`` if the fan-out limit would be exceeded. + + When returning ``False`` the decision is logged for audit so the + blocked call is traceable. + """ + allowed = self._ruleset.track_start(tool_name) + if not allowed: + self._logger.log_decision( + agent_id=self.agent_id, + tool_name=tool_name, + arguments={}, + allowed=False, + reason="fanout limit exceeded during track_start", + rule="", + ) + return allowed def track_end(self, tool_name: str) -> None: """Notify the rule set that *tool_name* has finished executing.""" diff --git a/coworker/guard/ruleset.py b/coworker/guard/ruleset.py index c8bdfc82..361ceb74 100644 --- a/coworker/guard/ruleset.py +++ b/coworker/guard/ruleset.py @@ -43,6 +43,7 @@ class GuardDecision: reason: str = "" rule: str = "" # name of the matching rule, if any needs_user: bool = False # True → surface should prompt the user for approval + is_fanout_block: bool = False # True when the denial was caused by max_concurrent class GuardRuleSet: @@ -158,6 +159,7 @@ def evaluate( allowed=False, reason=f"{rule.reason} (max {rule.max_concurrent} concurrent)", rule=rule.name, + is_fanout_block=True, ) # Within limits — this rule doesn't apply; continue to next. # A rule with max_concurrent but no other match patterns @@ -182,9 +184,21 @@ def evaluate( # Fan-out tracking # ------------------------------------------------------------------ - def track_start(self, tool_name: str) -> None: - """Increment the active-count for *tool_name* (tool is starting).""" + def track_start(self, tool_name: str) -> bool: + """Try to start tracking *tool_name*. + + Returns ``True`` if the tool was allowed to start (counter incremented), + or ``False`` if the fan-out limit would be exceeded (no change). + + Delegates to :meth:`evaluate` so the fan-out check logic stays in one + place — only fan-out denials ("concurrent" in reason) cause a rejection; + other deny rules are already enforced at authorization time. + """ + decision = self.evaluate(tool_name, {}) + if not decision.allowed and decision.is_fanout_block: + return False self._fanout_counters[tool_name] = self._fanout_counters.get(tool_name, 0) + 1 + return True def track_end(self, tool_name: str) -> None: """Decrement the active-count for *tool_name* (tool finished).""" diff --git a/tests/test_engine.py b/tests/test_engine.py index 90e38fdb..8a4c343b 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -9,7 +9,7 @@ import aisuite as ai from coworker.engine import ApprovalOutcome, PermissionRequest, TurnEngine from coworker.events import EventType -from coworker.permissions import PermissionEngine +from coworker.permissions import Mode, PermissionEngine from coworker.providers import ( AssistantTurn, ModelCapabilities, @@ -20,6 +20,16 @@ from coworker.tools import ToolRegistry +def _multi_tool_turn(calls): + return AssistantTurn( + tool_calls=[ + ToolCall(id=f"call_{i}", name=name, arguments=args) + for i, (name, args) in enumerate(calls) + ], + finish_reason="tool_calls", + ) + + def _text_turn(text): return AssistantTurn(text=text, finish_reason="stop") @@ -47,11 +57,12 @@ def capabilities(self, model): return ModelCapabilities() -def _engine(tmp_path, turns, *, approver=None, loop=False, max_iterations=12): +def _engine(tmp_path, turns, *, approver=None, loop=False, max_iterations=12, + mode=None, guard_middleware=None): provider = ScriptedProvider(turns, loop=loop) registry = ToolRegistry() registry.register_all(ai.toolkits.files(root=str(tmp_path), allow_write=True)) - permissions = PermissionEngine(workspace_root=tmp_path) + permissions = PermissionEngine(workspace_root=tmp_path, mode=mode or Mode.INTERACTIVE) engine = TurnEngine( provider=provider, registry=registry, @@ -59,6 +70,7 @@ def _engine(tmp_path, turns, *, approver=None, loop=False, max_iterations=12): model="gpt-5.5", approver=approver, max_iterations=max_iterations, + guard_middleware=guard_middleware, ) return engine, provider @@ -201,25 +213,16 @@ def test_steering_injects_next_turn(tmp_path): # -- parallel tool execution ------------------------------------------------------ -def _multi_tool_turn(calls): - return AssistantTurn( - tool_calls=[ - ToolCall(id=f"call_{i}", name=name, arguments=args) - for i, (name, args) in enumerate(calls) - ], - finish_reason="tool_calls", - ) - - -def _bare_engine(tmp_path, turns): +def _bare_engine(tmp_path, turns, *, mode=None, guard_middleware=None): provider = ScriptedProvider(turns) registry = ToolRegistry() - permissions = PermissionEngine(workspace_root=tmp_path) + permissions = PermissionEngine(workspace_root=tmp_path, mode=mode or Mode.INTERACTIVE) engine = TurnEngine( provider=provider, registry=registry, permissions=permissions, model="gpt-5.5", + guard_middleware=guard_middleware, ) return engine, registry @@ -419,6 +422,376 @@ def capabilities(self, model): assert "images" in text # degradation is called out in the marker +# -- guard middleware integration ------------------------------------------------ + + +def _guard_engine(tmp_path, turns, rules, *, mode=None, loop=False): + """Build an engine wired with GuardMiddleware for integration tests. + Delegates to ``_engine()`` then attaches the guard middleware.""" + from coworker.guard.middleware import GuardMiddleware + from coworker.guard.ruleset import GuardRuleSet + + engine, provider = _engine( + tmp_path, turns, mode=mode or Mode.AUTO, loop=loop, + ) + guard = GuardMiddleware( + engine.permissions, + ruleset=GuardRuleSet(rules), + log_path=str(tmp_path / "guard.log"), + ) + engine.guard_middleware = guard + return engine, provider + + +def test_guard_blocks_disallowed_tool(tmp_path): + """GuardMiddleware denies a tool that matches a deny rule — the engine + surfaces it as a denied tool with a guard: prefixed reason, and the file + is NOT written.""" + from coworker.guard.ruleset import GuardRule + + rules = [ + GuardRule( + name="block-writes", + tool="write_file", + action="deny", + reason="All writes blocked by guard rule", + ), + ] + engine, _ = _guard_engine( + tmp_path, + [ + _tool_turn("write_file", {"path": "x.py", "content": "print(1)\n"}), + _text_turn("blocked by guard"), + ], + rules=rules, + ) + events = _collect(engine, "write x.py") + + finished = [e for e in events if e.type == EventType.TOOL_FINISHED] + assert len(finished) == 1 + assert finished[0].data["status"] == "denied" + assert "guard:" in finished[0].data.get("reason", "") + + # Permission was AUTO so the permission engine itself did NOT block; + # only the guard layer stopped it. Verify no file was written. + assert not (tmp_path / "x.py").exists() + + # The error message in history should cite the guard rule. + assert any( + m.get("role") == "tool" + and "guard:" in m.get("content", "") + and "not executed" in m.get("content", "") + for m in engine.messages + ) + + +def test_guard_blocks_in_read_only_mode_via_permission_engine(tmp_path): + """In PLAN mode, the guard middleware respects the permission engine's + read-only gate — a consequential tool is blocked even without a matching + guard rule.""" + + # Empty ruleset — the guard has no opinion; the permission engine should + # still block the write in plan mode. + engine, _ = _guard_engine( + tmp_path, + [ + _tool_turn("write_file", {"path": "x.py", "content": "x"}), + _text_turn("blocked"), + ], + rules=[], + mode=Mode.PLAN, + ) + events = _collect(engine, "write x.py") + + finished = [e for e in events if e.type == EventType.TOOL_FINISHED] + assert len(finished) == 1 + assert finished[0].data["status"] == "denied" + # The reason should cite plan mode being read-only, not the guard. + assert "read-only" in finished[0].data.get("reason", "").lower() + assert not (tmp_path / "x.py").exists() + + +def _guard_bare_engine(tmp_path, turns, rules): + """Build a bare engine with GuardMiddleware and no pre-registered tools. + Delegates to ``_bare_engine()`` then attaches the guard middleware. + Caller registers tools on the returned ``registry``.""" + from coworker.guard.middleware import GuardMiddleware + from coworker.guard.ruleset import GuardRuleSet + + engine, registry = _bare_engine(tmp_path, turns) + guard = GuardMiddleware( + engine.permissions, + ruleset=GuardRuleSet(rules), + log_path=str(tmp_path / "guard.log"), + ) + engine.guard_middleware = guard + return engine, registry + + +def test_guard_allows_non_matching_tool(tmp_path): + """GuardMiddleware lets through a tool call when no rule matches.""" + from coworker.guard.ruleset import GuardRule + + (tmp_path / "a.txt").write_text("hello") + + # A guard rule that only targets run_shell, not read_file. + rules = [ + GuardRule( + name="block-dangerous-shell", + tool="run_shell", + action="deny", + match_command_contains="rm -rf", + ), + ] + engine, _ = _guard_engine( + tmp_path, + [ + _tool_turn("read_file", {"path": "a.txt"}), + _text_turn("it says hello"), + ], + rules=rules, + ) + events = _collect(engine, "read a.txt") + + finished = [e for e in events if e.type == EventType.TOOL_FINISHED] + assert len(finished) == 1 + assert finished[0].data["status"] == "ok" + + # The result content made it through. + assert any( + m.get("role") == "tool" and "hello" in m.get("content", "") + for m in engine.messages + ) + + +def test_guard_concurrent_tools_pass_through_within_turn(tmp_path): + """Two concurrent low-risk tools both pass GuardMiddleware within a single + turn. The engine authorises ALL tools before executing ANY of them, so the + fan-out counter is always 0 during evaluation within one turn — the limit + can only fire *across* iterations in the current architecture. + + This test verifies the wiring: barrier-based concurrency + guard middleware + + track_start/track_end all work together, and that after the turn the + fan-out counter is back to 0. + """ + from coworker.guard.ruleset import GuardRule + + barrier = threading.Barrier(2, timeout=5) + low = ai.ToolMetadata( + category="search", risk_level="low", requires_approval=False + ) + + # Use an empty tool name (match any tool) so the rule fires against + # whatever the registered function is called. + rules = [ + GuardRule( + name="limit-concurrent", + tool="", + action="deny", + max_concurrent=1, + ), + ] + + def explore(task: str): + barrier.wait() + return {"report": task} + + engine, registry = _guard_bare_engine( + tmp_path, + [ + _multi_tool_turn( + [("explore", {"task": "x"}), ("explore", {"task": "y"})] + ), + _text_turn("done"), + ], + rules, + ) + registry.register(explore, metadata=low) + + events = _collect(engine, "research both") + finished = [e for e in events if e.type == EventType.TOOL_FINISHED] + assert len(finished) == 2 + # Both succeed because authorisation happens before track_start — + # the fan-out counter is 0 when both are evaluated. + assert all(e.data["status"] == "ok" for e in finished) + + # The fan-out counter must be back to 0 after the turn completes. + guard = engine.guard_middleware + assert guard._ruleset._fanout_counters.get("explore", 0) == 0 + + +def test_guard_serial_tools_track_correctly(tmp_path): + """Serial tools with max_concurrent=1 both pass because each finishes before + the next authorises. Verifies the engine calls track_start/track_end through + the guard middleware for non-concurrent (medium-risk) tools. + """ + from coworker.guard.ruleset import GuardRule + + medium = ai.ToolMetadata( + category="filesystem", risk_level="medium", requires_approval=False + ) + + # Empty tool name = match any tool. + rules = [ + GuardRule( + name="limit-serial", + tool="", + action="deny", + max_concurrent=1, + ), + ] + order: list[str] = [] + + def serial_tool(): + order.append("start") + time.sleep(0.1) + order.append("end") + return "ok" + + engine, registry = _guard_bare_engine( + tmp_path, + [ + _multi_tool_turn([("serial_tool", {}), ("serial_tool", {})]), + _text_turn("done"), + ], + rules, + ) + registry.register(serial_tool, metadata=medium) + + _collect(engine, "go") + + # Serial execution order: they run one at a time even when the turn + # requests two calls (medium-risk = not parallel_safe). + assert order == ["start", "end", "start", "end"] + + # Fan-out counters reset after turn. + guard = engine.guard_middleware + assert guard._ruleset._fanout_counters.get("serial_tool", 0) == 0 + + +def test_guard_concurrent_fanout_limit_denies_overflow(tmp_path): + """Three concurrent low-risk tools with max_concurrent=2 — the first two + succeed (proved concurrent via threading.Barrier) and the third is denied + by the guard middleware's eager ``track_start`` check. + + Unlike the earlier ``pass_through_within_turn`` test, the empty-tool + rule fires on every ``track_start`` call, so the moment counter reaches + *max_concurrent* the next tool is blocked before it ever executes. + """ + from coworker.guard.ruleset import GuardRule + + barrier = threading.Barrier(2, timeout=5) + low = ai.ToolMetadata( + category="search", risk_level="low", requires_approval=False + ) + + rules = [ + GuardRule( + name="limit-concurrent", + tool="", + action="deny", + max_concurrent=2, + ), + ] + + def explore(task: str): + barrier.wait() + return {"report": task} + + engine, registry = _guard_bare_engine( + tmp_path, + [ + _multi_tool_turn( + [ + ("explore", {"task": "a"}), + ("explore", {"task": "b"}), + ("explore", {"task": "c"}), + ] + ), + _text_turn("done"), + ], + rules, + ) + registry.register(explore, metadata=low) + + events = _collect(engine, "research three things") + finished = [e for e in events if e.type == EventType.TOOL_FINISHED] + assert len(finished) == 3 + + ok_tools = [e for e in finished if e.data["status"] == "ok"] + denied_tools = [e for e in finished if e.data["status"] == "denied"] + + # Two succeed (concurrent, barrier proves it), one is denied by fanout. + assert len(ok_tools) == 2, f"expected 2 ok, got {len(ok_tools)}" + assert len(denied_tools) == 1, f"expected 1 denied, got {len(denied_tools)}" + + # Denied tool cites the guard. + assert "guard:" in denied_tools[0].data.get("reason", "") + + # Fan-out counter reset after turn completes. + guard = engine.guard_middleware + assert guard._ruleset._fanout_counters.get("explore", 0) == 0 + + +def test_guard_middleware_exception_fails_safe_via_engine(tmp_path): + """When the guard middleware's ruleset raises during ``evaluate()``, the + middleware catches the exception and returns ``Decision(needs_user=True)``. + The engine emits ``PERMISSION_REQUIRED``, then the default approver + (``_deny_all``) denies the tool. + """ + + class _BrokenRuleset: + def evaluate(self, tool_name, arguments): + raise RuntimeError("guard ruleset crash") + + def track_start(self, tool_name): + return True + + def track_end(self, tool_name): + pass + + from coworker.guard.middleware import GuardMiddleware + + provider = ScriptedProvider( + [_tool_turn("read_file", {"path": "x.txt"}), _text_turn("done")] + ) + registry = ToolRegistry() + registry.register_all(ai.toolkits.files(root=str(tmp_path))) + permissions = PermissionEngine(workspace_root=tmp_path) + guard = GuardMiddleware( + permissions, + ruleset=_BrokenRuleset(), # type: ignore[arg-type] + log_path=str(tmp_path / "guard.log"), + ) + engine = TurnEngine( + provider=provider, + registry=registry, + permissions=permissions, + model="gpt-5.5", + guard_middleware=guard, + ) + + events = _collect(engine, "read x.txt") + + # The middleware exception triggers a permission prompt (fail safe). + assert EventType.PERMISSION_REQUIRED in _types(events) + + # The PERMISSION_REQUIRED event carries the guard error in its reason. + perm_events = [e for e in events if e.type == EventType.PERMISSION_REQUIRED] + assert any("guard" in e.data.get("reason", "") for e in perm_events) + + # The default approver (_deny_all) denies the tool. + finished = [e for e in events if e.type == EventType.TOOL_FINISHED] + assert len(finished) == 1 + assert finished[0].data["status"] == "denied" + + # The history carries an explanation. + assert any( + m.get("role") == "tool" and "not executed" in m.get("content", "") + for m in engine.messages + ) + + def test_outbound_replaces_images_for_non_vision_models(tmp_path): class NoVisionProvider(ScriptedProvider): def capabilities(self, model): diff --git a/tests/test_guard.py b/tests/test_guard.py index a5f4f12f..240802c5 100644 --- a/tests/test_guard.py +++ b/tests/test_guard.py @@ -9,6 +9,7 @@ import logging import os +import time from pathlib import Path from types import SimpleNamespace @@ -329,6 +330,7 @@ def test_has_changed_true_after_file_modification(self, tmp_path): yaml_path.write_text(yaml.dump({"rules": []})) loader = GuardConfigLoader(yaml_path) loader.load_config() + time.sleep(0.15) # ensure mtime differs (NTFS resolution ~100ms) # Modify the file yaml_path.write_text(yaml.dump({"rules": [{"name": "new", "tool": "x", "action": "deny"}]})) assert loader.has_changed() @@ -338,6 +340,7 @@ def test_cached_config_returned_without_change(self, tmp_path): yaml_path.write_text(yaml.dump({"rules": [{"name": "original"}]})) loader = GuardConfigLoader(yaml_path) first = loader.load_config() + time.sleep(0.15) # ensure mtime differs (NTFS resolution ~100ms) # Modify file behind our back yaml_path.write_text(yaml.dump({"rules": [{"name": "modified"}]})) # Without calling load_config again, has_changed is True @@ -421,13 +424,7 @@ def evaluate(self, tool_name, arguments): raise RuntimeError("something went wrong") def track_start(self, tool_name): - pass - - def track_end(self, tool_name): - pass - - def track_start(self, tool_name): - pass + return True def track_end(self, tool_name): pass