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/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/engine.py b/coworker/engine.py index 78a88316..aeed8838 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 ------------------------------------------------------ @@ -483,22 +486,77 @@ 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: + 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): - 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: + 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) + 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 +590,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..12de7bd8 --- /dev/null +++ b/coworker/guard/middleware.py @@ -0,0 +1,181 @@ +"""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) -> 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.""" + 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..361ceb74 --- /dev/null +++ b/coworker/guard/ruleset.py @@ -0,0 +1,207 @@ +"""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 + is_fanout_block: bool = False # True when the denial was caused by max_concurrent + + +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, + is_fanout_block=True, + ) + # 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) -> 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).""" + current = self._fanout_counters.get(tool_name, 0) + if current > 0: + self._fanout_counters[tool_name] = current - 1 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_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 new file mode 100644 index 00000000..240802c5 --- /dev/null +++ b/tests/test_guard.py @@ -0,0 +1,491 @@ +"""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 +import time +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() + 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() + + 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() + 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 + 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): + return True + + 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 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