diff --git a/.gitignore b/.gitignore index f3f70cfbb..a46e2eb22 100644 --- a/.gitignore +++ b/.gitignore @@ -43,6 +43,7 @@ MANIFEST *.spec release.sh build.sh +scripts/tb21-bench.sh *.html # Installer logs diff --git a/ms_agent/a2a/agent_card.py b/ms_agent/a2a/agent_card.py index b7b3af288..29751084e 100644 --- a/ms_agent/a2a/agent_card.py +++ b/ms_agent/a2a/agent_card.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os from typing import Any, Dict, List diff --git a/ms_agent/a2a/client.py b/ms_agent/a2a/client.py index c5a5d3943..ea00f0186 100644 --- a/ms_agent/a2a/client.py +++ b/ms_agent/a2a/client.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os from typing import Any, Dict, List, Optional diff --git a/ms_agent/a2a/errors.py b/ms_agent/a2a/errors.py index 659aa700e..560fcdc54 100644 --- a/ms_agent/a2a/errors.py +++ b/ms_agent/a2a/errors.py @@ -1,3 +1,4 @@ +from __future__ import annotations from ms_agent.utils.logger import get_logger logger = get_logger() diff --git a/ms_agent/a2a/executor.py b/ms_agent/a2a/executor.py index b67c37bad..2535bb316 100644 --- a/ms_agent/a2a/executor.py +++ b/ms_agent/a2a/executor.py @@ -1,3 +1,4 @@ +from __future__ import annotations import logging import sys from typing import Any diff --git a/ms_agent/acp/client.py b/ms_agent/acp/client.py index 02ad4b0e2..ad9c12c86 100644 --- a/ms_agent/acp/client.py +++ b/ms_agent/acp/client.py @@ -1,3 +1,4 @@ +from __future__ import annotations import asyncio from contextlib import asynccontextmanager from typing import Any, Dict, List, Optional diff --git a/ms_agent/acp/errors.py b/ms_agent/acp/errors.py index 92b4dd0b8..a74bf0dbb 100644 --- a/ms_agent/acp/errors.py +++ b/ms_agent/acp/errors.py @@ -1,3 +1,4 @@ +from __future__ import annotations from acp import RequestError diff --git a/ms_agent/acp/http_adapter.py b/ms_agent/acp/http_adapter.py index 7ef2a5cb2..8905aef51 100644 --- a/ms_agent/acp/http_adapter.py +++ b/ms_agent/acp/http_adapter.py @@ -1,3 +1,4 @@ +from __future__ import annotations import asyncio import os import secrets diff --git a/ms_agent/acp/proxy.py b/ms_agent/acp/proxy.py index f857cad28..e0186aa30 100644 --- a/ms_agent/acp/proxy.py +++ b/ms_agent/acp/proxy.py @@ -1,3 +1,4 @@ +from __future__ import annotations import logging import os import sys diff --git a/ms_agent/acp/registry.py b/ms_agent/acp/registry.py index 96719d29a..ff9030cb8 100644 --- a/ms_agent/acp/registry.py +++ b/ms_agent/acp/registry.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os import sys from typing import Any, Dict diff --git a/ms_agent/acp/server.py b/ms_agent/acp/server.py index d7f74e973..07217a0c6 100644 --- a/ms_agent/acp/server.py +++ b/ms_agent/acp/server.py @@ -1,3 +1,4 @@ +from __future__ import annotations import logging import sys from typing import Any diff --git a/ms_agent/acp/session_store.py b/ms_agent/acp/session_store.py index c659e2a09..e05c9f1c7 100644 --- a/ms_agent/acp/session_store.py +++ b/ms_agent/acp/session_store.py @@ -1,3 +1,4 @@ +from __future__ import annotations import asyncio import os import uuid diff --git a/ms_agent/agent/agent.yaml b/ms_agent/agent/agent.yaml index a8a65d440..aeca4dcda 100644 --- a/ms_agent/agent/agent.yaml +++ b/ms_agent/agent/agent.yaml @@ -56,6 +56,21 @@ callbacks: - input_callback tools: + file_system: + mcp: false + include: + - write_file + - read_file + - edit_file + - grep + - glob + code_executor: + mcp: false + implementation: python_env + include: + - shell_executor + - python_executor + - notebook_executor help: | A commonly use config, try whatever you want! diff --git a/ms_agent/agent/base.py b/ms_agent/agent/base.py index 0c6ee3fc3..790536f98 100644 --- a/ms_agent/agent/base.py +++ b/ms_agent/agent/base.py @@ -1,3 +1,4 @@ +from __future__ import annotations # Copyright (c) ModelScope Contributors. All rights reserved. import os from abc import ABC, abstractmethod diff --git a/ms_agent/agent/llm_agent.py b/ms_agent/agent/llm_agent.py index 44e34e250..7ba210378 100644 --- a/ms_agent/agent/llm_agent.py +++ b/ms_agent/agent/llm_agent.py @@ -1,3 +1,4 @@ +from __future__ import annotations # Copyright (c) ModelScope Contributors. All rights reserved. import asyncio import importlib @@ -15,7 +16,6 @@ from ms_agent.agent.runtime import Runtime from ms_agent.callbacks import Callback, callbacks_mapping -from ms_agent.knowledge_search import SirchmunkSearch from ms_agent.llm.llm import LLM from ms_agent.llm.utils import Message, ToolResult from ms_agent.memory import Memory, get_memory_meta_safe, memory_mapping diff --git a/ms_agent/callbacks/repetition_guard.py b/ms_agent/callbacks/repetition_guard.py new file mode 100644 index 000000000..c6ca65b1e --- /dev/null +++ b/ms_agent/callbacks/repetition_guard.py @@ -0,0 +1,189 @@ +"""Callback that detects stuck loops and injects strategy-variation prompts. + +When the agent repeats the same tool call (same name + same arguments) multiple +times without making progress, this callback injects a user message encouraging +a different approach. After ``max_warnings`` injections, the callback forces +the agent to stop to avoid wasting remaining API budget. +""" + +from __future__ import annotations + +import hashlib +import json +from collections import Counter +from dataclasses import dataclass +from typing import List + +from omegaconf import DictConfig + +from ms_agent.agent.runtime import Runtime +from ms_agent.callbacks.base import Callback +from ms_agent.llm.utils import Message +from ms_agent.utils.logger import get_logger + +logger = get_logger() + +_DEFAULT_THRESHOLD = 3 +_DEFAULT_LOOKBACK = 8 +_DEFAULT_MAX_WARNINGS = 2 + +_STRATEGY_PROMPT = ( + "IMPORTANT: You have been repeating the same approach multiple times " + "without making progress. The exact same tool call or command has been " + "attempted {count} times with the same arguments.\n\n" + "Repeated action: {description}\n\n" + "Please try a fundamentally different strategy. Consider:\n" + "1. Check whether the prerequisites for your approach are actually met.\n" + "2. Break the problem into smaller, verifiable sub-steps.\n" + "3. Use a different tool or a different command to achieve the same goal.\n" + "4. Read error messages carefully — they often suggest specific fixes.\n" + "5. If a dependency is missing, install it before retrying.\n\n" + "Do NOT repeat the same command again." +) + +_FORCE_STOP_PROMPT = ( + "You have been warned multiple times about repeating the same approach, " + "but continue to retry the same failing action. " + "To conserve resources, execution is being stopped. " + "Please review the errors above and formulate a new plan before " + "the next attempt." +) + + +@dataclass(frozen=True) +class _Repetition: + key: str + tool_name: str + count: int + description: str + + +def _args_hash(arguments: str) -> str: + try: + parsed = json.loads(arguments) if isinstance(arguments, str) else arguments + except (json.JSONDecodeError, TypeError): + parsed = arguments + canonical = json.dumps(parsed, sort_keys=True, ensure_ascii=False) + return hashlib.md5(canonical.encode("utf-8")).hexdigest()[:12] + + +class RepetitionGuardCallback(Callback): + """Detects stuck loops and injects strategy-variation prompts.""" + + def __init__(self, config: DictConfig) -> None: + super().__init__(config) + guard_cfg = getattr(config, "repetition_guard", None) or {} + self.threshold: int = int(guard_cfg.get("threshold", _DEFAULT_THRESHOLD)) + self.lookback: int = int(guard_cfg.get("lookback_rounds", _DEFAULT_LOOKBACK)) + self.max_warnings: int = int(guard_cfg.get("max_warnings", _DEFAULT_MAX_WARNINGS)) + self._warnings_given: int = 0 + self._warned_keys: set[str] = set() + + async def after_tool_call(self, runtime: Runtime, messages: List[Message]) -> None: + if self._warnings_given >= self.max_warnings: + logger.info( + "[RepetitionGuard] Max warnings (%d) reached — forcing stop.", + self.max_warnings, + ) + messages.append(Message(role="user", content=_FORCE_STOP_PROMPT)) + runtime.should_stop = True + return + + recent = _extract_recent_tool_calls(messages, self.lookback) + repetition = _detect_repetition(recent, self.threshold) + + if repetition is None: + return + if repetition.key in self._warned_keys: + return + + self._warned_keys.add(repetition.key) + self._warnings_given += 1 + logger.info( + "[RepetitionGuard] Stuck loop detected (%s, %dx). Injecting strategy prompt. " + "Warning %d/%d.", + repetition.tool_name, + repetition.count, + self._warnings_given, + self.max_warnings, + ) + runtime.should_stop = False + prompt = _STRATEGY_PROMPT.format( + count=repetition.count, + description=repetition.description, + ) + messages.append(Message(role="user", content=prompt)) + + +def _extract_recent_tool_calls( + messages: List[Message], + lookback: int, +) -> list[tuple[str, str, str]]: + """Return ``(tool_name, args_hash, description)`` for recent rounds. + + Walks backwards through *messages* collecting assistant tool-call entries. + Stops after scanning *lookback* assistant-with-tool-calls messages. + """ + calls: list[tuple[str, str, str]] = [] + rounds_seen = 0 + + for msg in reversed(messages): + if rounds_seen >= lookback: + break + if msg.role == "assistant" and msg.tool_calls: + rounds_seen += 1 + for tc in msg.tool_calls: + name = tc.get("tool_name", "") + raw_args = tc.get("arguments", "{}") + ah = _args_hash(raw_args) + desc = _summarize_call(name, raw_args) + calls.append((name, ah, desc)) + + return calls + + +def _detect_repetition( + calls: list[tuple[str, str, str]], + threshold: int, +) -> _Repetition | None: + if not calls: + return None + + key_counts: Counter[str] = Counter() + key_to_info: dict[str, tuple[str, str]] = {} + + for name, ah, desc in calls: + key = f"{name}:{ah}" + key_counts[key] += 1 + if key not in key_to_info: + key_to_info[key] = (name, desc) + + most_common_key, count = key_counts.most_common(1)[0] + if count < threshold: + return None + + tool_name, description = key_to_info[most_common_key] + return _Repetition( + key=most_common_key, + tool_name=tool_name, + count=count, + description=description, + ) + + +def _summarize_call(tool_name: str, raw_args: str) -> str: + try: + args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args + except (json.JSONDecodeError, TypeError): + args = raw_args + + if isinstance(args, dict) and "command" in args: + cmd = str(args["command"]) + if len(cmd) > 120: + cmd = cmd[:120] + "..." + return f"{tool_name}(command={cmd})" + + summary = json.dumps(args, ensure_ascii=False) + if len(summary) > 120: + summary = summary[:120] + "..." + return f"{tool_name}({summary})" diff --git a/ms_agent/callbacks/utils.py b/ms_agent/callbacks/utils.py index 6a58286f4..2571c6047 100644 --- a/ms_agent/callbacks/utils.py +++ b/ms_agent/callbacks/utils.py @@ -1,4 +1,8 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from ms_agent.callbacks.input_callback import InputCallback +from ms_agent.callbacks.repetition_guard import RepetitionGuardCallback -callbacks_mapping = {'input_callback': InputCallback} +callbacks_mapping = { + 'input_callback': InputCallback, + 'repetition_guard': RepetitionGuardCallback, +} diff --git a/ms_agent/capabilities/async_task.py b/ms_agent/capabilities/async_task.py index 74ac90fbe..f8eee2232 100644 --- a/ms_agent/capabilities/async_task.py +++ b/ms_agent/capabilities/async_task.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. import asyncio import logging diff --git a/ms_agent/capabilities/descriptor.py b/ms_agent/capabilities/descriptor.py index 3df1bab74..17f6651b9 100644 --- a/ms_agent/capabilities/descriptor.py +++ b/ms_agent/capabilities/descriptor.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. import os import shutil diff --git a/ms_agent/capabilities/mcp_server.py b/ms_agent/capabilities/mcp_server.py index d8891edbd..8bef4ef4d 100644 --- a/ms_agent/capabilities/mcp_server.py +++ b/ms_agent/capabilities/mcp_server.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. import argparse import json diff --git a/ms_agent/capabilities/registry.py b/ms_agent/capabilities/registry.py index 75c8f8609..c4f7eb52e 100644 --- a/ms_agent/capabilities/registry.py +++ b/ms_agent/capabilities/registry.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. import logging from typing import Any, Awaitable, Callable diff --git a/ms_agent/capabilities/wrappers/agent_delegate.py b/ms_agent/capabilities/wrappers/agent_delegate.py index 75376dc33..a158ee107 100644 --- a/ms_agent/capabilities/wrappers/agent_delegate.py +++ b/ms_agent/capabilities/wrappers/agent_delegate.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. import contextlib import logging diff --git a/ms_agent/capabilities/wrappers/code_genesis.py b/ms_agent/capabilities/wrappers/code_genesis.py index 596e96970..02d67eba4 100644 --- a/ms_agent/capabilities/wrappers/code_genesis.py +++ b/ms_agent/capabilities/wrappers/code_genesis.py @@ -1,3 +1,4 @@ +from __future__ import annotations # Copyright (c) ModelScope Contributors. All rights reserved. import asyncio import os diff --git a/ms_agent/capabilities/wrappers/deep_research.py b/ms_agent/capabilities/wrappers/deep_research.py index 50bf13061..d1a24e5cf 100644 --- a/ms_agent/capabilities/wrappers/deep_research.py +++ b/ms_agent/capabilities/wrappers/deep_research.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. import asyncio import os diff --git a/ms_agent/capabilities/wrappers/doc_research.py b/ms_agent/capabilities/wrappers/doc_research.py index 00c2c016a..09231f204 100644 --- a/ms_agent/capabilities/wrappers/doc_research.py +++ b/ms_agent/capabilities/wrappers/doc_research.py @@ -1,3 +1,4 @@ +from __future__ import annotations # Copyright (c) ModelScope Contributors. All rights reserved. import asyncio import logging diff --git a/ms_agent/capabilities/wrappers/filesystem.py b/ms_agent/capabilities/wrappers/filesystem.py index 5becb5615..a11a99edf 100644 --- a/ms_agent/capabilities/wrappers/filesystem.py +++ b/ms_agent/capabilities/wrappers/filesystem.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. import os from typing import Any diff --git a/ms_agent/capabilities/wrappers/fin_research.py b/ms_agent/capabilities/wrappers/fin_research.py index 4e860fb75..d1902b387 100644 --- a/ms_agent/capabilities/wrappers/fin_research.py +++ b/ms_agent/capabilities/wrappers/fin_research.py @@ -1,3 +1,4 @@ +from __future__ import annotations # Copyright (c) ModelScope Contributors. All rights reserved. import asyncio import os diff --git a/ms_agent/capabilities/wrappers/singularity_cinema.py b/ms_agent/capabilities/wrappers/singularity_cinema.py index 469867b01..5cfbca865 100644 --- a/ms_agent/capabilities/wrappers/singularity_cinema.py +++ b/ms_agent/capabilities/wrappers/singularity_cinema.py @@ -1,3 +1,4 @@ +from __future__ import annotations # Copyright (c) ModelScope Contributors. All rights reserved. import asyncio import os diff --git a/ms_agent/command/builtin/context_cmds.py b/ms_agent/command/builtin/context_cmds.py index 0fd7a2a6c..0fe97c419 100644 --- a/ms_agent/command/builtin/context_cmds.py +++ b/ms_agent/command/builtin/context_cmds.py @@ -1,3 +1,4 @@ +from __future__ import annotations from ms_agent.command.router import CommandRouter from ms_agent.command.types import ( CommandContext, diff --git a/ms_agent/command/types.py b/ms_agent/command/types.py index fed234724..d9261c2dc 100644 --- a/ms_agent/command/types.py +++ b/ms_agent/command/types.py @@ -1,3 +1,4 @@ +from __future__ import annotations from dataclasses import dataclass, field from enum import Enum from typing import Any, Awaitable, Callable, Optional diff --git a/ms_agent/llm/anthropic_llm.py b/ms_agent/llm/anthropic_llm.py index eafe965bf..978d13306 100644 --- a/ms_agent/llm/anthropic_llm.py +++ b/ms_agent/llm/anthropic_llm.py @@ -390,7 +390,7 @@ def _call_llm(self, else: return self.client.messages.create(**params) - @retry(max_attempts=LLM.retry_count, delay=1.0) + @retry(max_attempts=LLM.retry_count, delay=3.0) def generate(self, messages: List[Message], tools: Optional[List[Tool]] = None, diff --git a/ms_agent/llm/openai_llm.py b/ms_agent/llm/openai_llm.py index d9a9ef927..6de24bfd8 100644 --- a/ms_agent/llm/openai_llm.py +++ b/ms_agent/llm/openai_llm.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. import httpx import inspect @@ -211,7 +213,7 @@ def format_tools(self, tools = None return tools - @retry(max_attempts=LLM.retry_count, delay=1.0) + @retry(max_attempts=LLM.retry_count, delay=3.0) def generate(self, messages: List[Message], tools: Optional[List[Tool]] = None, diff --git a/ms_agent/memory/diversity.py b/ms_agent/memory/diversity.py index 1d3e5af8d..e1e2a37f6 100644 --- a/ms_agent/memory/diversity.py +++ b/ms_agent/memory/diversity.py @@ -1,5 +1,7 @@ +import asyncio import re from copy import deepcopy +from omegaconf import DictConfig from typing import List from omegaconf import DictConfig diff --git a/ms_agent/memory/memory_manager.py b/ms_agent/memory/memory_manager.py index 4d31e2f92..e8589a3ab 100644 --- a/ms_agent/memory/memory_manager.py +++ b/ms_agent/memory/memory_manager.py @@ -1,4 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +from omegaconf import DictConfig from typing import Dict from omegaconf import DictConfig, OmegaConf diff --git a/ms_agent/memory/utils.py b/ms_agent/memory/utils.py index 0c939c4fe..8f86d34df 100644 --- a/ms_agent/memory/utils.py +++ b/ms_agent/memory/utils.py @@ -1,4 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + from omegaconf import DictConfig, OmegaConf from .condenser.code_condenser import CodeCondenser diff --git a/ms_agent/permission/config.py b/ms_agent/permission/config.py index 5394fc55d..243ea14d4 100644 --- a/ms_agent/permission/config.py +++ b/ms_agent/permission/config.py @@ -116,9 +116,13 @@ def from_dict(cls, d: dict[str, Any], project_root: str | None = None) -> Permis whitelist = tuple(d.get('whitelist', ())) ask_rules = tuple(d.get('ask_rules', ())) user_blacklist = tuple(d.get('blacklist', ())) - blacklist = _DEFAULT_BLACKLIST + tuple( - p for p in user_blacklist if p not in _DEFAULT_BLACKLIST - ) + if d.get('no_default_blacklist', False): + # Container / sandbox mode: strip built-in network blocks + blacklist = user_blacklist + else: + blacklist = _DEFAULT_BLACKLIST + tuple( + p for p in user_blacklist if p not in _DEFAULT_BLACKLIST + ) safety_raw = d.get('safety_rules', {}) # Merge directory configs from top level into safety config diff --git a/ms_agent/permission/path_extractors.py b/ms_agent/permission/path_extractors.py index bd1f3c9b9..fb64447a4 100644 --- a/ms_agent/permission/path_extractors.py +++ b/ms_agent/permission/path_extractors.py @@ -14,10 +14,10 @@ import os import re from dataclasses import dataclass -from typing import Callable, Literal +from typing import Callable, Literal, Optional CommandExtractor = Callable[[list[str]], list[str]] -CommandValidator = Callable[[list[str]], str | None] +CommandValidator = Callable[[list[str]], Optional[str]] @dataclass(frozen=True) diff --git a/ms_agent/permission/path_validator.py b/ms_agent/permission/path_validator.py index f2fc4a7b5..4a55f48e5 100644 --- a/ms_agent/permission/path_validator.py +++ b/ms_agent/permission/path_validator.py @@ -43,13 +43,17 @@ def _expand_tilde(path: str, home_dir: str) -> tuple[str, str | None]: return path, f'Unsupported tilde expansion: {path}' +_SHELL_VAR = re.compile(r'(? str | None: - if '$' in path: + if _SHELL_VAR.search(path): return f'Path contains shell variable expansion: {path}' if '%' in path: return f'Path contains Windows variable expansion: {path}' - if path.startswith('='): - return f'Path starts with = (Zsh expansion): {path}' + # Only flag Zsh =(…) process substitution, not bare = or =value + if path.startswith('=('): + return f'Path starts with =( (Zsh process substitution): {path}' return None @@ -114,7 +118,10 @@ def validate_path( ) if _has_glob(path): - if op_type in ('write', 'create'): + # 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: return PathValidationResult( allowed=False, resolved_path=path, action='deny', reason=f'Glob patterns not allowed in {op_type} operations: {path}', diff --git a/ms_agent/permission/safety.py b/ms_agent/permission/safety.py index 81a06206a..b9bb532b6 100644 --- a/ms_agent/permission/safety.py +++ b/ms_agent/permission/safety.py @@ -60,7 +60,16 @@ def check(self, tool_name: str, tool_args: dict[str, Any]) -> SafetyDecision: return self._check_file_path(tool_args.get('path', ''), 'write') if tool_name.endswith('---read_file'): - return self._check_file_path(tool_args.get('path', ''), 'read') + paths = tool_args.get('paths') + path = tool_args.get('path', '') + if isinstance(paths, list) and paths: + for p in paths: + if isinstance(p, str) and p.strip(): + result = self._check_file_path(p.strip(), 'read') + if result.action != 'allow': + return result + return SafetyDecision(action='allow', reason='All read paths validated') + return self._check_file_path(path, 'read') if tool_name.endswith('---grep') or tool_name.endswith('---glob'): return self._check_file_path(tool_args.get('path', '.'), 'read') diff --git a/ms_agent/permission/shell_validator.py b/ms_agent/permission/shell_validator.py index e4de48ac5..10e7f6df1 100644 --- a/ms_agent/permission/shell_validator.py +++ b/ms_agent/permission/shell_validator.py @@ -14,6 +14,7 @@ import re import shlex from dataclasses import dataclass +from pathlib import Path from typing import Literal, Sequence from .path_extractors import ( @@ -109,14 +110,22 @@ def check(self, command: str, *, _depth: int = 0) -> SafetyDecision: # 3. Split compound commands sub_commands = _split_compound(command) - # Track cd presence for cd+write detection - has_cd = False - has_write_or_create = False + # Track cwd through cd commands for accurate path validation + _current_cwd = self._workspace_root for sub_cmd in sub_commands: + # Strip comment lines (starting with #) before parsing + stripped = sub_cmd.strip() + if stripped.startswith('#'): + continue # Skip pure comment + try: tokens = shlex.split(sub_cmd) except ValueError: + # If parsing fails, check if it's a comment with unmatched quotes + if stripped.startswith('#'): + continue + # Otherwise, ask for confirmation (could be command injection) return SafetyDecision(action='ask', reason=f'Failed to parse command: {sub_cmd}', category='parse_failure') if not tokens: @@ -135,27 +144,23 @@ def check(self, command: str, *, _depth: int = 0) -> SafetyDecision: base_cmd = os.path.basename(tokens[0]) args = tokens[1:] + # Track cd: resolve target and update _current_cwd if base_cmd == 'cd': - has_cd = True + cd_entry = self._extractors.get('cd') + if cd_entry: + cd_targets = cd_entry.extractor(args) + if cd_targets: + target = cd_targets[0] + if os.path.isabs(target): + _current_cwd = str(Path(target).resolve()) + else: + _current_cwd = str((Path(_current_cwd) / target).resolve()) # 5. Command path extraction and validation - result = self._check_command(base_cmd, args, _depth=_depth) + result = self._check_command(base_cmd, args, _depth=_depth, cwd=_current_cwd) if result.action != 'allow': return result - entry = self._extractors.get(base_cmd) - if entry and entry.op_type in ('write', 'create'): - has_write_or_create = True - - # 6. cd + write/create compound → ask - if has_cd and has_write_or_create: - return SafetyDecision( - action='ask', - reason='Command combines cd with write/create operations — ' - 'path validation may not reflect runtime working directory', - category='cd_write_compound', - ) - return SafetyDecision(action='allow', reason='Shell command passed all checks') def _check_command_substitutions( @@ -184,6 +189,7 @@ def _check_command( args: list[str], *, _depth: int = 0, + cwd: str | None = None, ) -> SafetyDecision: entry = self._extractors.get(base_cmd) if entry is None: @@ -197,18 +203,18 @@ def _check_command( # sed special handling if base_cmd == 'sed': - return self._check_sed(args, entry) + return self._check_sed(args, entry, cwd=cwd) if base_cmd == 'find': - return self._check_find(args, entry, _depth=_depth) + return self._check_find(args, entry, _depth=_depth, cwd=cwd) paths = entry.extractor(args) if not paths: return SafetyDecision(action='allow', reason=f'{base_cmd}: no paths to validate') - return self._validate_paths(paths, entry.op_type, base_cmd) + return self._validate_paths(paths, entry.op_type, base_cmd, cwd=cwd) - def _check_sed(self, args: list[str], entry: ExtractorEntry) -> SafetyDecision: + def _check_sed(self, args: list[str], entry: ExtractorEntry, *, cwd: str | None = None) -> SafetyDecision: op_type = entry.op_type if is_sed_read_only(args): op_type = 'read' @@ -218,13 +224,13 @@ def _check_sed(self, args: list[str], entry: ExtractorEntry) -> SafetyDecision: for expr in expressions: result = check_sed_expression_safety(expr) if not result.safe: - return SafetyDecision(action='deny', reason=result.reason) + return SafetyDecision(action='ask', reason=result.reason) paths = entry.extractor(args) if not paths: return SafetyDecision(action='allow', reason='sed: no file paths') - return self._validate_paths(paths, op_type, 'sed') + return self._validate_paths(paths, op_type, 'sed', cwd=cwd) def _check_find( self, @@ -232,6 +238,7 @@ def _check_find( entry: ExtractorEntry, *, _depth: int, + cwd: str | None = None, ) -> SafetyDecision: for exec_cmd in extract_find_exec_commands(args): result = self.check(exec_cmd, _depth=_depth + 1) @@ -248,7 +255,7 @@ def _check_find( if not paths: return SafetyDecision(action='allow', reason='find: no paths to validate') - return self._validate_paths(paths, op_type, 'find') + return self._validate_paths(paths, op_type, 'find', cwd=cwd) @staticmethod def _collect_sed_expressions(args: list[str]) -> list[str]: @@ -282,8 +289,10 @@ def _validate_paths( paths: list[str], op_type: Literal['read', 'write', 'create'], cmd_name: str, + *, + cwd: str | None = None, ) -> SafetyDecision: - cwd = self._workspace_root + effective_cwd = cwd or self._workspace_root for path in paths: # Dangerous removal check for rm/rmdir @@ -293,7 +302,7 @@ def _validate_paths( reason=f'Dangerous removal path: {path}', ) - result = validate_path(path, cwd, self._allowed_dirs, op_type, read_only_dirs=self._read_only_dirs) + result = validate_path(path, effective_cwd, self._allowed_dirs, op_type, read_only_dirs=self._read_only_dirs) if not result.allowed: return SafetyDecision(action=result.action, reason=result.reason, category=result.category) diff --git a/ms_agent/project/manager.py b/ms_agent/project/manager.py index 3dcbec892..d1a6700ce 100644 --- a/ms_agent/project/manager.py +++ b/ms_agent/project/manager.py @@ -1,3 +1,4 @@ +from __future__ import annotations import os import shutil from dataclasses import asdict, replace diff --git a/ms_agent/skill/loader.py b/ms_agent/skill/loader.py index fcadb37ff..918587efc 100644 --- a/ms_agent/skill/loader.py +++ b/ms_agent/skill/loader.py @@ -1,3 +1,4 @@ +from __future__ import annotations # Copyright (c) ModelScope Contributors. All rights reserved. import os import re diff --git a/ms_agent/tools/a2a_agent_tool.py b/ms_agent/tools/a2a_agent_tool.py index 5dd6270e7..cda6290e2 100644 --- a/ms_agent/tools/a2a_agent_tool.py +++ b/ms_agent/tools/a2a_agent_tool.py @@ -1,3 +1,4 @@ +from __future__ import annotations from typing import Any, Dict, List from ms_agent.llm.utils import Tool diff --git a/ms_agent/tools/acp_agent_tool.py b/ms_agent/tools/acp_agent_tool.py index cd77e82c9..a72e0a3b2 100644 --- a/ms_agent/tools/acp_agent_tool.py +++ b/ms_agent/tools/acp_agent_tool.py @@ -1,3 +1,4 @@ +from __future__ import annotations from typing import Any, Dict, List from ms_agent.acp.client import ACPClientManager diff --git a/ms_agent/tools/agent_tool.py b/ms_agent/tools/agent_tool.py index 51fbfd121..fe4d9cb52 100644 --- a/ms_agent/tools/agent_tool.py +++ b/ms_agent/tools/agent_tool.py @@ -20,7 +20,7 @@ from ms_agent.utils.stats import (append_stats, build_timing_record, get_stats_path, monotonic, now_iso, summarize_usage) -from ms_agent.utils.stream_writer import SubAgentStreamWriter +from ms_agent.utils.stream_writer import SubAgentStreamWriter, _msg_to_dict logger = get_logger() @@ -1289,10 +1289,9 @@ def _save_transcript(self, messages: Any, path = os.path.join(subagents_dir, f'{agent_tag}.jsonl') with open(path, 'w', encoding='utf-8') as f: for msg in messages: - if hasattr(msg, 'to_dict'): - f.write( - json.dumps(msg.to_dict(), ensure_ascii=False) - + '\n') + f.write( + json.dumps(_msg_to_dict(msg), ensure_ascii=False) + + '\n') except Exception as exc: logger.warning( f'Failed to save sub-agent transcript for {agent_tag}: {exc}') diff --git a/ms_agent/tools/docling/chunker.py b/ms_agent/tools/docling/chunker.py index 5a7c894c6..b054b299d 100644 --- a/ms_agent/tools/docling/chunker.py +++ b/ms_agent/tools/docling/chunker.py @@ -1,3 +1,4 @@ +from __future__ import annotations from docling_core.transforms.chunker import BaseChunk, DocChunk from docling_core.transforms.chunker.hierarchical_chunker import ( ChunkingDocSerializer, ChunkingSerializerProvider) diff --git a/ms_agent/tools/docling/doc_loader.py b/ms_agent/tools/docling/doc_loader.py index daac553e5..c36b92321 100644 --- a/ms_agent/tools/docling/doc_loader.py +++ b/ms_agent/tools/docling/doc_loader.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # flake8: noqa # yapf: disable import ast diff --git a/ms_agent/tools/filesystem_tool.py b/ms_agent/tools/filesystem_tool.py index 5640317c6..a130a15cc 100644 --- a/ms_agent/tools/filesystem_tool.py +++ b/ms_agent/tools/filesystem_tool.py @@ -1,3 +1,4 @@ +from __future__ import annotations # Copyright (c) Alibaba, Inc. and its affiliates. import asyncio import base64 @@ -485,11 +486,11 @@ async def _grep_rg_file( if glob_pat: args.extend(['--glob', glob_pat]) if output_mode == 'files_with_matches': - args.extend(['-l', pattern, str(file_path)]) + args.extend(['-l', '-e', pattern, str(file_path)]) elif output_mode == 'count': - args.extend(['-c', pattern, str(file_path)]) + args.extend(['-c', '-e', pattern, str(file_path)]) else: - args.extend(['-n', pattern, str(file_path)]) + args.extend(['-n', '-e', pattern, str(file_path)]) proc = await asyncio.create_subprocess_exec( *args, stdout=asyncio.subprocess.PIPE, @@ -522,11 +523,11 @@ async def _grep_rg_dir( if glob_pat: args.extend(['--glob', glob_pat]) if output_mode == 'files_with_matches': - args.extend(['-l', pattern, str(root)]) + args.extend(['-l', '-e', pattern, str(root)]) elif output_mode == 'count': - args.extend(['--count-matches', pattern, str(root)]) + args.extend(['--count-matches', '-e', pattern, str(root)]) else: - args.extend(['-n', pattern, str(root)]) + args.extend(['-n', '-e', pattern, str(root)]) proc = await asyncio.create_subprocess_exec( *args, stdout=asyncio.subprocess.PIPE, @@ -975,6 +976,8 @@ async def edit_file(self, Success or error message. """ try: + if not path: + return 'Error: `path` is required.' if old_string is None: return 'Error: `old_string` is required.' if new_string is None: diff --git a/ms_agent/tools/mcp_client.py b/ms_agent/tools/mcp_client.py index 0bd6e95a3..7f2da5ee6 100644 --- a/ms_agent/tools/mcp_client.py +++ b/ms_agent/tools/mcp_client.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. import copy import os diff --git a/ms_agent/tools/search/sirchmunk_search.py b/ms_agent/tools/search/sirchmunk_search.py index cd8aaacf5..70c609f48 100644 --- a/ms_agent/tools/search/sirchmunk_search.py +++ b/ms_agent/tools/search/sirchmunk_search.py @@ -1,6 +1,7 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. """Sirchmunk backend for the ``localsearch`` tool. - Configuration lives under ``tools.localsearch`` (same namespace as other tools). Legacy top-level ``knowledge_search`` is still accepted for backward compatibility. """ diff --git a/ms_agent/tools/tool_manager.py b/ms_agent/tools/tool_manager.py index 41175fb8f..1ebce8cad 100644 --- a/ms_agent/tools/tool_manager.py +++ b/ms_agent/tools/tool_manager.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. import asyncio import importlib @@ -223,6 +225,17 @@ def __init__(self, # Find cls which base class is `ToolBase` if issubclass(cls, ToolBase) and cls.__module__ == _plugin: self.register_tool(cls(self.config)) + # Used temporarily during async initialization; the actual client is managed in self.servers + self.mcp_client = mcp_client + self.mcp_config = mcp_config + self.servers = None + self._managed_client = mcp_client is None + + # Initialize concurrency limiter (will be set in connect) + self._concurrent_limiter = None + self._init_lock = None + self._sync_lock = asyncio.Lock() + self._tool_index = {} self._mcp_index_keys: set[str] = set() self._skip_mcp_reindex = False @@ -242,17 +255,6 @@ def ensure_plugin_agent_tools(self, registry) -> None: self.extra_tools.append(agent_tool) agent_tool.sync_plugin_agents(registry) - # Used temporarily during async initialization; the actual client is managed in self.servers - self.mcp_client = mcp_client - self.mcp_config = mcp_config - self.servers = None - self._managed_client = mcp_client is None - - # Initialize concurrency limiter (will be set in connect) - self._concurrent_limiter = None - self._init_lock = None - self._sync_lock = asyncio.Lock() - def register_tool(self, tool: ToolBase): self.extra_tools.append(tool) @@ -268,8 +270,18 @@ async def connect(self): elif MCPClient is not None: self.servers = MCPClient(self.mcp_config, self.config) await self.servers.connect() + elif MCPClient is not None and not MCP_AVAILABLE: + logger.warning_once( + 'mcp package not installed; MCP tools disabled for this run' + ) for tool in self.extra_tools: - await tool.connect() + try: + await tool.connect() + except Exception as e: + logger.warning( + f'Tool {getattr(tool, "name", type(tool).__name__)} ' + f'failed to connect: {e}; disabling.' + ) if not self._skip_mcp_reindex: await self.reindex_tool() @@ -534,6 +546,16 @@ async def single_call_tool(self, tool_info: ToolCall): tool_args=call_args), timeout=wait_sec) + # Truncate excessively long tool outputs to prevent context window explosion + max_len = int(os.getenv('MAX_TOOL_OUTPUT_LEN', 20000)) + if isinstance(response, str) and len(response) > max_len: + half = max_len // 2 + trunc_notice = ( + f"\n\n...[SYSTEM: Output truncated, {len(response)} chars total, " + f"showing first and last {half} chars]...\n\n" + ) + response = response[:half] + trunc_notice + response[-half:] + if (self.mcp_success_handler is not None and tool_ins is self.servers): await self.mcp_success_handler(server_name) @@ -581,10 +603,26 @@ async def single_call_tool(self, tool_info: ToolCall): tool_info.get('tool_name', ''), self.TOOL_SPLITER), exc=asyncio.TimeoutError(timeout_msg), ) - return timeout_msg + return json.dumps({ + 'success': False, + 'error': 'timeout', + 'tool_name': tn, + 'message': timeout_msg, + 'recovery_hint': ( + 'The tool took too long. Try: (1) add "timeout" field with a larger value in seconds ' + f'(max {self.tool_call_timeout_max:.0f}s), (2) break the task into smaller steps, ' + 'or (3) simplify the command/input.' + ), + }, ensure_ascii=False) except Exception as e: import traceback - logger.warning(traceback.format_exc()) + tb_str = traceback.format_exc() + logger.warning(tb_str) + exc_type_name = type(e).__name__ + exc_msg = str(e) or '(no error message)' + tn = tool_info.get('tool_name', '(unknown)') + tb_lines = tb_str.strip().splitlines() + tb_tail = '\n'.join(tb_lines[-6:]) if len(tb_lines) > 6 else '\n'.join(tb_lines) if tool_ins is not None and tool_ins is self.servers: await self._report_mcp_failure( server_name, @@ -594,7 +632,20 @@ async def single_call_tool(self, tool_info: ToolCall): tool_info.get('tool_name', ''), self.TOOL_SPLITER), exc=e, ) - return f'Tool calling failed: {brief_info}, details: {str(e)}' + return json.dumps({ + 'success': False, + 'error': exc_type_name, + 'tool_name': tn, + 'message': f'{exc_type_name}: {exc_msg}', + 'call_info': brief_info, + 'traceback_tail': tb_tail, + 'recovery_hint': ( + f'The tool "{tn}" raised {exc_type_name}. ' + 'Review the error and traceback above, check your input parameters, ' + 'verify prerequisites (files, dependencies, running services), ' + 'and retry with corrected arguments or a different approach.' + ), + }, ensure_ascii=False, default=str) async def parallel_call_tool(self, tool_list: List[ToolCall]): tasks = [self.single_call_tool(tool) for tool in tool_list] diff --git a/ms_agent/utils/constants.py b/ms_agent/utils/constants.py index cbd37e271..3def00c6b 100644 --- a/ms_agent/utils/constants.py +++ b/ms_agent/utils/constants.py @@ -35,7 +35,7 @@ DEFAULT_OUTPUT_WRAPPER = ['', ''] -DEFAULT_RETRY_COUNT = 3 +DEFAULT_RETRY_COUNT = 5 DEFAULT_SEARCH_LIMIT = 3 diff --git a/ms_agent/utils/parser_utils.py b/ms_agent/utils/parser_utils.py index af22034c6..fc6d7c423 100644 --- a/ms_agent/utils/parser_utils.py +++ b/ms_agent/utils/parser_utils.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import json import os import re diff --git a/ms_agent/utils/snapshot.py b/ms_agent/utils/snapshot.py index d7362405a..fd9637534 100644 --- a/ms_agent/utils/snapshot.py +++ b/ms_agent/utils/snapshot.py @@ -1,3 +1,4 @@ +from __future__ import annotations # Copyright (c) Alibaba, Inc. and its affiliates. """ Lightweight snapshot utility for ms-agent output directories. @@ -165,7 +166,8 @@ def take_snapshot(output_dir: str, logger.warning_once('[snapshot] git not found — snapshots disabled.') return None except subprocess.CalledProcessError as e: - logger.warning(f'[snapshot] git error: {e.stderr.strip()}') + stderr = (e.stderr or '').strip() + logger.warning(f'[snapshot] git error: {stderr or e}') return None except Exception as e: logger.warning(f'[snapshot] unexpected error: {e}') @@ -229,6 +231,13 @@ def restore_snapshot(output_dir: str, commit_hash: str) -> tuple[bool, int]: meta = _load_meta(output_dir) message_count = meta.get(commit_hash, {}).get('message_count', 0) return True, message_count + except FileNotFoundError: + logger.warning_once('[snapshot] git not found — snapshots disabled.') + return False, 0 except subprocess.CalledProcessError as e: - logger.warning(f'[snapshot] restore failed: {e.stderr.strip()}') + stderr = (e.stderr or '').strip() + logger.warning(f'[snapshot] restore failed: {stderr or e}') + return False, 0 + except Exception as e: + logger.warning(f'[snapshot] unexpected restore error: {e}') return False, 0 diff --git a/ms_agent/utils/utils.py b/ms_agent/utils/utils.py index 91d9c6d42..93605819b 100644 --- a/ms_agent/utils/utils.py +++ b/ms_agent/utils/utils.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) ModelScope Contributors. All rights reserved. import base64 import glob diff --git a/ms_agent/utils/workspace_policy.py b/ms_agent/utils/workspace_policy.py new file mode 100644 index 000000000..c11edf4db --- /dev/null +++ b/ms_agent/utils/workspace_policy.py @@ -0,0 +1,207 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Workspace path policy: allow-roots (default output_dir) and optional deny globs.""" + +from __future__ import annotations + +import fnmatch +import os +import re +from pathlib import Path +from typing import Iterable, Sequence + + +class WorkspacePolicyError(ValueError): + """Raised when a path or command violates workspace policy.""" + + +class WorkspacePolicyKernel: + """Resolve user paths under allowed workspace roots; optional shell read-only rules.""" + + def __init__( + self, + output_dir: Path | str, + *, + extra_allow_roots: Sequence[str | Path] | None = None, + deny_globs: Sequence[str] | None = None, + shell_default_mode: str = 'workspace_write', + shell_network_enabled: bool = False, + max_command_chars: int = 8192, + ) -> None: + self._output = Path(output_dir).expanduser().resolve() + self._roots: list[Path] = [self._output] + if extra_allow_roots: + for r in extra_allow_roots: + p = Path(r).expanduser().resolve() + if p not in self._roots: + self._roots.append(p) + if deny_globs is None or len(tuple(deny_globs)) == 0: + self._deny_globs: tuple[str, ...] = ('**/.git/**', ) + else: + self._deny_globs = tuple(deny_globs) + self.shell_default_mode = shell_default_mode + self.shell_network_enabled = shell_network_enabled + self.max_command_chars = max_command_chars + + @property + def workspace_root(self) -> Path: + return self._output + + @property + def allow_roots(self) -> tuple[Path, ...]: + return tuple(self._roots) + + @property + def deny_globs(self) -> tuple[str, ...]: + return self._deny_globs + + def resolve_under_roots(self, user_path: str | Path) -> Path: + """Resolve *user_path* to an absolute path that must lie under one allow root.""" + raw = Path(user_path).expanduser() + if raw.is_absolute(): + resolved = raw.resolve() + else: + resolved = (self._output / raw).resolve() + for root in self._roots: + try: + resolved.relative_to(root) + break + except ValueError: + continue + else: + raise WorkspacePolicyError( + f'Path is outside allowed workspace roots: {resolved}') + if self._is_denied(resolved): + raise WorkspacePolicyError( + f'Path matches a deny_globs pattern: {resolved}') + return resolved + + def _is_denied(self, path: Path) -> bool: + if not self._deny_globs: + return False + rel = None + try: + rel = path.relative_to(self._output) + except ValueError: + rel = path + rel_s = rel.as_posix() + for pat in self._deny_globs: + if fnmatch.fnmatch(rel_s, pat) or fnmatch.fnmatch(path.name, pat): + return True + if fnmatch.fnmatch(str(path), pat): + return True + return False + + def path_is_allowed(self, path: Path) -> bool: + path = path.expanduser().resolve() + for root in self._roots: + try: + path.relative_to(root) + break + except ValueError: + continue + else: + return False + return not self._is_denied(path) + + def assert_shell_command_allowed(self, command: str) -> None: + """Length and mode-based checks before executing shell.""" + if not command or not command.strip(): + raise WorkspacePolicyError('Empty shell command') + if len(command) > self.max_command_chars: + raise WorkspacePolicyError( + f'Shell command exceeds max length ({self.max_command_chars})') + + mode = self.shell_default_mode + if mode == 'read_only': + if _shell_looks_mutating_or_network(command, allow_network=False): + raise WorkspacePolicyError( + 'Shell is in read_only mode: mutating or network commands are not allowed' + ) + elif mode == 'workspace_write': + if not self.shell_network_enabled and _shell_looks_network( + command): + raise WorkspacePolicyError( + 'Network commands are disabled for shell (enable tools.code_executor.shell.network_enabled)' + ) + # future: explicit 'network' mode could allow curl etc. + + +def _shell_looks_network(command: str) -> bool: + lowered = command.lower() + tokens = ( + 'curl ', + 'wget ', + 'ssh ', + 'scp ', + 'rsync ', + 'ftp ', + 'nc ', + 'netcat ', + 'pip install', + 'pip3 install', + 'npm install', + 'yarn add', + 'pnpm add', + ) + return any(t in lowered for t in tokens) + + +def _shell_looks_mutating_or_network(command: str, *, + allow_network: bool) -> bool: + if not allow_network and _shell_looks_network(command): + return True + # redirection that creates/overwrites files + if re.search(r'[>]{1,2}\s*[^\s]', command): + return True + if re.search(r'\b(rm|rmdir|mv|cp|chmod|chown|chgrp|mkdir|touch|tee)\b', + command): + return True + return False + + +def iter_files_under( + root: Path, + *, + deny_globs: Iterable[str] = (), + max_files: int = 100_000, +) -> Iterable[Path]: + """Yield files under *root* (depth-first), skipping directories matching deny globs.""" + deny = tuple(deny_globs) + count = 0 + root = root.resolve() + + def dir_skipped(dirpath: Path) -> bool: + try: + rel = dirpath.relative_to(root).as_posix() + except ValueError: + return True + for pat in deny: + if fnmatch.fnmatch(rel, pat) or fnmatch.fnmatch(rel + '/', pat): + return True + parts = rel.split('/') + for i in range(len(parts)): + sub = '/'.join(parts[:i + 1]) + if fnmatch.fnmatch(sub, pat.rstrip('/')) or fnmatch.fnmatch( + sub + '/', pat): + return True + return False + + for dirpath, dirnames, filenames in os.walk( + root, topdown=True, followlinks=False): + dp = Path(dirpath) + if dir_skipped(dp): + dirnames[:] = [] + continue + # prune skipped subdirs + keep: list[str] = [] + for d in dirnames: + child = dp / d + if dir_skipped(child): + continue + keep.append(d) + dirnames[:] = keep + for name in filenames: + count += 1 + if count > max_files: + return + yield dp / name diff --git a/tests/permission/test_path_validator.py b/tests/permission/test_path_validator.py index f92a605b3..f81760066 100644 --- a/tests/permission/test_path_validator.py +++ b/tests/permission/test_path_validator.py @@ -58,7 +58,13 @@ def test_windows_variable_rejected(self): assert not r.allowed def test_zsh_equals_rejected(self): + # Bare =ls is allowed (not Zsh process substitution) r = validate_path('=ls', '/tmp', ['/tmp'], 'read') + assert r.allowed + + def test_zsh_process_substitution_rejected(self): + # =(…) is Zsh process substitution and should be rejected + r = validate_path('=(ls)', '/tmp', ['/tmp'], 'read') assert not r.allowed def test_glob_in_write_rejected(self): diff --git a/tests/permission/test_security_regression.py b/tests/permission/test_security_regression.py index ac28f4727..4179ab186 100644 --- a/tests/permission/test_security_regression.py +++ b/tests/permission/test_security_regression.py @@ -43,12 +43,12 @@ def test_redirect_to_etc(self, guard): assert r.action in ('deny', 'ask') def test_cd_plus_mv(self, guard, tmp_path): - """cd dir && mv a b → ask (cd + write compound)""" + """cd dir && mv a b → allow (cd target resolved, paths validated)""" r = guard.check( 'code_executor---shell_executor', {'command': f'cd {tmp_path} && mv {tmp_path}/a {tmp_path}/b'}, ) - assert r.action == 'ask' + assert r.action == 'allow' def test_rm_dollar_home(self, guard): """rm $HOME/.ssh/* → ask/deny (shell expansion in path)""" @@ -123,12 +123,12 @@ def test_mv_target_directory(self, guard, tmp_path): assert r.action == 'ask' def test_sed_write_expression(self, guard, tmp_path): - """sed -e 's/x/y/w /etc/passwd' file → deny (sed expression safety)""" + """sed -e 's/x/y/w /etc/passwd' file → ask (sed expression safety, relaxed from deny)""" r = guard.check( 'code_executor---shell_executor', {'command': f"sed -e 's/x/y/w /etc/passwd' {tmp_path}/file"}, ) - assert r.action == 'deny' + assert r.action == 'ask' class TestSensitivePathWrites: diff --git a/tests/permission/test_shell_validator.py b/tests/permission/test_shell_validator.py index 3e6fcedd6..2f214b3ab 100644 --- a/tests/permission/test_shell_validator.py +++ b/tests/permission/test_shell_validator.py @@ -62,7 +62,7 @@ def test_nohup_rm(self, validator): class TestCompoundCommands: def test_cd_plus_write(self, validator, tmp_path): r = validator.check(f'cd {tmp_path} && rm {tmp_path}/test.txt') - assert r.action == 'ask' + assert r.action == 'allow' # cd target resolved, paths validated against resolved cwd def test_multiple_safe(self, validator, tmp_path): r = validator.check(f'ls {tmp_path} && cat {tmp_path}/f') diff --git a/tests/test_repetition_guard.py b/tests/test_repetition_guard.py new file mode 100644 index 000000000..40eeb1816 --- /dev/null +++ b/tests/test_repetition_guard.py @@ -0,0 +1,260 @@ +"""Tests for RepetitionGuardCallback.""" + +from __future__ import annotations + +import json + +import pytest +from omegaconf import DictConfig + +from ms_agent.agent.runtime import Runtime +from ms_agent.callbacks.repetition_guard import ( + RepetitionGuardCallback, + _args_hash, + _detect_repetition, + _extract_recent_tool_calls, + _summarize_call, +) +from ms_agent.llm.utils import Message + + +def _make_config(**overrides) -> DictConfig: + guard = {"threshold": 3, "lookback_rounds": 8, "max_warnings": 2} + guard.update(overrides) + return DictConfig({"repetition_guard": guard}) + + +def _tool_call(tool_name: str, arguments: dict | str, call_id: str = "c1") -> dict: + args = json.dumps(arguments) if isinstance(arguments, dict) else arguments + return { + "id": call_id, + "type": "function", + "tool_name": tool_name, + "arguments": args, + } + + +def _assistant_with_tools(*tool_calls) -> Message: + return Message(role="assistant", content="Let me do a tool calling.", tool_calls=list(tool_calls)) + + +def _tool_result(call_id: str, content: str = "ok") -> Message: + return Message(role="tool", content=content, tool_call_id=call_id) + + +class TestArgsHash: + def test_identical_dicts(self): + a = _args_hash('{"command": "ls", "timeout": 30}') + b = _args_hash('{"timeout": 30, "command": "ls"}') + assert a == b + + def test_different_args(self): + a = _args_hash('{"command": "ls"}') + b = _args_hash('{"command": "pwd"}') + assert a != b + + def test_invalid_json(self): + h = _args_hash("not json") + assert isinstance(h, str) and len(h) == 12 + + +class TestExtractRecentToolCalls: + def test_empty_messages(self): + assert _extract_recent_tool_calls([], 8) == [] + + def test_collects_from_assistant_messages(self): + msgs = [ + Message(role="system", content="sys"), + Message(role="user", content="do something"), + _assistant_with_tools(_tool_call("shell", {"command": "ls"})), + _tool_result("c1", "file1 file2"), + _assistant_with_tools(_tool_call("shell", {"command": "pwd"})), + _tool_result("c1", "/app"), + ] + calls = _extract_recent_tool_calls(msgs, 8) + assert len(calls) == 2 + assert calls[0][0] == "shell" + assert calls[1][0] == "shell" + + def test_respects_lookback(self): + msgs = [Message(role="system", content="sys")] + for i in range(10): + msgs.append(_assistant_with_tools(_tool_call("shell", {"command": f"cmd{i}"}, f"c{i}"))) + msgs.append(_tool_result(f"c{i}")) + + calls = _extract_recent_tool_calls(msgs, 3) + assert len(calls) == 3 + + +class TestDetectRepetition: + def test_no_repetition(self): + calls = [ + ("shell", "hash1", "shell(ls)"), + ("shell", "hash2", "shell(pwd)"), + ("shell", "hash3", "shell(cat)"), + ] + assert _detect_repetition(calls, 3) is None + + def test_detects_at_threshold(self): + calls = [ + ("shell", "abc123", "shell(ls)"), + ("shell", "abc123", "shell(ls)"), + ("shell", "abc123", "shell(ls)"), + ] + rep = _detect_repetition(calls, 3) + assert rep is not None + assert rep.count == 3 + assert rep.tool_name == "shell" + + def test_below_threshold(self): + calls = [ + ("shell", "abc123", "shell(ls)"), + ("shell", "abc123", "shell(ls)"), + ] + assert _detect_repetition(calls, 3) is None + + def test_empty_calls(self): + assert _detect_repetition([], 3) is None + + +class TestSummarizeCall: + def test_command_arg(self): + s = _summarize_call("shell", '{"command": "ls -la"}') + assert "shell" in s + assert "ls -la" in s + + def test_truncates_long_command(self): + long_cmd = "x" * 200 + s = _summarize_call("shell", json.dumps({"command": long_cmd})) + assert len(s) < 200 + assert "..." in s + + +class TestRepetitionGuardCallback: + @pytest.mark.asyncio + async def test_no_trigger_on_varied_calls(self): + cb = RepetitionGuardCallback(_make_config()) + runtime = Runtime() + msgs = [ + Message(role="system", content="sys"), + Message(role="user", content="do something"), + ] + for i in range(5): + msgs.append(_assistant_with_tools(_tool_call("shell", {"command": f"cmd{i}"}, f"c{i}"))) + msgs.append(_tool_result(f"c{i}")) + + await cb.after_tool_call(runtime, msgs) + user_msgs = [m for m in msgs if m.role == "user"] + assert len(user_msgs) == 1 + + @pytest.mark.asyncio + async def test_triggers_on_repeated_calls(self): + cb = RepetitionGuardCallback(_make_config()) + runtime = Runtime() + msgs = [ + Message(role="system", content="sys"), + Message(role="user", content="do something"), + ] + for i in range(3): + msgs.append( + _assistant_with_tools( + _tool_call("shell", {"command": "make build"}, f"c{i}") + ) + ) + msgs.append(_tool_result(f"c{i}", "error: build failed")) + + await cb.after_tool_call(runtime, msgs) + user_msgs = [m for m in msgs if m.role == "user"] + assert len(user_msgs) == 2 + injected = user_msgs[-1].content + assert "repeating the same approach" in injected + assert "make build" in injected + assert runtime.should_stop is False + + @pytest.mark.asyncio + async def test_force_stop_after_max_warnings(self): + cb = RepetitionGuardCallback(_make_config(max_warnings=1)) + runtime = Runtime() + msgs = [ + Message(role="system", content="sys"), + Message(role="user", content="do something"), + ] + for i in range(3): + msgs.append( + _assistant_with_tools( + _tool_call("shell", {"command": "make build"}, f"c{i}") + ) + ) + msgs.append(_tool_result(f"c{i}", "error")) + + await cb.after_tool_call(runtime, msgs) + assert cb._warnings_given == 1 + + msgs.append( + _assistant_with_tools( + _tool_call("shell", {"command": "make build"}, "c99") + ) + ) + msgs.append(_tool_result("c99", "error")) + + await cb.after_tool_call(runtime, msgs) + assert runtime.should_stop is True + + @pytest.mark.asyncio + async def test_does_not_warn_same_key_twice(self): + cb = RepetitionGuardCallback(_make_config()) + runtime = Runtime() + msgs = [ + Message(role="system", content="sys"), + Message(role="user", content="do something"), + ] + for i in range(4): + msgs.append( + _assistant_with_tools( + _tool_call("shell", {"command": "ls"}, f"c{i}") + ) + ) + msgs.append(_tool_result(f"c{i}")) + + await cb.after_tool_call(runtime, msgs) + assert cb._warnings_given == 1 + + await cb.after_tool_call(runtime, msgs) + assert cb._warnings_given == 1 + + @pytest.mark.asyncio + async def test_warns_on_new_pattern(self): + cb = RepetitionGuardCallback(_make_config()) + runtime = Runtime() + msgs = [ + Message(role="system", content="sys"), + Message(role="user", content="do something"), + ] + for i in range(3): + msgs.append( + _assistant_with_tools( + _tool_call("shell", {"command": "ls"}, f"c{i}") + ) + ) + msgs.append(_tool_result(f"c{i}")) + + await cb.after_tool_call(runtime, msgs) + assert cb._warnings_given == 1 + + for i in range(3): + msgs.append( + _assistant_with_tools( + _tool_call("shell", {"command": "pwd"}, f"d{i}") + ) + ) + msgs.append(_tool_result(f"d{i}")) + + await cb.after_tool_call(runtime, msgs) + assert cb._warnings_given == 2 + + @pytest.mark.asyncio + async def test_default_config(self): + cb = RepetitionGuardCallback(DictConfig({})) + assert cb.threshold == 3 + assert cb.lookback == 8 + assert cb.max_warnings == 2 diff --git a/tests/tools/test_local_shell_executor_smoke.py b/tests/tools/test_local_shell_executor_smoke.py new file mode 100644 index 000000000..2e09e73bc --- /dev/null +++ b/tests/tools/test_local_shell_executor_smoke.py @@ -0,0 +1,56 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Regression: LocalCodeExecutionTool.shell_executor via ToolManager (no LLM / network).""" +import json +import os +import shutil +import sys +import tempfile +import unittest + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +from omegaconf import OmegaConf + +from ms_agent.agent.llm_agent import LLMAgent # noqa: F401 — breaks tools import cycle +from ms_agent.llm.utils import ToolCall +from ms_agent.tools.tool_manager import ToolManager + + +class TestLocalShellExecutorSmoke(unittest.IsolatedAsyncioTestCase): + + async def test_shell_executor_via_tool_manager(self): + td = tempfile.mkdtemp() + try: + cfg = OmegaConf.create({ + 'output_dir': td, + 'tools': { + 'code_executor': { + 'mcp': False, + 'implementation': 'python_env', + 'include': ['shell_executor'], + }, + }, + 'tool_call_timeout': 60, + }) + tm = ToolManager(cfg) + await tm.connect() + self.assertIn('code_executor---shell_executor', tm._tool_index) + + tc = ToolCall( + tool_name='code_executor---shell_executor', + arguments={'command': 'echo ms_agent_shell_ok'}, + id='regression-call', + ) + raw = await tm.single_call_tool(tc) + data = json.loads(raw) + self.assertTrue(data.get('success'), raw) + self.assertIn('ms_agent_shell_ok', data.get('output', '')) + await tm.cleanup() + finally: + shutil.rmtree(td, ignore_errors=True) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/tools/test_tool_manager_timeout.py b/tests/tools/test_tool_manager_timeout.py new file mode 100644 index 000000000..d77b11fcb --- /dev/null +++ b/tests/tools/test_tool_manager_timeout.py @@ -0,0 +1,64 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import sys +import unittest + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +from ms_agent.agent.llm_agent import LLMAgent # noqa: F401 — breaks tools import cycle + +from ms_agent.tools.tool_manager import ( + effective_tool_wait_seconds, + parse_timeout_from_tool_args, +) + + +class TestToolManagerTimeout(unittest.TestCase): + + def test_parse_timeout(self): + self.assertIsNone(parse_timeout_from_tool_args(None)) + self.assertIsNone(parse_timeout_from_tool_args({})) + self.assertIsNone(parse_timeout_from_tool_args({'timeout': None})) + self.assertEqual(parse_timeout_from_tool_args({'timeout': 45}), 45.0) + self.assertEqual(parse_timeout_from_tool_args({'timeout': '90'}), 90.0) + self.assertIsNone(parse_timeout_from_tool_args({'timeout': True})) + + def test_effective_wait(self): + self.assertEqual( + effective_tool_wait_seconds( + {}, + default_sec=120, + max_sec=600, + ), + 120.0, + ) + self.assertEqual( + effective_tool_wait_seconds( + {'timeout': 30}, + default_sec=120, + max_sec=600, + ), + 30.0, + ) + self.assertEqual( + effective_tool_wait_seconds( + {'timeout': 9999}, + default_sec=120, + max_sec=600, + ), + 600.0, + ) + self.assertEqual( + effective_tool_wait_seconds( + {}, + default_sec=900, + max_sec=600, + ), + 600.0, + ) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/utils/test_filesystem_tool_config.py b/tests/utils/test_filesystem_tool_config.py index 5bc8310c9..2f68b63b6 100644 --- a/tests/utils/test_filesystem_tool_config.py +++ b/tests/utils/test_filesystem_tool_config.py @@ -2,8 +2,12 @@ """FileSystemTool config: include aliases and grep/glob registration.""" import asyncio +import json +import os +import shutil import tempfile +import pytest from omegaconf import OmegaConf from ms_agent.tools.filesystem_tool import FileSystemTool @@ -52,3 +56,55 @@ async def _run(): assert names == ['grep', 'glob'] asyncio.run(_run()) + + +@pytest.mark.skipif(not shutil.which('rg'), reason='rg not installed') +def test_grep_hyphen_prefixed_pattern_uses_rg_e_flag(): + async def _run(): + with tempfile.TemporaryDirectory() as td: + sample = os.path.join(td, 'sample.js') + with open(sample, 'w', encoding='utf-8') as f: + f.write('-import foo from "bar"\n') + + cfg = OmegaConf.create({ + 'output_dir': td, + 'tools': { + 'file_system': { + 'mcp': False, + 'include': ['grep'], + }, + }, + }) + fs = FileSystemTool(cfg) + raw = await fs.grep( + pattern='-import', + path='sample.js', + output_mode='content', + ) + result = json.loads(raw) + assert result['success'] is True + assert '-import' in result['output'] + + asyncio.run(_run()) + + +def test_edit_file_requires_path(): + async def _run(): + with tempfile.TemporaryDirectory() as td: + cfg = OmegaConf.create({ + 'output_dir': td, + 'tools': { + 'file_system': { + 'mcp': False, + 'include': ['edit'], + }, + }, + }) + fs = FileSystemTool(cfg) + result = await fs.edit_file( + old_string='a', + new_string='b', + ) + assert result == 'Error: `path` is required.' + + asyncio.run(_run()) diff --git a/tests/utils/test_snapshot_smoke.py b/tests/utils/test_snapshot_smoke.py index e4cce6509..5c8330954 100644 --- a/tests/utils/test_snapshot_smoke.py +++ b/tests/utils/test_snapshot_smoke.py @@ -132,6 +132,15 @@ def test_no_repo_returns_false(self): self.assertFalse(ok) self.assertEqual(mc, 0) + def test_restore_without_git_returns_false(self): + with tempfile.TemporaryDirectory() as td: + os.makedirs(os.path.join(td, '.ms_agent_snapshots')) + with patch('ms_agent.utils.snapshot._git', + side_effect=FileNotFoundError('git')): + ok, mc = restore_snapshot(td, 'abc1234') + self.assertFalse(ok) + self.assertEqual(mc, 0) + def test_restore_reverts_file_content(self): with tempfile.TemporaryDirectory() as td: path = os.path.join(td, 'data.txt') diff --git a/tests/utils/test_task_manager_smoke.py b/tests/utils/test_task_manager_smoke.py index f29f57321..874172b3a 100644 --- a/tests/utils/test_task_manager_smoke.py +++ b/tests/utils/test_task_manager_smoke.py @@ -359,5 +359,39 @@ async def fake_launch_background(payload, spec, call_id): self.assertEqual(data['status'], 'async_launched') +class TestAgentToolTranscript(unittest.TestCase): + + def test_save_transcript_writes_message_objects_and_dicts(self): + import json + import tempfile + + from omegaconf import OmegaConf + + from ms_agent.llm.utils import Message + from ms_agent.tools.agent_tool import AgentTool + + with tempfile.TemporaryDirectory() as td: + config = OmegaConf.create({ + 'tag': 'test', + 'output_dir': td, + 'tools': {}, + }) + tool = AgentTool(config) + messages = [ + Message(role='user', content='hello'), + {'role': 'assistant', 'content': 'hi'}, + ] + tool._save_transcript(messages, 'worker-1') + path = os.path.join(td, 'subagents', 'worker-1.jsonl') + self.assertTrue(os.path.isfile(path)) + lines = open(path, encoding='utf-8').read().strip().split('\n') + self.assertEqual(len(lines), 2) + parsed = [json.loads(line) for line in lines] + self.assertEqual(parsed[0]['role'], 'user') + self.assertEqual(parsed[0]['content'], 'hello') + self.assertEqual(parsed[1]['role'], 'assistant') + self.assertEqual(parsed[1]['content'], 'hi') + + if __name__ == '__main__': unittest.main() diff --git a/tests/utils/test_workspace_policy.py b/tests/utils/test_workspace_policy.py new file mode 100644 index 000000000..012888533 --- /dev/null +++ b/tests/utils/test_workspace_policy.py @@ -0,0 +1,78 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Tests for WorkspacePolicyKernel.""" + +import tempfile +from pathlib import Path + +import pytest + +from ms_agent.utils.workspace_policy import WorkspacePolicyError, WorkspacePolicyKernel + + +def test_default_root_is_output_dir(): + with tempfile.TemporaryDirectory() as td: + out = Path(td) / 'out' + out.mkdir() + k = WorkspacePolicyKernel(out) + p = k.resolve_under_roots('foo/bar') + assert p == (out / 'foo' / 'bar').resolve() + + +def test_rejects_escape(): + with tempfile.TemporaryDirectory() as td: + out = Path(td) / 'out' + out.mkdir() + k = WorkspacePolicyKernel(out) + with pytest.raises(WorkspacePolicyError): + k.resolve_under_roots('../../etc/passwd') + + +def test_extra_allow_root(): + with tempfile.TemporaryDirectory() as td: + out = Path(td) / 'out' + other = Path(td) / 'other' + out.mkdir() + other.mkdir() + k = WorkspacePolicyKernel(out, extra_allow_roots=[str(other)]) + assert k.resolve_under_roots(str(other / 'x')) == (other / 'x').resolve() + + +def test_read_only_blocks_redirect(): + with tempfile.TemporaryDirectory() as td: + out = Path(td) / 'out' + out.mkdir() + k = WorkspacePolicyKernel( + out, + shell_default_mode='read_only', + ) + with pytest.raises(WorkspacePolicyError): + k.assert_shell_command_allowed('echo x > file.txt') + + +def test_workspace_write_allows_redirect_but_blocks_network(): + with tempfile.TemporaryDirectory() as td: + out = Path(td) / 'out' + out.mkdir() + k = WorkspacePolicyKernel( + out, + shell_default_mode='workspace_write', + shell_network_enabled=False, + ) + k.assert_shell_command_allowed('echo x > file.txt') + with pytest.raises(WorkspacePolicyError): + k.assert_shell_command_allowed('curl https://example.com') + + +def test_artifact_manager_spill(tmp_path): + from ms_agent.utils.artifact_manager import ArtifactManager + + am = ArtifactManager(tmp_path, max_combined_bytes=32) + big = 'a' * 100 + packed = am.pack_text_result( + tool_name='t', + call_id='c1', + stdout=big, + stderr='', + ) + assert packed.get('truncated') is True + assert 'artifact_path' in packed