diff --git a/.github/workflows/python-ci.yml b/.github/workflows/python-ci.yml index 73fe5a01..e06960d8 100644 --- a/.github/workflows/python-ci.yml +++ b/.github/workflows/python-ci.yml @@ -66,12 +66,13 @@ jobs: tests/application/test_session_deletion_service.py tests/application/test_execution_coordinator.py - # These suites carry Windows-gated cases (the Job Object backend) that - # the ubuntu job can only skip. This is the sole place they actually - # execute. - - name: Verify Job Object sandbox + # These suites carry Windows-gated cases (NTFS ACLs, the Job Object + # backend) that the ubuntu job can only skip. This is the sole place + # they actually execute. + - name: Verify Windows ACLs and Job Object sandbox run: >- python -m pytest -q + tests/test_private_storage_windows.py tests/test_harness_sandbox.py tests/test_exec_sandbox_wiring.py diff --git a/core/agent_runtime/runner.py b/core/agent_runtime/runner.py index d4ed5ccb..093285ef 100644 --- a/core/agent_runtime/runner.py +++ b/core/agent_runtime/runner.py @@ -88,6 +88,49 @@ # it survives across turns and is not re-summarized every step. _BACKFILL_CONTENT = "[Tool result unavailable — call was interrupted or lost]" +# PreCompact checkpoint re-injection (bounded, provider-safe). A PreCompact +# hook may attach ``additional_contexts`` that must survive a successful +# compaction so the post-compaction model can restore working context. The +# re-injection is a plain ``role: user`` message (provider-agnostic) carrying a +# clearly delimited prefix, with hard limits per context and in total — a +# runaway hook can never blow the post-compaction window back open. +_PRECOMPACT_CHECKPOINT_PREFIX = "[PreCompact checkpoint]" +_PRECOMPACT_CONTEXT_LIMIT = 2000 # chars per additional context +_PRECOMPACT_TOTAL_LIMIT = 8000 # chars for the whole checkpoint block + + +def _build_precompact_checkpoint( + contexts: list[str], + *, + total_limit: int = _PRECOMPACT_TOTAL_LIMIT, +) -> str | None: + """Bounded, delimited representation of PreCompact hook context. + + Each context is stripped, truncated to ``_PRECOMPACT_CONTEXT_LIMIT`` chars + and the combined block capped at ``_PRECOMPACT_TOTAL_LIMIT``. Returns + ``None`` when nothing survives (empty input or all contexts blank). + """ + prefix = _PRECOMPACT_CHECKPOINT_PREFIX + "\n" + total_limit = min(max(total_limit, 0), _PRECOMPACT_TOTAL_LIMIT) + content_limit = total_limit - len(prefix) + if not contexts or content_limit <= 0: + return None + parts: list[str] = [] + used = 0 + for ctx in contexts: + text = (ctx or "").strip() + if not text: + continue + text = text[:_PRECOMPACT_CONTEXT_LIMIT] + room = content_limit - used + if room <= 0: + break + parts.append(text[:room]) + used += min(len(text), room) + 1 # +1 for the newline separator + if not parts: + return None + return prefix + "\n".join(parts) + @dataclass(slots=True) class AgentRunSpec: @@ -1688,15 +1731,16 @@ async def _maybe_compact( if estimate is None or estimate <= trigger: return messages + pre_contexts: list[str] = [] signature = history_signature(messages) if signature == self._refused_compaction: # Already tried on exactly this history and it did not shrink. return messages - if spec.pre_compact_hook is not None: pre = await self._call_tool_hook(spec.pre_compact_hook, "auto") if pre is not None and getattr(pre, "block", False): return messages # a PreCompact hook aborted compaction this turn + pre_contexts = list(getattr(pre, "additional_contexts", None) or []) summary = await self._summarize( spec, @@ -1723,6 +1767,28 @@ async def _maybe_compact( spec.session_key or "default", ) return messages + # Bounded checkpoint re-injection: the PreCompact hook's + # ``additional_contexts`` survive a successful compaction as a single + # provider-agnostic user message. ``_build_precompact_checkpoint`` + # caps each context and the total block, so a runaway hook can never + # blow the post-compaction window back open. When the hook blocks or + # summarization fails we returned above — the checkpoint only ever + # appears after a successful compaction. + # A checkpoint must not turn successful compaction back into growth. + # Bound it to the actual character reduction in addition to the + # absolute hook-context cap. + checkpoint_room = ( + self._history_chars(messages) - self._history_chars(compacted) - 1 + ) + checkpoint = _build_precompact_checkpoint( + pre_contexts, + total_limit=checkpoint_room, + ) + if checkpoint: + self._append_injected_messages( + compacted, + [{"role": "user", "content": checkpoint}], + ) self._refused_compaction = None if spec.post_compact_hook is not None: await self._call_tool_hook(spec.post_compact_hook, "auto") diff --git a/core/events/session.py b/core/events/session.py index 1d754c01..37a91663 100644 --- a/core/events/session.py +++ b/core/events/session.py @@ -418,6 +418,8 @@ def __init__( # no hooks are configured, so the whole feature is dormant at zero cost. self._hooks_engine = hooks_engine self._session_started = False + self._session_end_fired = False + self._session_end_lock = asyncio.Lock() # When this session is a spawned sub-agent, (agent_id, agent_type) — its # lifecycle fires SubagentStart/SubagentStop instead of SessionStart/Stop. self._agent_context = agent_context @@ -622,6 +624,10 @@ async def submit(self, op: Op) -> None: if task is not None and not task.done() and task.cancelling() == 0: task.cancel() elif isinstance(op, Shutdown): + # SessionEnd fires exactly once, here — at the real session + # termination boundary — never per turn. Per-turn notifications + # are the Stop event's job (see _EVENTS_WITHOUT_MATCHER). + await self._run_end_hook(reason="other") self._emit(ShutdownComplete()) else: # pragma: no cover - exhaustive guard self._emit(ErrorEvent(message=f"unknown op: {op!r}")) @@ -641,6 +647,7 @@ async def aclose(self) -> None: if task.cancelling() == 0: task.cancel() await asyncio.gather(task, return_exceptions=True) + await self._run_end_hook(reason="other") control = getattr(self, "_agent_control", None) if control is not None: await control.close() @@ -690,6 +697,30 @@ async def _run_start_hook(self): logger.exception("start hook failed") return None + async def _run_end_hook(self, reason: str = "other") -> None: + """Run SessionEnd hooks when the session itself terminates. + + Notification-only: a failure is logged and never crashes the + shutdown. Fired exactly once from ``submit(Shutdown)`` or the actual + resource teardown path, whichever comes first. ``other`` is the + compatible exit reason for a DeepCode runtime shutdown; per-turn + notifications belong to the Stop event, not SessionEnd. + """ + async with self._session_end_lock: + if self._session_end_fired: + return + self._session_end_fired = True + # Spawned agents use their dedicated SubagentStop lifecycle. + if self._agent_context is not None: + return + engine = self._hooks_engine + if engine is None or not engine.has_event("SessionEnd"): + return + try: + await engine.run_session_end(reason=reason) + except Exception: # noqa: BLE001 - hooks never crash a shutdown + logger.exception("session end hook failed") + async def _run_prompt_hooks( self, text: str, hook_contexts: list[str] ) -> str | None: @@ -810,6 +841,10 @@ async def _run_user_input(self, op: UserInput | str) -> None: self._active_turn_task = None if terminal is not None: self._emit(terminal) + # SessionEnd is NOT fired here: this finally block runs after + # every turn, and SessionEnd must fire exactly once at session + # termination (submit(Shutdown)), not per turn. Per-turn + # notifications are the Stop event's responsibility. async def _execute_turn( self, diff --git a/core/harness/hooks/discovery.py b/core/harness/hooks/discovery.py index 45bf4a20..0ef40663 100644 --- a/core/harness/hooks/discovery.py +++ b/core/harness/hooks/discovery.py @@ -38,6 +38,8 @@ ) _DEFAULT_TIMEOUT_SEC = 600 +_SESSION_END_DEFAULT_TIMEOUT_SEC = 2 +_SESSION_END_MAX_TIMEOUT_SEC = 60 @dataclass(slots=True) @@ -150,12 +152,19 @@ def _append_group( warnings.append(f"skipping empty hook command in {path}") continue timeout = handler.get("timeout") + default_timeout = ( + _SESSION_END_DEFAULT_TIMEOUT_SEC + if event_name == "SessionEnd" + else _DEFAULT_TIMEOUT_SEC + ) try: timeout_sec = ( - max(1, int(timeout)) if timeout is not None else _DEFAULT_TIMEOUT_SEC + max(1, int(timeout)) if timeout is not None else default_timeout ) except (TypeError, ValueError): - timeout_sec = _DEFAULT_TIMEOUT_SEC + timeout_sec = default_timeout + if event_name == "SessionEnd": + timeout_sec = min(timeout_sec, _SESSION_END_MAX_TIMEOUT_SEC) status_message = handler.get("statusMessage") or handler.get("status_message") handlers.append( Handler( diff --git a/core/harness/hooks/engine.py b/core/harness/hooks/engine.py index 26f66a11..71bdc19c 100644 --- a/core/harness/hooks/engine.py +++ b/core/harness/hooks/engine.py @@ -176,6 +176,22 @@ async def run_session_start(self, source: str = "startup") -> ContextOutcome: additional_contexts=folded.additional_contexts, ) + async def run_session_end(self, reason: str = "other") -> ContextOutcome: + """Session lifecycle end — fires exactly once when the session terminates. + + Called from ``AgentSession.submit(Shutdown)``, never per turn. The + reason doubles as the matcher input. DeepCode runtime shutdown maps to + the compatible ``other`` reason. The caller logs failures so a hook + can never crash the session close. + """ + payload = {"hook_event_name": "SessionEnd", "reason": reason} + folded = await self._dispatch("SessionEnd", reason, payload) + return ContextOutcome( + block=folded.block, + block_reason=folded.block_reason, + additional_contexts=folded.additional_contexts, + ) + async def run_user_prompt_submit(self, prompt: str) -> ContextOutcome: payload = {"hook_event_name": "UserPromptSubmit", "prompt": prompt} folded = await self._dispatch("UserPromptSubmit", None, payload) @@ -192,10 +208,18 @@ async def run_stop(self, stop_hook_active: bool = False) -> StopOutcome: async def run_pre_compact(self, trigger: str = "auto") -> ContextOutcome: """Before a summarization pass. A ``block`` (continue:false) asks to skip - compaction this turn; the matcher runs against ``trigger`` (auto/manual).""" + compaction this turn; the matcher runs against ``trigger`` (auto/manual). + + ``additional_contexts`` from hook ``hookSpecificOutput.additionalContext`` + are passed through so a PreCompact hook can inject a checkpoint summary + (memento-style) that survives the compaction.""" payload = {"hook_event_name": "PreCompact", "trigger": trigger} folded = await self._dispatch("PreCompact", trigger, payload) - return ContextOutcome(block=folded.block, block_reason=folded.block_reason) + return ContextOutcome( + block=folded.block, + block_reason=folded.block_reason, + additional_contexts=folded.additional_contexts, + ) async def run_post_compact(self, trigger: str = "auto") -> ContextOutcome: """After a summarization pass — a notification hook (state saved, etc.).""" diff --git a/core/harness/hooks/events.py b/core/harness/hooks/events.py index ed393156..61668754 100644 --- a/core/harness/hooks/events.py +++ b/core/harness/hooks/events.py @@ -15,8 +15,8 @@ - ``Stop`` — when the agent would end its turn (may force-continue) The remaining reference events (``PermissionRequest``, ``PreCompact``, -``PostCompact``, ``SubagentStart``, ``SubagentStop``) are recognized by the -engine but wired opportunistically as DeepCode grows the matching kernel seams. +``PostCompact``, ``SessionEnd``, ``SubagentStart``, ``SubagentStop``) are +recognized by the engine and wired at their matching lifecycle seams. """ from __future__ import annotations @@ -31,14 +31,18 @@ "PreCompact", "PostCompact", "SessionStart", + "SessionEnd", "UserPromptSubmit", "SubagentStart", "SubagentStop", "Stop", ) -# Events whose ``matcher`` field is meaningful. ``UserPromptSubmit`` and ``Stop`` -# fire unconditionally, so their matchers are ignored (mirrors the reference). +# Events whose ``matcher`` field is meaningful. ``UserPromptSubmit`` and +# ``Stop`` fire unconditionally, so their matchers are ignored (mirrors the +# reference). ``SessionEnd`` DOES honour its matcher: the session-exit reason +# is matched against the ``matcher`` field so hooks can target a specific exit +# path. DeepCode runtime shutdown uses the compatible ``other`` reason. _EVENTS_WITHOUT_MATCHER: frozenset[str] = frozenset({"UserPromptSubmit", "Stop"}) diff --git a/core/harness/hooks/execution.py b/core/harness/hooks/execution.py index b11977ef..ff772a9d 100644 --- a/core/harness/hooks/execution.py +++ b/core/harness/hooks/execution.py @@ -250,7 +250,9 @@ def _decode_permission_request(obj: dict) -> HandlerDecision: "Stop": lambda o: _block_from_decision(o, "Stop"), "SubagentStop": lambda o: _block_from_decision(o, "SubagentStop"), "SessionStart": _decode_additional_context_only, + "SessionEnd": _decode_additional_context_only, "SubagentStart": _decode_additional_context_only, + "PreCompact": lambda o: _block_from_decision(o, "PreCompact"), "PermissionRequest": _decode_permission_request, } diff --git a/core/harness/memory.py b/core/harness/memory.py index 64a4fd55..144c2f15 100644 --- a/core/harness/memory.py +++ b/core/harness/memory.py @@ -23,6 +23,9 @@ from __future__ import annotations +import os +import re +from functools import lru_cache from pathlib import Path from typing import Any @@ -41,6 +44,69 @@ _REMINDER_OPEN = "" _REMINDER_CLOSE = "" _REMINDER_CLOSE_ESCAPED = "</system-reminder>" +# Comma-separated glob patterns for instruction files that must not be loaded, +# for example ``code/CLAUDE.md,**/vendor/**``. +_INSTRUCTION_EXCLUDE_ENV = "DEEPCODE_INSTRUCTION_EXCLUDES" + + +@lru_cache(maxsize=256) +def _glob_to_re(pattern: str) -> re.Pattern[str]: + """Compile a path glob where ``**`` crosses directory boundaries.""" + + parts = [] + i, n = 0, len(pattern) + while i < n: + c = pattern[i] + if c == "*": + if i + 1 < n and pattern[i + 1] == "*": + if i + 2 < n and pattern[i + 2] in "/\\": + parts.append(r"(?:.*/)?") + i += 3 + else: + parts.append(".*") + i += 2 + else: + parts.append(r"[^/\\]*") + i += 1 + elif c == "?": + parts.append(r"[^/\\]") + i += 1 + else: + parts.append(re.escape(c)) + i += 1 + flags = re.IGNORECASE if os.name == "nt" else 0 + return re.compile("^" + "".join(parts) + "$", flags) + + +def _instruction_excluded(candidate: Path, *, root: Path | None = None) -> bool: + """Whether the candidate instruction file is excluded by pattern. + + Patterns containing a separator match both the absolute path and, when + available, the path relative to the repository root. A bare filename such + as ``CLAUDE.md`` matches that filename at any searched level. + """ + patterns = [ + p.strip() + for p in os.environ.get(_INSTRUCTION_EXCLUDE_ENV, "").split(",") + if p.strip() + ] + if not patterns: + return False + candidates = {str(candidate).replace("\\", "/"), candidate.name} + if root is not None: + try: + candidates.add(candidate.relative_to(root).as_posix()) + except ValueError: + pass + for pat in patterns: + normalized = pat.replace("\\", "/") + try: + compiled = _glob_to_re(normalized) + if any(compiled.fullmatch(value) for value in candidates): + return True + except re.error: + continue + return False def memory_dir(workspace: str | Path) -> Path: @@ -135,7 +201,10 @@ def project_instructions(workspace: str | Path) -> str: for directory in search_dirs: for name in _PROJECT_FILES: candidate = directory / name - if candidate.is_file(): + if candidate.is_file() and not _instruction_excluded( + candidate, + root=root or workspace, + ): try: body = candidate.read_text( encoding="utf-8", errors="replace" diff --git a/core/mcp/models.py b/core/mcp/models.py index 67c9a812..a16f4d95 100644 --- a/core/mcp/models.py +++ b/core/mcp/models.py @@ -154,6 +154,10 @@ class McpServerDefinition(_ConfigModel): enabled: bool = True required: bool = False + defer_loading: bool = Field( + default=False, + validation_alias=AliasChoices("deferLoading", "defer_loading"), + ) supports_parallel_tool_calls: bool = False startup_timeout_seconds: Annotated[float, Field(gt=0, le=300)] = Field( default=10.0, @@ -266,6 +270,9 @@ def _valid_tool_names( @model_validator(mode="after") def _transport_contract(self) -> McpServerDefinition: + if self.required and self.defer_loading: + raise ValueError("required MCP servers cannot defer loading") + if self.bearer_token_env_var is not None and not _ENV_NAME.fullmatch( self.bearer_token_env_var ): diff --git a/core/mcp/runtime.py b/core/mcp/runtime.py index 23527967..97596685 100644 --- a/core/mcp/runtime.py +++ b/core/mcp/runtime.py @@ -9,6 +9,7 @@ from loguru import logger +from core.agent_runtime.tools.base import Tool, ToolResult from core.agent_runtime.tools.registry import ToolRegistry from core.mcp.connection import CredentialResolver, McpConnection, OAuthProviderFactory from core.mcp.models import McpRuntimePlan, McpStartupError @@ -26,6 +27,74 @@ class McpServerRuntimeStatus: error: str | None = None +class _ActivateMcpServerTool(Tool): + """Model-visible bridge that makes deferred servers reachable.""" + + def __init__( + self, + runtime: McpSessionRuntime, + *, + name: str, + server_ids: tuple[str, ...], + ) -> None: + self._runtime = runtime + self._name = name + self._server_ids = server_ids + + @property + def name(self) -> str: + return self._name + + @property + def description(self) -> str: + available = ", ".join(self._server_ids) + return ( + "Start one configured deferred MCP server so its tools become " + f"available to this session. Deferred server ids: {available}." + ) + + @property + def parameters(self) -> dict[str, Any]: + return { + "type": "object", + "properties": { + "server_id": { + "type": "string", + "enum": list(self._server_ids), + "description": "Configured MCP server id to activate.", + } + }, + "required": ["server_id"], + "additionalProperties": False, + } + + def presentation_detail(self, arguments: dict[str, Any]) -> str | None: + value = arguments.get("server_id") + return value if isinstance(value, str) else "" + + async def execute(self, *, server_id: str) -> ToolResult: + activated = await self._runtime.activate_server(server_id) + if not activated: + return ToolResult( + f"Error: MCP server '{server_id}' could not be activated.", + is_error=True, + metadata={"serverId": server_id, "activated": False}, + ) + tools = self._runtime.skill_capabilities.get(server_id, ()) + suffix = f" Available tools: {', '.join(tools)}." if tools else "" + instructions = self._runtime.server_instruction(server_id) + if instructions: + suffix += f"\n\nServer instructions:\n{instructions}" + return ToolResult( + f"MCP server '{server_id}' activated.{suffix}", + metadata={ + "serverId": server_id, + "activated": True, + "tools": list(tools), + }, + ) + + class McpSessionRuntime: """Materialize one immutable MCP plan into a DeepCode ToolRegistry.""" @@ -95,6 +164,14 @@ def instruction_context(self) -> str | None: return None return "\n\n".join(sections)[:8_000] + def server_instruction(self, server_id: str) -> str | None: + """Return bounded instructions for one active server.""" + + connection = self._connections.get(server_id) + if connection is None or not connection.instructions: + return None + return str(connection.instructions)[:4_000] + async def ensure_started(self) -> None: async with self._lock: if self._closed: @@ -104,6 +181,11 @@ async def ensure_started(self) -> None: if not self.plan.servers: self._started = True return + deferred_ids = { + server.server_id + for server in self.plan.servers + if server.definition.defer_loading + } connections = { server.server_id: McpConnection( server, @@ -111,13 +193,14 @@ async def ensure_started(self) -> None: oauth_provider_factory=self._oauth_provider_factory, ) for server in self.plan.servers + if server.server_id not in deferred_ids } for server_id, status in tuple(self._statuses.items()): self._statuses[server_id] = McpServerRuntimeStatus( status.server_id, status.name, status.source, - "starting", + "deferred" if server_id in deferred_ids else "starting", 0, ) self._publish_statuses() @@ -165,54 +248,158 @@ async def ensure_started(self) -> None: used = set(self.registry.tool_names) registered: list[str] = [] for connection, definitions in ready: - server = connection.server - seen_raw: set[str] = set() - exposed_count = 0 - server_tools: list[str] = [] - for definition in sorted(definitions, key=lambda item: str(item.name)): - raw_name = str(definition.name) - if raw_name in seen_raw: - logger.warning( - "MCP server '{}' returned duplicate tool name '{}'", - server.name, - raw_name, - ) - continue - seen_raw.add(raw_name) - if not server.definition.exposes(raw_name): - continue - name = visible_tool_name(server.server_id, raw_name, used=used) - adapter = McpToolAdapter( - connection, - definition, - visible_name=name, + registered.extend( + self._register_server_tools(connection, definitions, used=used) + ) + if deferred_ids: + activation_name = visible_tool_name( + "deepcode_runtime", + "activate_server", + used=used, + ) + self.registry.register( + _ActivateMcpServerTool( + self, + name=activation_name, + server_ids=tuple(sorted(deferred_ids)), ) - self.registry.register(adapter) - registered.append(name) - server_tools.append(name) - exposed_count += 1 - self._warn_unmatched_filters(server, seen_raw) - tools = tuple(server_tools) - self._capabilities[server.server_id] = tools - if server.policy_key is not None: - self._capabilities[server.policy_key] = tools - if ( - server.plugin_id is not None - and server.plugin_server_name is not None - ): - self._capabilities[ - f"{server.plugin_id}:{server.plugin_server_name}" - ] = tools - self._statuses[server.server_id] = McpServerRuntimeStatus( + ) + registered.append(activation_name) + self._registered_tools = tuple(registered) + self._started = True + self._publish_statuses() + + async def activate_server(self, server_id: str) -> bool: + """Start one MCP server on demand and register its tools. + + Deferred servers (``deferLoading: true``) are skipped by + :meth:`ensure_started`; call this to bring one up when it is actually + needed. Idempotent: returns ``True`` immediately when the server is + already ready. Returns ``False`` for unknown servers or when startup + fails (status is set to ``failed``). + """ + if not self._started: + if self._closed: + return False + try: + await self.ensure_started() + except RuntimeError: + if self._closed: + return False + raise + async with self._lock: + if self._closed: + return False + existing = self._connections.get(server_id) + if existing is not None and existing.ready: + return True + server = next( + (s for s in self.plan.servers if s.server_id == server_id), + None, + ) + if server is None: + return False + self._statuses[server_id] = McpServerRuntimeStatus( + server.server_id, + server.name, + server.source.value, + "starting", + 0, + ) + connection = McpConnection( + server, + credential_resolver=self._credential_resolver, + oauth_provider_factory=self._oauth_provider_factory, + ) + try: + definitions = await connection.start() + except asyncio.CancelledError: + await connection.close() + self._statuses[server_id] = McpServerRuntimeStatus( server.server_id, server.name, server.source.value, - "ready", - exposed_count, + "deferred" if server.definition.defer_loading else "failed", + 0, ) - self._registered_tools = tuple(registered) - self._started = True + self._publish_statuses() + raise + except BaseException as exc: # noqa: BLE001 - startup boundary + error = f"{type(exc).__name__}: {exc}" + self._statuses[server_id] = McpServerRuntimeStatus( + server.server_id, + server.name, + server.source.value, + "failed", + 0, + error, + ) + logger.warning( + "MCP server '{}' activation failed: {}", server.name, error + ) + await connection.close() + self._publish_statuses() + return False + self._connections[server_id] = connection + registered = self._register_server_tools( + connection, + definitions, + used=set(self.registry.tool_names), + ) + self._registered_tools = (*self._registered_tools, *registered) self._publish_statuses() + return True + + def _register_server_tools( + self, + connection: McpConnection, + definitions: tuple[Any, ...], + *, + used: set[str], + ) -> list[str]: + """Register one ready server's tools; returns the visible names.""" + server = connection.server + seen_raw: set[str] = set() + exposed_count = 0 + server_tools: list[str] = [] + for definition in sorted(definitions, key=lambda item: str(item.name)): + raw_name = str(definition.name) + if raw_name in seen_raw: + logger.warning( + "MCP server '{}' returned duplicate tool name '{}'", + server.name, + raw_name, + ) + continue + seen_raw.add(raw_name) + if not server.definition.exposes(raw_name): + continue + name = visible_tool_name(server.server_id, raw_name, used=used) + adapter = McpToolAdapter( + connection, + definition, + visible_name=name, + ) + self.registry.register(adapter) + server_tools.append(name) + exposed_count += 1 + self._warn_unmatched_filters(server, seen_raw) + tools = tuple(server_tools) + self._capabilities[server.server_id] = tools + if server.policy_key is not None: + self._capabilities[server.policy_key] = tools + if server.plugin_id is not None and server.plugin_server_name is not None: + self._capabilities[ + f"{server.plugin_id}:{server.plugin_server_name}" + ] = tools + self._statuses[server.server_id] = McpServerRuntimeStatus( + server.server_id, + server.name, + server.source.value, + "ready", + exposed_count, + ) + return server_tools async def aclose(self) -> None: async with self._lock: diff --git a/core/private_storage.py b/core/private_storage.py index e4e662ac..d6789952 100644 --- a/core/private_storage.py +++ b/core/private_storage.py @@ -5,14 +5,18 @@ the process umask, which is commonly permissive on desktop systems. POSIX permissions are repaired to ``0700`` for directories and ``0600`` for -regular files. Windows access control is inherited from the user's profile; -the mode arguments are still supplied at creation time where supported. +regular files. On Windows the current user is granted full control and the +inherited access entries are then stripped; the restriction is applied in a +fail-safe order so a failed grant leaves the inherited ACLs untouched and the +path stays accessible. """ from __future__ import annotations import os import stat +import subprocess +from functools import lru_cache from pathlib import Path PRIVATE_DIRECTORY_MODE = 0o700 @@ -23,6 +27,145 @@ class UnsafePrivateFileError(OSError): """A private-state path is not a regular file owned by this path entry.""" +_WINDOWS_BROAD_ACCESS_SIDS = ( + "*S-1-1-0", # Everyone + "*S-1-5-11", # Authenticated Users + "*S-1-5-32-545", # BUILTIN\\Users +) + + +@lru_cache(maxsize=1) +def _windows_identity() -> str | None: + """Return the fully-qualified current user (``DOMAIN\\user``) on Windows.""" + + if os.name != "nt": + return None + try: + import ctypes + from ctypes import wintypes + + get_name = ctypes.WinDLL("secur32", use_last_error=True).GetUserNameExW + get_name.argtypes = ( + wintypes.ULONG, + wintypes.LPWSTR, + ctypes.POINTER(wintypes.ULONG), + ) + get_name.restype = wintypes.BOOL + size = wintypes.ULONG(0) + # NameSamCompatible (2) yields DOMAIN\user, the form icacls accepts. + get_name(2, None, ctypes.byref(size)) + if size.value <= 1: + return None + buffer = ctypes.create_unicode_buffer(size.value) + if not get_name(2, buffer, ctypes.byref(size)): + return None + principal = buffer.value.strip() + except (AttributeError, OSError, TypeError, ValueError): + return None + return principal or None + + +@lru_cache(maxsize=1) +def _windows_icacls() -> str | None: + """Resolve the trusted system icacls executable without consulting PATH.""" + + if os.name != "nt": + return None + try: + import ctypes + from ctypes import wintypes + + get_system_directory = ctypes.WinDLL( + "kernel32", use_last_error=True + ).GetSystemDirectoryW + get_system_directory.argtypes = (wintypes.LPWSTR, wintypes.UINT) + get_system_directory.restype = wintypes.UINT + buffer = ctypes.create_unicode_buffer(32_768) + length = get_system_directory(buffer, len(buffer)) + if length <= 0 or length >= len(buffer): + return None + executable = Path(buffer.value) / "icacls.exe" + return os.fspath(executable) if executable.is_file() else None + except (AttributeError, OSError, TypeError, ValueError): + return None + + +def _run_icacls(executable: str, path: Path, *arguments: str) -> bool: + """Run one bounded icacls operation and report whether it succeeded.""" + + try: + subprocess.run( + [executable, os.fspath(path), *arguments], + capture_output=True, + text=True, + encoding="mbcs", + errors="replace", + timeout=15, + check=True, + ) + except (OSError, subprocess.SubprocessError): + return False + return True + + +def _restrict_windows_acl(path: Path) -> None: + """Restrict ``path`` to the current user, failing safe. + + The current user is granted full control **before** inherited access + entries are stripped. If the grant fails (service account, transient + timeout, ...) the inherited ACLs are left untouched so the path stays + accessible to the caller; the previous strip-first order could leave a + path with no usable ACE and make it unopenable. + """ + + identity = _windows_identity() + executable = _windows_icacls() + if identity is None or executable is None: + return + grant = f"{identity}:{'(OI)(CI)' if path.is_dir() else ''}F" + if not _run_icacls(executable, path, "/grant:r", grant): + # Fail safe: keep the inherited ACLs; the path stays accessible. + return + if not _run_icacls(executable, path, "/inheritance:r"): + # Strip failed: the path is merely less restricted, still usable. + return + # ``/inheritance:r`` removes inherited ACEs but deliberately leaves + # explicit entries alone. Remove the broad built-in principals as a final + # best-effort step so hardening also repairs legacy explicit grants. + _run_icacls(executable, path, "/remove", *_WINDOWS_BROAD_ACCESS_SIDS) + + +def _open_private_descriptor(target: Path, flags: int) -> tuple[int, bool]: + """Open ``target`` and atomically report whether this call created it.""" + + open_flags = flags | getattr(os, "O_NOFOLLOW", 0) + if not flags & os.O_CREAT or flags & os.O_EXCL: + descriptor = os.open(target, open_flags, PRIVATE_FILE_MODE) + return descriptor, bool(flags & os.O_CREAT) + + # A pre-open exists() check races with another creator. First attempt an + # exclusive creation; if the path already exists, open it without O_CREAT + # so a concurrent delete is observable and can retry the decision. + while True: + try: + descriptor = os.open( + target, + open_flags | os.O_EXCL, + PRIVATE_FILE_MODE, + ) + return descriptor, True + except FileExistsError: + try: + descriptor = os.open( + target, + open_flags & ~os.O_CREAT, + PRIVATE_FILE_MODE, + ) + except FileNotFoundError: + continue + return descriptor, False + + def ensure_private_directory(path: Path | str) -> Path: """Create ``path`` and make every newly created component user-private.""" @@ -35,10 +178,14 @@ def ensure_private_directory(path: Path | str) -> Path: for component in reversed(missing): component.mkdir(mode=PRIVATE_DIRECTORY_MODE, exist_ok=True) - _chmod(component, PRIVATE_DIRECTORY_MODE) + _chmod(component, PRIVATE_DIRECTORY_MODE, force=True) + was_missing = not directory.exists() directory.mkdir(parents=True, exist_ok=True, mode=PRIVATE_DIRECTORY_MODE) - _chmod(directory, PRIVATE_DIRECTORY_MODE) + # Restrict only what this call actually created. An existing directory was + # restricted at its own creation; re-running icacls on it on every open + # costs two subprocesses without changing the ACL. + _chmod(directory, PRIVATE_DIRECTORY_MODE, force=was_missing) return directory @@ -47,11 +194,7 @@ def open_private_file(path: Path | str, flags: int) -> int: target = Path(path) ensure_private_directory(target.parent) - descriptor = os.open( - target, - flags | getattr(os, "O_NOFOLLOW", 0), - PRIVATE_FILE_MODE, - ) + descriptor, created = _open_private_descriptor(target, flags) try: metadata = target.lstat() opened = os.fstat(descriptor) @@ -63,6 +206,8 @@ def open_private_file(path: Path | str, flags: int) -> int: ) if os.name != "nt": os.fchmod(descriptor, PRIVATE_FILE_MODE) + if created: + _restrict_windows_acl(target) return descriptor except BaseException: os.close(descriptor) @@ -112,31 +257,50 @@ def ensure_private_file(path: Path | str) -> None: except OSError: return if stat.S_ISREG(metadata.st_mode): - _chmod(target, PRIVATE_FILE_MODE) + _chmod(target, PRIVATE_FILE_MODE, force=True) def harden_private_tree(root: Path | str) -> Path: """Repair a DeepCode-owned tree while refusing to traverse symlinks.""" base = ensure_private_directory(root) - if os.name == "nt": - return base + if _is_directory_link(base): + raise UnsafePrivateFileError("private storage root must not be a link") for current, directories, files in os.walk(base, followlinks=False): current_path = Path(current) - _chmod(current_path, PRIVATE_DIRECTORY_MODE) + _chmod(current_path, PRIVATE_DIRECTORY_MODE, force=True) directories[:] = [ - name for name in directories if not (current_path / name).is_symlink() + name for name in directories if not _is_directory_link(current_path / name) ] for name in directories: - _chmod(current_path / name, PRIVATE_DIRECTORY_MODE) + _chmod(current_path / name, PRIVATE_DIRECTORY_MODE, force=True) for name in files: ensure_private_file(current_path / name) return base -def _chmod(path: Path, mode: int) -> None: +def _is_directory_link(path: Path) -> bool: + """Treat symlinks and Windows junctions as traversal boundaries.""" + + try: + if path.is_symlink(): + return True + is_junction = getattr(path, "is_junction", None) + return bool(callable(is_junction) and is_junction()) + except OSError: + # An unreadable/replaced entry is unsafe to traverse. + return True + + +def _chmod(path: Path, mode: int, *, force: bool = False) -> None: if os.name == "nt": + # harden_private_tree (force=True) deliberately re-applies the + # restriction even to existing paths (it repairs legacy trees whose + # ACLs may be absent or permissive). Default callers pass force=False + # so an already-restricted path is not re-churned on every open. + if force: + _restrict_windows_acl(path) return try: os.chmod(path, mode, follow_symlinks=False) diff --git a/desktop/sidecar-requirements.in b/desktop/sidecar-requirements.in index d4605f06..a9ca2641 100644 --- a/desktop/sidecar-requirements.in +++ b/desktop/sidecar-requirements.in @@ -17,6 +17,6 @@ PyYAML==6.0.3 tiktoken==0.12.0 # Packaging toolchain. -pip==26.1.2 +pip==26.2 pyinstaller==6.21.0 pyinstaller-hooks-contrib==2026.6 diff --git a/desktop/sidecar-requirements.lock b/desktop/sidecar-requirements.lock index 8e5e0ef9..5e0caa59 100644 --- a/desktop/sidecar-requirements.lock +++ b/desktop/sidecar-requirements.lock @@ -105,7 +105,7 @@ packaging==26.2 # pyinstaller-hooks-contrib pefile==2024.8.26 ; sys_platform == 'win32' # via pyinstaller -pip==26.1.2 +pip==26.2 # via -r sidecar-requirements.in prompt-toolkit==3.0.52 # via -r sidecar-requirements.in diff --git a/tests/test_mcp_runtime_lazy.py b/tests/test_mcp_runtime_lazy.py new file mode 100644 index 00000000..452c5fdb --- /dev/null +++ b/tests/test_mcp_runtime_lazy.py @@ -0,0 +1,215 @@ +"""Lazy MCP server activation (deferLoading) tests. + +``McpServerDefinition.defer_loading`` skips a server at ``ensure_started``; +``McpSessionRuntime.activate_server`` brings it up on demand and registers its +tools. This mirrors the CLI-side lazy-connect design (CodeWhale-derived). +""" + +from __future__ import annotations + +import asyncio +import sys +import tempfile +import textwrap +import unittest +from pathlib import Path +from unittest.mock import patch + +import pytest +from pydantic import ValidationError + +from core.agent_runtime.tools.registry import ToolRegistry +from core.mcp.models import ( + McpRuntimePlan, + McpServerDefinition, + McpServerSource, + ResolvedMcpServer, +) +from core.mcp.runtime import McpSessionRuntime + +FAKE_SERVER_SRC = textwrap.dedent( + """ + import json, sys + for line in sys.stdin: + msg = json.loads(line) + method = msg.get("method") + if method == "initialize": + json.dump({"jsonrpc": "2.0", "id": msg["id"], "result": { + "protocolVersion": "2025-03-26", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "fake", "version": "1.0"}}}, sys.stdout) + print(flush=True) + elif method == "tools/list": + json.dump({"jsonrpc": "2.0", "id": msg["id"], "result": {"tools": [ + {"name": "echo", "description": "echo back", + "inputSchema": {"type": "object", "properties": {}}}]}}, + sys.stdout) + print(flush=True) + """ +) + + +class LazyActivationTests(unittest.IsolatedAsyncioTestCase): + def setUp(self) -> None: + self.tmp = tempfile.mkdtemp(prefix="deepcode-mcp-lazy-") + self.fake_server = Path(self.tmp) / "fake_mcp_server.py" + self.fake_server.write_text(FAKE_SERVER_SRC, encoding="utf-8") + + def _server(self, server_id: str, defer_loading: bool) -> ResolvedMcpServer: + return ResolvedMcpServer( + server_id=server_id, + name=server_id, + source=McpServerSource.USER, + definition=McpServerDefinition( + type="stdio", + command=sys.executable, + args=(str(self.fake_server),), + startup_timeout_seconds=10.0, + defer_loading=defer_loading, + ), + config_dir=Path(self.tmp), + workspace=Path(self.tmp), + ) + + def _plan(self, *servers: ResolvedMcpServer) -> McpRuntimePlan: + return McpRuntimePlan( + workspace=Path(self.tmp), + servers=tuple(servers), + revision="test", + ) + + async def test_deferred_server_not_started_at_ensure(self) -> None: + registry = ToolRegistry() + runtime = McpSessionRuntime( + self._plan(self._server("eager1", False), self._server("lazy1", True)), + registry, + ) + try: + await runtime.ensure_started() + statuses = {s.server_id: s for s in runtime.statuses} + self.assertEqual(statuses["eager1"].state, "ready") + self.assertEqual(statuses["lazy1"].state, "deferred") + self.assertNotIn("lazy1", runtime.available_server_ids) + # Only the eager server's remote tool is registered, plus the + # activation bridge that makes deferred servers reachable. + self.assertIn("eager1", runtime.skill_capabilities) + self.assertNotIn("lazy1", runtime.skill_capabilities) + self.assertIn("mcp__deepcode_runtime__activate_server", registry.tool_names) + finally: + await runtime.aclose() + self.assertNotIn("mcp__deepcode_runtime__activate_server", registry.tool_names) + + async def test_activate_server_registers_tools(self) -> None: + registry = ToolRegistry() + runtime = McpSessionRuntime( + self._plan(self._server("eager1", False), self._server("lazy1", True)), + registry, + ) + try: + await runtime.ensure_started() + result = await registry.execute( + "mcp__deepcode_runtime__activate_server", {"server_id": "lazy1"} + ) + self.assertIn("activated", str(result)) + statuses = {s.server_id: s for s in runtime.statuses} + self.assertEqual(statuses["lazy1"].state, "ready") + self.assertIn("lazy1", runtime.available_server_ids) + self.assertIn("lazy1", runtime.skill_capabilities) + self.assertGreater(len(runtime.skill_capabilities["lazy1"]), 0) + self.assertIn("mcp__lazy1__echo", registry.tool_names) + finally: + await runtime.aclose() + + async def test_activate_is_idempotent(self) -> None: + runtime = McpSessionRuntime( + self._plan(self._server("lazy1", True)), + ToolRegistry(), + ) + await runtime.ensure_started() + try: + self.assertTrue(await runtime.activate_server("lazy1")) + tools_after_first = set(runtime.skill_capabilities.get("lazy1", ())) + self.assertTrue(await runtime.activate_server("lazy1")) + tools_after_second = set(runtime.skill_capabilities.get("lazy1", ())) + self.assertEqual(tools_after_first, tools_after_second) + finally: + await runtime.aclose() + + async def test_activate_unknown_server_returns_false(self) -> None: + runtime = McpSessionRuntime( + self._plan(self._server("lazy1", True)), + ToolRegistry(), + ) + await runtime.ensure_started() + try: + self.assertFalse(await runtime.activate_server("no-such-server")) + finally: + await runtime.aclose() + + async def test_activate_after_close_returns_false(self) -> None: + runtime = McpSessionRuntime( + self._plan(self._server("lazy1", True)), + ToolRegistry(), + ) + await runtime.aclose() + self.assertFalse(await runtime.activate_server("lazy1")) + + async def test_activation_cancellation_propagates_and_restores_state(self) -> None: + runtime = McpSessionRuntime( + self._plan(self._server("lazy1", True)), + ToolRegistry(), + ) + await runtime.ensure_started() + + async def cancel_start(_connection): + raise asyncio.CancelledError + + try: + with ( + patch("core.mcp.runtime.McpConnection.start", cancel_start), + self.assertRaises(asyncio.CancelledError), + ): + await runtime.activate_server("lazy1") + statuses = {s.server_id: s for s in runtime.statuses} + self.assertEqual(statuses["lazy1"].state, "deferred") + finally: + await runtime.aclose() + + async def test_activate_failure_marks_failed(self) -> None: + broken = ResolvedMcpServer( + server_id="broken1", + name="broken1", + source=McpServerSource.USER, + definition=McpServerDefinition( + type="stdio", + command=sys.executable, + args=("C:/definitely_missing_mcp_server_xyz.py",), + startup_timeout_seconds=5.0, + defer_loading=True, + ), + config_dir=Path(self.tmp), + workspace=Path(self.tmp), + ) + runtime = McpSessionRuntime(self._plan(broken), ToolRegistry()) + await runtime.ensure_started() + try: + self.assertFalse(await runtime.activate_server("broken1")) + statuses = {s.server_id: s for s in runtime.statuses} + self.assertEqual(statuses["broken1"].state, "failed") + self.assertIsNotNone(statuses["broken1"].error) + finally: + await runtime.aclose() + + +def test_required_server_cannot_be_deferred() -> None: + with pytest.raises(ValidationError, match="required MCP servers cannot defer"): + McpServerDefinition( + type="stdio", + command=sys.executable, + required=True, + defer_loading=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_memory.py b/tests/test_memory.py index c1562ee1..fd114374 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -10,9 +10,11 @@ if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) -from core.harness.memory import ( # noqa: E402 +from core.harness.memory import ( + _INSTRUCTION_EXCLUDE_ENV, _MAX_INJECT_CHARS, MemoryTool, + _instruction_excluded, memory_dir, project_instructions, system_preamble, @@ -36,6 +38,48 @@ def test_project_instructions_prefers_agents_md(tmp_path): assert "spaces" not in out # AGENTS.md wins over CLAUDE.md +def test_instruction_excluded_matches_globs(monkeypatch): + monkeypatch.setenv(_INSTRUCTION_EXCLUDE_ENV, "**/code/CLAUDE.md,**/vendor/**") + # 匹配: 任意层级前缀 + 目录段精确匹配 + assert _instruction_excluded(Path("repo/code/CLAUDE.md")) + assert _instruction_excluded(Path("repo/vendor/x/AGENTS.md")) + assert _instruction_excluded(Path("repo/vendor/AGENTS.md")) + assert _instruction_excluded(Path("code/CLAUDE.md")) # 零层前缀 + # 反例: 前缀同名目录不误匹配 (vendorized ≠ vendor/) + assert not _instruction_excluded(Path("repo/CLAUDE.md")) + assert not _instruction_excluded(Path("repo/vendorized/AGENTS.md")) + assert not _instruction_excluded(Path("repo/vendorized/x/CLAUDE.md")) + + +def test_instruction_excluded_matches_repo_relative_and_bare_names( + tmp_path, monkeypatch +): + repo = tmp_path / "repo" + candidate = repo / "code" / "CLAUDE.md" + monkeypatch.setenv(_INSTRUCTION_EXCLUDE_ENV, "code/CLAUDE.md") + assert _instruction_excluded(candidate, root=repo) + + monkeypatch.setenv(_INSTRUCTION_EXCLUDE_ENV, "CLAUDE.md") + assert _instruction_excluded(candidate, root=repo) + + +def test_project_instructions_skips_excluded_file(tmp_path, monkeypatch): + repo = tmp_path / "repo" + (repo / ".git").mkdir(parents=True) + (repo / "code").mkdir() + (repo / "CLAUDE.md").write_text("root instructions") + (repo / "code" / "CLAUDE.md").write_text("subdir instructions") + monkeypatch.setenv(_INSTRUCTION_EXCLUDE_ENV, "**/code/CLAUDE.md") + out = project_instructions(repo / "code") + assert "root instructions" in out + assert "subdir instructions" not in out + + +def test_instruction_excluded_treats_regex_metacharacters_literally(monkeypatch): + monkeypatch.setenv(_INSTRUCTION_EXCLUDE_ENV, "**/[code/CLAUDE.md") + assert not _instruction_excluded(Path("code/CLAUDE.md")) + + def test_project_instructions_absent(tmp_path): assert project_instructions(tmp_path) == "" diff --git a/tests/test_private_storage.py b/tests/test_private_storage.py index 08b03165..78f3986a 100644 --- a/tests/test_private_storage.py +++ b/tests/test_private_storage.py @@ -154,6 +154,16 @@ def test_private_tree_repair_does_not_follow_symlinks(tmp_path: Path) -> None: assert _mode(external) == 0o644 +def test_private_tree_rejects_a_link_as_its_root(tmp_path: Path) -> None: + external = tmp_path / "external" + external.mkdir() + link = tmp_path / "private-link" + link.symlink_to(external, target_is_directory=True) + + with pytest.raises(OSError, match="must not be a link"): + harden_private_tree(link) + + def test_suite_redirects_all_user_state_to_the_test_directory( tmp_path: Path, ) -> None: diff --git a/tests/test_private_storage_acl_once.py b/tests/test_private_storage_acl_once.py new file mode 100644 index 00000000..2ea778d9 --- /dev/null +++ b/tests/test_private_storage_acl_once.py @@ -0,0 +1,173 @@ +"""Cross-platform tests for the per-open ACL optimization in private_storage. + +The Windows ACL restriction is applied at file *creation*; opening an +existing private file must not re-run icacls (maintainer feedback on the +earlier ACL PR: "open_private_file() currently calls it on each call; once at +creation is enough"). These tests mock `_restrict_windows_acl` to count calls, +so they run on any platform. +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from core.private_storage import open_private_file + + +def _file_calls(calls, target: Path) -> int: + """Count restrictions applied to the target file itself (excludes the + parent-directory restriction that ensure_private_directory performs).""" + return sum(1 for p in calls if Path(p) == target) + + +def test_open_existing_file_does_not_rerun_acl(monkeypatch, tmp_path: Path) -> None: + import core.private_storage as ps + + calls = [] + monkeypatch.setattr(ps, "_restrict_windows_acl", lambda p: calls.append(p)) + + target = tmp_path / "existing.jsonl" + # First open: file does not exist → new → restrict once. + fd = open_private_file(target, os.O_CREAT | os.O_RDWR) + os.close(fd) + assert _file_calls(calls, target) == 1, "new file must be restricted exactly once" + + # Second open: file exists → must NOT re-run the ACL restriction. + fd = open_private_file(target, os.O_RDWR) + os.close(fd) + assert _file_calls(calls, target) == 1, ( + "existing file must not re-run the ACL restriction" + ) + + +def test_open_created_file_restricts_once(monkeypatch, tmp_path: Path) -> None: + import core.private_storage as ps + + calls = [] + monkeypatch.setattr(ps, "_restrict_windows_acl", lambda p: calls.append(p)) + + target = tmp_path / "fresh.jsonl" + for _ in range(3): + fd = open_private_file(target, os.O_CREAT | os.O_RDWR) + os.close(fd) + assert _file_calls(calls, target) == 1, "created once, restricted once, never again" + + +def test_open_without_creat_never_restricts(monkeypatch, tmp_path: Path) -> None: + import core.private_storage as ps + + calls = [] + monkeypatch.setattr(ps, "_restrict_windows_acl", lambda p: calls.append(p)) + + target = tmp_path / "pre.jsonl" + target.write_text("x", encoding="utf-8") + # O_RDONLY (no O_CREAT) on an existing file → no new file → no restriction. + fd = open_private_file(target, os.O_RDONLY) + os.close(fd) + assert _file_calls(calls, target) == 0, ( + "read-only open of an existing file must not restrict" + ) + + +def test_concurrent_creator_is_not_mistaken_for_our_new_file( + monkeypatch, tmp_path: Path +) -> None: + import core.private_storage as ps + + target = tmp_path / "raced.jsonl" + original_open = os.open + restrictions = [] + raced = False + + def racing_open(path, flags, mode=0o777): + nonlocal raced + if Path(path) == target and flags & os.O_EXCL and not raced: + raced = True + other = original_open(target, os.O_CREAT | os.O_EXCL | os.O_WRONLY, mode) + os.close(other) + raise FileExistsError(os.fspath(target)) + return original_open(path, flags, mode) + + monkeypatch.setattr(ps.os, "open", racing_open) + monkeypatch.setattr( + ps, "_restrict_windows_acl", lambda path: restrictions.append(Path(path)) + ) + + descriptor = ps.open_private_file(target, os.O_CREAT | os.O_RDWR) + os.close(descriptor) + + assert raced is True + assert target not in restrictions + + +def test_acl_restriction_orders_grant_strip_and_broad_principal_removal( + monkeypatch, tmp_path: Path +) -> None: + import core.private_storage as ps + + calls = [] + monkeypatch.setattr(ps, "_windows_identity", lambda: "DOMAIN\\user") + monkeypatch.setattr(ps, "_windows_icacls", lambda: "trusted-icacls.exe") + + def record(executable, path, *arguments): + calls.append((executable, Path(path), arguments)) + return True + + monkeypatch.setattr(ps, "_run_icacls", record) + + directory = tmp_path / "private" + directory.mkdir() + ps._restrict_windows_acl(directory) + + assert calls[0][2] == ("/grant:r", "DOMAIN\\user:(OI)(CI)F") + assert calls[1][2] == ("/inheritance:r",) + assert calls[2][2] == ("/remove", *ps._WINDOWS_BROAD_ACCESS_SIDS) + + +def test_acl_restriction_stops_before_strip_when_grant_fails( + monkeypatch, tmp_path: Path +) -> None: + import core.private_storage as ps + + calls = [] + monkeypatch.setattr(ps, "_windows_identity", lambda: "DOMAIN\\user") + monkeypatch.setattr(ps, "_windows_icacls", lambda: "trusted-icacls.exe") + + def fail_grant(executable, path, *arguments): + calls.append(arguments) + return False + + monkeypatch.setattr(ps, "_run_icacls", fail_grant) + + ps._restrict_windows_acl(tmp_path / "private.json") + + assert calls == [("/grant:r", "DOMAIN\\user:F")] + + +def test_acl_restriction_does_not_remove_entries_when_strip_fails( + monkeypatch, tmp_path: Path +) -> None: + import core.private_storage as ps + + calls = [] + monkeypatch.setattr(ps, "_windows_identity", lambda: "DOMAIN\\user") + monkeypatch.setattr(ps, "_windows_icacls", lambda: "trusted-icacls.exe") + + def fail_strip(executable, path, *arguments): + calls.append(arguments) + return arguments != ("/inheritance:r",) + + monkeypatch.setattr(ps, "_run_icacls", fail_strip) + + ps._restrict_windows_acl(tmp_path / "private.json") + + assert calls == [ + ("/grant:r", "DOMAIN\\user:F"), + ("/inheritance:r",), + ] diff --git a/tests/test_private_storage_windows.py b/tests/test_private_storage_windows.py new file mode 100644 index 00000000..83d73311 --- /dev/null +++ b/tests/test_private_storage_windows.py @@ -0,0 +1,174 @@ +"""Windows NTFS ACL restriction tests for core.private_storage. + +These tests assert that private directories and files are restricted to the +current user with full control and that dangerous well-known ACEs (Everyone, +Authenticated Users, BUILTIN\\Users) are removed after the restriction runs. +They are skipped on non-Windows platforms. +""" + +from __future__ import annotations + +import os +import subprocess +from pathlib import Path + +import pytest + +from core.private_storage import ( + ensure_private_directory, + harden_private_tree, + open_private_file, +) + +pytestmark = pytest.mark.skipif( + os.name != "nt", + reason="NTFS ACL restriction applies on Windows only", +) + +_DANGEROUS_ACES = ("Authenticated Users", "BUILTIN\\Users", "Everyone") + + +def _windows_identity() -> str: + completed = subprocess.run( + ["whoami"], + capture_output=True, + text=True, + encoding="mbcs", + errors="replace", + timeout=5, + check=True, + ) + return (completed.stdout or "").strip() + + +def _acl_lines(path: Path) -> list[str]: + completed = subprocess.run( + ["icacls", os.fspath(path)], + capture_output=True, + text=True, + encoding="mbcs", + errors="replace", + timeout=15, + check=True, + ) + return [ + line.strip() for line in (completed.stdout or "").splitlines() if ":" in line + ] + + +def _assert_no_dangerous_aces(path: Path) -> None: + lines = _acl_lines(path) + joined = "\n".join(lines).lower() + for ace in _DANGEROUS_ACES: + assert ace.lower() not in joined, ( + f"{path} still exposes dangerous ACE {ace!r}:\n{joined}" + ) + + +def _assert_current_user_has_full_control(path: Path) -> None: + identity = _windows_identity().lower() + # ``whoami`` may return either ``domain\user`` or a bare ``user`` depending + # on which binary is on PATH, while ``icacls`` always prints the fully + # qualified principal. Compare the last path segment so both match. + short_name = identity.rsplit("\\", 1)[-1] + lines = _acl_lines(path) + for line in lines: + # Split from the right: icacls lines start with a Windows path that + # contains a drive-letter colon (``C:\\...``), so the first colon is + # not the principal/rights separator. + principal, rights = line.rsplit(":", 1) + principal_short = principal.strip().lower().rsplit("\\", 1)[-1] + if principal_short == short_name: + assert "(f)" in rights.lower(), ( + f"{path} does not grant the current user full control:\n{line}" + ) + return + raise AssertionError( + f"{path} has no ACE for the current user {identity!r}:\n" + "\n".join(lines) + ) + + +def test_windows_private_directory_is_restricted(tmp_path: Path) -> None: + directory = ensure_private_directory(tmp_path / "private" / "nested") + + _assert_no_dangerous_aces(directory) + _assert_current_user_has_full_control(directory) + + # The directory grant must be inheritable so files created by libraries + # that do not call open_private_file (SQLite WAL files, lock files, etc.) + # remain accessible without inheriting broad user-group access. + child = directory / "library-created.tmp" + child.write_text("x", encoding="utf-8") + _assert_no_dangerous_aces(child) + _assert_current_user_has_full_control(child) + + +def test_windows_private_file_is_restricted(tmp_path: Path) -> None: + target = tmp_path / "private" / "credentials.json" + descriptor = open_private_file(target, os.O_WRONLY | os.O_CREAT) + try: + os.write(descriptor, b"secret") + finally: + os.close(descriptor) + + _assert_no_dangerous_aces(target) + _assert_current_user_has_full_control(target) + assert target.read_bytes() == b"secret" + + +def test_windows_harden_private_tree_restricts_every_entry(tmp_path: Path) -> None: + root = tmp_path / "legacy-private" + session = root / "session-1" + session.mkdir(parents=True) + (session / "session.jsonl").write_text("legacy\n", encoding="utf-8") + (root / "settings.json").write_text("{}", encoding="utf-8") + subprocess.run( + ["icacls", os.fspath(root / "settings.json"), "/grant", "*S-1-1-0:R"], + capture_output=True, + text=True, + encoding="mbcs", + errors="replace", + timeout=15, + check=True, + ) + + harden_private_tree(root) + + for path in (root, session, session / "session.jsonl", root / "settings.json"): + _assert_no_dangerous_aces(path) + _assert_current_user_has_full_control(path) + + +def test_windows_open_existing_file_does_not_rerun_acl( + monkeypatch, tmp_path: Path +) -> None: + """The per-open ACL re-run is gone: opening an existing private file does + not call _restrict_windows_acl again (the ACL was applied at creation).""" + import core.private_storage as ps + + calls = [] + monkeypatch.setattr(ps, "_restrict_windows_acl", lambda p: calls.append(p)) + + target = tmp_path / "existing.jsonl" + # First open: file does not exist → new → restrict once. + fd = ps.open_private_file(target, os.O_CREAT | os.O_RDWR) + os.close(fd) + assert len(calls) == 1, "new file must be restricted exactly once" + + # Second open: file exists → must NOT re-run the ACL restriction. + fd = ps.open_private_file(target, os.O_RDWR) + os.close(fd) + assert len(calls) == 1, "existing file must not re-run the ACL restriction" + + +def test_windows_open_created_file_restricts_once(monkeypatch, tmp_path: Path) -> None: + """A file created via open_private_file is restricted exactly once.""" + import core.private_storage as ps + + calls = [] + monkeypatch.setattr(ps, "_restrict_windows_acl", lambda p: calls.append(p)) + + for _ in range(3): + fd = ps.open_private_file(tmp_path / "fresh.jsonl", os.O_CREAT | os.O_RDWR) + os.close(fd) + assert len(calls) == 1, "created once, restricted once, never again" diff --git a/tests/test_session_end_lifecycle.py b/tests/test_session_end_lifecycle.py new file mode 100644 index 00000000..6b375f9e --- /dev/null +++ b/tests/test_session_end_lifecycle.py @@ -0,0 +1,353 @@ +"""SessionEnd lifecycle and PreCompact checkpoint end-to-end tests. + +Covers the lifecycle contracts introduced by the SessionEnd / PreCompact work: + +- ``SessionEnd`` fires exactly once at real session termination + (``AgentSession.submit(Shutdown)``), never per turn; the session-exit reason + doubles as the matcher input (DeepCode shutdown maps to ``other``); a failing + hook is non-fatal and never blocks ``ShutdownComplete``. +- The ``PreCompact`` hook's ``additional_contexts`` survive a successful + compaction as a single bounded, provider-agnostic user message, and are + absent when the hook blocks or summarization fails. + +Hooks are exercised as REAL subprocesses (``sh -lc`` commands that echo JSON or +exit with a code), matching ``test_hooks.py`` so we test the true execution +path, not a mock of it. +""" + +from __future__ import annotations + +import asyncio +import json +import shutil +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from core.events.protocol import Shutdown, ShutdownComplete, UserInput # noqa: E402 +from core.harness.hooks.discovery import Handler # noqa: E402 +from core.harness.hooks.engine import HooksEngine # noqa: E402 + +pytestmark = pytest.mark.skipif( + shutil.which("sh") is None, reason="POSIX shell required" +) + + +def _handler(event, command, *, matcher=None, order=0, timeout=30): + return Handler( + event_name=event, + matcher=matcher, + command=command, + timeout_sec=timeout, + source="project", + source_path="/tmp/hooks.json", + display_order=order, + ) + + +def _engine(handlers, cwd="/tmp"): + return HooksEngine(handlers, cwd, session_id="sess-1") + + +def _session(hooks_engine): + from core.events.session import AgentSession + + return AgentSession( + provider=None, + tools=_FakeTools(), + model="m", + hooks_engine=hooks_engine, + context_window_tokens=8000, + ) + + +class _FakeTools: + """Minimal tool registry: records the params each tool ran with.""" + + def __init__(self): + self.calls = [] + + def get_definitions(self): + return [] + + async def execute(self, name, params): + self.calls.append((name, params)) + return f"ran {name} with {params}" + + async def aclose(self): + return None + + +# --------------------------------------------------------------------------- +# SessionEnd lifecycle — fires exactly once at real session termination +# --------------------------------------------------------------------------- + + +def test_session_end_fires_exactly_once_on_shutdown(tmp_path): + count = tmp_path / "count.txt" + eng = _engine([_handler("SessionEnd", f"echo x >> {count}")]) + session = _session(eng) + + asyncio.run(session.submit(Shutdown())) + + assert count.read_text().count("x") == 1 + event = asyncio.run(session.next_event()) + assert isinstance(event.msg, ShutdownComplete) + + asyncio.run(session.submit(Shutdown())) + asyncio.run(session.aclose()) + assert count.read_text().count("x") == 1 + + +def test_session_end_fires_from_real_close_path(tmp_path): + count = tmp_path / "count.txt" + eng = _engine([_handler("SessionEnd", f"echo x >> {count}")]) + session = _session(eng) + + asyncio.run(session.aclose()) + asyncio.run(session.aclose()) + + assert count.read_text().count("x") == 1 + + +def test_session_end_reason_matcher(tmp_path): + shutdown_hits = tmp_path / "shutdown.txt" + other_hits = tmp_path / "other.txt" + eng = _engine( + [ + _handler("SessionEnd", f"echo x >> {shutdown_hits}", matcher="other"), + _handler("SessionEnd", f"echo x >> {other_hits}", matcher="complete"), + ] + ) + session = _session(eng) + + asyncio.run(session.submit(Shutdown())) + + assert shutdown_hits.read_text().count("x") == 1 + assert not other_hits.exists() + + +def test_session_end_hook_failure_non_fatal(): + eng = _engine([_handler("SessionEnd", "exit 3")]) + session = _session(eng) + + # A failing SessionEnd hook must never crash the session close. + asyncio.run(session.submit(Shutdown())) + event = asyncio.run(session.next_event()) + assert isinstance(event.msg, ShutdownComplete) + + +def test_session_end_is_not_emitted_for_subagent_session(tmp_path): + count = tmp_path / "count.txt" + eng = _engine([_handler("SessionEnd", f"echo x >> {count}")]) + session = _session(eng) + session._agent_context = ("child-1", "subagent") + + asyncio.run(session.aclose()) + + assert not count.exists() + + +def test_session_end_discovery_uses_bounded_exit_timeout(tmp_path): + from core.harness.hooks.discovery import discover_hooks + + home = tmp_path / "home" + workspace = tmp_path / "workspace" + (home / ".deepcode").mkdir(parents=True) + workspace.mkdir() + (home / ".deepcode" / "hooks.json").write_text( + json.dumps( + { + "hooks": { + "SessionEnd": [ + {"matcher": "*", "hooks": [{"command": "echo default"}]}, + { + "matcher": "*", + "hooks": [{"command": "echo capped", "timeout": 999}], + }, + ] + } + } + ), + encoding="utf-8", + ) + + result = discover_hooks(str(workspace), str(home)) + + assert [handler.timeout_sec for handler in result.handlers] == [2, 60] + + +def test_normal_turn_does_not_trigger_session_end(tmp_path): + count = tmp_path / "count.txt" + eng = _engine([_handler("SessionEnd", f"echo x >> {count}")]) + session = _session(eng) + + async def _noop_user_input(op): + return None + + session._run_user_input = _noop_user_input # type: ignore[assignment] + asyncio.run(session.submit(UserInput("hello"))) + assert not count.exists() + + asyncio.run(session.submit(Shutdown())) + assert count.read_text().count("x") == 1 + + +# --------------------------------------------------------------------------- +# PreCompact checkpoint — bounded, provider-safe re-injection +# --------------------------------------------------------------------------- + + +def test_build_precompact_checkpoint_empty(): + from core.agent_runtime.runner import _build_precompact_checkpoint + + assert _build_precompact_checkpoint([]) is None + assert _build_precompact_checkpoint([" "]) is None + + +def test_build_precompact_checkpoint_limits(): + from core.agent_runtime.runner import ( + _PRECOMPACT_CHECKPOINT_PREFIX, + _PRECOMPACT_CONTEXT_LIMIT, + _PRECOMPACT_TOTAL_LIMIT, + _build_precompact_checkpoint, + ) + + long_ctx = "y" * (_PRECOMPACT_CONTEXT_LIMIT * 2) + checkpoint = _build_precompact_checkpoint([long_ctx]) + assert checkpoint is not None + body = checkpoint[len(_PRECOMPACT_CHECKPOINT_PREFIX) + 1 :] + assert len(body) == _PRECOMPACT_CONTEXT_LIMIT + + many = ["z" * 3000] * 10 + checkpoint = _build_precompact_checkpoint(many) + assert checkpoint is not None + body = checkpoint[len(_PRECOMPACT_CHECKPOINT_PREFIX) + 1 :] + assert len(checkpoint) <= _PRECOMPACT_TOTAL_LIMIT + + +def test_build_precompact_checkpoint_honors_smaller_dynamic_limit(): + from core.agent_runtime.runner import _build_precompact_checkpoint + + checkpoint = _build_precompact_checkpoint(["x" * 500], total_limit=100) + assert checkpoint is not None + assert len(checkpoint) <= 100 + + +def test_maybe_compact_checkpoint_injected_after_success(monkeypatch): + from types import SimpleNamespace + + from core.agent_runtime.runner import AgentRunner, AgentRunSpec + + runner = AgentRunner(provider=object()) + monkeypatch.setattr(runner, "_estimate_prompt", lambda spec, messages: 999_999) + + async def fake_summarize(spec, messages, *, response_observer=None): + return "handoff summary" + + monkeypatch.setattr(runner, "_summarize", fake_summarize) + + async def pre_compact_hook(trigger): + return SimpleNamespace(block=False, additional_contexts=["checkpoint ctx"]) + + spec = AgentRunSpec( + initial_messages=[], + tools=_FakeTools(), + model="m", + max_iterations=1, + max_tool_result_chars=100000, + context_window_tokens=8000, + pre_compact_hook=pre_compact_hook, + ) + messages = [ + {"role": "user", "content": "turn 1 " + "context " * 50}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "turn 2 " + "context " * 50}, + {"role": "assistant", "content": "a2"}, + {"role": "user", "content": "turn 3 " + "context " * 50}, + ] + compacted = asyncio.run(runner._maybe_compact(spec, messages)) + checkpoint_msgs = [ + m + for m in compacted + if m.get("role") == "user" and "PreCompact checkpoint" in str(m.get("content")) + ] + assert checkpoint_msgs, "checkpoint must survive a successful compaction" + assert "checkpoint ctx" in checkpoint_msgs[0]["content"] + assert sum(len(str(m.get("content", ""))) for m in compacted) < sum( + len(str(m.get("content", ""))) for m in messages + ) + + +def test_maybe_compact_block_skips_checkpoint(monkeypatch): + from types import SimpleNamespace + + from core.agent_runtime.runner import AgentRunner, AgentRunSpec + + runner = AgentRunner(provider=object()) + monkeypatch.setattr(runner, "_estimate_prompt", lambda spec, messages: 999_999) + + async def pre_compact_hook(trigger): + return SimpleNamespace(block=True, additional_contexts=["should not appear"]) + + spec = AgentRunSpec( + initial_messages=[], + tools=_FakeTools(), + model="m", + max_iterations=1, + max_tool_result_chars=100000, + context_window_tokens=8000, + pre_compact_hook=pre_compact_hook, + ) + messages = [ + {"role": "user", "content": "turn 1 " + "context " * 50}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "turn 2 " + "context " * 50}, + {"role": "assistant", "content": "a2"}, + {"role": "user", "content": "turn 3 " + "context " * 50}, + ] + compacted = asyncio.run(runner._maybe_compact(spec, messages)) + assert compacted is messages # compaction aborted this turn + assert "PreCompact checkpoint" not in json.dumps(compacted) + + +def test_maybe_compact_summarize_failure_no_checkpoint(monkeypatch): + from types import SimpleNamespace + + from core.agent_runtime.runner import AgentRunner, AgentRunSpec + + runner = AgentRunner(provider=object()) + monkeypatch.setattr(runner, "_estimate_prompt", lambda spec, messages: 999_999) + + async def fake_summarize_fails(spec, messages, *, response_observer=None): + return None # summarization failed + + monkeypatch.setattr(runner, "_summarize", fake_summarize_fails) + + async def pre_compact_hook(trigger): + return SimpleNamespace(block=False, additional_contexts=["should not appear"]) + + spec = AgentRunSpec( + initial_messages=[], + tools=_FakeTools(), + model="m", + max_iterations=1, + max_tool_result_chars=100000, + context_window_tokens=8000, + pre_compact_hook=pre_compact_hook, + ) + messages = [ + {"role": "user", "content": "turn 1 " + "context " * 50}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "turn 2 " + "context " * 50}, + {"role": "assistant", "content": "a2"}, + {"role": "user", "content": "turn 3 " + "context " * 50}, + ] + compacted = asyncio.run(runner._maybe_compact(spec, messages)) + assert compacted is messages + assert "PreCompact checkpoint" not in json.dumps(compacted)