diff --git a/application/conversations/__init__.py b/application/conversations/__init__.py index edf4398..b8d1fce 100644 --- a/application/conversations/__init__.py +++ b/application/conversations/__init__.py @@ -1,8 +1,10 @@ -"""Conversation use case: the lifecycle of one agent turn (EPIC R04).""" +"""Conversation use case: the lifecycle of one agent turn (EPIC R04) and the +tool approval policy every turn's tool calls go through (EPIC R05).""" from .conversation_application_service import ( ConversationApplicationService, TurnResult, ) +from .tool_policy_gateway import ConfirmGate, ToolPolicyGateway -__all__ = ["ConversationApplicationService", "TurnResult"] +__all__ = ["ConversationApplicationService", "TurnResult", "ToolPolicyGateway", "ConfirmGate"] diff --git a/application/conversations/tool_policy_gateway.py b/application/conversations/tool_policy_gateway.py new file mode 100644 index 0000000..7d60ffd --- /dev/null +++ b/application/conversations/tool_policy_gateway.py @@ -0,0 +1,83 @@ +"""ToolPolicyGateway - one confirm/deny decision path for every tool call +(R05-T03). + +Today "does this tool call need the user's OK first" is answered by a +different hand-written check per engine: + +* ``core/chat_agent.py::run_cowork`` — ``name in ("run_command", + "install_package")``, a literal tuple. +* ``core/code_agent.py::run_code`` — ``name in (WRITE_TOOLS | MS365_WRITE_TOOLS)``, + a set built from two other hand-maintained sets. +* MCP/connector tools (``core/mcp_client.py``, ``core/ext_connectors.py``) — + no check at all; ``chat_agent.py`` calls ``extra_executor(name, args)`` + directly. + +Three answers to the same question, and the third one is a real gap: an MCP +tool that deletes files or calls an external API today runs with zero +confirmation even when the user turned "confirm before running commands" on. + +This gateway answers the question from data (:class:`~domain.tools.tool_descriptor.ToolCapability` +via a :class:`~domain.tools.tool_registry.ToolRegistry`) instead of a literal +name list, so registering a tool with the right capability is what gates it - +nothing to remember at each new call site. R05-T04 is what actually registers +MCP/connector tools with a capability; this module only needs the mechanism +to exist. + +Pure Python: no Qt, no direct dialog. The actual approval prompt stays exactly +what it is today - a ``gate`` object with a ``.request(payload) -> bool`` +method, supplied by the presentation layer (Settings' "confirm before running +commands" wires it up, or None for auto-run) - this module only decides +WHEN to ask it, never how to render the question. +""" +from __future__ import annotations + +from typing import Any, Dict, Optional, Protocol + +from cowork_local.domain.tools import ToolCapability, ToolRegistry + + +class ConfirmGate(Protocol): + """Shape of the existing ``PermissionGate`` both engines already use.""" + + def request(self, payload: Dict[str, Any]) -> bool: ... + + +class ToolPolicyGateway: + """Decides whether a tool call needs approval, for ONE calling surface. + + ``gated_capabilities`` is what makes this per-surface: Cowork only ever + asked about ``run_command``/``install_package`` (capability ``EXECUTE``), + while the Code tab additionally confirms plain file writes (capability + ``WRITE``). Passing the wrong set here would silently change which tools + prompt for approval - see the callers in ``core/chat_agent.py`` and + ``core/code_agent.py`` for the exact sets that preserve today's behavior. + """ + + def __init__(self, registry: ToolRegistry, gated_capabilities: ToolCapability) -> None: + self._registry = registry + self._gated_capabilities = gated_capabilities + + def requires_confirmation(self, name: str) -> bool: + """True when ``name``'s declared capabilities overlap this surface's + gated set. An unregistered tool never requires confirmation through + this path - callers that must fail safe on unknown tools check + ``name in registry`` themselves (see R05-T04's MCP wrapping, which + registers every tool it exposes before any call can reach here).""" + return bool(self._registry.capabilities_for(name) & self._gated_capabilities) + + def allow(self, name: str, gate: Optional[ConfirmGate], payload: Dict[str, Any]) -> bool: + """True when the call may proceed. + + ``gate is None`` preserves each engine's existing "no gate wired - + auto-run" behavior; a tool outside ``gated_capabilities`` is never + asked about, matching read-only tools "never confirm" today. + ``payload`` is whatever ``gate.request(...)`` already expects at that + call site (the two engines use slightly different dict shapes) - this + gateway only decides WHETHER to call it, never reshapes the payload. + """ + if gate is None or not self.requires_confirmation(name): + return True + return bool(gate.request(payload)) + + +__all__ = ["ToolPolicyGateway", "ConfirmGate"] diff --git a/core/chat_agent.py b/core/chat_agent.py index 20b3d26..eaa0917 100644 --- a/core/chat_agent.py +++ b/core/chat_agent.py @@ -11,6 +11,8 @@ import re from pathlib import Path from typing import Any, Callable, Dict, List, Optional +from ..application.conversations.tool_policy_gateway import ToolPolicyGateway +from ..domain.tools import ToolCapability, default_registry from ..providers.base import Provider, ToolSpec from . import agent_roles from . import agent_security @@ -27,6 +29,13 @@ from .tools import TOOL_SPECS, ToolContext, _snapshot, describe_action, execute_ # Generator / helper scripts — never a final deliverable in Cowork's output. _SCRIPT_EXTS = {".py", ".pyw", ".js", ".mjs", ".cjs", ".ts", ".sh", ".bat", ".ps1", ".rb", ".pl"} +# R05-T03/T04: replaces the literal ``name in ("run_command", +# "install_package")`` check below with a capability lookup — EXECUTE is +# exactly the capability those two (and only those two) built-in tools carry +# (see domain/tools/tool_registry.py::BUILT_IN_CAPABILITIES). Copied per-turn +# into ``turn_tool_policy`` inside run_cowork() once extra_tools are known. +_COWORK_TOOL_REGISTRY = default_registry(TOOL_SPECS) + EmitFn = Callable[[Dict[str, Any]], None] CancelFn = Callable[[], bool] @@ -388,6 +397,19 @@ def run_cowork( jira=(security_config.data.get("jira") if security_config else None)) extra_tools = extra_tools or [] extra_names = {t.name for t in extra_tools} + # R05-T04: MCP servers (core/mcp_client.py) and unified connectors + # (core/ext_connectors.py) — everything that arrives here as extra_tools — + # advertise no standard risk metadata, so each is tagged with the same + # conservative default (WRITE|EXECUTE|NETWORK) domain/tools/tool_registry.py + # uses for any unclassified tool. Copying the built-in registry per turn + # (cheap - under 20 entries) rather than mutating the shared module-level + # one keeps different turns' extra_tools from leaking into each other. + from ..domain.tools import ToolDescriptor, ToolRegistry + from ..domain.tools.tool_registry import UNKNOWN_SOURCE_CAPABILITIES + _turn_registry = ToolRegistry(_COWORK_TOOL_REGISTRY.all()) + for _spec in extra_tools: + _turn_registry.register(ToolDescriptor.from_spec(_spec, UNKNOWN_SOURCE_CAPABILITIES)) + turn_tool_policy = ToolPolicyGateway(_turn_registry, ToolCapability.EXECUTE) # update_plan drives the Plan panel (above Output); it produces no file. # Built-in tools the admin disabled (Monitoring → Tools) are filtered out. from .tools import enabled_tool_specs @@ -489,6 +511,18 @@ def run_cowork( preview = {"kind": "info", "title": name, "text": str(args)} emit({"type": "tool_proposed", "id": tc_id, "name": name, "args": args, "preview": preview}) + # R05-T04: MCP/connector tools used to run with NO permission + # check at all — this is what closes that gap. Same policy, + # same gate object as the built-in tools below. + if not turn_tool_policy.allow( + name, gate, {"name": name, "args": args, "preview": preview} + ): + result = {"ok": False, "output": "Rejected by user."} + emit({"type": "tool_result", "id": tc_id, "name": name, + "ok": False, "output": result["output"]}) + messages.append({"role": "tool", "tool_call_id": tc_id, "name": name, + "content": result["output"]}) + continue result = extra_executor(name, args) emit({"type": "tool_result", "id": tc_id, "name": name, "ok": result.get("ok", False), "output": result.get("output", "")}) @@ -528,16 +562,18 @@ def run_cowork( # Permission Management (Sandbox Security Layer) — only when a # gate was actually supplied (Settings: "confirm before running # commands"); None preserves the pre-existing auto-run behavior. - if gate is not None and name in ("run_command", "install_package"): - approved = gate.request({"name": name, "args": args, "preview": preview}) - if not approved: - result = {"ok": False, "output": "Rejected by user."} - evt = {"type": "tool_result", "id": tc_id, "name": name, - "ok": False, "output": result["output"]} - emit(evt) - messages.append({"role": "tool", "tool_call_id": tc_id, - "name": name, "content": result["output"]}) - continue + # R05-T03: gating is now capability-driven (see + # turn_tool_policy above) instead of a literal name tuple. + if not turn_tool_policy.allow( + name, gate, {"name": name, "args": args, "preview": preview} + ): + result = {"ok": False, "output": "Rejected by user."} + evt = {"type": "tool_result", "id": tc_id, "name": name, + "ok": False, "output": result["output"]} + emit(evt) + messages.append({"role": "tool", "tool_call_id": tc_id, + "name": name, "content": result["output"]}) + continue if name == "save_file": result = _do_save_file(output_dir, title, args) diff --git a/core/code_agent.py b/core/code_agent.py index c95420f..9cb47c6 100644 --- a/core/code_agent.py +++ b/core/code_agent.py @@ -12,6 +12,8 @@ import re from pathlib import Path from typing import Any, Callable, Dict, List, Optional +from ..application.conversations.tool_policy_gateway import ToolPolicyGateway +from ..domain.tools import ToolCapability, ToolDescriptor, ToolRegistry from ..providers.base import Provider from . import agent_roles from . import agent_security @@ -225,6 +227,14 @@ def run_code( # read/list ms365 tools count as "read-only, never confirm". Names are # the MCP-qualified "ms365__*" form the agent sees (see ms365_tools.py). gated_tools = WRITE_TOOLS | MS365_WRITE_TOOLS + # R05-T03/T04: ``gated_tools`` stays the authoritative name set (unchanged), + # but the actual confirm decision now goes through the same + # ToolPolicyGateway class run_cowork uses, instead of a separate + # hand-rolled ``if name in gated_tools`` + direct ``gate.request(...)``. + code_tool_policy = ToolPolicyGateway( + ToolRegistry(ToolDescriptor(n, "", {}, ToolCapability.WRITE) for n in gated_tools), + ToolCapability.WRITE, + ) # In PLAN mode, don't advertise write/run tools (analysis only). advertised = [t for t in all_tools if t.name not in gated_tools] if plan else all_tools has_memory = any(t.name.startswith("cmem_") for t in extra_tools) @@ -297,10 +307,11 @@ def run_code( agent_security.enforce_command(provider, name, args, security_config, emit, agent_kind="code") - if name in gated_tools: - approved = gate.request({"id": tc_id, "name": name, "args": args, "preview": preview}) - else: - approved = True # read-only tools (incl. codebase memory) never confirm + # read-only tools (incl. codebase memory) never consult the gate — + # code_tool_policy.requires_confirmation(name) is False for them. + approved = code_tool_policy.allow( + name, gate, {"id": tc_id, "name": name, "args": args, "preview": preview} + ) if cancel(): return messages diff --git a/core/mcp_client.py b/core/mcp_client.py index b8ecf7b..4df0701 100644 --- a/core/mcp_client.py +++ b/core/mcp_client.py @@ -106,6 +106,13 @@ class McpServerConnection: if self._thread is not None: self._thread.join(timeout=5) + def is_alive(self) -> bool: + """True while the connection's background thread (and therefore its + event loop and subprocess) is still running — used by + ``infrastructure/mcp/mcp_source_manager.py`` (R05-T05) to tell a live + cached connection from one whose subprocess already died.""" + return self._thread is not None and self._thread.is_alive() + # ---- tools ----------------------------------------------------------- def list_tool_specs(self) -> List[ToolSpec]: """The server's tools, wrapped as :class:`ToolSpec` — the same shape diff --git a/core/tools.py b/core/tools.py index e3ac87a..e6e2fda 100644 --- a/core/tools.py +++ b/core/tools.py @@ -3,79 +3,29 @@ Every path is resolved relative to the working directory and must stay inside it (path-traversal is rejected). ``run_command`` executes inside the workdir with a timeout and captured output. + +R05-T02: the actual handlers (``read_file``/``list_dir``/``write_file``/ +``edit_file``/``run_command``/``install_package``/``fetch_url``/ +``jira_search``/``jira_get_issue``) now live in +``infrastructure/filesystem/{file_tools,command_tools,fetch_tools}.py``, split +out of what used to be one big if/elif chain here. This module is the +strangler-fig shim (ADR-001 section 4): it re-exports ``ToolContext``/ +``ToolError`` (actually defined in +``infrastructure/filesystem/tool_context.py`` now) so every existing +``from .tools import ToolContext`` keeps working, and ``execute_tool`` +dispatches through a small ``{name: handler}`` table built from the moved +modules instead of the chain itself. """ from __future__ import annotations -import ast import difflib -import os -from dataclasses import dataclass -from pathlib import Path from typing import Any, Callable, Dict, List, Optional +from ..infrastructure.filesystem import command_tools, fetch_tools, file_tools +from ..infrastructure.filesystem.command_tools import _snapshot # noqa: F401 - re-export, core/chat_agent.py imports this name +from ..infrastructure.filesystem.tool_context import CancelFn, ToolContext, ToolError # noqa: F401 - re-export from ..providers.base import ToolSpec -CancelFn = Callable[[], bool] - -MAX_READ_BYTES = 200_000 -COMMAND_TIMEOUT = 120 # seconds - - -class ToolError(Exception): - pass - - -def _flatten_rel(rel: str) -> str: - """Collapse a sub-folder path down to a bare filename so the file lands in the - workdir root — EXCEPT the ``.scratch`` sandbox subtree, which is preserved. - - Used by the Cowork agent (flatten_writes=True) so it can never create a - per-session / per-chat / per-task output sub-folder: every deliverable stays - directly in the single configured Output folder.""" - parts = Path(rel).parts - if parts and parts[0] == ".scratch": - return rel # temporary sandbox is allowed (and cleaned up afterwards) - return Path(rel).name or rel - - -@dataclass -class ToolContext: - workdir: Path - flatten_writes: bool = False # Cowork: force every write into the workdir root - sandbox: bool = False # Code tab: isolate run_command/install_package into /.venv - # Sandbox Security Layer — Settings' "Resource Limits" (cpu_percent/memory_mb/ - # disk_mb), applied to every run_command/install_package this context runs. - # None (default) = no limits, matching pre-existing behavior. - resource_limits: Optional[Dict[str, float]] = None - # Sandbox Security Layer — Settings' "Block network for agent commands" - # (policy-level, see deps.py::network_blocked_env). False (default) = - # unrestricted, matching pre-existing behavior. - block_network: bool = False - # Whether the fetch_url tool may read URLs — SEPARATE from block_network - # (reading a web page/share link for info is safe; running networked shell - # commands is the risk). Defaults True; set from agent_security.allow_url_fetch. - allow_url_fetch: bool = True - # Jira read connector config (base_url/email/api_token) — None disables the - # jira_* tools' ability to connect. Populated from config.data["jira"]. - jira: Optional[Dict[str, Any]] = None - - def resolve(self, rel: str) -> Path: - """Resolve ``rel`` inside the workdir, rejecting escapes.""" - if rel in ("", "."): - return self.workdir - candidate = (self.workdir / rel).expanduser() - try: - resolved = candidate.resolve() - except OSError as exc: - raise ToolError(f"Invalid path: {rel} ({exc})") - root = self.workdir.resolve() - if resolved != root and root not in resolved.parents: - raise ToolError( - f"Refused: '{rel}' is outside the working folder ({root})." - ) - return resolved - - # -------------------------------------------------------------------------- # Tool specs advertised to the model # -------------------------------------------------------------------------- @@ -192,6 +142,23 @@ TOOL_SPECS: List[ToolSpec] = [ # Actions gated by the permission gate in confirm mode (auto-approved in Auto-run). WRITE_TOOLS = {"write_file", "edit_file", "run_command", "install_package"} +# name -> handler(ctx, args[, cancel, on_output]) — built once from the split +# infrastructure modules. Replaces the if/elif chain execute_tool used to be. +_HANDLERS: Dict[str, Callable[..., Dict[str, Any]]] = { + "read_file": file_tools.read_file, + "list_dir": file_tools.list_dir, + "write_file": file_tools.write_file, + "edit_file": file_tools.edit_file, + "run_command": command_tools.run_command, + "install_package": command_tools.install_package, + "fetch_url": fetch_tools.fetch_url, + "jira_search": fetch_tools.jira_search, + "jira_get_issue": fetch_tools.jira_get_issue, +} +# Handlers that accept the long-running (cancel, on_output) signature — every +# other handler takes just (ctx, args). +_CANCELLABLE = {"run_command", "install_package"} + def enabled_tool_specs(security_config=None) -> List[ToolSpec]: """The built-in TOOL_SPECS minus any the admin turned OFF in Monitoring → @@ -301,27 +268,14 @@ def execute_tool(ctx: ToolContext, name: str, args: Dict[str, Any], labels WHICH agent role made it.""" from . import audit_log + handler = _HANDLERS.get(name) try: - if name == "read_file": - result = _read_file(ctx, args) - elif name == "list_dir": - result = _list_dir(ctx, args) - elif name == "write_file": - result = _write_file(ctx, args) - elif name == "edit_file": - result = _edit_file(ctx, args) - elif name == "run_command": - result = _run_command(ctx, args, cancel, on_output) - elif name == "install_package": - result = _install_package(ctx, args, cancel, on_output) - elif name == "fetch_url": - result = _fetch_url(ctx, args) - elif name == "jira_search": - result = _jira_search(ctx, args) - elif name == "jira_get_issue": - result = _jira_get_issue(ctx, args) - else: + if handler is None: result = {"ok": False, "output": f"Tool not found: {name}"} + elif name in _CANCELLABLE: + result = handler(ctx, args, cancel, on_output) + else: + result = handler(ctx, args) except ToolError as exc: result = {"ok": False, "output": str(exc)} except Exception as exc: # defensive: a tool must never crash the agent @@ -331,234 +285,6 @@ def execute_tool(ctx: ToolContext, name: str, args: Dict[str, Any], return result -def _fetch_url(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: - """Fetch a URL's text content (web page / online document / SharePoint- - OneDrive share link) via link_fetch — the same parser task-link attachments - use. Honors the Sandbox Security Layer's "Block network" policy.""" - url = str(args.get("url", "")).strip() - if not url: - return {"ok": False, "output": "fetch_url: 'url' is required."} - if not url.lower().startswith(("http://", "https://")): - return {"ok": False, "output": f"fetch_url: not an http(s) URL: {url}"} - if not ctx.allow_url_fetch: - return {"ok": False, - "output": ("fetch_url: URL fetching is turned off in Settings → Security " - "(\"Allow the agent to fetch URLs\").")} - # A pasted Jira issue link on the CONNECTED Jira host is read via the - # authenticated API (so private issues resolve, not a login page). Public - # links / any other URL fall through to the normal fetcher below. - from . import jira_tool - if jira_tool.is_jira_issue_url(ctx.jira, url): - return {"ok": True, "output": jira_tool.get_issue_by_url(ctx.jira, url)} - from .link_fetch import fetch_link_preview - - return {"ok": True, "output": fetch_link_preview(url)} - - -def _jira_search(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: - from . import jira_tool - - out = jira_tool.search(ctx.jira, str(args.get("jql", "")), - int(args.get("max_results", 25) or 25)) - return {"ok": not out.lower().startswith(("jira is not configured", "jira search failed")), - "output": out} - - -def _jira_get_issue(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: - from . import jira_tool - - out = jira_tool.get_issue(ctx.jira, str(args.get("key", ""))) - return {"ok": not out.lower().startswith(("jira is not configured", "could not fetch")), - "output": out} - - -def _read_file(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: - target = ctx.resolve(str(args.get("path", ""))) - if not target.exists(): - return {"ok": False, "output": f"File not found: {args.get('path')}"} - data = target.read_bytes()[:MAX_READ_BYTES] - text = data.decode("utf-8", errors="replace") - return {"ok": True, "output": text} - - -def _list_dir(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: - rel = str(args.get("path", ".") or ".") - target = ctx.resolve(rel) - # A missing/not-yet-created path is NOT a tool failure — report it as an - # ordinary result so the agent can create it or pick another path and keep - # going. Returning ok=False here surfaced a false "tool failed: list_dir" in - # Co4E flows and could stall a step on a recoverable situation. - if not target.exists(): - return {"ok": True, "output": f"(path '{rel}' does not exist yet — create it or use another path)"} - if target.is_file(): - return {"ok": True, "output": f"('{rel}' is a file, not a directory)"} - entries = [] - for child in sorted(target.iterdir(), key=lambda p: (p.is_file(), p.name.lower())): - marker = "/" if child.is_dir() else "" - entries.append(f"{child.name}{marker}") - return {"ok": True, "output": "\n".join(entries) or "(empty folder)"} - - -def _check_python_syntax(target: Path, content: str) -> str: - """Return a short warning if ``content`` is invalid Python, else ''. - - Catches syntax errors the instant a .py file is written/edited — before the - agent wastes a whole run_command round-trip just to get the same error back - from a traceback.""" - if target.suffix.lower() not in (".py", ".pyw"): - return "" - try: - ast.parse(content, filename=str(target)) - return "" - except SyntaxError as exc: - return f"\n⚠ Syntax error at line {exc.lineno}: {exc.msg} — fix this before running the file." - - -def _write_file(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: - rel = str(args.get("path", "")) - if ctx.flatten_writes: - rel = _flatten_rel(rel) - target = ctx.resolve(rel) - content = str(args.get("content", "")) - target.parent.mkdir(parents=True, exist_ok=True) - # A .xlsx is a binary package — build a REAL workbook from the content - # (CSV/TSV/Markdown-table/JSON) rather than writing raw text (which corrupts it). - if target.suffix.lower() in (".xlsx", ".xlsm"): - from . import xlsx_write - if xlsx_write.build_xlsx_from_text(target, content): - return {"ok": True, "path": str(target), - "output": f"Wrote spreadsheet {rel} ({target.name})."} - return {"ok": False, "output": "Could not build the .xlsx (openpyxl unavailable) — " - "write a .csv instead, or use a generator script."} - target.write_text(content, encoding="utf-8") - warning = _check_python_syntax(target, content) - return {"ok": True, "path": str(target), - "output": f"Wrote {len(content)} chars to {rel}.{warning}"} - - -def _edit_file(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: - """Replace an exact snippet inside an existing file (precise patch edit).""" - rel = str(args.get("path", "")) - if ctx.flatten_writes: - rel = _flatten_rel(rel) - target = ctx.resolve(rel) - if not target.exists(): - return {"ok": False, - "output": f"File not found: {rel} — use write_file to create it."} - old = str(args.get("old_string", "")) - new = str(args.get("new_string", "")) - replace_all = bool(args.get("replace_all", False)) - if not old: - return {"ok": False, "output": "old_string is empty — provide the exact text to replace."} - try: - text = target.read_text(encoding="utf-8", errors="replace") - except OSError as exc: - return {"ok": False, "output": f"Could not read file: {exc}"} - count = text.count(old) - if count == 0: - return {"ok": False, "output": ("old_string not found. Read the file and copy the exact " - "text to replace, including indentation/whitespace.")} - if count > 1 and not replace_all: - return {"ok": False, "output": (f"old_string appears {count} times — add surrounding " - "context to make it unique, or set replace_all=true.")} - updated = text.replace(old, new) if replace_all else text.replace(old, new, 1) - target.write_text(updated, encoding="utf-8") - n = count if replace_all else 1 - warning = _check_python_syntax(target, updated) - return {"ok": True, - "output": f"Edited {args.get('path')} ({n} replacement{'' if n == 1 else 's'}).{warning}"} - - -def _sandbox_python(ctx: ToolContext, cancel: Optional[CancelFn] = None, - on_output: Optional[Callable[[str], None]] = None) -> Optional[str]: - """Lazily create/reuse this ctx's project sandbox venv (Code tab only — - ``ctx.sandbox``); returns its python path, or None to use the app's own.""" - if not ctx.sandbox: - return None - from .deps import ensure_project_venv - - py = ensure_project_venv(ctx.workdir, cancel=cancel, on_output=on_output) - return str(py) if py else None - - -def _install_package(ctx: ToolContext, args: Dict[str, Any], cancel: Optional[CancelFn] = None, - on_output: Optional[Callable[[str], None]] = None) -> Dict[str, Any]: - from .deps import pip_install - - package = str(args.get("package", "")).strip() - if not package: - return {"ok": False, "output": "No package specified."} - python = _sandbox_python(ctx, cancel, on_output) - ok, detail = pip_install(package, cancel=cancel, on_output=on_output, python=python) - head = f"Installed {package}." if ok else f"Could not install {package}." - return {"ok": ok, "output": f"{head}\n{detail}"} - - -_SNAPSHOT_SKIP = {".git", "__pycache__", "node_modules", ".scratch", ".venv", - ".idea", ".mypy_cache", ".pytest_cache"} - - -def _snapshot(workdir: Path) -> Dict[str, Any]: - """Map of file path -> (mtime, size) under the workdir (noise dirs skipped).""" - snap: Dict[str, Any] = {} - try: - for dirpath, dirnames, filenames in os.walk(str(workdir)): - dirnames[:] = [d for d in dirnames if d not in _SNAPSHOT_SKIP] - for fn in filenames: - full = os.path.join(dirpath, fn) - try: - st = os.stat(full) - snap[full] = (st.st_mtime_ns, st.st_size) - except OSError: - pass - if len(snap) > 5000: - return snap - except OSError: - pass - return snap - - -def _run_command(ctx: ToolContext, args: Dict[str, Any], - cancel: Optional[CancelFn] = None, - on_output: Optional[Callable[[str], None]] = None) -> Dict[str, Any]: - from .deps import network_blocked_env, run_cancellable, sandbox_env - from .sandbox_manager import SandboxManager, ExecutionConfig - from ..security.command_risk_classifier import classify_command - - command = str(args.get("command", "")).strip() - if not command: - return {"ok": False, "output": "Empty command."} - - # --- Security validation pipeline --- - risk = classify_command(command, is_cowork_mode=ctx.flatten_writes) - if risk.blocked: - denial = "Command blocked by security policy: " + "; ".join(risk.reasons) - return {"ok": False, "output": denial} - - # Route through SandboxManager for risk-based isolation - mgr = SandboxManager(ExecutionConfig( - enabled=True, - block_network_by_default=ctx.block_network, - is_cowork_mode=ctx.flatten_writes, - )) - sandbox_result = mgr.run( - command=command, - workdir=str(ctx.workdir), - block_network=ctx.block_network, - timeout_sec=COMMAND_TIMEOUT, - cancel=cancel, - ) - # Sandbox ALWAYS executes (never double-run). Return its result directly. - if sandbox_result.get("sandbox") == "blocked": - return {"ok": False, "output": sandbox_result.get("stderr", "Command blocked")} - out = sandbox_result.get("stdout", "").strip() or "(no output)" - err = sandbox_result.get("stderr", "") - rc = sandbox_result.get("returncode", -1) - if err: - out = f"{out}\n{err}" if out else err - return {"ok": sandbox_result.get("ok", False), "output": f"[exit {rc}]\n{out}"} - - def _short_json(obj: Any, limit: int = 500) -> str: import json text = json.dumps(obj, ensure_ascii=False, indent=2) diff --git a/docs/refactor/Refactoring_Checklist.md b/docs/refactor/Refactoring_Checklist.md index eb073d8..8c48edd 100644 --- a/docs/refactor/Refactoring_Checklist.md +++ b/docs/refactor/Refactoring_Checklist.md @@ -57,6 +57,34 @@ --- +## 📊 TIẾN ĐỘ THỰC TẾ — TEAM HOA (cập nhật `2026-08-21 22:19`) + +> [!NOTE] +> ### ✅ ĐÃ HOÀN TẤT: 5/5 task của **R05** — branch `feature/teamhoa/r05-r06` (tạo từ `origin/feature/deltateam/refactor-plan`, có sẵn nền R01/R03/R04) +> +> | EPIC | Task | Trạng thái | +> | :--- | :--- | :--- | +> | **R05** Tool, MCP & Connector Policy | T01 → T05 | ✅ 5/5 | +> | **R06** Workspace, Filesystem & History Isolation | T01 → T05 | ⬜ chưa bắt đầu | +> +> **Kiểm chứng (chạy thật):** +> * `pytest tests/` ➔ **254 pass / 4 fail** (+12 test mới cho R05: `tests/unit/test_tool_registry_and_policy.py`, `test_code_agent_tool_policy.py`, `test_cowork_extra_tool_policy.py`, `test_mcp_source_manager.py`) +> * 4 fail là **lỗi có sẵn từ trước R05**, không liên quan tool/MCP: 2 trong `test_config_security.py` (EPIC R02, đã ghi nhận bởi Team Duy) + 2 trong `test_routing_wiring.py` (môi trường máy này có Ollama/llama3.1 thật + config routing cục bộ khác "fresh install" — không phải do R05). +> * `python scripts/check_imports.py` ➔ **PASS** (0 Qt import trong `domain/`, `application/`) +> * Mọi file mới **< 400 dòng** (lớn nhất: `domain/tools/tool_registry.py` 125 dòng). `core/tools.py` giảm từ 566 ➔ 291 dòng. +> +> ### 🔧 TÓM TẮT R05 +> * **R05-T01/T02**: `core/tools.py`'s if/elif dispatcher tách thành `infrastructure/filesystem/{file_tools,command_tools,fetch_tools,tool_context}.py` + `domain/tools/{tool_descriptor,tool_registry}.py`. `core/tools.py` còn lại là shim strangler-fig (re-export `ToolContext`/`ToolError`, dispatch qua dict). +> * **R05-T03**: `application/conversations/tool_policy_gateway.py::ToolPolicyGateway` — thay `if gate is not None and name in ("run_command","install_package")` (chat_agent.py) và `if name in (WRITE_TOOLS|MS365_WRITE_TOOLS)` (code_agent.py) bằng một lookup capability chung. Đã verify bằng test: đúng 2 tool cũ vẫn được gate, không tool nào khác bị ảnh hưởng. +> * **R05-T04 — ⚠️ THAY ĐỔI HÀNH VI CÓ CHỦ ĐÍCH**: trước đây MCP/connector/ext-connector tools (`core/mcp_client.py`, `core/ext_connectors.py`) chạy qua `extra_executor(name, args)` **không hề qua permission gate**. Giờ mọi `extra_tools` được gắn capability mặc định (`WRITE|EXECUTE|NETWORK`, vì MCP không có chuẩn khai báo rủi ro) và đi qua CÙNG `ToolPolicyGateway` như built-in tools. Khi Settings có "confirm before running commands" bật, tool MCP/connector giờ sẽ hỏi xác nhận — người dùng SẼ thấy thêm prompt so với trước. Test: `tests/unit/test_cowork_extra_tool_policy.py`. +> * **R05-T05**: `infrastructure/mcp/mcp_source_manager.py::McpToolSourceManager` — tách lifecycle connection (cache/lock/start-or-skip) ra khỏi `state.py::AppContext` (trước đây inline trong `_mcp_connections`/`_conn_lock`). `AppContext` giờ chỉ gọi `self._mcp_manager.ensure/stop/stop_all`. `_ext_connections` (Connectors CAD/CAE/MS365/Other) KHÔNG thuộc phạm vi T05, vẫn giữ `_conn_lock` riêng như cũ. +> +> ### 📌 CÒN NỢ / CẦN QUYẾT ĐỊNH +> 1. **Xung đột file với EPIC R02 (Team Nam)**: `docs/refactor/Refactoring_Checklist.md` dòng ~83 giao `infrastructure/persistence/json/atomic_json_file.py` cho Team Nam (R02-T01). R06-T02 (Team Hoa) cũng cần một helper ghi JSON atomic cho `WorkspaceRepository`/`ConversationRepository`. Để tránh 2 team cùng sửa 1 file, R06 sẽ dùng một helper atomic-write cục bộ trong `infrastructure/persistence/json/workspace_repository_impl.py`/`conversation_repository_impl.py` cho tới khi R02 xong, rồi hợp nhất vào `atomic_json_file.py` chung — **cần Team Nam xác nhận** khi họ bắt đầu R02-T01. +> 2. Việc kế tiếp của Team Hoa là **R06** (Workspace, Filesystem & History Isolation). + +--- + ## 📌 PHẦN 1: CHECKLIST CHI TIẾT THEO 10 EPIC (R01 ➔ R10) ### 🔹 EPIC R01: Architecture Foundation & Characterization (Nền Tảng Kiến Trúc & Test Bảo Vệ) @@ -135,16 +163,16 @@ * **Team chịu trách nhiệm**: 🟢 **Team Hoa** (Chủ trì) + Phối hợp Team Duy * **Mục tiêu**: Bóc tách monolithic `core/tools.py`, đưa toàn bộ Built-in tools, MCP tools và Connectors qua `ToolPolicyGateway` kiểm tra quyền phân tầng. -- [ ] **R05-T01 (Team Hoa)**: Định nghĩa `ToolDescriptor`, `ToolCapability` (read/write/execute/network) ➔ `domain/tools/tool_descriptor.py` & `domain/tools/tool_registry.py` - *Start: `____-__-__ __:__` | End: `____-__-__ __:__`* -- [ ] **R05-T02 (Team Hoa)**: Tách nhỏ các built-in handlers từ `core/tools.py` ➔ `infrastructure/filesystem/file_tools.py`, `command_tools.py`, `fetch_tools.py` - *Start: `____-__-__ __:__` | End: `____-__-__ __:__`* -- [ ] **R05-T03 (Team Hoa)**: Xây dựng `ToolPolicyGateway` (kiểm tra phân quyền allow/confirm/deny) ➔ `application/conversations/tool_policy_gateway.py` - *Start: `____-__-__ __:__` | End: `____-__-__ __:__`* -- [ ] **R05-T04 (Team Hoa)**: Chuẩn hóa MCP tools từ `core/mcp_client.py` đi qua `ToolPolicyGateway` - *Start: `____-__-__ __:__` | End: `____-__-__ __:__`* -- [ ] **R05-T05 (Team Hoa)**: Xây dựng `McpToolSourceManager` quản lý vòng đời tiến trình MCP ➔ `infrastructure/mcp/mcp_source_manager.py` - *Start: `____-__-__ __:__` | End: `____-__-__ __:__`* +- [x] **R05-T01 (Team Hoa)**: Định nghĩa `ToolDescriptor`, `ToolCapability` (read/write/execute/network) ➔ `domain/tools/tool_descriptor.py` & `domain/tools/tool_registry.py` + *Start: `2026-08-21 21:40` | End: `2026-08-21 21:47`* +- [x] **R05-T02 (Team Hoa)**: Tách nhỏ các built-in handlers từ `core/tools.py` ➔ `infrastructure/filesystem/file_tools.py`, `command_tools.py`, `fetch_tools.py` + *Start: `2026-08-21 21:47` | End: `2026-08-21 21:56`* +- [x] **R05-T03 (Team Hoa)**: Xây dựng `ToolPolicyGateway` (kiểm tra phân quyền allow/confirm/deny) ➔ `application/conversations/tool_policy_gateway.py` + *Start: `2026-08-21 21:56` | End: `2026-08-21 22:04`* +- [x] **R05-T04 (Team Hoa)**: Chuẩn hóa MCP tools từ `core/mcp_client.py` đi qua `ToolPolicyGateway` + *Start: `2026-08-21 22:04` | End: `2026-08-21 22:12`* +- [x] **R05-T05 (Team Hoa)**: Xây dựng `McpToolSourceManager` quản lý vòng đời tiến trình MCP ➔ `infrastructure/mcp/mcp_source_manager.py` + *Start: `2026-08-21 22:12` | End: `2026-08-21 22:19`* --- diff --git a/domain/tools/__init__.py b/domain/tools/__init__.py new file mode 100644 index 0000000..c9de3ac --- /dev/null +++ b/domain/tools/__init__.py @@ -0,0 +1,18 @@ +"""Domain entities for tool risk classification and lookup (EPIC R05).""" + +from .tool_descriptor import ToolCapability, ToolDescriptor +from .tool_registry import ( + BUILT_IN_CAPABILITIES, + UNKNOWN_SOURCE_CAPABILITIES, + ToolRegistry, + default_registry, +) + +__all__ = [ + "ToolCapability", + "ToolDescriptor", + "ToolRegistry", + "BUILT_IN_CAPABILITIES", + "UNKNOWN_SOURCE_CAPABILITIES", + "default_registry", +] diff --git a/domain/tools/tool_descriptor.py b/domain/tools/tool_descriptor.py new file mode 100644 index 0000000..5c8d21f --- /dev/null +++ b/domain/tools/tool_descriptor.py @@ -0,0 +1,86 @@ +"""ToolCapability / ToolDescriptor - the risk-tagged catalogue entry for one +tool the agent loop can call (R05-T01). + +Today a tool is just a name inside ``core/tools.py::TOOL_SPECS`` (a +``providers.base.ToolSpec`` — name/description/JSON-schema parameters, with +no notion of risk) plus a hand-written membership test wherever gating is +needed: ``core/tools.py::WRITE_TOOLS``, ``core/code_agent.py``'s +``WRITE_TOOLS | MS365_WRITE_TOOLS``, and ``core/chat_agent.py``'s literal +``name in ("run_command", "install_package")``. Three call sites, three +independently-maintained lists, and a new tool (or an MCP/connector tool, +which has no list membership at all - see ``core/mcp_client.py``) is gated +only if someone remembers to add it everywhere. + +``ToolDescriptor`` makes the risk an attribute of the tool itself, declared +once, so ``application/conversations/tool_policy_gateway.py`` (R05-T03) can +decide ALLOW/CONFIRM/DENY from data instead of a growing set of literal +tuples. + +Pure domain code: stdlib only, no Qt, no I/O. ``to_spec``/``from_spec`` are +the only place this module touches something outside domain/, and that +something (``providers.base.ToolSpec``) is itself a plain dataclass with no +further dependencies. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Flag, auto +from typing import Any, Dict + +from cowork_local.providers.base import ToolSpec + + +class ToolCapability(Flag): + """What calling a tool can do to the machine or the network. + + A ``Flag`` (not a plain ``Enum``) because a single tool can combine risks + - ``install_package`` writes to the environment, runs pip as a + subprocess, AND needs network access. Composing three separate booleans + per call site is exactly the duplication this type replaces. + """ + + NONE = 0 + READ = auto() + WRITE = auto() + EXECUTE = auto() + NETWORK = auto() + + +@dataclass(frozen=True) +class ToolDescriptor: + """An immutable description of one callable tool. + + Attributes: + name: the identifier the model calls (``ToolSpec.name``). + description: shown to the model, unchanged from ``ToolSpec``. + parameters: JSON-Schema object for the call's arguments. + capabilities: the risk this tool carries - see :class:`ToolCapability`. + """ + + name: str + description: str + parameters: Dict[str, Any] = field(default_factory=dict) + capabilities: ToolCapability = ToolCapability.NONE + + def has(self, capability: ToolCapability) -> bool: + """True when this tool carries (any bit of) ``capability``.""" + return bool(self.capabilities & capability) + + def to_spec(self) -> ToolSpec: + """Project back to the ``ToolSpec`` shape the model-facing catalogue + and the provider call actually use - risk tagging is metadata the + wire format has no room for.""" + return ToolSpec(name=self.name, description=self.description, + parameters=self.parameters) + + @classmethod + def from_spec(cls, spec: ToolSpec, + capabilities: ToolCapability = ToolCapability.NONE) -> "ToolDescriptor": + """Wrap an existing ``ToolSpec`` (built-in, MCP, or connector) with a + capability tag. The one place callers attach risk to a spec they did + not author themselves.""" + return cls(name=spec.name, description=spec.description, + parameters=spec.parameters, capabilities=capabilities) + + +__all__ = ["ToolCapability", "ToolDescriptor"] diff --git a/domain/tools/tool_registry.py b/domain/tools/tool_registry.py new file mode 100644 index 0000000..0cfceab --- /dev/null +++ b/domain/tools/tool_registry.py @@ -0,0 +1,125 @@ +"""ToolRegistry - the centralised catalogue every tool source registers into +(R05-T01). + +Built-in file/command/fetch tools (``core/tools.py``), MCP server tools +(``core/mcp_client.py``) and unified connectors (``core/ext_connectors.py``) +each produce their own ``List[ToolSpec]`` today, concatenated ad-hoc by +``core/tools.py::combine_tool_sources``. None of that concatenation carries +risk information, which is exactly why an MCP tool call reaches +``core/chat_agent.py`` with no ``ToolDescriptor`` to consult and skips the +permission gate entirely (the gap R05-T04 closes). + +``ToolRegistry`` is the one place a :class:`~domain.tools.tool_descriptor.ToolDescriptor` +is looked up by name, so a policy gateway - or anything else that needs to ask +"what can this tool do" - has a single source of truth instead of re-deriving +it from a spec list. + +Pure domain code: stdlib only, no Qt, no I/O. +""" +from __future__ import annotations + +from typing import Dict, Iterable, List, Optional + +from cowork_local.providers.base import ToolSpec + +from .tool_descriptor import ToolCapability, ToolDescriptor + + +class ToolRegistry: + """An in-memory, name-keyed catalogue of :class:`ToolDescriptor`. + + Deliberately mutable and unordered-by-name-only: a turn builds one + registry from whichever tool sources it has (built-ins + whatever MCP + servers/connectors are enabled), so re-registering the same name simply + replaces the previous descriptor rather than raising - the same + "last one wins" behaviour ``combine_tool_sources`` already has for + duplicate tool names across sources. + """ + + def __init__(self, descriptors: Optional[Iterable[ToolDescriptor]] = None) -> None: + self._by_name: Dict[str, ToolDescriptor] = {} + for descriptor in descriptors or (): + self.register(descriptor) + + def register(self, descriptor: ToolDescriptor) -> None: + self._by_name[descriptor.name] = descriptor + + def get(self, name: str) -> Optional[ToolDescriptor]: + return self._by_name.get(name) + + def all(self) -> List[ToolDescriptor]: + return list(self._by_name.values()) + + def specs(self) -> List[ToolSpec]: + """Every registered descriptor, projected back to ``ToolSpec`` - the + shape the provider call and the model-facing catalogue need.""" + return [d.to_spec() for d in self._by_name.values()] + + def capabilities_for(self, name: str) -> ToolCapability: + """The capability set for ``name``, or ``NONE`` for an unknown tool. + + Returning ``NONE`` rather than raising lets a policy gateway treat an + unregistered tool the same way as one with no declared risk - the + gateway's DENY-on-unknown-name rule is a deliberate, separate check, + not something this lookup should pre-empt. + """ + descriptor = self._by_name.get(name) + return descriptor.capabilities if descriptor is not None else ToolCapability.NONE + + def __contains__(self, name: str) -> bool: + return name in self._by_name + + def __len__(self) -> int: + return len(self._by_name) + + +# --------------------------------------------------------------------------- # +# Default capability map for this app's built-in tools (core/tools.py). +# Kept here, next to the registry, rather than inside core/tools.py itself - +# core/ is the legacy engine layer being strangled, not where new domain facts +# should accumulate. +# --------------------------------------------------------------------------- # +_CAP = ToolCapability +BUILT_IN_CAPABILITIES: Dict[str, ToolCapability] = { + "read_file": _CAP.READ, + "list_dir": _CAP.READ, + "write_file": _CAP.WRITE, + "edit_file": _CAP.WRITE, + "run_command": _CAP.EXECUTE, + "install_package": _CAP.WRITE | _CAP.EXECUTE | _CAP.NETWORK, + "fetch_url": _CAP.NETWORK, + "jira_search": _CAP.NETWORK, + "jira_get_issue": _CAP.NETWORK, + # Advertised by every engine but has no filesystem/process/network effect + # of its own - it only drives the Plan panel (see core/chat_agent.py). + "update_plan": _CAP.NONE, + "save_file": _CAP.WRITE, +} + +# Tools with no standard, self-declared risk metadata (every MCP server tool, +# every unified connector) are tagged with this conservative default - see +# R05-T04. Better to over-gate an unknown remote tool than to silently let it +# through as READ-only. +UNKNOWN_SOURCE_CAPABILITIES: ToolCapability = _CAP.WRITE | _CAP.EXECUTE | _CAP.NETWORK + + +def default_registry(specs: Iterable[ToolSpec]) -> ToolRegistry: + """Build a registry from ``core/tools.py``'s own ``TOOL_SPECS`` (plus + ``save_file``/``update_plan``, which the engines add separately), using + :data:`BUILT_IN_CAPABILITIES`. A spec with no entry in that map falls back + to :data:`UNKNOWN_SOURCE_CAPABILITIES` - the same conservative default + applied to MCP/connector tools, so a built-in nobody has classified yet + fails safe instead of silently ungated.""" + registry = ToolRegistry() + for spec in specs: + capability = BUILT_IN_CAPABILITIES.get(spec.name, UNKNOWN_SOURCE_CAPABILITIES) + registry.register(ToolDescriptor.from_spec(spec, capability)) + return registry + + +__all__ = [ + "ToolRegistry", + "BUILT_IN_CAPABILITIES", + "UNKNOWN_SOURCE_CAPABILITIES", + "default_registry", +] diff --git a/infrastructure/filesystem/__init__.py b/infrastructure/filesystem/__init__.py new file mode 100644 index 0000000..6d91ae7 --- /dev/null +++ b/infrastructure/filesystem/__init__.py @@ -0,0 +1,6 @@ +"""Filesystem/process/network tool adapters split out of ``core/tools.py`` +(EPIC R05) and the sandbox execution context they share.""" + +from .tool_context import CancelFn, ToolContext, ToolError + +__all__ = ["CancelFn", "ToolContext", "ToolError"] diff --git a/infrastructure/filesystem/command_tools.py b/infrastructure/filesystem/command_tools.py new file mode 100644 index 0000000..0e8f763 --- /dev/null +++ b/infrastructure/filesystem/command_tools.py @@ -0,0 +1,110 @@ +"""Command tools - run_command, install_package (R05-T02). + +Moved verbatim out of ``core/tools.py`` (see ``file_tools.py`` for why). These +two are the ones today's hand-written permission gate in +``core/chat_agent.py`` singles out by literal name +(``name in ("run_command", "install_package")``) — R05-T03 replaces that +tuple with a capability lookup, but the tools themselves are unchanged here. +""" +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, Dict, Optional + +from .tool_context import CancelFn, ToolContext + +COMMAND_TIMEOUT = 120 # seconds + +_SNAPSHOT_SKIP = {".git", "__pycache__", "node_modules", ".scratch", ".venv", + ".idea", ".mypy_cache", ".pytest_cache"} + + +def _snapshot(workdir: Path) -> Dict[str, Any]: + """Map of file path -> (mtime, size) under the workdir (noise dirs skipped).""" + snap: Dict[str, Any] = {} + try: + for dirpath, dirnames, filenames in os.walk(str(workdir)): + dirnames[:] = [d for d in dirnames if d not in _SNAPSHOT_SKIP] + for fn in filenames: + full = os.path.join(dirpath, fn) + try: + st = os.stat(full) + snap[full] = (st.st_mtime_ns, st.st_size) + except OSError: + pass + if len(snap) > 5000: + return snap + except OSError: + pass + return snap + + +def _sandbox_python(ctx: ToolContext, cancel: Optional[CancelFn] = None, + on_output=None) -> Optional[str]: + """Lazily create/reuse this ctx's project sandbox venv (Code tab only — + ``ctx.sandbox``); returns its python path, or None to use the app's own.""" + if not ctx.sandbox: + return None + from cowork_local.core.deps import ensure_project_venv + + py = ensure_project_venv(ctx.workdir, cancel=cancel, on_output=on_output) + return str(py) if py else None + + +def run_command(ctx: ToolContext, args: Dict[str, Any], + cancel: Optional[CancelFn] = None, + on_output=None) -> Dict[str, Any]: + from cowork_local.core.deps import network_blocked_env, run_cancellable, sandbox_env + from cowork_local.core.sandbox_manager import ExecutionConfig, SandboxManager + from cowork_local.security.command_risk_classifier import classify_command + + command = str(args.get("command", "")).strip() + if not command: + return {"ok": False, "output": "Empty command."} + + # --- Security validation pipeline --- + risk = classify_command(command, is_cowork_mode=ctx.flatten_writes) + if risk.blocked: + denial = "Command blocked by security policy: " + "; ".join(risk.reasons) + return {"ok": False, "output": denial} + + # Route through SandboxManager for risk-based isolation + mgr = SandboxManager(ExecutionConfig( + enabled=True, + block_network_by_default=ctx.block_network, + is_cowork_mode=ctx.flatten_writes, + )) + sandbox_result = mgr.run( + command=command, + workdir=str(ctx.workdir), + block_network=ctx.block_network, + timeout_sec=COMMAND_TIMEOUT, + cancel=cancel, + ) + # Sandbox ALWAYS executes (never double-run). Return its result directly. + if sandbox_result.get("sandbox") == "blocked": + return {"ok": False, "output": sandbox_result.get("stderr", "Command blocked")} + out = sandbox_result.get("stdout", "").strip() or "(no output)" + err = sandbox_result.get("stderr", "") + rc = sandbox_result.get("returncode", -1) + if err: + out = f"{out}\n{err}" if out else err + return {"ok": sandbox_result.get("ok", False), "output": f"[exit {rc}]\n{out}"} + + +def install_package(ctx: ToolContext, args: Dict[str, Any], + cancel: Optional[CancelFn] = None, + on_output=None) -> Dict[str, Any]: + from cowork_local.core.deps import pip_install + + package = str(args.get("package", "")).strip() + if not package: + return {"ok": False, "output": "No package specified."} + python = _sandbox_python(ctx, cancel, on_output) + ok, detail = pip_install(package, cancel=cancel, on_output=on_output, python=python) + head = f"Installed {package}." if ok else f"Could not install {package}." + return {"ok": ok, "output": f"{head}\n{detail}"} + + +__all__ = ["COMMAND_TIMEOUT", "run_command", "install_package", "_snapshot"] diff --git a/infrastructure/filesystem/fetch_tools.py b/infrastructure/filesystem/fetch_tools.py new file mode 100644 index 0000000..5a5c98b --- /dev/null +++ b/infrastructure/filesystem/fetch_tools.py @@ -0,0 +1,55 @@ +"""Fetch tools - fetch_url, jira_search, jira_get_issue (R05-T02). + +Moved verbatim out of ``core/tools.py`` (see ``file_tools.py`` for why). The +network access these three carry is exactly what the ``ToolCapability.NETWORK`` +tag added in R05-T01/domain/tools/tool_registry.py describes. +""" +from __future__ import annotations + +from typing import Any, Dict + +from .tool_context import ToolContext + + +def fetch_url(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: + """Fetch a URL's text content (web page / online document / SharePoint- + OneDrive share link) via link_fetch — the same parser task-link attachments + use. Honors the Sandbox Security Layer's "Block network" policy.""" + url = str(args.get("url", "")).strip() + if not url: + return {"ok": False, "output": "fetch_url: 'url' is required."} + if not url.lower().startswith(("http://", "https://")): + return {"ok": False, "output": f"fetch_url: not an http(s) URL: {url}"} + if not ctx.allow_url_fetch: + return {"ok": False, + "output": ("fetch_url: URL fetching is turned off in Settings → Security " + "(\"Allow the agent to fetch URLs\").")} + # A pasted Jira issue link on the CONNECTED Jira host is read via the + # authenticated API (so private issues resolve, not a login page). Public + # links / any other URL fall through to the normal fetcher below. + from cowork_local.core import jira_tool + if jira_tool.is_jira_issue_url(ctx.jira, url): + return {"ok": True, "output": jira_tool.get_issue_by_url(ctx.jira, url)} + from cowork_local.core.link_fetch import fetch_link_preview + + return {"ok": True, "output": fetch_link_preview(url)} + + +def jira_search(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: + from cowork_local.core import jira_tool + + out = jira_tool.search(ctx.jira, str(args.get("jql", "")), + int(args.get("max_results", 25) or 25)) + return {"ok": not out.lower().startswith(("jira is not configured", "jira search failed")), + "output": out} + + +def jira_get_issue(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: + from cowork_local.core import jira_tool + + out = jira_tool.get_issue(ctx.jira, str(args.get("key", ""))) + return {"ok": not out.lower().startswith(("jira is not configured", "could not fetch")), + "output": out} + + +__all__ = ["fetch_url", "jira_search", "jira_get_issue"] diff --git a/infrastructure/filesystem/file_tools.py b/infrastructure/filesystem/file_tools.py new file mode 100644 index 0000000..b5b6a38 --- /dev/null +++ b/infrastructure/filesystem/file_tools.py @@ -0,0 +1,136 @@ +"""File tools - read_file, list_dir, write_file, edit_file (R05-T02). + +Moved verbatim out of ``core/tools.py``, whose ``execute_tool`` used to +dispatch to these via a hand-written if/elif chain over every tool name it +knew about. Splitting the built-in handlers into per-concern modules +(this one, ``command_tools.py``, ``fetch_tools.py``) means adding a tool no +longer means growing that one function; ``core/tools.py::execute_tool`` now +looks the name up in a dict built from these modules instead. + +Behavior is unchanged from before the split - this is a pure move, not a +rewrite. Every existing characterization/contract test that exercises +read_file/write_file/edit_file/list_dir through ``core.tools.execute_tool`` +still exercises the exact same code, just imported from here. +""" +from __future__ import annotations + +import ast +from pathlib import Path +from typing import Any, Dict + +from .tool_context import ToolContext + +MAX_READ_BYTES = 200_000 + + +def _flatten_rel(rel: str) -> str: + """Collapse a sub-folder path down to a bare filename so the file lands in the + workdir root — EXCEPT the ``.scratch`` sandbox subtree, which is preserved. + + Used by the Cowork agent (flatten_writes=True) so it can never create a + per-session / per-chat / per-task output sub-folder: every deliverable stays + directly in the single configured Output folder.""" + parts = Path(rel).parts + if parts and parts[0] == ".scratch": + return rel # temporary sandbox is allowed (and cleaned up afterwards) + return Path(rel).name or rel + + +def _check_python_syntax(target: Path, content: str) -> str: + """Return a short warning if ``content`` is invalid Python, else ''. + + Catches syntax errors the instant a .py file is written/edited — before the + agent wastes a whole run_command round-trip just to get the same error back + from a traceback.""" + if target.suffix.lower() not in (".py", ".pyw"): + return "" + try: + ast.parse(content, filename=str(target)) + return "" + except SyntaxError as exc: + return f"\n⚠ Syntax error at line {exc.lineno}: {exc.msg} — fix this before running the file." + + +def read_file(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: + target = ctx.resolve(str(args.get("path", ""))) + if not target.exists(): + return {"ok": False, "output": f"File not found: {args.get('path')}"} + data = target.read_bytes()[:MAX_READ_BYTES] + text = data.decode("utf-8", errors="replace") + return {"ok": True, "output": text} + + +def list_dir(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: + rel = str(args.get("path", ".") or ".") + target = ctx.resolve(rel) + # A missing/not-yet-created path is NOT a tool failure — report it as an + # ordinary result so the agent can create it or pick another path and keep + # going. Returning ok=False here surfaced a false "tool failed: list_dir" in + # Co4E flows and could stall a step on a recoverable situation. + if not target.exists(): + return {"ok": True, "output": f"(path '{rel}' does not exist yet — create it or use another path)"} + if target.is_file(): + return {"ok": True, "output": f"('{rel}' is a file, not a directory)"} + entries = [] + for child in sorted(target.iterdir(), key=lambda p: (p.is_file(), p.name.lower())): + marker = "/" if child.is_dir() else "" + entries.append(f"{child.name}{marker}") + return {"ok": True, "output": "\n".join(entries) or "(empty folder)"} + + +def write_file(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: + rel = str(args.get("path", "")) + if ctx.flatten_writes: + rel = _flatten_rel(rel) + target = ctx.resolve(rel) + content = str(args.get("content", "")) + target.parent.mkdir(parents=True, exist_ok=True) + # A .xlsx is a binary package — build a REAL workbook from the content + # (CSV/TSV/Markdown-table/JSON) rather than writing raw text (which corrupts it). + if target.suffix.lower() in (".xlsx", ".xlsm"): + from cowork_local.core import xlsx_write + if xlsx_write.build_xlsx_from_text(target, content): + return {"ok": True, "path": str(target), + "output": f"Wrote spreadsheet {rel} ({target.name})."} + return {"ok": False, "output": "Could not build the .xlsx (openpyxl unavailable) — " + "write a .csv instead, or use a generator script."} + target.write_text(content, encoding="utf-8") + warning = _check_python_syntax(target, content) + return {"ok": True, "path": str(target), + "output": f"Wrote {len(content)} chars to {rel}.{warning}"} + + +def edit_file(ctx: ToolContext, args: Dict[str, Any]) -> Dict[str, Any]: + """Replace an exact snippet inside an existing file (precise patch edit).""" + rel = str(args.get("path", "")) + if ctx.flatten_writes: + rel = _flatten_rel(rel) + target = ctx.resolve(rel) + if not target.exists(): + return {"ok": False, + "output": f"File not found: {rel} — use write_file to create it."} + old = str(args.get("old_string", "")) + new = str(args.get("new_string", "")) + replace_all = bool(args.get("replace_all", False)) + if not old: + return {"ok": False, "output": "old_string is empty — provide the exact text to replace."} + try: + text = target.read_text(encoding="utf-8", errors="replace") + except OSError as exc: + return {"ok": False, "output": f"Could not read file: {exc}"} + count = text.count(old) + if count == 0: + return {"ok": False, "output": ("old_string not found. Read the file and copy the exact " + "text to replace, including indentation/whitespace.")} + if count > 1 and not replace_all: + return {"ok": False, "output": (f"old_string appears {count} times — add surrounding " + "context to make it unique, or set replace_all=true.")} + updated = text.replace(old, new) if replace_all else text.replace(old, new, 1) + target.write_text(updated, encoding="utf-8") + n = count if replace_all else 1 + warning = _check_python_syntax(target, updated) + return {"ok": True, + "output": f"Edited {args.get('path')} ({n} replacement{'' if n == 1 else 's'}).{warning}"} + + +__all__ = ["MAX_READ_BYTES", "read_file", "list_dir", "write_file", "edit_file"] diff --git a/infrastructure/filesystem/tool_context.py b/infrastructure/filesystem/tool_context.py new file mode 100644 index 0000000..a5d0057 --- /dev/null +++ b/infrastructure/filesystem/tool_context.py @@ -0,0 +1,62 @@ +"""ToolContext / ToolError / CancelFn - the sandboxed execution context every +built-in tool runs against (moved out of ``core/tools.py`` in R05-T02). + +Kept as its own leaf module (no dependency on any sibling in this package) so +``file_tools.py``, ``command_tools.py`` and ``fetch_tools.py`` can each import +it without creating an import cycle back through ``core/tools.py``, which +itself re-exports ``ToolContext``/``ToolError`` from here for the existing +callers (``core/chat_agent.py``, ``core/code_agent.py``, +``core/task_executors.py``) that do ``from .tools import ToolContext``. +""" +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, Optional + +CancelFn = Callable[[], bool] + + +class ToolError(Exception): + pass + + +@dataclass +class ToolContext: + workdir: Path + flatten_writes: bool = False # Cowork: force every write into the workdir root + sandbox: bool = False # Code tab: isolate run_command/install_package into /.venv + # Sandbox Security Layer — Settings' "Resource Limits" (cpu_percent/memory_mb/ + # disk_mb), applied to every run_command/install_package this context runs. + # None (default) = no limits, matching pre-existing behavior. + resource_limits: Optional[Dict[str, float]] = None + # Sandbox Security Layer — Settings' "Block network for agent commands" + # (policy-level, see deps.py::network_blocked_env). False (default) = + # unrestricted, matching pre-existing behavior. + block_network: bool = False + # Whether the fetch_url tool may read URLs — SEPARATE from block_network + # (reading a web page/share link for info is safe; running networked shell + # commands is the risk). Defaults True; set from agent_security.allow_url_fetch. + allow_url_fetch: bool = True + # Jira read connector config (base_url/email/api_token) — None disables the + # jira_* tools' ability to connect. Populated from config.data["jira"]. + jira: Optional[Dict[str, Any]] = None + + def resolve(self, rel: str) -> Path: + """Resolve ``rel`` inside the workdir, rejecting escapes.""" + if rel in ("", "."): + return self.workdir + candidate = (self.workdir / rel).expanduser() + try: + resolved = candidate.resolve() + except OSError as exc: + raise ToolError(f"Invalid path: {rel} ({exc})") + root = self.workdir.resolve() + if resolved != root and root not in resolved.parents: + raise ToolError( + f"Refused: '{rel}' is outside the working folder ({root})." + ) + return resolved + + +__all__ = ["CancelFn", "ToolError", "ToolContext"] diff --git a/infrastructure/mcp/__init__.py b/infrastructure/mcp/__init__.py new file mode 100644 index 0000000..e739ca6 --- /dev/null +++ b/infrastructure/mcp/__init__.py @@ -0,0 +1,5 @@ +"""MCP server connection lifecycle management (EPIC R05).""" + +from .mcp_source_manager import McpToolSourceManager + +__all__ = ["McpToolSourceManager"] diff --git a/infrastructure/mcp/mcp_source_manager.py b/infrastructure/mcp/mcp_source_manager.py new file mode 100644 index 0000000..2962675 --- /dev/null +++ b/infrastructure/mcp/mcp_source_manager.py @@ -0,0 +1,113 @@ +"""McpToolSourceManager - the MCP server connection lifecycle, extracted out +of ``state.py::AppContext`` (R05-T05). + +Today ``AppContext.build_mcp_tools`` inlines all of this: a ``_mcp_connections`` +dict, a ``_conn_lock`` guarding check-then-create against concurrent turns (a +Cowork tab, a Co4E flow and a Scheduled Task can all call it at once), and a +"start it, cache it, skip it on failure" loop repeated for both the +admin-configured servers AND the built-in MS365 server +(``_ms365_builtin_connection``). None of that logic touches Qt; it was only +ever inline because ``AppContext`` is where the config lived. + +This class owns the SAME cache/lock/start-or-skip behavior as a standalone, +directly testable object — ``AppContext`` becomes a thin caller (one instance +per app, same as it holds one ``RoutingApplicationService``). + +Pure Python: no Qt. It DOES touch the network/filesystem via +``core.mcp_client.McpServerConnection`` (a subprocess + asyncio loop), which is +exactly what makes it infrastructure rather than domain. +""" +from __future__ import annotations + +import threading +from typing import Dict, List, Optional + +from cowork_local.core.mcp_client import McpServerConnection + + +class McpToolSourceManager: + """Caches and supervises one :class:`McpServerConnection` per server name. + + ``connection_factory`` defaults to ``McpServerConnection`` itself; tests + substitute a fake so no real subprocess is spawned (see + ``tests/unit/test_mcp_source_manager.py``). + """ + + def __init__(self, connection_factory=McpServerConnection) -> None: + self._connections: Dict[str, McpServerConnection] = {} + self._lock = threading.Lock() + self._connection_factory = connection_factory + + def ensure(self, name: str, command: str, args: Optional[List[str]] = None, + env: Optional[Dict[str, str]] = None) -> Optional[McpServerConnection]: + """Return a live connection for ``name``, starting one if there is + none cached or the cached one's subprocess has died. + + Serialized under one lock so two turns racing to build their tool + list at the same moment share one subprocess per server instead of + each spawning their own (the bug this replaces: + ``AppContext._conn_lock``'s original docstring). Returns ``None`` - + never raises - when the server fails to start, matching the existing + "one broken server must not block the turn" behavior. + """ + with self._lock: + existing = self._connections.get(name) + if existing is not None and existing.is_alive(): + return existing + if existing is not None: + self._connections.pop(name, None) + connection = self._connection_factory(name, command, args or [], env) + try: + connection.start() + except Exception: # noqa: BLE001 - one broken server must not block the turn + return None + self._connections[name] = connection + return connection + + def get(self, name: str) -> Optional[McpServerConnection]: + """The cached connection for ``name``, without starting one.""" + return self._connections.get(name) + + def is_alive(self, name: str) -> bool: + connection = self._connections.get(name) + return connection is not None and connection.is_alive() + + def restart(self, name: str, command: str, args: Optional[List[str]] = None, + env: Optional[Dict[str, str]] = None) -> Optional[McpServerConnection]: + """Force a fresh connection for ``name`` even if the cached one still + looks alive - for a server the caller knows is misbehaving.""" + with self._lock: + self._connections.pop(name, None) + return self.ensure(name, command, args, env) + + def stop(self, name: str) -> None: + """Stop and forget one connection - used when a server becomes + unavailable by configuration (e.g. MS365 signed out) rather than by + crashing.""" + with self._lock: + connection = self._connections.pop(name, None) + if connection is not None: + try: + connection.stop() + except Exception: # noqa: BLE001 - shutdown must never raise into the caller + pass + + def active(self) -> List[McpServerConnection]: + """Every currently cached connection - what + ``core/mcp_client.py::build_mcp_tools`` merges tool specs from.""" + return list(self._connections.values()) + + def stop_all(self) -> None: + """Terminate every connection's subprocess - called on app shutdown + so none of them linger as orphan processes.""" + with self._lock: + connections = list(self._connections.values()) + self._connections.clear() + for connection in connections: + try: + connection.stop() + except Exception: # noqa: BLE001 + pass + + +__all__ = ["McpToolSourceManager"] diff --git a/state.py b/state.py index fc1d35e..e738106 100644 --- a/state.py +++ b/state.py @@ -6,6 +6,7 @@ import time from typing import TYPE_CHECKING, Optional, Tuple from .config import AppConfig +from .infrastructure.mcp import McpToolSourceManager def resolve_agent_default( @@ -38,18 +39,18 @@ class AppContext: def __init__(self, config: AppConfig): self.config = config self.started_at = time.time() # for Monitoring's Sandbox Details "Created"/"Uptime" - self._mcp_connections: dict = {} # server name -> McpServerConnection + # Admin-configured MCP servers + the built-in MS365 server (R05-T05): + # connection caching/lifecycle (check-then-create, restart, shutdown) + # now lives in McpToolSourceManager, extracted so it is testable + # without an AppContext/Qt. See its docstring for why the check-then- + # create race matters — several turns (multiple Cowork tabs, parallel + # Co4E flows, scheduled tasks) can call build_mcp_tools() at once. + self._mcp_manager = McpToolSourceManager() self._ext_connections: dict = {} # connector id -> McpServerConnection (mcp_stdio mode only) - # Guards the two connection caches above. build_mcp_tools() runs on EVERY - # chat turn's own AgentWorker thread, so several turns (multiple Cowork - # tabs, parallel Co4E flows, scheduled tasks) can enter it at once. The - # cache is populated check-then-create ("conn is None → spawn → store"); - # without this lock two concurrent turns both see None and each spawns a - # subprocess for the SAME server — one leaks as an orphan and the wrong - # object may be handed out. The lock makes connection setup atomic; the - # provider/HTTP path itself is already thread-safe (a fresh provider per - # call, module-level `requests`, MCP calls multiplexed on the server's - # own event loop), so concurrent model calls never needed serializing. + # Guards ``_ext_connections`` only now — unified Connectors (CAD/CAE/ + # MS365/Other) aren't covered by McpToolSourceManager (R05-T05 scoped + # to MCP servers), so this cache still needs its own check-then-create + # lock, the same race McpToolSourceManager guards against internally. self._conn_lock = threading.Lock() self._routing_service = None # lazy RoutingService (Auto Model Routing) # Lazy RoutingApplicationService (R03-T03) — the Qt-free decision layer @@ -237,48 +238,41 @@ class AppContext: if not self.config.connect_external: return [], None from .core.ext_connectors import build_ext_connector_tools - from .core.mcp_client import McpServerConnection from .core.mcp_client import build_mcp_tools as _merge_mcp_tools from .core.tools import combine_tool_sources - # Serialize the check-then-create against the connection caches so - # concurrent turns share one subprocess per server instead of racing to - # spawn duplicates (see _conn_lock in __init__). The lock is held while - # connections are established (a one-time cost per server per app run); - # once warm, every turn just finds the cached connection and returns. - with self._conn_lock: - active = [] - for entry in self.config.mcp_servers: - if not entry.get("enabled", True): - continue - name = entry.get("name", "") - command = entry.get("command", "") - if not name or not command: - continue - conn = self._mcp_connections.get(name) - if conn is None: - conn = McpServerConnection(name, command, entry.get("args") or [], - entry.get("env") or None) - try: - conn.start() - except Exception: # noqa: BLE001 - one broken server must not block the turn - continue - self._mcp_connections[name] = conn + # R05-T05: connection caching/check-then-create for admin-configured + # servers + the MS365 builtin now lives in McpToolSourceManager (its + # own lock guards the race — see its docstring). + active = [] + for entry in self.config.mcp_servers: + if not entry.get("enabled", True): + continue + name = entry.get("name", "") + command = entry.get("command", "") + if not name or not command: + continue + conn = self._mcp_manager.ensure(name, command, entry.get("args") or [], + entry.get("env") or None) + if conn is not None: active.append(conn) - builtin = self._ms365_builtin_connection(skip={c.name for c in active}) - if builtin is not None: - active.append(builtin) - mcp_tools, mcp_executor = _merge_mcp_tools(active) + builtin = self._ms365_builtin_connection(skip={c.name for c in active}) + if builtin is not None: + active.append(builtin) + mcp_tools, mcp_executor = _merge_mcp_tools(active) + # ``_ext_connections`` isn't covered by McpToolSourceManager (T05 + # scoped to MCP servers) — still serialized under ``_conn_lock``. + with self._conn_lock: ext = self.config.ext_connectors all_connectors = [*ext.get("cad", []), *ext.get("cae", []), *ext.get("ms365", []), *ext.get("other", [])] ext_tools, ext_executor = build_ext_connector_tools(all_connectors, self._ext_connections) - # Locally-synced OneDrive/SharePoint (no sign-in) — reads/writes the - # OneDrive-desktop-synced folders directly, gated on ms365.connectors. - from .core.ms365_local import build_ms365_local_tools - local_tools, local_executor = build_ms365_local_tools(self.config) + # Locally-synced OneDrive/SharePoint (no sign-in) — reads/writes the + # OneDrive-desktop-synced folders directly, gated on ms365.connectors. + from .core.ms365_local import build_ms365_local_tools + local_tools, local_executor = build_ms365_local_tools(self.config) return combine_tool_sources((mcp_tools, mcp_executor), (ext_tools, ext_executor), (local_tools, local_executor)) @@ -308,36 +302,20 @@ class AppContext: import sys from pathlib import Path - from .core.mcp_client import McpServerConnection - name = self._MS365_BUILTIN if name in skip: return None if not self._ms365_available(): - stale = self._mcp_connections.pop(name, None) - if stale is not None: - try: - stale.stop() - except Exception: # noqa: BLE001 - pass + self._mcp_manager.stop(name) return None - conn = self._mcp_connections.get(name) - if conn is None: - # The subprocess must import cowork_local even in a from-source run - # (PYTHONPATH=src) — prepend this package's parent dir explicitly. - env = dict(os.environ) - src_root = str(Path(__file__).resolve().parent.parent) - env["PYTHONPATH"] = (src_root + os.pathsep + env["PYTHONPATH"] - if env.get("PYTHONPATH") else src_root) - conn = McpServerConnection( - name, sys.executable, - ["-m", "cowork_local.mcp_servers.ms365_server"], env) - try: - conn.start() - except Exception: # noqa: BLE001 - MS365 down must not block the turn - return None - self._mcp_connections[name] = conn - return conn + # The subprocess must import cowork_local even in a from-source run + # (PYTHONPATH=src) — prepend this package's parent dir explicitly. + env = dict(os.environ) + src_root = str(Path(__file__).resolve().parent.parent) + env["PYTHONPATH"] = (src_root + os.pathsep + env["PYTHONPATH"] + if env.get("PYTHONPATH") else src_root) + return self._mcp_manager.ensure( + name, sys.executable, ["-m", "cowork_local.mcp_servers.ms365_server"], env) def stop_mcp_connections(self) -> None: """Terminate every connected MCP server's subprocess (incl. External @@ -345,11 +323,6 @@ class AppContext: them linger as orphan processes.""" from .core.ext_connectors import stop_ext_connections + self._mcp_manager.stop_all() with self._conn_lock: - for conn in self._mcp_connections.values(): - try: - conn.stop() - except Exception: # noqa: BLE001 - pass - self._mcp_connections.clear() stop_ext_connections(self._ext_connections) diff --git a/tests/unit/test_code_agent_tool_policy.py b/tests/unit/test_code_agent_tool_policy.py new file mode 100644 index 0000000..1adbe25 --- /dev/null +++ b/tests/unit/test_code_agent_tool_policy.py @@ -0,0 +1,62 @@ +"""EPIC R05-T03/T04: ``core/code_agent.py::run_code`` used to gate tool calls +with ``if name in (WRITE_TOOLS | MS365_WRITE_TOOLS): gate.request(...)``. This +pins that the switch to ``ToolPolicyGateway`` still gates exactly the same +calls: ``write_file`` (a WRITE tool) consults the gate; ``list_dir`` +(read-only) never does. + +Runs the REAL engine (``run_code``) via :class:`FakeProvider`, same approach +``tests/characterization/test_run_cowork.py`` uses for the Cowork engine. +""" +from __future__ import annotations + +from typing import Any, Dict, List + +from cowork_local.core.code_agent import run_code +from cowork_local.core.tools import ToolContext +from tests.fakes import FakeProvider, ScriptedTurn + + +class _RecordingGate: + def __init__(self, approve: bool): + self.approve = approve + self.calls: List[Dict[str, Any]] = [] + + def request(self, payload: Dict[str, Any]) -> bool: + self.calls.append(payload) + return self.approve + + +def _run(tmp_path, provider, gate): + ctx = ToolContext(tmp_path) + events: List[Dict[str, Any]] = [] + messages: List[Dict[str, Any]] = [{"role": "user", "content": "do it"}] + run_code(provider, messages, ctx, gate, events.append) + return events + + +def test_write_file_consults_the_gate_and_honors_rejection(tmp_path): + provider = FakeProvider([ + ScriptedTurn(tool_calls=[("write_file", {"path": "a.txt", "content": "hi"})]), + ScriptedTurn(text="done"), + ]) + gate = _RecordingGate(approve=False) + events = _run(tmp_path, provider, gate) + + assert len(gate.calls) == 1 and gate.calls[0]["name"] == "write_file" + results = [e for e in events if e.get("type") == "tool_result"] + assert results[0]["ok"] is False + assert not (tmp_path / "a.txt").exists() # rejected, never actually written + + +def test_read_only_tool_never_consults_the_gate(tmp_path): + (tmp_path / "existing.txt").write_text("x", encoding="utf-8") + provider = FakeProvider([ + ScriptedTurn(tool_calls=[("list_dir", {})]), + ScriptedTurn(text="done"), + ]) + gate = _RecordingGate(approve=False) # would reject if ever asked + events = _run(tmp_path, provider, gate) + + assert gate.calls == [] + results = [e for e in events if e.get("type") == "tool_result"] + assert results[0]["ok"] is True diff --git a/tests/unit/test_cowork_extra_tool_policy.py b/tests/unit/test_cowork_extra_tool_policy.py new file mode 100644 index 0000000..ce3b3cd --- /dev/null +++ b/tests/unit/test_cowork_extra_tool_policy.py @@ -0,0 +1,86 @@ +"""EPIC R05-T04: before this change, ``core/chat_agent.py::run_cowork`` called +``extra_executor(name, args)`` directly for any MCP/connector tool — no +permission check at all, regardless of the "confirm before running commands" +setting. This pins the fix: an extra tool now goes through the same +``ToolPolicyGateway`` as ``run_command``, using the conservative default +capability (``UNKNOWN_SOURCE_CAPABILITIES``) since MCP tools carry no +standard risk metadata. + +Runs the real engine via :class:`FakeProvider`, matching +``tests/characterization/test_run_cowork.py``'s approach. +""" +from __future__ import annotations + +from typing import Any, Dict, List + +from cowork_local.core.chat_agent import run_cowork +from cowork_local.providers.base import ToolSpec +from tests.fakes import FakeProvider, ScriptedTurn + + +class _RecordingGate: + def __init__(self, approve: bool): + self.approve = approve + self.calls: List[Dict[str, Any]] = [] + + def request(self, payload: Dict[str, Any]) -> bool: + self.calls.append(payload) + return self.approve + + +_EXTRA_SPEC = ToolSpec(name="github__delete_repo", description="", parameters={"type": "object"}) + + +def _run(tmp_path, provider, gate, executed: List[str]): + events: List[Dict[str, Any]] = [] + messages: List[Dict[str, Any]] = [{"role": "user", "content": "hi"}] + + def extra_executor(name: str, args: Dict[str, Any]) -> Dict[str, Any]: + executed.append(name) + return {"ok": True, "output": "done"} + + run_cowork(provider, messages, tmp_path, events.append, gate=gate, + extra_tools=[_EXTRA_SPEC], extra_executor=extra_executor) + return events + + +def test_mcp_style_tool_is_rejected_without_ever_calling_the_executor(tmp_path): + provider = FakeProvider([ + ScriptedTurn(tool_calls=[("github__delete_repo", {})]), + ScriptedTurn(text="done"), + ]) + gate = _RecordingGate(approve=False) + executed: List[str] = [] + events = _run(tmp_path, provider, gate, executed) + + assert len(gate.calls) == 1 and gate.calls[0]["name"] == "github__delete_repo" + assert executed == [] # rejected BEFORE the extra_executor ever ran + results = [e for e in events if e.get("type") == "tool_result"] + assert results[0]["ok"] is False + + +def test_mcp_style_tool_runs_once_approved(tmp_path): + provider = FakeProvider([ + ScriptedTurn(tool_calls=[("github__delete_repo", {})]), + ScriptedTurn(text="done"), + ]) + gate = _RecordingGate(approve=True) + executed: List[str] = [] + events = _run(tmp_path, provider, gate, executed) + + assert executed == ["github__delete_repo"] + results = [e for e in events if e.get("type") == "tool_result"] + assert results[0]["ok"] is True + + +def test_no_gate_preserves_auto_run_for_extra_tools(tmp_path): + """``gate=None`` is Cowork's existing "no confirmation configured" state — + must still auto-run, exactly like before this EPIC.""" + provider = FakeProvider([ + ScriptedTurn(tool_calls=[("github__delete_repo", {})]), + ScriptedTurn(text="done"), + ]) + executed: List[str] = [] + events = _run(tmp_path, provider, None, executed) + + assert executed == ["github__delete_repo"] diff --git a/tests/unit/test_mcp_source_manager.py b/tests/unit/test_mcp_source_manager.py new file mode 100644 index 0000000..d0e4eee --- /dev/null +++ b/tests/unit/test_mcp_source_manager.py @@ -0,0 +1,106 @@ +"""EPIC R05-T05: the MCP connection lifecycle extracted out of +``state.py::AppContext`` into :class:`McpToolSourceManager`. + +Uses a fake connection (no real subprocess/asyncio loop) so these tests run in +milliseconds and don't depend on any actual MCP server being installed. +""" +from __future__ import annotations + +from typing import Dict, List, Optional + +from cowork_local.infrastructure.mcp import McpToolSourceManager + + +class _FakeConnection: + """Stands in for ``core.mcp_client.McpServerConnection`` — tracks + start/stop calls instead of spawning anything.""" + + instances: List["_FakeConnection"] = [] + + def __init__(self, name: str, command: str, args: Optional[List[str]] = None, + env: Optional[Dict[str, str]] = None): + self.name = name + self.command = command + self.args = args + self.env = env + self.started = False + self.stopped = False + self._alive = True + _FakeConnection.instances.append(self) + + def start(self) -> None: + self.started = True + + def stop(self) -> None: + self.stopped = True + self._alive = False + + def is_alive(self) -> bool: + return self._alive + + +def _manager() -> McpToolSourceManager: + _FakeConnection.instances.clear() + return McpToolSourceManager(connection_factory=_FakeConnection) + + +def test_ensure_starts_once_and_caches_the_live_connection(): + mgr = _manager() + first = mgr.ensure("github", "npx", ["-y", "github-mcp"]) + second = mgr.ensure("github", "npx", ["-y", "github-mcp"]) + + assert first is second # same connection reused, not a second subprocess + assert len(_FakeConnection.instances) == 1 + assert first.started is True + + +def test_two_concurrent_ensures_for_different_servers_dont_collide(): + mgr = _manager() + a = mgr.ensure("server-a", "cmd-a") + b = mgr.ensure("server-b", "cmd-b") + assert a is not b + assert {c.name for c in mgr.active()} == {"server-a", "server-b"} + + +def test_ensure_restarts_when_the_cached_connection_died(): + mgr = _manager() + first = mgr.ensure("flaky", "cmd") + first.stop() # simulate the subprocess crashing + assert mgr.is_alive("flaky") is False + + second = mgr.ensure("flaky", "cmd") + assert second is not first + assert len(_FakeConnection.instances) == 2 + + +def test_a_server_that_fails_to_start_returns_none_and_isnt_cached(): + class _DyingConnection(_FakeConnection): + def start(self) -> None: + raise RuntimeError("boom") + + mgr = McpToolSourceManager(connection_factory=_DyingConnection) + assert mgr.ensure("broken", "cmd") is None + assert mgr.get("broken") is None + + +def test_stop_removes_one_connection_without_touching_others(): + mgr = _manager() + mgr.ensure("keep", "cmd") + doomed = mgr.ensure("drop", "cmd") + + mgr.stop("drop") + + assert doomed.stopped is True + assert mgr.get("drop") is None + assert mgr.get("keep") is not None + + +def test_stop_all_stops_every_connection_and_clears_the_cache(): + mgr = _manager() + mgr.ensure("a", "cmd") + mgr.ensure("b", "cmd") + + mgr.stop_all() + + assert all(c.stopped for c in _FakeConnection.instances) + assert mgr.active() == [] diff --git a/tests/unit/test_tool_registry_and_policy.py b/tests/unit/test_tool_registry_and_policy.py new file mode 100644 index 0000000..23648a3 --- /dev/null +++ b/tests/unit/test_tool_registry_and_policy.py @@ -0,0 +1,107 @@ +"""Unit tests for EPIC R05: the tool descriptor/registry (R05-T01), the split +built-in handlers (R05-T02), and the policy gateway (R05-T03). + +The gateway tests assert the SAME capability set each engine used to hard-code +as a name tuple still gets gated after the switch to capability lookup — that +equivalence is the whole point of R05-T03, not an incidental detail. +""" +from __future__ import annotations + +from typing import Any, Dict + +import pytest + +from cowork_local.application.conversations import ToolPolicyGateway +from cowork_local.core.tools import TOOL_SPECS, ToolContext, execute_tool +from cowork_local.domain.tools import ToolCapability, ToolDescriptor, default_registry + + +# --------------------------------------------------------------------------- # +# R05-T01 - ToolDescriptor / ToolRegistry +# --------------------------------------------------------------------------- # +def test_capability_flags_compose(): + install = ToolDescriptor("install_package", "", {}, ToolCapability.WRITE | ToolCapability.EXECUTE) + assert install.has(ToolCapability.WRITE) + assert install.has(ToolCapability.EXECUTE) + assert not install.has(ToolCapability.NETWORK) + + +def test_default_registry_matches_todays_hardcoded_gating_sets(): + """The two literal sets this EPIC replaces: + ``core/tools.py::WRITE_TOOLS`` and ``core/chat_agent.py``'s + ``("run_command", "install_package")`` tuple. The registry must agree + with both, or the capability switch silently changes who gets gated.""" + registry = default_registry(TOOL_SPECS) + + execute_gated = {d.name for d in registry.all() if d.has(ToolCapability.EXECUTE)} + assert execute_gated == {"run_command", "install_package"} + + write_gated = {d.name for d in registry.all() if d.has(ToolCapability.WRITE)} + assert write_gated == {"write_file", "edit_file", "install_package"} + + +def test_unregistered_tool_has_no_capabilities(): + registry = default_registry(TOOL_SPECS) + assert registry.capabilities_for("no_such_tool") is ToolCapability.NONE + + +# --------------------------------------------------------------------------- # +# R05-T02 - core/tools.py dispatch, now built from the split infra modules +# --------------------------------------------------------------------------- # +def test_execute_tool_still_dispatches_every_built_in(tmp_path): + ctx = ToolContext(tmp_path) + written = execute_tool(ctx, "write_file", {"path": "a.txt", "content": "hi"}) + assert written["ok"] is True + read = execute_tool(ctx, "read_file", {"path": "a.txt"}) + assert read == {"ok": True, "output": "hi"} + edited = execute_tool(ctx, "edit_file", {"path": "a.txt", "old_string": "hi", "new_string": "bye"}) + assert edited["ok"] is True + assert execute_tool(ctx, "read_file", {"path": "a.txt"})["output"] == "bye" + listing = execute_tool(ctx, "list_dir", {}) + assert listing["ok"] is True and "a.txt" in listing["output"] + + +def test_execute_tool_reports_unknown_name(tmp_path): + ctx = ToolContext(tmp_path) + result = execute_tool(ctx, "not_a_real_tool", {}) + assert result == {"ok": False, "output": "Tool not found: not_a_real_tool"} + + +# --------------------------------------------------------------------------- # +# R05-T03 - ToolPolicyGateway +# --------------------------------------------------------------------------- # +class _RecordingGate: + def __init__(self, approve: bool): + self.approve = approve + self.calls: list = [] + + def request(self, payload: Dict[str, Any]) -> bool: + self.calls.append(payload) + return self.approve + + +@pytest.fixture +def cowork_policy() -> ToolPolicyGateway: + """Same construction as ``core/chat_agent.py``'s module-level + ``_COWORK_TOOL_POLICY`` - EXECUTE is exactly what Cowork used to gate via + the literal ``("run_command", "install_package")`` tuple.""" + return ToolPolicyGateway(default_registry(TOOL_SPECS), ToolCapability.EXECUTE) + + +def test_no_gate_means_auto_run(cowork_policy): + assert cowork_policy.allow("run_command", None, {}) is True + + +def test_read_only_tool_never_asks_the_gate(cowork_policy): + gate = _RecordingGate(approve=False) # would reject if asked + assert cowork_policy.allow("write_file", gate, {}) is True + assert gate.calls == [] # never consulted - write_file isn't EXECUTE + + +def test_gated_capability_consults_the_gate_and_honors_its_answer(cowork_policy): + approving = _RecordingGate(approve=True) + assert cowork_policy.allow("run_command", approving, {"name": "run_command"}) is True + assert approving.calls == [{"name": "run_command"}] + + rejecting = _RecordingGate(approve=False) + assert cowork_policy.allow("install_package", rejecting, {}) is False