diff --git a/.gitignore b/.gitignore index d602ca38a..06df41f8a 100644 --- a/.gitignore +++ b/.gitignore @@ -132,6 +132,7 @@ venv.bak/ .firecrawl .claude/ .claude-trace/ +.qoder # custom *.pkl diff --git a/ms_agent/agent/llm_agent.py b/ms_agent/agent/llm_agent.py index f2b3d11bd..542f61f77 100644 --- a/ms_agent/agent/llm_agent.py +++ b/ms_agent/agent/llm_agent.py @@ -444,6 +444,10 @@ def _build_system_content(self) -> str: content += '\n\n' + LIVE_FILES_HINT.format( home=str(global_home())) + internals = self._build_workspace_internals_section() + if internals: + content += '\n\n' + internals + if self._memory_guidance: content += '\n\n' + self._memory_guidance @@ -461,6 +465,46 @@ def _build_system_content(self) -> str: return content + def _build_workspace_internals_section(self) -> str: + """Describe the framework's own directories, when they sit in the + working directory. + + Only for the layout where they do. A project opened from an existing + folder keeps its records elsewhere, and telling that agent to watch out + for a ``sessions/`` directory it will never encounter would be a + fabricated warning. + """ + from ms_agent.prompting.builtin import WORKSPACE_INTERNALS_HINT + from ms_agent.utils.workspace_context import resolve_workspace_root + + try: + workspace_root = Path(resolve_workspace_root(self.config)) + except Exception: # noqa: BLE001 - never break prompt assembly + return '' + if not (workspace_root / 'sessions').is_dir(): + return '' + + # The directory the log is actually writing to, not the agent's tag: + # naming a path that does not exist is worse than naming none, since + # the model will go looking for it. + session_dir = None + log = getattr(self, 'session_log', None) + directory = getattr(log, 'directory', None) + if directory is not None: + session_dir = Path(directory).name + if not session_dir: + session_dir = getattr(self.runtime, 'session_id', None) + + try: + home = str(global_home()) + except Exception: # noqa: BLE001 + home = '~/.ms_agent' + hint = WORKSPACE_INTERNALS_HINT.format( + session_line=(f' This conversation is `sessions/{session_dir}/`.' + if session_dir else ''), + home=home) + return hint + def _check_skill_tool_dependencies(self): """Warn if skills are enabled but essential tools are missing.""" if (not self._skill_catalog @@ -858,10 +902,8 @@ def _build_permission_objects(self): perm_config = PermissionConfig.from_dict( raw, project_root=workspace_root) - allowed_dirs = [workspace_root] - for directory in perm_config.safety.allowed_directories: - if directory not in allowed_dirs: - allowed_dirs.append(directory) + allowed_dirs = list( + perm_config.safety.effective_allowed_directories(workspace_root)) read_only_dirs = list(perm_config.safety.read_only_directories) safety_guard = SafetyGuard( config=perm_config.safety, diff --git a/ms_agent/permission/ask_resolver.py b/ms_agent/permission/ask_resolver.py index ca223c486..2b58e1e53 100644 --- a/ms_agent/permission/ask_resolver.py +++ b/ms_agent/permission/ask_resolver.py @@ -19,9 +19,31 @@ 'command_validator': 'deny', 'shell_expansion': 'deny', 'read_outside_dirs': 'deny', + # Running code is what full access is FOR. Refusing every `python3 -c` in + # the mode whose whole meaning is "stop asking me" would make the mode + # useless, and the confirmation this category exists for is the one an + # interactive user gets. What auto mode cannot honestly claim is that it + # inspected the code — hence the message on the ask, and the setting-page + # copy that says so. + 'interpreter_exec': 'allow', + # A private key does not become less private because the user is not + # watching. There is no path here that reads it without someone deciding. + 'sensitive_read': 'deny', } +#: Safety confirmations a standing answer may satisfy. +#: +#: The rest are deliberately not here. "This reads a private key" has to be +#: decided each time, because what makes it risky is the specific file, and a +#: pattern broad enough to remember would cover files the user never saw. +#: "This runs code I cannot analyse" is different: every inline `python3 -c` is +#: the same decision, it comes up constantly, and an ask the user cannot settle +#: is one they answer by turning confirmations off — which is the outcome the +#: confirmation existed to prevent. +REMEMBERABLE_ASK_CATEGORIES: frozenset = frozenset({'interpreter_exec'}) + + def resolve_ask( decision: SafetyDecision, mode: str, diff --git a/ms_agent/permission/config.py b/ms_agent/permission/config.py index dde6a17aa..3e164d622 100644 --- a/ms_agent/permission/config.py +++ b/ms_agent/permission/config.py @@ -7,6 +7,7 @@ from __future__ import annotations import os +from pathlib import Path from dataclasses import dataclass from typing import Any, Literal @@ -33,6 +34,30 @@ '**/.git/**', ) +#: Locations whose CONTENTS are a credential. Kept apart from +#: ``sensitive_paths`` because the two lists answer different questions: +#: that one protects things from being CHANGED (``.git/config``, +#: ``~/.bashrc``), which is no reason to refuse reading them — an agent that +#: cannot run ``cat .git/config`` to find a remote is just broken. These are +#: things that must not be COPIED, most of all into a transcript that gets +#: written to disk and replayed to a model. +#: +#: Patterns are fnmatch, where ``*`` crosses ``/``, so a leading ``*/`` covers +#: any user's home rather than only the one the process happens to run as. +_DEFAULT_SENSITIVE_READ_PATHS: tuple[str, ...] = ( + '*/.ssh/*', + '*/.gnupg/*', + '*/.aws/*', + '*/.kube/config', + '*/.docker/config.json', + '*/.netrc', + '*/id_rsa', + '*/id_dsa', + '*/id_ecdsa', + '*/id_ed25519', + '*/*.pem', +) + _DEFAULT_DANGEROUS_REMOVAL: tuple[str, ...] = ( '*', '/*', @@ -41,16 +66,45 @@ ) +def default_temp_directories() -> tuple[str, ...]: + """Scratch directories the agent may write to besides its workspace. + + Refusing these buys nothing: the OS temp directory is world-writable by + design and holds nothing to protect. It costs plenty, though — every tool + that stages output through a temp file, and every ``python3 - < SafetyConfig: patterns = tuple(d.get('patterns', _DEFAULT_SAFETY_PATTERNS)) sensitive = tuple(d.get('sensitive_paths', _DEFAULT_SENSITIVE_PATHS)) + sensitive_read = tuple( + d.get('sensitive_read_paths', _DEFAULT_SENSITIVE_READ_PATHS)) dangerous = tuple( d.get('dangerous_removal_paths', _DEFAULT_DANGEROUS_REMOVAL)) @@ -80,13 +136,28 @@ def _expand_dirs(raw: list[str]) -> tuple[str, ...]: return cls( patterns=patterns, sensitive_paths=sensitive, + sensitive_read_paths=sensitive_read, dangerous_removal_paths=dangerous, read_policy=read_policy, max_command_chars=max_chars, allowed_directories=allowed, read_only_directories=read_only, + allow_temp_dir=bool(d.get('allow_temp_dir', True)), ) + def effective_allowed_directories( + self, workspace_root: str) -> tuple[str, ...]: + """Every directory writes are permitted in, workspace root first.""" + out = [workspace_root] + for directory in self.allowed_directories: + if directory not in out: + out.append(directory) + if self.allow_temp_dir: + for directory in default_temp_directories(): + if directory not in out: + out.append(directory) + return tuple(out) + #: Nothing by default. A blacklist entry can never be overridden — not by the #: mode, not by a whitelist, not by the user answering a prompt — so it is the diff --git a/ms_agent/permission/enforcer.py b/ms_agent/permission/enforcer.py index c8d59880e..f94b149b9 100644 --- a/ms_agent/permission/enforcer.py +++ b/ms_agent/permission/enforcer.py @@ -17,6 +17,9 @@ from .matcher import CONTENT_SEP, PermissionMatcher from .memory import PermissionMemory from .suggestions import generate_suggestions +from ms_agent.utils import get_logger + +logger = get_logger() @dataclass(frozen=True) @@ -24,6 +27,17 @@ class PermissionDecision: action: Literal['allow', 'deny', 'ask'] reason: str updated_args: dict[str, Any] | None = None + #: Whether a standing answer may satisfy this confirmation next time. + #: + #: Safety confirmations default to False — a remembered answer must not + #: stand in for looking at THIS call. But that is not true of all of them + #: equally. "This command runs code I cannot analyse" is a thing a user can + #: reasonably decide once for a project, the way they decide about `git`; + #: "this reads a private key" is not. Marking the first kind rememberable + #: is what keeps the confirmation useful — an ask that reappears no matter + #: how the user answers it does not make anyone safer, it just teaches them + #: to turn confirmations off entirely. + rememberable: bool = False class PermissionEnforcer: @@ -77,6 +91,11 @@ async def _ask_user(self, if 'call_id' in kwargs and not self._handler_accepts('call_id'): kwargs.pop('call_id') if getattr(self._handler, 'supports_concurrent_asks', False): + # A handler that services asks concurrently needs to know which of + # them are safety confirmations, so a remembered answer is never + # applied to one. Older handlers with a fixed signature don't. + if self._handler_accepts('forced'): + return await self._handler.ask(forced=forced, **kwargs) return await self._handler.ask(**kwargs) async with self._ask_lock_for_loop(): if not forced and self._memory.matches(kwargs['tool_name'], @@ -125,9 +144,15 @@ async def check( ) if force_decision and force_decision.action == 'ask': + rememberable = getattr(force_decision, 'rememberable', False) + if rememberable and self._memory.matches(tool_name, tool_args): + return PermissionDecision( + action='allow', + reason='Allowed by remembered permission', + ) suggestions = generate_suggestions(tool_name, tool_args) response = await self._ask_user( - forced=True, + forced=not rememberable, tool_name=tool_name, tool_args=tool_args, context=force_decision.reason or '', @@ -209,6 +234,30 @@ def _remember_pattern(self, response: PermissionResponse, tool_name: str, return s return tool_name + def _release_asks_covered_by_memory(self, pattern: str) -> int: + """Apply a just-remembered answer to the other cards still on screen. + + Only reaches handlers that show several cards at once. There, "always + allow" is a statement about a pattern, so re-asking about a sibling the + pattern covers asks a question the user has answered. It also stops + mattering only in one direction: since a wait has no deadline, an + unanswered sibling now holds the turn open instead of being denied + after a couple of minutes, so leaving them up turns a mis-set + expectation into a stuck conversation. + """ + resolver = getattr(self._handler, 'resolve_matching', None) + if resolver is None: + return 0 + released = resolver( + lambda name, args: self._memory.matches(name, args), + PermissionResponse(action=PermissionAction.ALLOW_ONCE), + ) + if released: + logger.info( + 'permission pattern %r also released %d waiting request(s)', + pattern, released) + return released + def _process_response( self, response: PermissionResponse | None, @@ -230,6 +279,7 @@ def _process_response( if response.action == PermissionAction.ALLOW_SESSION: pattern = self._remember_pattern(response, tool_name, tool_args) self._memory.add_session(pattern) + self._release_asks_covered_by_memory(pattern) return PermissionDecision( action='allow', reason=f'User allowed for session (pattern: {pattern})', @@ -238,6 +288,7 @@ def _process_response( if response.action == PermissionAction.ALLOW_ALWAYS: pattern = self._remember_pattern(response, tool_name, tool_args) self._memory.add(pattern, scope='project', source='user') + self._release_asks_covered_by_memory(pattern) return PermissionDecision( action='allow', reason=f'User allowed always (pattern: {pattern})', diff --git a/ms_agent/permission/handler.py b/ms_agent/permission/handler.py index 9e9802c57..68f3db034 100644 --- a/ms_agent/permission/handler.py +++ b/ms_agent/permission/handler.py @@ -13,7 +13,7 @@ import sys from dataclasses import dataclass from enum import Enum -from typing import Any, Protocol +from typing import Any, Callable, Protocol from uuid import uuid4 @@ -152,6 +152,15 @@ def emit(self, event: dict[str, Any]) -> None: ... +@dataclass +class _PendingAsk: + """One card the user has not answered yet, and what it was about.""" + future: 'asyncio.Future[PermissionResponse]' + tool_name: str + tool_args: dict + forced: bool = False + + class WebPermissionHandler: """Async handler that suspends on a Future until the frontend responds.""" @@ -164,9 +173,17 @@ class WebPermissionHandler: def __init__( self, event_emitter: EventEmitter, - timeout: float = 120.0, + timeout: float | None = None, ) -> None: - self._pending: dict[str, asyncio.Future[PermissionResponse]] = {} + """``timeout=None`` waits indefinitely for an answer. + + That is the right default for a handler whose whole purpose is to ask + a person something: expiring the question answers it on their behalf, + with the one answer they cannot undo. The host sets a bound where one + makes sense — a full-access session, where the human may not be at the + keyboard at all. + """ + self._pending: dict[str, _PendingAsk] = {} self._event_emitter = event_emitter self._timeout = timeout @@ -177,11 +194,19 @@ async def ask( context: str, suggestions: list[str] | None = None, call_id: str = '', + forced: bool = False, ) -> PermissionResponse: request_id = uuid4().hex loop = asyncio.get_running_loop() future: asyncio.Future[PermissionResponse] = loop.create_future() - self._pending[request_id] = future + # What was asked is kept beside the future so an answer to ONE card can + # be applied to the others it covers (see resolve_matching). + self._pending[request_id] = _PendingAsk( + future=future, + tool_name=tool_name, + tool_args=dict(tool_args or {}), + forced=forced, + ) self._event_emitter.emit({ 'type': @@ -202,16 +227,97 @@ async def ask( }) try: + if self._timeout is None: + return await future return await asyncio.wait_for(future, timeout=self._timeout) except asyncio.TimeoutError: + # Said plainly, and said to the MODEL: a timeout is not a person + # declining. Without the distinction the agent reads an ordinary + # refusal, tries a variation, and waits out the whole timeout + # again — one unattended prompt costing several times what the + # limit says it should. return PermissionResponse( action=PermissionAction.DENY, - feedback='Permission request timed out', + feedback=( + f'No response within {self._timeout:.0f}s, so this call ' + 'was not run. This is a TIMEOUT, not a refusal by the ' + 'user — nobody saw the request. Do not re-request the ' + 'same approval; finish what you can without it and say ' + 'plainly what is left waiting on approval.'), ) finally: self._pending.pop(request_id, None) + def awaiting_request_ids(self) -> set: + """Every request still open for an answer. + + A host replaying a reconnected turn needs this to tell a card that is + still live from one that was already decided. + """ + return { + request_id + for request_id, pending in self._pending.items() + if not pending.future.done() + } + + def is_awaiting(self, request_id: str) -> bool: + """Whether this request is still open for an answer. + + Public because a host has to ask before routing a click, and reaching + into ``_pending`` to ask makes the host's code depend on how pending + asks happen to be stored — which is how adding a field to that record + turned every approval click into a 500. + """ + pending = self._pending.get(request_id) + return pending is not None and not pending.future.done() + def resolve(self, request_id: str, response: PermissionResponse) -> None: - future = self._pending.get(request_id) - if future and not future.done(): - future.set_result(response) + pending = self._pending.get(request_id) + if pending and not pending.future.done(): + pending.future.set_result(response) + + def resolve_matching( + self, + covers: Callable[[str, dict[str, Any]], bool], + response: PermissionResponse, + ) -> int: + """Answer the still-open asks that a decision just made unnecessary. + + A round can put several cards up at once, and answering one of them + with "always allow" is a statement about a PATTERN, not about that one + call. Leaving its siblings up asks the user the question they just + answered — and since a wait has no deadline, an unanswered sibling + holds the turn open indefinitely rather than being quietly denied. + + Safety confirmations are skipped: those exist precisely so a remembered + answer cannot stand in for looking at this one. + """ + resolved = 0 + for request_id, pending in list(self._pending.items()): + if pending.forced or pending.future.done(): + continue + if not covers(pending.tool_name, pending.tool_args): + continue + pending.future.set_result(response) + self._pending.pop(request_id, None) + resolved += 1 + return resolved + + def cancel_pending(self, feedback: str = 'Session closed') -> int: + """Answer every outstanding ask so nothing is left waiting on a person + who has gone. Returns how many were resolved. + + Needed once waits can be unbounded: a suspended ask holds its turn, and + a held turn is exempt from idle reclamation, so an abandoned prompt + would otherwise pin its session for the life of the process. + """ + resolved = 0 + for request_id, pending in list(self._pending.items()): + if pending.future.done(): + continue + pending.future.set_result( + PermissionResponse( + action=PermissionAction.DENY, feedback=feedback)) + resolved += 1 + self._pending.pop(request_id, None) + return resolved diff --git a/ms_agent/permission/path_extractors.py b/ms_agent/permission/path_extractors.py index cb721bf86..fc7df87cf 100644 --- a/ms_agent/permission/path_extractors.py +++ b/ms_agent/permission/path_extractors.py @@ -414,6 +414,78 @@ def _make_filter_entry( ) +#: Commands that run code supplied as an argument, a script, or on stdin. +#: +#: Path extraction is a claim about what a command touches, and for these that +#: claim cannot be made: the argument to ``python3 -c`` is a program, and a +#: program can open, write or delete anything. So they are not extracted from — +#: they are surfaced, letting the mode decide (see ``ask_resolver``). Without +#: this they fell to the "unregistered command" branch and were allowed +#: outright, which is how ``python3 -c "import os; os.remove(...)"`` walked +#: past a policy that stops a plain ``rm``. +INTERPRETER_COMMANDS: frozenset[str] = frozenset({ + 'python', + 'python2', + 'python3', + 'pypy', + 'pypy3', + 'uv', + 'uvx', + 'pipx', + 'node', + 'nodejs', + 'deno', + 'bun', + 'npx', + 'ruby', + 'perl', + 'php', + 'lua', + 'Rscript', + 'sh', + 'bash', + 'zsh', + 'ksh', + 'dash', + 'fish', + 'osascript', + 'eval', + 'exec', +}) + +#: Interpreter flags that take code inline rather than a file path. +_INLINE_CODE_FLAGS = frozenset({'-c', '-e', '--eval', '--exec', '-E'}) + + +def interpreter_runs_inline_code(args: list[str]) -> bool: + """Whether an interpreter invocation carries its program in the argv. + + ``python3 script.py`` names a file the workspace rules can judge; + ``python3 -c '…'`` and a bare ``python3`` reading stdin do not. + """ + if not args: + return True # bare REPL / stdin + for arg in args: + if arg == '--': + break + if arg in _INLINE_CODE_FLAGS or arg.startswith('--eval='): + return True + if arg == '-': + return True # explicit stdin, e.g. `python3 - < list[str]: + """The script path an interpreter was pointed at, if it was pointed at one.""" + for arg in args: + if arg == '--': + continue + if arg.startswith('-'): + continue + return [arg] + return [] + + def build_extractor_registry() -> dict[str, ExtractorEntry]: """Build the full 36-command extractor registry.""" registry: dict[str, ExtractorEntry] = {} diff --git a/ms_agent/permission/path_validator.py b/ms_agent/permission/path_validator.py index 74f465f8c..fdb2f85c8 100644 --- a/ms_agent/permission/path_validator.py +++ b/ms_agent/permission/path_validator.py @@ -89,6 +89,27 @@ def _is_under_allowed(resolved: Path, allowed_dirs: Sequence[str]) -> bool: return False +def matches_sensitive(resolved: str, patterns: Sequence[str]) -> bool: + """Whether a resolved path is one of the configured sensitive locations. + + Compared against both the literal pattern and its ``~``-expanded form, and + a directory pattern (``~/.ssh/*``) also matches the directory itself, so + listing it is judged the same as reading a file inside it. + """ + import fnmatch + + for raw in patterns or (): + pattern = os.path.expanduser(str(raw)) + if fnmatch.fnmatch(resolved, pattern): + return True + # `~/.ssh/*` should also cover `~/.ssh` and anything nested deeper. + if pattern.endswith('/*'): + parent = pattern[:-2] + if resolved == parent or resolved.startswith(parent + os.sep): + return True + return False + + def validate_path( path: str, cwd: str, @@ -97,6 +118,7 @@ def validate_path( *, read_only_dirs: Sequence[str] = (), home_dir: str | None = None, + sensitive_paths: Sequence[str] = (), ) -> PathValidationResult: """Validate a single filesystem path for a given operation type. @@ -128,10 +150,7 @@ def validate_path( ) if _has_glob(path): - # Heuristic: if path contains parentheses or is very long, it's likely - # code content (e.g., from heredoc), not a real file path. Skip glob check. - looks_like_code = '(' in path or ')' in path or len(path) > 200 - if op_type in ('write', 'create') and not looks_like_code: + if op_type in ('write', 'create'): return PathValidationResult( allowed=False, resolved_path=path, @@ -148,6 +167,19 @@ def validate_path( resolved_str = str(resolved) + # Credential stores are gated wherever they sit. The list existed already + # but only the write path consulted it, so `cat ~/.ssh/id_rsa` was judged + # purely on "is it in the workspace" — and under the default loose read + # policy that resolved to allow, putting the key in the transcript. + if op_type == 'read' and matches_sensitive(resolved_str, sensitive_paths): + return PathValidationResult( + allowed=False, + resolved_path=resolved_str, + action='ask', + reason=f'Read of sensitive path: {resolved_str}', + category='sensitive_read', + ) + if not _is_under_allowed(resolved, allowed_dirs): if op_type == 'read': if _is_under_allowed(resolved, read_only_dirs): diff --git a/ms_agent/permission/safety.py b/ms_agent/permission/safety.py index 0380f679e..a66a3656f 100644 --- a/ms_agent/permission/safety.py +++ b/ms_agent/permission/safety.py @@ -33,6 +33,7 @@ def __init__( self._allowed_dirs = list(allowed_dirs) self._read_only_dirs = list(read_only_dirs) self._sensitive_paths = list(config.sensitive_paths) + self._sensitive_read_paths = list(config.sensitive_read_paths) self._workspace_root = workspace_root path_safety_cfg = PathSafetyConfig( @@ -41,6 +42,7 @@ def __init__( read_only_directories=tuple(self._read_only_dirs), workspace_root=workspace_root, dangerous_removal_paths=tuple(config.dangerous_removal_paths), + sensitive_read_paths=tuple(config.sensitive_read_paths), ) self._shell_validator = ShellPathValidator( allowed_dirs=self._allowed_dirs, @@ -105,7 +107,8 @@ def _check_file_path(self, path: str, cwd, self._allowed_dirs, op_type, - read_only_dirs=self._read_only_dirs) + read_only_dirs=self._read_only_dirs, + sensitive_paths=self._sensitive_read_paths) if not result.allowed: return SafetyDecision( action=result.action, diff --git a/ms_agent/permission/shell_validator.py b/ms_agent/permission/shell_validator.py index 506a8d073..2912132c1 100644 --- a/ms_agent/permission/shell_validator.py +++ b/ms_agent/permission/shell_validator.py @@ -17,8 +17,11 @@ from pathlib import Path from typing import Literal, Sequence -from .path_extractors import (ExtractorEntry, build_extractor_registry, - extract_find_exec_commands, find_uses_delete) +from .path_extractors import (INTERPRETER_COMMANDS, ExtractorEntry, + build_extractor_registry, + extract_find_exec_commands, + extract_interpreter_script, find_uses_delete, + interpreter_runs_inline_code) from .path_validator import (PathValidationResult, is_dangerous_removal_path, validate_path) from .sed_validator import check_sed_expression_safety, is_sed_read_only @@ -26,12 +29,217 @@ _PROCESS_INPUT_SUB = re.compile(r'<\s*\(') _PROCESS_OUTPUT_SUB = re.compile(r'>\s*\(') -_REDIRECT_PATTERN = re.compile(r'(?:&>>|&>|>>|>\||>)' - r'\s*' - r'(\S+)') _FD_REDIRECT = re.compile(r'^\d*>&\d+$') _MAX_SUBSTITUTION_DEPTH = 16 +#: Character devices every shell uses as a sink or a source. They are not +#: files the workspace policy has anything to say about. +_REDIRECT_DEVICE_ALLOWLIST = frozenset({ + '/dev/null', + '/dev/stdout', + '/dev/stderr', + '/dev/stdin', + '/dev/tty', + '/dev/zero', +}) + +#: Ends a word in shell source. Used to know where a redirect target stops — +#: notably at the `)` closing a subshell, which is not part of the filename. +_WORD_TERMINATORS = frozenset(' \t\n\r;|&()<>') + + +def _extract_redirect_targets(command: str) -> list[str]: + """Find the target of every output redirection in one command. + + Walks the source tracking quote state instead of pattern-matching it. A + regex reading ``\\S+`` after the operator cannot see either boundary that + matters: it takes ``>`` inside a quoted argument for a redirection + (``git commit -m "a > b"``), and it swallows whatever punctuation abuts the + filename, so ``(… 2>/dev/null)`` yields the target ``/dev/null)`` — which + matches no allowlist entry and gets refused as a write outside the + workspace. + """ + targets: list[str] = [] + i = 0 + n = len(command) + in_single = False + in_double = False + + while i < n: + c = command[i] + + if c == '\\' and not in_single and i + 1 < n: + i += 2 + continue + if c == "'" and not in_double: + in_single = not in_single + i += 1 + continue + if c == '"' and not in_single: + in_double = not in_double + i += 1 + continue + if in_single or in_double: + i += 1 + continue + + if c != '>': + i += 1 + continue + + # Step back over an fd prefix (``2>``) and ``&`` of ``&>``. + start = i + i += 1 + if i < n and command[i] == '>': # >> + i += 1 + elif i < n and command[i] == '|': # >| + i += 1 + elif i < n and command[i] == '&': + # ``>&2`` / ``2>&1`` duplicates a descriptor; no file involved. + j = i + 1 + while j < n and command[j].isdigit(): + j += 1 + if j > i + 1 or (j < n and command[j] == '-'): + i = j + continue + + while i < n and command[i] in ' \t': + i += 1 + + word_start = i + word: list[str] = [] + w_single = False + w_double = False + while i < n: + ch = command[i] + if ch == '\\' and not w_single and i + 1 < n: + word.append(command[i + 1]) + i += 2 + continue + if ch == "'" and not w_double: + w_single = not w_single + i += 1 + continue + if ch == '"' and not w_single: + w_double = not w_double + i += 1 + continue + if not w_single and not w_double and ch in _WORD_TERMINATORS: + break + word.append(ch) + i += 1 + + if i == word_start and not word: + # A bare ``>`` with nothing after it — malformed, but not ours to + # interpret; leave it to the shell. + del start + continue + targets.append(''.join(word)) + + return targets + + +_REDIRECT_TOKEN = re.compile(r'^\d*(?:&?>>?\|?|<&|<&)$') + + +def _strip_redirections(tokens: list[str]) -> list[str]: + """Remove redirection operators and their targets from a token list. + + ``shlex.split`` has no idea ``>`` means anything, so ``ls >> /dev/null`` + arrives as three plain words and the extractor reads ``/dev/null`` as a + file ``ls`` was asked to list. Redirect targets are checked separately, by + :func:`_extract_redirect_targets`, against the rules for the write they + actually are. + """ + out: list[str] = [] + i = 0 + while i < len(tokens): + token = tokens[i] + if _FD_REDIRECT.match(token): + i += 1 + continue + if _REDIRECT_TOKEN.match(token): + # Operator and, unless it is a descriptor duplication, its target. + i += 1 + if i < len(tokens) and not _REDIRECT_TOKEN.match(tokens[i]): + i += 1 + continue + # Glued forms such as ``2>/dev/null`` or ``>out.txt``. + glued = re.match(r'^\d*(?:&?>>?\|?|< str: + """Peel subshell/command-group punctuation off one sub-command. + + Splitting a compound command on its operators leaves the delimiters of any + group attached to the neighbouring word — ``(cd sub && ls)`` yields + ``(cd sub`` and ``ls)``. Only unquoted leading/trailing delimiters are + removed, so a real argument like ``echo "(hi)"`` is untouched. + """ + out = sub_cmd.strip() + while out and out[0] in '({': + out = out[1:].lstrip() + while out and out[-1] in ')}': + out = out[:-1].rstrip() + if out.endswith(';'): + out = out[:-1].rstrip() + return out + + +_HEREDOC_START = re.compile(r'<<-?\s*([\'"]?)([A-Za-z_][A-Za-z0-9_]*)\1') + + +def _strip_heredoc_bodies(command: str) -> str: + """Drop heredoc payloads, keeping the command lines that introduce them. + + A heredoc body is DATA. Left in place it is split on newlines like any + compound command and each line analysed as if it were one, so a line of + Python such as ``result = [x * 2 for x in data if x > 1]`` reads as a + redirection into a file named ``1]`` — a glob in a create position, and a + refusal for a script that touches nothing. + """ + if '<<' not in command: + return command + + lines = command.splitlines() + kept: list[str] = [] + pending: list[tuple] = [] + opener_at: int = -1 # where the still-open heredoc began, on kept[-1] + + for line in lines: + if pending: + delimiter, strip_tabs = pending[0] + candidate = line.lstrip('\t') if strip_tabs else line + if candidate.strip() == delimiter: + pending.pop(0) + if not pending: + opener_at = -1 + continue + + kept.append(line) + for match in _HEREDOC_START.finditer(line): + if not pending: + opener_at = match.start() + pending.append((match.group(2), match.group(0).startswith('<<-'))) + + if pending and kept and opener_at >= 0: + # Opened and never closed. The shell will reject this outright, so + # there is nothing here to judge — and judging it anyway means reading + # the payload as filenames and refusing for a reason that is not true. + # Seen when a multi-line command reaches us flattened onto one line + # (a paste that lost its newlines): the refusal claimed a "glob in a + # create operation" for a script that creates nothing, hiding the + # actual problem, which is a syntax error. + kept[-1] = kept[-1][:opener_at].rstrip() + + return '\n'.join(kept) + @dataclass(frozen=True) class SafetyDecision: @@ -47,6 +255,7 @@ class PathSafetyConfig: read_only_directories: tuple[str, ...] = () workspace_root: str | None = None dangerous_removal_paths: tuple[str, ...] = () + sensitive_read_paths: tuple[str, ...] = () class ShellPathValidator: @@ -61,6 +270,7 @@ def __init__( self._config = safety_config or PathSafetyConfig() self._read_only_dirs = list(self._config.read_only_directories) self._workspace_root = self._config.workspace_root or os.getcwd() + self._sensitive_read_paths = list(self._config.sensitive_read_paths) self._extractors = build_extractor_registry() def check(self, command: str, *, _depth: int = 0) -> SafetyDecision: @@ -81,6 +291,11 @@ def check(self, command: str, *, _depth: int = 0) -> SafetyDecision: f'Command exceeds max length ({self._config.max_command_chars})', ) + # Heredoc payloads are data, not commands. Analysing them as commands + # is how an inline Python script gets refused for a "glob in a create + # operation" it never contained. + command = _strip_heredoc_bodies(command) + # 1. Process substitution if _PROCESS_OUTPUT_SUB.search(command): return SafetyDecision( @@ -115,6 +330,15 @@ def check(self, command: str, *, _depth: int = 0) -> SafetyDecision: if stripped.startswith('#'): continue # Skip pure comment + # Unwrap subshell / group punctuation so the first token is the + # command. Without this `(cd sub && …)` presents `(cd` as the + # command name, which matches no extractor and silently skips the + # cd tracking the rest of the loop depends on. + stripped = _unwrap_group(stripped) + if not stripped or stripped.startswith('#'): + continue + sub_cmd = stripped + try: tokens = shlex.split(sub_cmd) except ValueError: @@ -131,7 +355,8 @@ def check(self, command: str, *, _depth: int = 0) -> SafetyDecision: if redirect_result.action != 'allow': return redirect_result - # 4. Strip safe wrappers + # 4. Strip safe wrappers and redirection syntax + tokens = _strip_redirections(tokens) tokens = strip_safe_wrappers(tokens) if not tokens: continue @@ -189,6 +414,9 @@ def _check_command( _depth: int = 0, cwd: str | None = None, ) -> SafetyDecision: + if base_cmd in INTERPRETER_COMMANDS: + return self._check_interpreter(base_cmd, args, cwd=cwd) + entry = self._extractors.get(base_cmd) if entry is None: return SafetyDecision( @@ -215,6 +443,39 @@ def _check_command( return self._validate_paths(paths, entry.op_type, base_cmd, cwd=cwd) + def _check_interpreter(self, + base_cmd: str, + args: list[str], + *, + cwd: str | None = None) -> SafetyDecision: + """Judge a command that runs code rather than touching named files. + + Pointing an interpreter at a script inside the workspace is the same + kind of act as reading that script, and is treated that way — otherwise + every ``python3 build.py`` in a normal project would need confirming, + and a prompt nobody can act on is a prompt everybody clicks through. + Code arriving inline or on stdin has no path to judge, so it is raised + for the mode to resolve. + """ + if not interpreter_runs_inline_code(args): + script = extract_interpreter_script(args) + if script: + result = self._validate_paths( + script, 'read', base_cmd, cwd=cwd) + if result.action != 'allow': + return result + return SafetyDecision( + action='allow', + reason=f'{base_cmd}: runs a script inside allowed dirs', + ) + + return SafetyDecision( + action='ask', + reason=(f'`{base_cmd}` runs code given inline or on stdin; what it ' + 'touches cannot be determined from the command'), + category='interpreter_exec', + ) + def _check_sed(self, args: list[str], entry: ExtractorEntry, @@ -315,7 +576,8 @@ def _validate_paths( effective_cwd, self._allowed_dirs, op_type, - read_only_dirs=self._read_only_dirs) + read_only_dirs=self._read_only_dirs, + sensitive_paths=self._sensitive_read_paths) if not result.allowed: return SafetyDecision( action=result.action, @@ -326,11 +588,12 @@ def _validate_paths( action='allow', reason=f'{cmd_name}: all paths validated') def _check_redirects(self, sub_cmd: str) -> SafetyDecision: - for match in _REDIRECT_PATTERN.finditer(sub_cmd): - target = match.group(1) + for target in _extract_redirect_targets(sub_cmd): + if not target: + continue if _FD_REDIRECT.match(target): continue - if target == '/dev/null': + if target in _REDIRECT_DEVICE_ALLOWLIST: continue if '$' in target or '%' in target: return SafetyDecision( @@ -345,6 +608,7 @@ def _check_redirects(self, sub_cmd: str) -> SafetyDecision: self._allowed_dirs, 'create', read_only_dirs=self._read_only_dirs, + sensitive_paths=self._sensitive_read_paths, ) if not result.allowed: return SafetyDecision( diff --git a/ms_agent/prompting/builtin.py b/ms_agent/prompting/builtin.py index a3ed88ccc..e262491f5 100644 --- a/ms_agent/prompting/builtin.py +++ b/ms_agent/prompting/builtin.py @@ -176,6 +176,45 @@ files actually live in {home}; project AGENTS.md files live in the project \ directory.""" +#: Injected when the framework keeps its own records INSIDE the working +#: directory, which is the layout of a managed project. +#: +#: Without it the agent has no way to tell its own bookkeeping apart from the +#: user's material, and the confusion is not hypothetical: searching the +#: workspace for a phrase finds that phrase in the transcript of the very +#: request being served, because the prompt was written there moments earlier. +#: Every hit is real, every hit is worthless, and the model has no reason to +#: suspect it. Naming the directories, and saying what changing them does, is +#: cheaper and less brittle than hiding them — hidden, they would also be +#: unavailable when the user genuinely asks about history or configuration. +WORKSPACE_INTERNALS_HINT = """\ +## Framework files in your working directory + +Two things under your working directory are maintained by the framework rather \ +than written by the user: + +- `sessions/` — a full transcript of every conversation in this project, \ +including the user's messages verbatim.{session_line} +- `.ms_agent/` — this project's state: `memory/` (what is remembered across \ +conversations), `snapshots/` (a git repository of previous workspace \ +versions), `permission_memory.json` (approvals the user chose to keep), \ +`web_search/` (cached search results), `mcp.json` and `project.json`. + +Settings that apply to every project live separately, in {home} — the location \ +is configurable, so a machine may have several and this conversation is using \ +that one. + +When you search the workspace, matches inside those two directories are the \ +framework's record of this and earlier conversations, not the user's content. \ +Anything you were just asked is already written to `sessions/`, so searching \ +for a phrase from the request will match your own transcript. Exclude them \ +unless the user is asking about history or configuration, and never cite such \ +a match as if it were something you found in their material. + +You may read these files, and edit them when asked. Be aware that editing \ +`memory/` or `permission_memory.json` changes how later conversations behave, \ +and that `snapshots/` is what makes reverting possible.""" + #: Filename -> template registry used by workspace_files.ensure logic. HOME_FILE_TEMPLATES = { 'SOUL.md': SOUL_TEMPLATE, diff --git a/ms_agent/tools/arg_coercion.py b/ms_agent/tools/arg_coercion.py new file mode 100644 index 000000000..674350ef0 --- /dev/null +++ b/ms_agent/tools/arg_coercion.py @@ -0,0 +1,162 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Reconcile the JSON a model emitted with the JSON a tool's schema asks for. + +Models routinely quote scalars — ``{"a": "19.5"}`` where the schema says +``number`` — and a tool validating strictly (Pydantic ``StrictFloat``, say) +rejects that. The model cannot see why: it gets a validator's internal +complaint, tries again, and produces the same quoting, so the round repeats +until something gives up. + +Only rewrites that cannot change meaning are made. ``"19.5"`` is 19.5 in every +reading; ``"nineteen"`` is left exactly as it is, for the tool to refuse on its +own terms. Where the schema says ``string``, nothing is touched at all — the +quotes are the answer there. +""" +from __future__ import annotations + +import json +from typing import Any, Dict, Optional + +__all__ = ['coerce_arguments'] + +_TRUE = frozenset({'true', 'True', 'TRUE', 'yes', 'on', '1'}) +_FALSE = frozenset({'false', 'False', 'FALSE', 'no', 'off', '0'}) + + +def _schema_types(schema: Any) -> set: + """The JSON Schema types a value may take, flattened across combinators.""" + if not isinstance(schema, dict): + return set() + out: set = set() + declared = schema.get('type') + if isinstance(declared, str): + out.add(declared) + elif isinstance(declared, list): + out.update(t for t in declared if isinstance(t, str)) + for key in ('anyOf', 'oneOf', 'allOf'): + for branch in schema.get(key) or (): + out |= _schema_types(branch) + return out + + +def _to_number(text: str, types: set) -> Any: + stripped = text.strip() + if not stripped: + return None + if 'integer' in types: + try: + return int(stripped, 10) + except ValueError: + pass + if 'number' in types: + try: + value = float(stripped) + except ValueError: + return None + # An integer-valued float where the schema also allows int keeps the + # narrower type, which is what a StrictInt branch is waiting for. + if 'integer' in types and value.is_integer(): + return int(value) + return value + return None + + +def _coerce_value(value: Any, schema: Any) -> Any: + types = _schema_types(schema) + if not types: + return value + + if isinstance(value, str): + # A schema that accepts a string means the string IS the answer. + if 'string' in types: + return value + if {'number', 'integer'} & types: + number = _to_number(value, types) + if number is not None: + return number + if 'boolean' in types: + token = value.strip() + if token in _TRUE: + return True + if token in _FALSE: + return False + if 'null' in types and value.strip() in ('null', 'None', ''): + return None + if {'object', 'array'} & types: + # Some models serialize a whole nested argument as JSON text. + try: + parsed = json.loads(value) + except (ValueError, TypeError): + return value + if isinstance(parsed, dict) and 'object' in types: + return _coerce_value(parsed, schema) + if isinstance(parsed, list) and 'array' in types: + return _coerce_value(parsed, schema) + return value + + if isinstance(value, dict) and 'object' in types: + properties = schema.get('properties') if isinstance(schema, + dict) else None + if not isinstance(properties, dict): + for key in ('anyOf', 'oneOf', 'allOf'): + for branch in (schema.get(key) or ()) if isinstance( + schema, dict) else (): + if isinstance(branch, dict) and isinstance( + branch.get('properties'), dict): + properties = branch['properties'] + break + if properties: + break + if isinstance(properties, dict): + return { + k: (_coerce_value(v, properties[k]) if k in properties else v) + for k, v in value.items() + } + return value + + if isinstance(value, list) and 'array' in types: + items = schema.get('items') if isinstance(schema, dict) else None + if isinstance(items, dict): + return [_coerce_value(item, items) for item in value] + return value + + # An int where the schema wants a float is already valid JSON-schema-wise, + # but a StrictFloat branch disagrees; widen only when int is not accepted. + if isinstance(value, bool): + return value + if isinstance(value, int) and 'number' in types and 'integer' not in types: + return float(value) + return value + + +def coerce_arguments( + arguments: Dict[str, Any], + schema: Optional[Dict[str, Any]], +) -> Dict[str, Any]: + """Return ``arguments`` with values reshaped to fit ``schema``. + + Returns the original object when there is nothing to do, so callers can + tell "unchanged" by identity. A malformed schema is not an error — it just + means no rewriting is possible. + """ + if not isinstance(arguments, dict) or not arguments: + return arguments + if not isinstance(schema, dict): + return arguments + properties = schema.get('properties') + if not isinstance(properties, dict): + return arguments + + out: Dict[str, Any] = {} + changed = False + for key, value in arguments.items(): + if key in properties: + new_value = _coerce_value(value, properties[key]) + if new_value is not value and new_value != value: + changed = True + elif type(new_value) is not type(value): + changed = True + out[key] = new_value + else: + out[key] = value + return out if changed else arguments diff --git a/ms_agent/tools/code/local_code_executor.py b/ms_agent/tools/code/local_code_executor.py index e46a6911f..576d6b9c6 100644 --- a/ms_agent/tools/code/local_code_executor.py +++ b/ms_agent/tools/code/local_code_executor.py @@ -4,7 +4,6 @@ import io import json import os -import shlex import shutil import time from contextlib import redirect_stderr, redirect_stdout @@ -20,6 +19,81 @@ logger = get_logger() +#: Ambient variables a POSIX command may read without them carrying anything +#: the agent should not have. The environment is built as an allow-list rather +#: than inherited, so anything missing here is simply absent for every command +#: the agent runs — and an absent one is not a neutral default: +#: +#: * ``USER``/``LOGNAME`` unset makes ``git commit`` and anything deriving an +#: identity from the environment fail or record the wrong author. +#: * ``TMPDIR`` unset sends every tool that wants scratch space to ``/tmp``, +#: which is outside the workspace and therefore refused by the path policy — +#: leaving the agent with nowhere to write a temporary file at all. +#: * ``SSL_CERT_FILE``/``REQUESTS_CA_BUNDLE`` unset breaks TLS for interpreters +#: with no system trust store, so ``pip install`` fails to verify PyPI. +#: * The proxy variables are how a developer behind one reaches the network; +#: dropping them turns every fetch into a timeout. +#: +#: Credentials are still NOT forwarded: no ``*_API_KEY``, ``*_TOKEN``, +#: ``AWS_*``, ``GITHUB_*`` or similar appears here, and the tool's own +#: ``shell_env`` config remains the way to pass one deliberately. +_POSIX_ENV_PASSTHROUGH = ( + 'PATH', + 'HOME', + 'USER', + 'LOGNAME', + 'SHELL', + 'TMPDIR', + 'LANG', + 'LC_ALL', + 'LC_CTYPE', + 'TERM', + 'TZ', + 'SSL_CERT_FILE', + 'SSL_CERT_DIR', + 'REQUESTS_CA_BUNDLE', + 'CURL_CA_BUNDLE', + 'NODE_EXTRA_CA_CERTS', + 'HTTP_PROXY', + 'HTTPS_PROXY', + 'ALL_PROXY', + 'NO_PROXY', + 'http_proxy', + 'https_proxy', + 'all_proxy', + 'no_proxy', +) + + +#: Set on every command the agent runs, unless the environment already says +#: otherwise. Nothing here changes what a command DOES — it changes what the +#: command assumes about who is reading. +#: +#: A pager is the sharp edge: ``git log`` or ``systemctl status`` hands its +#: output to ``less``, which waits for a keypress that will never come, and the +#: tool call sits there until it times out. Progress bars and colour codes are +#: milder, but they fill the result with control characters that cost context +#: and read as noise. ``AI_AGENT`` is a cross-vendor convention some tools +#: (``gh``) report in their User-Agent. +_AGENT_FRIENDLY_ENV = { + 'PAGER': 'cat', + 'GIT_PAGER': 'cat', + 'MANPAGER': 'cat', + # -F quit if it all fits on one screen, -R keep colour, -X no alternate + # screen. Needed because PAGER=cat only covers programs that CONSULT it; + # anything invoking `less` directly would still sit waiting for a key. + 'LESS': '-FRX', + 'GIT_TERMINAL_PROMPT': '0', + 'GIT_MERGE_AUTOEDIT': 'no', + 'PIP_DISABLE_PIP_VERSION_CHECK': '1', + 'PIP_PROGRESS_BAR': 'off', + 'PYTHONUNBUFFERED': '1', + 'TQDM_DISABLE': '1', + 'DEBIAN_FRONTEND': 'noninteractive', + 'NO_COLOR': '1', + 'AI_AGENT': 'ms_agent', +} + def _is_relative_to(path: Path, base: Path) -> bool: try: @@ -328,10 +402,19 @@ def _build_env(self, field: str, inherit: bool = False) -> Dict[str, str]: else: env: Dict[str, str] = { 'INHERITED_FROM_LOCAL': 'False', - 'PATH': os.environ.get('PATH', ''), - 'HOME': os.environ.get('HOME', ''), - 'LANG': os.environ.get('LANG', ''), } + for key in _POSIX_ENV_PASSTHROUGH: + value = os.environ.get(key) + if value is not None: + env[key] = value + env.setdefault('PATH', '') + # Set outright, NOT deferring to whatever the parent had. These + # exist to undo settings made for a human at a terminal, so + # inheriting them defeats the entire point: a developer's + # `PAGER=less` is exactly the value that leaves `git log` waiting + # for a keypress nobody will press. A deliberate choice still wins + # — `tools.code_executor.shell_env` is applied after this. + env.update(_AGENT_FRIENDLY_ENV) if os.name == 'nt': # ``create_subprocess_shell`` uses the native Windows command # processor. Keep the non-secret OS/user variables that cmd @@ -459,12 +542,30 @@ async def _get_tools_inner(self) -> Dict[str, Any]: Tool( tool_name='shell_executor', server_name='code_executor', - description= - ('Execute shell commands locally under the workspace output directory (cwd). ' - 'Subject to policy (read_only vs workspace_write, network toggle). ' - 'Large stdout/stderr may be spilled to .ms_agent_artifacts. ' - 'Use run_in_background=true to return immediately with task_id; poll via task notifications.' - ), + description=( + 'Run a shell command.\n' + '\n' + 'Each call starts a NEW shell in the workspace root, ' + f'which is {self._ws.root}. Nothing carries over ' + 'between calls: not the working directory, not ' + 'environment variables, and not an activated ' + 'virtualenv or conda environment.\n' + '\n' + 'So do each of those in the SAME call as the work ' + 'that needs it, chained with && — for example ' + '`cd sub && ls`, or ' + '`. .venv/bin/activate && pip install requests && ' + 'python -c "import requests"`. Prefer absolute paths ' + 'or a leading cd over assuming a directory from an ' + 'earlier call.\n' + '\n' + 'Paths are restricted to the workspace (plus the OS ' + 'temp directory); reading credential files and running ' + 'code inline (python -c, heredocs) may require ' + 'approval. Large output is spilled to ' + '.ms_agent_artifacts and the result says where. Use ' + 'run_in_background=true for a long command: it returns ' + 'a task_id immediately.'), parameters={ 'type': 'object', 'properties': { @@ -600,21 +701,6 @@ async def call_tool(self, server_name: str, *, tool_name: str, ensure_ascii=False, indent=2) - def _prepare_shell_command(self, command: str) -> str: - """Use the native Windows shell; wrap composite POSIX input when needed.""" - if os.name == 'nt': - # asyncio.create_subprocess_shell delegates to cmd.exe on Windows; - # wrapping native syntax in a usually unavailable POSIX shell - # makes otherwise valid compound commands fail. - return command - shell_meta = ('&&', '||', '|', ';', '>', '<', '`', '$(', 'cd ', - 'export ') - already_wrapped = command.lstrip().startswith( - ('sh ', 'bash ', '/bin/sh ', '/bin/bash ')) - if not already_wrapped and any(meta in command for meta in shell_meta): - return f'sh -lc {shlex.quote(command)}' - return command - async def notebook_executor(self, code: str, description: str = '', @@ -718,7 +804,16 @@ async def shell_executor(self, exec_timeout = timeout or self._shell_timeout call_id = call_id or f'shell-{os.urandom(4).hex()}' - shell_cmd = self._prepare_shell_command(command) + # Handed to the shell verbatim. ``create_subprocess_shell`` already runs + # it under ``/bin/sh -c`` (``cmd.exe`` on Windows), which understands + # every compound form on its own, so there is nothing to pre-wrap. + # Wrapping used to be conditional on the command CONTAINING a + # metacharacter, and the wrapper was a LOGIN shell: adding a `;` to a + # command re-ran the profile, which on macOS reorders PATH via + # path_helper, so `python3 -V` and `python3 -V ; true` resolved to + # different interpreters. Installing with one form and importing with + # the other is then a ModuleNotFoundError with no visible cause. + shell_cmd = command if run_in_background: if self._task_manager is None: diff --git a/ms_agent/tools/filesystem_tool.py b/ms_agent/tools/filesystem_tool.py index b8497efc6..d5031cc0c 100644 --- a/ms_agent/tools/filesystem_tool.py +++ b/ms_agent/tools/filesystem_tool.py @@ -928,17 +928,23 @@ async def read_file(self, f'Use offset and limit to read specific portions.') continue - # Dedup: return stub if file unchanged since last read + # Whether this path was read before, and looks the same. Used + # only to ANNOTATE the content — never to withhold it. + # + # Withholding assumed the earlier read was still in the model's + # context, which the tool cannot know: after a truncation or a + # compaction it is not, and the model is then unable to get the + # file back at all. Observed consequence — it fell back to the + # shell, hit an approval it could not clear, and finally edited + # the file to reconstruct what it could no longer see. The + # mtime comparison is also not proof of sameness: two writes + # within one filesystem timestamp tick are indistinguishable, + # so this could withhold content that HAD changed. mtime = os.path.getmtime(target_path_real) cached = self._read_cache.get(target_path_real) - if (cached and cached['mtime'] == mtime - and cached['offset'] == offset - and cached['limit'] == limit): - results[path] = { - 'type': 'file_unchanged', - 'message': 'File has not changed since last read.', - } - continue + unchanged = bool(cached and cached['mtime'] == mtime + and cached['offset'] == offset + and cached['limit'] == limit) with open(target_path_real, 'rb') as f: raw_bytes = f.read() @@ -970,7 +976,14 @@ async def read_file(self, else: selected = lines - results[path] = ''.join(selected) + body = ''.join(selected) + if unchanged: + # Says the same thing the old stub said, without taking the + # content away to say it. The model can now decide for + # itself whether it still had this in context. + body = ('[note: unchanged since your last read of this ' + 'file]\n' + body) + results[path] = body # Update dedup cache self._read_cache[target_path_real] = { diff --git a/ms_agent/tools/mcp_client.py b/ms_agent/tools/mcp_client.py index 77dd76c79..d29ed403a 100644 --- a/ms_agent/tools/mcp_client.py +++ b/ms_agent/tools/mcp_client.py @@ -1,9 +1,11 @@ from __future__ import annotations # Copyright (c) ModelScope Contributors. All rights reserved. +import asyncio import copy import os -from contextlib import AsyncExitStack +import re +from contextlib import AsyncExitStack, suppress from datetime import timedelta from mcp import ClientSession, ListToolsResult, StdioServerParameters from mcp.client.sse import sse_client @@ -27,12 +29,71 @@ DEFAULT_HTTP_TIMEOUT = 5 DEFAULT_SSE_READ_TIMEOUT = 60 * 5 -CONNECTION_TIMEOUT = os.getenv('CONNECTION_TIMEOUT', 120) + + +def _int_env(name: str, default: int) -> int: + """``os.getenv`` hands back a STRING once the variable is set, and these + values end up in ``timedelta(seconds=...)``, which rejects one.""" + try: + return int(os.getenv(name, default)) + except (TypeError, ValueError): + return default + + +CONNECTION_TIMEOUT = _int_env('CONNECTION_TIMEOUT', 120) + +#: How many servers may be started at once. Bounded because each stdio server +#: is a child process and a large config would otherwise fork all of them +#: simultaneously. +MAX_PARALLEL_CONNECTS = _int_env('MCP_MAX_PARALLEL_CONNECTS', 8) + +#: How long to wait for a server's owner task to close its transport before +#: cancelling it outright. +SERVER_STOP_TIMEOUT = _int_env('MCP_STOP_TIMEOUT', 15) + +#: Wall-clock ceiling on bringing ONE server up (spawn/dial + initialize). +#: Without it a stdio server that never answers — a cold ``uvx`` still +#: resolving, a wrapper whose upstream 403s — parks the whole connect, +#: and with servers brought up together that is a session that never +#: finishes preparing its tools. +DEFAULT_SERVER_STARTUP_TIMEOUT = _int_env('MCP_STARTUP_TIMEOUT', 60) DEFAULT_STREAMABLE_HTTP_TIMEOUT = timedelta(seconds=30) DEFAULT_STREAMABLE_HTTP_SSE_READ_TIMEOUT = timedelta(seconds=60 * 5) +_PYDANTIC_DOC_LINK = re.compile(r'\s*For further information visit \S+') +_PYDANTIC_FIELD_LINE = re.compile(r'^(\S+?)(?:\.\w+)?\n\s+(.+?)\s*\[type=', + re.MULTILINE) + + +def _summarize_tool_error(detail: str) -> str: + """Make a validator's complaint answerable. + + A union of strict types reports once per branch, so rejecting two quoted + numbers arrives as four near-identical paragraphs, each ending in a link to + the validation library's documentation. None of that tells the model what + to send instead, and its length buries the one line that does. The + per-field expectation is lifted to the front; the original follows for + anyone reading the transcript. + """ + if 'validation error' not in detail: + return detail + + seen: dict = {} + for field, message in _PYDANTIC_FIELD_LINE.findall(detail): + seen.setdefault(field, message) + if not seen: + return _PYDANTIC_DOC_LINK.sub('', detail).strip() + + lines = [f'- `{field}`: {message}' for field, message in seen.items()] + return ('the arguments were rejected by the tool:\n' + + '\n'.join(lines) + + '\nSend each value as its declared JSON type (a number ' + 'unquoted, not as a string) and call again.\n\nOriginal error:\n' + + _PYDANTIC_DOC_LINK.sub('', detail).strip()) + + class MCPClient(ToolBase): """MCP client for all mcp tools @@ -50,7 +111,11 @@ def __init__( ): super().__init__(config) self.sessions: Dict[str, ClientSession] = {} - self._server_stacks: Dict[str, AsyncExitStack] = {} + # One owner task per server, holding that server's transport open for + # its whole lifetime (see _own_server), plus the event that asks it to + # let go. + self._server_tasks: Dict[str, 'asyncio.Task'] = {} + self._server_shutdown: Dict[str, asyncio.Event] = {} self.exit_stack = AsyncExitStack() self.mcp_config: Dict[str, Dict[str, Any]] = {'mcpServers': {}} if config is not None: @@ -73,12 +138,22 @@ async def call_tool(self, server_name: str, tool_name: str, if response.isError: sep = '\n\n' if all(isinstance(item, str) for item in response.content): - return f'execute tool call error: [{server_name}]{tool_name}, {sep.join(response.content)}' + detail = sep.join(response.content) else: - item_list = [] - for item in response.content: - item_list.append(item.text) - return f'execute tool call error: [{server_name}]{tool_name}, {sep.join(item_list)}' + detail = sep.join( + getattr(item, 'text', str(item)) + for item in response.content) + # Marked as an error rather than returned as ordinary text. A + # failure reported only in prose reaches the model with no error + # flag and reaches the UI with no failed step, so a call that was + # refused looks like a call that answered. + return { + 'result': + (f'execute tool call error: [{server_name}]{tool_name}, ' + f'{_summarize_tool_error(detail)}'), + 'is_error': + True, + } for content in response.content: if content.type == 'text': texts.append(content.text) @@ -127,6 +202,10 @@ async def get_tools_for_server(self, server_name: str) -> List[Tool]: new_eg = enhance_error( e, f'MCP `{server_name}` list tool failed, details: ') raise new_eg from e + # Logged from here rather than from connect(): the same listing that + # registers the tools also reports them, so the log costs no extra + # round trip to the server. + self.print_tools(server_name, response) return self._filter_session_tools(server_name, response) async def get_tools(self) -> Dict: @@ -172,12 +251,9 @@ def is_connected(self, server_name: str) -> bool: async def disconnect_server(self, server_name: str) -> None: """Disconnect a single MCP server.""" - stack = self._server_stacks.pop(server_name, None) - self.sessions.pop(server_name, None) self.exclude_functions.pop(server_name, None) self.include_functions.pop(server_name, None) - if stack is not None: - await stack.aclose() + await self._stop_server(server_name) async def connect_single_server( self, @@ -207,15 +283,14 @@ async def connect_single_server( **server, ) - async def connect_to_server(self, - server_name: str, - timeout: int = CONNECTION_TIMEOUT, - **kwargs): - if self.is_connected(server_name): - return server_name - logger.info(f'connect to {server_name}') - stack = AsyncExitStack() - self._server_stacks[server_name] = stack + async def _open_session(self, stack: AsyncExitStack, server_name: str, + timeout: int, **kwargs) -> ClientSession: + """Open the transport for one server and return its initialized session. + + Every context entered here is registered on ``stack``, which the + caller must also close — see :meth:`_own_server` for why that has to + happen in the same task. + """ # transport: stdio, sse, streamable_http, websocket transport = kwargs.get('transport') or kwargs.get('type') command = kwargs.get('command') @@ -305,42 +380,205 @@ async def connect_to_server(self, stdio, write = await stack.enter_async_context( stdio_client(server_params)) + session_kwargs = session_kwargs or {} + read_timeout = max( + session_kwargs.pop('read_timeout_seconds', timeout), 1) + # The url branch above has always passed this; stdio never did, so + # a child that accepted the pipe and then went quiet left every + # request on it waiting forever. session = await stack.enter_async_context( - ClientSession(stdio, write)) + ClientSession( + stdio, + write, + read_timeout_seconds=timedelta(seconds=read_timeout), + **session_kwargs)) else: raise ValueError( "'url' or 'command' parameter is required for connection") await session.initialize() - # Store session - self.sessions[server_name] = session - self.print_tools(server_name, await session.list_tools()) + return session + + async def _own_server( + self, + server_name: str, + kwargs: Dict[str, Any], + ready: 'asyncio.Future', + shutdown: asyncio.Event, + ) -> None: + """Hold one server's transport open for as long as it is connected. + + The anyio streams underneath ``stdio_client`` / ``streamablehttp_client`` + carry a cancel scope that must be exited by the same task that entered + it. Handing the open context back to the caller therefore only works + while connect and cleanup happen in one task — the moment servers are + brought up concurrently, teardown raises "Attempted to exit cancel + scope in a different task". Keeping open, serve and close inside this + one task removes that coupling entirely. + """ + stack = AsyncExitStack() + try: + async with stack: + session = await self._open_session(stack, server_name, **kwargs) + self.sessions[server_name] = session + if not ready.done(): + ready.set_result(server_name) + await shutdown.wait() + except BaseException as exc: # noqa: BLE001 + if not ready.done(): + ready.set_exception(exc) + elif not isinstance(exc, asyncio.CancelledError): + logger.warning('MCP server %s dropped: %s', server_name, exc) + finally: + self.sessions.pop(server_name, None) + + async def connect_to_server(self, + server_name: str, + timeout: int = CONNECTION_TIMEOUT, + startup_timeout: Optional[int] = None, + **kwargs): + if self.is_connected(server_name): + return server_name + logger.info(f'connect to {server_name}') + ready: 'asyncio.Future' = asyncio.get_running_loop().create_future() + shutdown = asyncio.Event() + task = asyncio.create_task( + self._own_server(server_name, { + 'timeout': timeout, + **kwargs + }, ready, shutdown), + name=f'mcp-server:{server_name}') + self._server_tasks[server_name] = task + self._server_shutdown[server_name] = shutdown + try: + if startup_timeout: + # Shield so a timeout stops US waiting without cancelling the + # owner mid-open; _stop_server below unwinds it in its own task. + await asyncio.wait_for( + asyncio.shield(ready), timeout=startup_timeout) + else: + await ready + except BaseException: + # It never came up, so there is no established session to shut down + # politely — waiting out the grace period here would just add it to + # the startup timeout the caller already spent. + await self._stop_server(server_name, graceful=False) + raise return server_name - async def connect(self, timeout: int = CONNECTION_TIMEOUT): - assert self.mcp_config, 'MCP config is required' + async def _stop_server(self, + server_name: str, + graceful: bool = True) -> None: + """Ask a server's owner task to close its transport, and wait for it.""" + shutdown = self._server_shutdown.pop(server_name, None) + task = self._server_tasks.pop(server_name, None) + self.sessions.pop(server_name, None) + if shutdown is not None: + shutdown.set() + if task is None or task.done(): + return + if graceful: + try: + await asyncio.wait_for( + asyncio.shield(task), timeout=SERVER_STOP_TIMEOUT) + return + except (asyncio.TimeoutError, asyncio.CancelledError): + pass + except BaseException as exc: # noqa: BLE001 + logger.debug('MCP server %s stopped with %s', server_name, exc) + return + task.cancel() + with suppress(BaseException): + await task + + def _plan_server(self, name: str, server: Dict[str, Any], + default_timeout: int) -> Dict[str, Any]: + """Pull the keys ``connect_to_server`` does not take out of a server + block, and resolve its env against the ambient one.""" envs = Env.load_env() + env_dict = server.pop('env', {}) + env_dict = { + key: value if value else envs.get(key, '') + for key, value in env_dict.items() + } + if 'exclude' in server: + self.exclude_functions[name] = server.pop('exclude') + if 'include' in server: + self.include_functions[name] = server.pop('include') + assert (not self.include_functions.get(name)) or ( + not self.exclude_functions.get(name) + ), 'Set either `include` or `exclude` in tools config.' + # Bound to THIS server. Assigning it to the shared default (as the + # sequential loop used to) leaked one server's timeout onto every + # server configured after it. + return { + 'server_name': name, + 'env': env_dict, + 'timeout': server.pop('timeout', default_timeout), + **server, + } + + async def connect( + self, + timeout: int = CONNECTION_TIMEOUT, + startup_timeout: int = DEFAULT_SERVER_STARTUP_TIMEOUT, + ) -> List[tuple]: + """Bring every configured server up, and report which ones did not. + + Servers are started TOGETHER rather than one after another. Startup + cost is dominated by what is being started — a cold ``uvx`` resolving + and downloading its package runs to seconds — and paying that in + sequence multiplied it by the number of servers before the user's + first message could even be read. + + A server that fails or exceeds ``startup_timeout`` is left out instead + of taking the rest down with it; the returned ``(name, exc)`` list lets + a caller surface exactly what is missing. If NOTHING came up the + failure is raised, since that is indistinguishable from a broken + configuration and callers rely on hearing about it. + """ + assert self.mcp_config, 'MCP config is required' mcp_config = self.mcp_config['mcpServers'] - for name, server in mcp_config.items(): - try: - env_dict = server.pop('env', {}) - env_dict = { - key: value if value else envs.get(key, '') - for key, value in env_dict.items() - } - if 'exclude' in server: - self.exclude_functions[name] = server.pop('exclude') - if 'include' in server: - self.include_functions[name] = server.pop('include') - assert (not self.include_functions.get(name)) or ( - not self.exclude_functions.get(name) - ), 'Set either `include` or `exclude` in tools config.' - timeout = server.pop('timeout', timeout) + plans = [ + self._plan_server(name, server, timeout) + for name, server in mcp_config.items() + ] + if not plans: + return [] + + limiter = asyncio.Semaphore(MAX_PARALLEL_CONNECTS) + + async def _one(plan: Dict[str, Any]): + async with limiter: + # The bound is applied inside connect_to_server, which unwinds + # a server that overran through its own owner task. await self.connect_to_server( - server_name=name, env=env_dict, timeout=timeout, **server) - except Exception as e: - new_eg = enhance_error(e, f'Connect `{name}` failed, details:') - raise new_eg from e + startup_timeout=startup_timeout, **plan) + + results = await asyncio.gather( + *(_one(plan) for plan in plans), return_exceptions=True) + + failures: List[tuple] = [] + for plan, result in zip(plans, results): + if not isinstance(result, BaseException): + continue + name = plan['server_name'] + if isinstance(result, asyncio.TimeoutError): + result = TimeoutError( + f'MCP server `{name}` did not finish starting within ' + f'{startup_timeout}s (set MCP_STARTUP_TIMEOUT to change). ' + 'A launcher such as uvx/npx populating a cold cache, or an ' + 'unreachable upstream, is the usual cause.') + failures.append((name, result)) + + if failures and not self.sessions: + name, exc = failures[0] + new_eg = enhance_error(exc, f'Connect `{name}` failed, details:') + raise new_eg from exc + for name, exc in failures: + logger.warning('MCP server %s unavailable this session: %s', name, + exc) + return failures async def add_mcp_config(self, mcp_config: Dict[str, Dict[str, Any]]): if mcp_config is None: @@ -366,7 +604,7 @@ async def add_mcp_config(self, mcp_config: Dict[str, Dict[str, Any]]): async def cleanup(self): """Clean up resources""" - for name in list(self._server_stacks): + for name in list(self._server_tasks): await self.disconnect_server(name) await self.exit_stack.aclose() diff --git a/ms_agent/tools/tool_manager.py b/ms_agent/tools/tool_manager.py index df6b9a228..ca1148ce6 100644 --- a/ms_agent/tools/tool_manager.py +++ b/ms_agent/tools/tool_manager.py @@ -15,6 +15,7 @@ from typing import Any, Awaitable, Callable, Dict, List, Optional from ms_agent.llm.utils import Tool, ToolCall +from ms_agent.tools.arg_coercion import coerce_arguments from ms_agent.tools.agent_tool import AgentTool from ms_agent.tools.base import ToolBase from ms_agent.tools.code import CodeExecutionTool, LocalCodeExecutionTool @@ -423,6 +424,45 @@ def _build_index_key(self, server_name: str, tool_name: str) -> str: server_name = server_name[:max(0, max_server_len)] return f'{server_name}{self.TOOL_SPLITER}{tool_name}' + def _resolve_tool_name(self, tool_name: str) -> tuple: + """Map what the model asked for onto a registered tool. + + Some models reach for the server segment alone — ``code_executor`` + instead of ``code_executor---shell_executor`` — and used to get an + AssertionError wrapped in a traceback, which says nothing about what + the name should have been. When the shorthand picks out exactly one + tool there is no ambiguity to resolve, so the call proceeds and the + result carries a note; otherwise it is reported as not found, with + candidates. + + Returns ``(resolved_name, note)``; ``resolved_name`` is ``None`` when + nothing matched. + """ + if tool_name in self._tool_index: + return tool_name, '' + prefix = f'{tool_name}{self.TOOL_SPLITER}' + matches = [key for key in self._tool_index if key.startswith(prefix)] + if len(matches) == 1: + return matches[0], ( + f'Note: "{tool_name}" was resolved to "{matches[0]}". ' + 'Use the exact tool name in subsequent calls.') + return None, '' + + def _unknown_tool_message(self, tool_name: str) -> str: + """Say what to call instead, rather than how the lookup failed.""" + import difflib + + available = sorted(self._tool_index) + close = difflib.get_close_matches(tool_name, available, n=3, cutoff=0.4) + if not close: + prefix = f'{tool_name}{self.TOOL_SPLITER}' + close = [k for k in available if k.startswith(prefix)][:3] + detail = (f' Closest matches: {", ".join(close)}.' if close else + f' Available tools: {", ".join(available[:20])}' + f'{" …" if len(available) > 20 else ""}.') + return (f'Tool "{tool_name}" is not available.{detail} ' + 'Use the exact tool name from the tool list.') + def _extend_mcp_tool_index( self, tool_ins: ToolBase, @@ -557,9 +597,28 @@ async def single_call_tool(self, tool_info: ToolCall): tool_args = json.loads(tool_args) except Exception: # noqa return f'The input {tool_args} is not a valid JSON, fix your arguments and try again' - assert tool_name in self._tool_index, f'Tool name {tool_name} not found' + resolved_name, resolution_note = self._resolve_tool_name( + tool_name) + if resolved_name is None: + return { + 'result': self._unknown_tool_message(tool_name), + 'is_error': True, + } + tool_name = resolved_name + tool_info['tool_name'] = tool_name index_snapshot = self._tool_index[tool_name] - tool_ins, server_name, _ = index_snapshot + tool_ins, server_name, tool_spec = index_snapshot + + # Reconcile the model's JSON with the tool's declared schema + # before anyone validates it. The schema is right here in the + # index; leaving a quoted number quoted just moves the failure + # to the tool, which reports it in terms of its own validator. + if isinstance(tool_args, dict): + coerced = coerce_arguments( + tool_args, (tool_spec or {}).get('parameters')) + if coerced is not tool_args: + tool_args = coerced + tool_info['arguments'] = tool_args # --- MCP availability (before SafetyGuard / PreToolUse) --- if (tool_ins is self.servers @@ -609,11 +668,19 @@ async def single_call_tool(self, tool_info: ToolCall): } # interactive mode: force the enforcer to confirm with # the user; whitelist/memory must not bypass a safety - # ask (REVIEW P1-2). + # ask (REVIEW P1-2) — except for the categories a + # user can reasonably settle once for a project, + # where an ask that ignores their answer is just an + # ask they learn to click through. + from ms_agent.permission.ask_resolver import \ + REMEMBERABLE_ASK_CATEGORIES from ms_agent.permission.enforcer import \ PermissionDecision safety_force_decision = PermissionDecision( - action='ask', reason=resolved.reason) + action='ask', + reason=resolved.reason, + rememberable=(resolved.category + in REMEMBERABLE_ASK_CATEGORIES)) # --- PreToolUse hooks --- hook_result = None @@ -697,6 +764,17 @@ async def single_call_tool(self, tool_info: ToolCall): and tool_ins is self.servers): await self.mcp_success_handler(server_name) + if resolution_note: + # Carried on the successful result so the correction lands + # while the model is looking at what it just did, instead + # of costing it a failed round to discover. + if isinstance(response, dict): + response = dict(response) + response['result'] = ( + f'{resolution_note}\n{response.get("result", "")}') + else: + response = f'{resolution_note}\n{response}' + # --- PostToolUse hooks --- hook_attachments = list(pre_attachments) if self._hook_runtime is not None and not self._hook_runtime.is_empty: diff --git a/tests/mcp/test_mcp_runtime.py b/tests/mcp/test_mcp_runtime.py index 7f82c810d..27911909f 100644 --- a/tests/mcp/test_mcp_runtime.py +++ b/tests/mcp/test_mcp_runtime.py @@ -527,19 +527,34 @@ async def test_sync_mcp_tools_list_failure_records(): @pytest.mark.asyncio async def test_mcp_client_aexit_disconnects_server_stacks(): - from contextlib import AsyncExitStack + import asyncio from ms_agent.tools.mcp_client import MCPClient client = MCPClient({'mcpServers': {}}) + + # Stand in for a live server: an owner task parked on its shutdown event, + # exactly as _own_server parks once its transport is open. + shutdown = asyncio.Event() + closed = asyncio.Event() + + async def _owner(): + try: + await shutdown.wait() + finally: + client.sessions.pop('fake', None) + closed.set() + client.sessions['fake'] = object() - stack = AsyncExitStack() - await stack.__aenter__() - client._server_stacks['fake'] = stack + client._server_shutdown['fake'] = shutdown + client._server_tasks['fake'] = asyncio.create_task(_owner()) + await asyncio.sleep(0) await client.__aexit__(None, None, None) + assert closed.is_set() assert 'fake' not in client.sessions - assert 'fake' not in client._server_stacks + assert 'fake' not in client._server_tasks + assert 'fake' not in client._server_shutdown def _make_tool_manager(client, runtime, hook_runtime=None): diff --git a/tests/permission/test_ask_timeout_modes.py b/tests/permission/test_ask_timeout_modes.py new file mode 100644 index 000000000..99b2cc513 --- /dev/null +++ b/tests/permission/test_ask_timeout_modes.py @@ -0,0 +1,139 @@ +"""How long an approval waits, and what happens when it does not arrive. + +Expiring an approval answers it as a refusal — an answer the user never gave. +So the wait is unbounded wherever a person intends to answer, and the message +on the bounded path says plainly that nobody refused anything. +""" +import asyncio + +import pytest + +from ms_agent.permission.handler import (PermissionAction, PermissionResponse, + WebPermissionHandler) + + +class _Emitter: + + def __init__(self): + self.events = [] + + def emit(self, event): + self.events.append(event) + + +@pytest.mark.asyncio +async def test_default_wait_is_unbounded(): + handler = WebPermissionHandler(_Emitter()) + assert handler._timeout is None + + task = asyncio.create_task(handler.ask('tool', {}, '', call_id='c1')) + await asyncio.sleep(0.05) + assert not task.done(), 'an unbounded ask must not resolve itself' + + request_id = handler._event_emitter.events[0]['request_id'] + handler.resolve(request_id, + PermissionResponse(action=PermissionAction.ALLOW_ONCE)) + assert (await task).action is PermissionAction.ALLOW_ONCE + + +@pytest.mark.asyncio +async def test_bounded_wait_denies_but_says_it_was_a_timeout(): + handler = WebPermissionHandler(_Emitter(), timeout=0.05) + response = await handler.ask('tool', {}, '', call_id='c1') + + assert response.action is PermissionAction.DENY + # The model has to be able to tell these apart: a refusal is a decision to + # work around, a timeout is nobody having looked. + assert 'TIMEOUT' in response.feedback + assert 'not a refusal' in response.feedback + assert 'Do not re-request' in response.feedback + + +@pytest.mark.asyncio +async def test_pending_asks_are_answered_when_the_session_closes(): + """An unbounded ask holds its turn open; an abandoned prompt would pin its + session for the life of the process unless something answers it.""" + handler = WebPermissionHandler(_Emitter()) + first = asyncio.create_task(handler.ask('a', {}, '', call_id='c1')) + second = asyncio.create_task(handler.ask('b', {}, '', call_id='c2')) + await asyncio.sleep(0.05) + + assert handler.cancel_pending('Session closed') == 2 + assert handler.cancel_pending() == 0 # idempotent + + for task in (first, second): + response = await task + assert response.action is PermissionAction.DENY + assert response.feedback == 'Session closed' + assert handler._pending == {} + + +@pytest.mark.asyncio +async def test_answering_one_card_releases_the_ones_it_covers(): + """"Always allow" is a statement about a pattern; a sibling card the + pattern covers must not be left holding the turn open indefinitely.""" + from ms_agent.permission.config import PermissionConfig + from ms_agent.permission.enforcer import PermissionEnforcer + + emitter = _Emitter() + handler = WebPermissionHandler(emitter) + enforcer = PermissionEnforcer( + config=PermissionConfig(mode='interactive'), handler=handler) + tool = 'code_executor---shell_executor' + + calls = [ + asyncio.create_task( + enforcer.check(tool, {'command': f'git status {i}'})) + for i in range(3) + ] + await asyncio.sleep(0.05) + assert len(handler._pending) == 3 + + handler.resolve( + emitter.events[0]['request_id'], + PermissionResponse( + action=PermissionAction.ALLOW_ALWAYS, pattern=f'{tool}:git *')) + await asyncio.sleep(0.05) + + assert all(c.done() for c in calls), 'siblings left waiting' + assert handler._pending == {} + decisions = await asyncio.gather(*calls) + assert [d.action for d in decisions] == ['allow'] * 3 + + +@pytest.mark.asyncio +async def test_a_remembered_answer_never_releases_a_safety_confirmation(): + """The forced path exists so a remembered answer cannot stand in for + looking at this particular call. Releasing siblings must not create a way + around it.""" + from ms_agent.permission.config import PermissionConfig + from ms_agent.permission.enforcer import (PermissionDecision, + PermissionEnforcer) + + emitter = _Emitter() + handler = WebPermissionHandler(emitter) + enforcer = PermissionEnforcer( + config=PermissionConfig(mode='interactive'), handler=handler) + tool = 'code_executor---shell_executor' + forced = PermissionDecision(action='ask', reason='reads outside workspace') + + calls = [ + asyncio.create_task( + enforcer.check( + tool, {'command': f'cat /etc/hosts {i}'}, + force_decision=forced)) for i in range(3) + ] + await asyncio.sleep(0.05) + + handler.resolve( + emitter.events[0]['request_id'], + PermissionResponse( + action=PermissionAction.ALLOW_ALWAYS, pattern=f'{tool}:cat *')) + await asyncio.sleep(0.05) + + assert sum(c.done() for c in calls) == 1, 'a safety ask was auto-released' + assert len(handler._pending) == 2 + for call in calls: + if not call.done(): + call.cancel() + await asyncio.gather(*calls, return_exceptions=True) diff --git a/tests/permission/test_interpreter_and_secrets.py b/tests/permission/test_interpreter_and_secrets.py new file mode 100644 index 000000000..f57837031 --- /dev/null +++ b/tests/permission/test_interpreter_and_secrets.py @@ -0,0 +1,156 @@ +"""Two holes the path policy had: an interpreter's argument is a program, so +extracting paths from it proves nothing; and the sensitive-path list was only +ever consulted for writes, so the files it names could be read out.""" +import asyncio + +import pytest + +from ms_agent.permission.ask_resolver import (REMEMBERABLE_ASK_CATEGORIES, + resolve_ask) +from ms_agent.permission.config import PermissionConfig, SafetyConfig +from ms_agent.permission.enforcer import (PermissionDecision, + PermissionEnforcer) +from ms_agent.permission.handler import (PermissionAction, PermissionResponse, + WebPermissionHandler) +from ms_agent.permission.shell_validator import (PathSafetyConfig, + ShellPathValidator) + +WS = '/tmp/ms-agent-secret-tests' + + +class _Emitter: + + def __init__(self): + self.events = [] + + def emit(self, event): + self.events.append(event) + + +def _validator() -> ShellPathValidator: + cfg = SafetyConfig() + allowed = list(cfg.effective_allowed_directories(WS)) + return ShellPathValidator( + allowed_dirs=allowed, + safety_config=PathSafetyConfig( + workspace_root=WS, + allowed_directories=tuple(allowed), + sensitive_read_paths=cfg.sensitive_read_paths, + )) + + +def _enforcer(): + emitter = _Emitter() + handler = WebPermissionHandler(emitter) + return emitter, handler, PermissionEnforcer( + config=PermissionConfig(mode='interactive'), handler=handler) + + +@pytest.mark.parametrize('command', [ + 'python3 -c "import os; os.remove(\'/etc/hosts\')"', + 'node -e "require(\'fs\').unlinkSync(\'/etc/hosts\')"', + 'bash -c "rm -rf /"', + 'python3 -', +]) +def test_inline_code_is_surfaced_rather_than_waved_through(command): + decision = _validator().check(command) + assert decision.action == 'ask', command + assert decision.category == 'interpreter_exec', command + + +@pytest.mark.parametrize('command', ['python3 build.py', 'node server.js']) +def test_running_a_script_inside_the_workspace_stays_quiet(command): + assert _validator().check(command).action == 'allow', command + + +def test_script_outside_the_workspace_is_still_judged(): + assert _validator().check('python3 /etc/evil.py').action != 'allow' + + +@pytest.mark.parametrize('mode, expected', [ + ('auto', 'allow'), + ('interactive', 'ask'), + ('strict', 'deny'), +]) +def test_interpreter_resolution_follows_the_mode(mode, expected): + decision = _validator().check('python3 -c "print(1)"') + assert resolve_ask(decision, mode, 'loose').action == expected + + +@pytest.mark.parametrize('command', [ + 'cat ~/.ssh/id_rsa', + 'cat /home/other/.aws/credentials', + 'cat /srv/certs/server.pem', +]) +def test_credentials_are_never_read_unattended(command): + decision = _validator().check(command) + assert decision.category == 'sensitive_read', command + assert resolve_ask(decision, 'auto', 'loose').action == 'deny', command + assert resolve_ask(decision, 'interactive', 'loose').action == 'ask' + + +def test_write_protection_does_not_become_read_refusal(): + """``sensitive_paths`` stops `.git/config` being CHANGED. Reading it is + how an agent finds a git remote; refusing that would be a regression.""" + decision = _validator().check('cat .git/config') + assert decision.category != 'sensitive_read' + assert resolve_ask(decision, 'auto', 'loose').action == 'allow' + + +def test_only_the_settleable_category_is_rememberable(): + """Inline code is the same decision every time; a credential read is risky + because of the specific file, so it must be decided each time.""" + assert REMEMBERABLE_ASK_CATEGORIES == frozenset({'interpreter_exec'}) + + +@pytest.mark.asyncio +async def test_always_allow_actually_sticks_for_inline_code(): + """An `interpreter_exec` ask used to take the forced path, which skips + memory — "always run" stored a pattern that was never consulted.""" + emitter, handler, enforcer = _enforcer() + tool = 'code_executor---shell_executor' + ask = PermissionDecision( + action='ask', reason='runs code inline', rememberable=True) + + first = asyncio.create_task( + enforcer.check( + tool, {'command': 'python3 -c "print(1)"'}, force_decision=ask)) + await asyncio.sleep(0.05) + handler.resolve( + emitter.events[0]['request_id'], + PermissionResponse( + action=PermissionAction.ALLOW_ALWAYS, pattern=f'{tool}:python3 *')) + assert (await first).action == 'allow' + + # Second time: no card at all. + before = len(emitter.events) + again = await enforcer.check( + tool, {'command': 'python3 -c "print(2)"'}, force_decision=ask) + assert again.action == 'allow' + assert len(emitter.events) == before, 'asked again despite "always allow"' + + +@pytest.mark.asyncio +async def test_a_credential_read_still_asks_every_time(): + emitter, handler, enforcer = _enforcer() + tool = 'code_executor---shell_executor' + ask = PermissionDecision(action='ask', reason='reads a sensitive path') + + first = asyncio.create_task( + enforcer.check( + tool, {'command': 'cat ~/.ssh/id_rsa'}, force_decision=ask)) + await asyncio.sleep(0.05) + handler.resolve( + emitter.events[0]['request_id'], + PermissionResponse( + action=PermissionAction.ALLOW_ALWAYS, pattern=f'{tool}:cat *')) + await first + + before = len(emitter.events) + second = asyncio.create_task( + enforcer.check( + tool, {'command': 'cat ~/.ssh/id_ed25519'}, force_decision=ask)) + await asyncio.sleep(0.05) + assert len(emitter.events) == before + 1, 'a credential read was waved through' + second.cancel() + await asyncio.gather(second, return_exceptions=True) diff --git a/tests/permission/test_shell_policy_misfires.py b/tests/permission/test_shell_policy_misfires.py new file mode 100644 index 000000000..5ccfae820 --- /dev/null +++ b/tests/permission/test_shell_policy_misfires.py @@ -0,0 +1,104 @@ +"""Commands the policy used to refuse although they touch nothing outside the +workspace. Each case here is one that was reported from a real session.""" +import pytest + +from ms_agent.permission.config import SafetyConfig +from ms_agent.permission.shell_validator import (PathSafetyConfig, + ShellPathValidator) + +WS = '/tmp/ms-agent-policy-tests' + + +def _validator(**safety_kwargs) -> ShellPathValidator: + cfg = SafetyConfig(**safety_kwargs) + allowed = list(cfg.effective_allowed_directories(WS)) + return ShellPathValidator( + allowed_dirs=allowed, + safety_config=PathSafetyConfig( + workspace_root=WS, + allowed_directories=tuple(allowed), + dangerous_removal_paths=cfg.dangerous_removal_paths, + )) + + +HEREDOC_WITH_GLOB_CHARS = """python3 - < 1] +print(result) +EOF""" + +HEREDOC_QUOTED_DELIMITER = """python3 - <<'PY' +with open('out_*.json', 'w') as f: + f.write('a > b') +PY""" + +HEREDOC_TAB_STRIPPED = """cat <<-END +\tbody containing > and [brackets] +\tEND""" + + +@pytest.mark.parametrize( + 'command', + [ + # The redirect target used to be read with \S+, so the subshell's + # closing paren became part of the filename: "/dev/null)". + 'rg -F -n "ORANGE-RIVER-731" . 2>/dev/null', + 'echo hi 2>/dev/null)', + '(ls >> /dev/null)', + 'make build > log.txt 2>&1', + # A '>' inside a quoted argument is not a redirection. + 'git commit -m "perf: speed > /etc of before"', + # Heredoc bodies are data, not commands. + HEREDOC_TAB_STRIPPED, + ], +) +def test_allows_commands_that_touch_nothing_outside_the_workspace(command): + assert _validator().check(command).action == 'allow', command + + +@pytest.mark.parametrize('command', + [HEREDOC_WITH_GLOB_CHARS, HEREDOC_QUOTED_DELIMITER]) +def test_inline_script_is_judged_as_code_not_as_bogus_paths(command): + decision = _validator().check(command) + assert decision.category == 'interpreter_exec' + assert 'Glob' not in decision.reason + assert 'outside allowed directories' not in decision.reason + + +def test_os_temp_directory_is_writable_by_default(): + assert _validator().check('touch /tmp/probe.py').action == 'allow' + + +def test_os_temp_directory_can_be_taken_away(): + decision = _validator(allow_temp_dir=False).check('touch /tmp/probe.py') + assert decision.action != 'allow' + + +@pytest.mark.parametrize('command, fragment', [ + ('echo x > /etc/passwd', 'outside allowed directories'), + ('ls >> /etc/hosts', 'outside allowed directories'), + ('touch "out_*.json"', 'Glob patterns not allowed'), + ('echo x > $HOME/y.txt', 'variable expansion'), +]) +def test_real_violations_are_still_refused(command, fragment): + decision = _validator().check(command) + assert decision.action != 'allow', command + assert fragment in decision.reason + + +@pytest.mark.parametrize('command', [ + # A pasted heredoc that arrived flattened onto one line, and one that was + # opened but never closed — both real-world shapes. + 'python3 - < 1]' + 'print(result)EOF', + 'cat < text\n', +]) +def test_an_unterminated_heredoc_is_not_refused_for_an_invented_reason(command): + """Reading the payload as filenames produced refusals like "Glob patterns + not allowed in create operations: 1]print" — untrue, and hiding the real + problem. Let the shell reject it on its own terms.""" + decision = _validator().check(command) + assert 'Glob' not in decision.reason + assert 'outside allowed directories' not in decision.reason + if decision.action != 'allow': + assert decision.category == 'interpreter_exec', decision.reason diff --git a/tests/prompting/test_workspace_internals_hint.py b/tests/prompting/test_workspace_internals_hint.py new file mode 100644 index 000000000..6d1fe30af --- /dev/null +++ b/tests/prompting/test_workspace_internals_hint.py @@ -0,0 +1,84 @@ +"""The section that tells the agent which directories under its working +directory are the framework's own records.""" +from omegaconf import OmegaConf + +from ms_agent.agent.llm_agent import LLMAgent + + +def _agent_for(workspace, session_id='abc123', log_dir=None): + agent = LLMAgent.__new__(LLMAgent) + agent.config = OmegaConf.create({'output_dir': str(workspace)}) + + class _Runtime: + pass + + runtime = _Runtime() + runtime.session_id = session_id + agent.runtime = runtime + + if log_dir is not None: + + class _Log: + directory = log_dir + + agent.session_log = _Log() + return agent + + +def test_section_is_absent_when_records_live_elsewhere(tmp_path): + """A project opened from an existing folder keeps its transcripts outside + the working directory; there is nothing to warn about.""" + workspace = tmp_path / 'plain-project' + workspace.mkdir() + assert _agent_for(workspace)._build_workspace_internals_section() == '' + + +def test_section_names_the_directories_and_this_session(tmp_path): + workspace = tmp_path / 'managed-project' + (workspace / 'sessions' / 'abc123').mkdir(parents=True) + (workspace / '.ms_agent').mkdir() + + section = _agent_for(workspace)._build_workspace_internals_section() + + assert 'sessions/' in section + assert '.ms_agent/' in section + assert 'sessions/abc123/' in section + # The reason the agent needs this at all: its own prompt is already on + # disk, so searching for a phrase from the request matches the transcript. + assert 'searching' in section + + +def test_the_named_directory_is_the_one_being_written_to(tmp_path): + """The agent's tag is not its session directory; sending the model after + `sessions/Agent-default/` would send it somewhere that does not exist.""" + workspace = tmp_path / 'managed-project' + real = workspace / 'sessions' / 'de43058dc0b1' + real.mkdir(parents=True) + + agent = _agent_for(workspace, session_id='Agent-default', log_dir=real) + section = agent._build_workspace_internals_section() + assert 'sessions/de43058dc0b1/' in section + assert 'Agent-default' not in section + + +def test_it_reaches_the_system_prompt(tmp_path): + """Guards the wiring, not just the builder.""" + workspace = tmp_path / 'managed-project' + (workspace / 'sessions' / 'sid').mkdir(parents=True) + + agent = _agent_for(workspace, session_id='sid') + # ``system`` reads through the config; set it where it actually comes from. + agent.config = OmegaConf.create({ + 'output_dir': str(workspace), + 'prompt': { + 'system': 'BASE PROMPT' + }, + }) + agent._memory_guidance = None + agent._skill_injector = None + agent._skill_runtime = None + agent._personalization_enabled = lambda: False + + content = agent._build_system_content() + assert content.startswith('BASE PROMPT') + assert 'sessions/sid/' in content diff --git a/tests/tools/test_arg_coercion.py b/tests/tools/test_arg_coercion.py new file mode 100644 index 000000000..91e678f62 --- /dev/null +++ b/tests/tools/test_arg_coercion.py @@ -0,0 +1,111 @@ +"""Reconciling model-emitted JSON with a tool's declared schema. + +The rule throughout: rewrite only where the reading cannot change. Anything +ambiguous is left for the tool to judge on its own terms. +""" +import pytest + +from ms_agent.tools.arg_coercion import coerce_arguments + +NUMERIC = { + 'type': 'object', + 'properties': { + 'a': { + 'anyOf': [{ + 'type': 'number' + }, { + 'type': 'integer' + }] + }, + 'b': { + 'type': 'integer' + }, + 'flag': { + 'type': 'boolean' + }, + 'name': { + 'type': 'string' + }, + }, +} + + +@pytest.mark.parametrize('given, expected', [ + ({'a': '19.5'}, {'a': 19.5}), + ({'b': '2'}, {'b': 2}), + ({'flag': 'true'}, {'flag': True}), +]) +def test_quoted_scalars_are_unquoted_to_match_the_schema(given, expected): + assert coerce_arguments(given, NUMERIC) == expected + + +@pytest.mark.parametrize('given', [ + {'name': '19.5'}, # schema says string: the quotes ARE the value + {'a': 'nineteen'}, # not a number in any reading + {'a': ''}, + {'unknown': '5'}, # not in the schema; nothing is claimed about it +]) +def test_ambiguous_values_are_left_alone(given): + assert coerce_arguments(given, NUMERIC) == given + + +def test_unchanged_arguments_are_returned_by_identity(): + args = {'name': 'x', 'a': 1} + assert coerce_arguments(args, NUMERIC) is args + + +def test_booleans_are_not_read_as_numbers(): + schema = {'type': 'object', 'properties': {'n': {'type': 'integer'}}} + assert coerce_arguments({'n': True}, schema) == {'n': True} + + +def test_structured_text_is_parsed_and_walked(): + schema = { + 'type': 'object', + 'properties': { + 'items': { + 'type': 'array', + 'items': { + 'type': 'integer' + } + }, + 'cfg': { + 'type': 'object', + 'properties': { + 'depth': { + 'type': 'integer' + } + }, + }, + }, + } + given = {'items': '["1", "2"]', 'cfg': {'depth': '3'}} + assert coerce_arguments(given, schema) == { + 'items': [1, 2], + 'cfg': { + 'depth': 3 + }, + } + + +def test_int_widens_only_where_integer_is_not_accepted(): + schema = { + 'type': 'object', + 'properties': { + 'only_float': { + 'type': 'number' + }, + 'either': { + 'type': ['number', 'integer'] + }, + }, + } + out = coerce_arguments({'only_float': 3, 'either': 3}, schema) + assert isinstance(out['only_float'], float) + assert isinstance(out['either'], int) + + +@pytest.mark.parametrize('schema', [None, 'nonsense']) +def test_a_missing_or_broken_schema_is_not_an_error(schema): + args = {'a': '1'} + assert coerce_arguments(args, schema) is args diff --git a/tests/tools/test_local_code_executor_windows.py b/tests/tools/test_local_code_executor_windows.py index 8ba819c6e..593c34b7e 100644 --- a/tests/tools/test_local_code_executor_windows.py +++ b/tests/tools/test_local_code_executor_windows.py @@ -1,3 +1,4 @@ +import asyncio import os from unittest import mock @@ -11,10 +12,57 @@ def _bare_tool() -> LocalCodeExecutionTool: return tool +class _FakeProcess: + + def __init__(self): + self.returncode = 0 + + async def communicate(self): + return b'', b'' + + +def _run_shell(command: str): + """Run one command with the subprocess stubbed, returning what the shell + would have been handed.""" + tool = _bare_tool() + tool._shell_timeout = 30 + tool._task_manager = None + tool.shell_env = {'PATH': '/usr/bin'} + tool._ws = mock.Mock(root='/tmp') + tool._artifacts = mock.Mock() + tool._artifacts.pack_json_shell_result.side_effect = ( + lambda **kwargs: kwargs) + seen = {} + + async def _fake_exec(cmd, **kwargs): + seen['cmd'] = cmd + return _FakeProcess() + + with mock.patch( + 'ms_agent.tools.code.local_code_executor.asyncio.' + 'create_subprocess_shell', + new=_fake_exec): + asyncio.run(tool.shell_executor(command=command)) + return seen['cmd'] + + +def test_command_reaches_the_shell_verbatim(): + """No rewriting, whatever punctuation the command contains. Rewriting used + to switch to a LOGIN shell on metacharacters, so `cmd` and `cmd ; true` + ran under different PATHs and could resolve different binaries.""" + for command in ( + 'python3 -V', + 'python3 -V ; true', + 'pip install jsonschema && python3 -c "import jsonschema"', + 'grep -r x . 2>/dev/null | head -3', + ): + assert _run_shell(command) == command + + def test_composite_command_uses_native_windows_shell(): command = 'cd work && echo ok > result.txt' with mock.patch('ms_agent.tools.code.local_code_executor.os.name', 'nt'): - assert _bare_tool()._prepare_shell_command(command) == command + assert _run_shell(command) == command def test_sanitized_env_keeps_windows_runtime_variables(): @@ -45,3 +93,78 @@ def test_sanitized_env_keeps_windows_runtime_variables(): assert env[key] == windows_env[key] assert env['INHERITED_FROM_LOCAL'] == 'False' assert 'SECRET_TOKEN' not in env + + +def test_posix_env_carries_identity_tmpdir_tls_and_proxy(): + """Each of these being absent breaks a specific, observed thing — see + ``_POSIX_ENV_PASSTHROUGH``.""" + parent = { + 'PATH': '/usr/bin', + 'HOME': '/home/tester', + 'USER': 'tester', + 'LOGNAME': 'tester', + 'TMPDIR': '/var/folders/xy/T/', + 'SSL_CERT_FILE': '/etc/ssl/cert.pem', + 'HTTPS_PROXY': 'http://127.0.0.1:7890', + 'TERM': 'xterm-256color', + 'AWS_SECRET_ACCESS_KEY': 'must-not-leak', + 'OPENAI_API_KEY': 'must-not-leak', + } + with mock.patch.dict(os.environ, parent, clear=True), mock.patch( + 'ms_agent.tools.code.local_code_executor.os.name', 'posix'): + env = _bare_tool()._build_env('shell_env', inherit=False) + + for key in ('PATH', 'HOME', 'USER', 'LOGNAME', 'TMPDIR', 'SSL_CERT_FILE', + 'HTTPS_PROXY', 'TERM'): + assert env[key] == parent[key], key + assert 'AWS_SECRET_ACCESS_KEY' not in env + assert 'OPENAI_API_KEY' not in env + + # And a parent without a variable must not have one invented for it: + # passing an empty TMPDIR is worse than passing none. + with mock.patch.dict(os.environ, {'PATH': '/usr/bin'}, clear=True), \ + mock.patch('ms_agent.tools.code.local_code_executor.os.name', + 'posix'): + sparse = _bare_tool()._build_env('shell_env', inherit=False) + assert 'TMPDIR' not in sparse + assert 'USER' not in sparse + + +def test_agent_env_overrides_the_parents_interactive_settings(): + """Set outright, not inherited: these exist to undo settings made for a + human at a terminal. A developer's `PAGER=less` is precisely the value + that leaves `git log` waiting for a keypress — observed live.""" + from ms_agent.tools.code.local_code_executor import _AGENT_FRIENDLY_ENV + + hostile_parent = { + 'PATH': '/usr/bin', + 'PAGER': 'less', + 'GIT_PAGER': 'less', + 'LESS': '-R', + 'AI_AGENT': 'some-other-agent', + 'NO_COLOR': '', + 'GIT_TERMINAL_PROMPT': '1', + } + with mock.patch.dict(os.environ, hostile_parent, clear=True), mock.patch( + 'ms_agent.tools.code.local_code_executor.os.name', 'posix'): + env = _bare_tool()._build_env('shell_env', inherit=False) + + for key, expected in _AGENT_FRIENDLY_ENV.items(): + assert env[key] == expected, f'{key} inherited from the parent' + assert env['PAGER'] == 'cat' + assert env['AI_AGENT'] == 'ms_agent' + assert env['GIT_TERMINAL_PROMPT'] == '0' + + +def test_config_can_still_override_the_agent_defaults(): + """Forcing them must not take away the deliberate-choice channel.""" + from types import SimpleNamespace + + tool = _bare_tool() + tool.tool_config = SimpleNamespace(shell_env={'PAGER': 'bat'}) + with mock.patch.dict(os.environ, {'PATH': '/usr/bin'}, clear=True), \ + mock.patch('ms_agent.tools.code.local_code_executor.os.name', + 'posix'): + env = tool._build_env('shell_env', inherit=False) + assert env['PAGER'] == 'bat' + assert env['GIT_PAGER'] == 'cat' # untouched keys keep the safe default diff --git a/tests/tools/test_tool_name_and_reread.py b/tests/tools/test_tool_name_and_reread.py new file mode 100644 index 000000000..279ca24d7 --- /dev/null +++ b/tests/tools/test_tool_name_and_reread.py @@ -0,0 +1,99 @@ +"""Two ways a tool call used to dead-end: an unusable name error, and a read +that returned no content.""" +import json +import os + +import pytest +from omegaconf import OmegaConf + +from ms_agent.tools.tool_manager import ToolManager + + +def _fs_config(workspace): + return OmegaConf.create({ + 'output_dir': str(workspace), + 'tools': {'file_system': {}}, + }) + + +def _manager_with(index_keys) -> ToolManager: + manager = ToolManager.__new__(ToolManager) + manager._tool_index = {key: (None, 'srv', {}) for key in index_keys} + return manager + + +def test_unambiguous_shorthand_resolves_and_says_so(): + manager = _manager_with(['code_executor---shell_executor']) + resolved, note = manager._resolve_tool_name('code_executor') + assert resolved == 'code_executor---shell_executor' + assert 'exact tool name' in note + + +def test_ambiguous_shorthand_is_not_guessed(): + manager = _manager_with( + ['code_executor---shell_executor', 'code_executor---file_operation']) + resolved, _ = manager._resolve_tool_name('code_executor') + assert resolved is None + + +def test_unknown_name_reports_candidates_not_a_traceback(): + manager = _manager_with( + ['file_system---read_file', 'file_system---write_file']) + message = manager._unknown_tool_message('file_system---raed_file') + assert 'file_system---read_file' in message + assert 'exact tool name' in message + assert 'AssertionError' not in message + assert 'Traceback' not in message + + +@pytest.mark.asyncio +async def test_reading_the_same_file_twice_returns_content_both_times(tmp_path): + """A second read must never come back empty-handed: after a context + truncation the first read is gone, and withholding content on the grounds + that it "has not changed" leaves the model with no way to obtain the file. + """ + from ms_agent.tools.filesystem_tool import FileSystemTool + + workspace = tmp_path / 'ws' + workspace.mkdir() + (workspace / 'rules.md').write_text('v1 rules\n', encoding='utf-8') + + tool = FileSystemTool(_fs_config(workspace)) + + first = await tool.call_tool( + 'file_system', tool_name='read_file', tool_args={'path': 'rules.md'}) + second = await tool.call_tool( + 'file_system', tool_name='read_file', tool_args={'path': 'rules.md'}) + + first_payload = json.loads(first)['rules.md'] + second_payload = json.loads(second)['rules.md'] + + assert 'v1 rules' in first_payload + assert 'v1 rules' in second_payload, 'content withheld on re-read' + assert 'unchanged since your last read' in second_payload + + +@pytest.mark.asyncio +async def test_rereading_after_a_change_reports_the_new_content(tmp_path): + from ms_agent.tools.filesystem_tool import FileSystemTool + + workspace = tmp_path / 'ws' + workspace.mkdir() + target = workspace / 'rules.md' + target.write_text('v1\n', encoding='utf-8') + + tool = FileSystemTool(_fs_config(workspace)) + await tool.call_tool( + 'file_system', tool_name='read_file', tool_args={'path': 'rules.md'}) + + # Push mtime forward so the change is visible whatever the clock's + # granularity is. + target.write_text('v2\n', encoding='utf-8') + future = os.path.getmtime(target) + 2 + os.utime(target, (future, future)) + + again = await tool.call_tool( + 'file_system', tool_name='read_file', tool_args={'path': 'rules.md'}) + payload = json.loads(again)['rules.md'] + assert 'v2' in payload + assert 'unchanged' not in payload