返回 ViMax
tools.py
根目录 / agent_runtime / tools.py
1 from __future__ import annotations
2
3 import asyncio
4 import glob
5 import inspect
6 import json
7 import os
8 import subprocess
9 from dataclasses import dataclass, field
10 from pathlib import Path
11 from threading import Event
12 from typing import Any, Awaitable, Callable
13
14 from .image_tools import ViewImageHandler
15 from .models import ToolCall, ToolResult
16
17 ToolHandler = Callable[..., Awaitable[ToolResult] | ToolResult]
18 ProgressCallback = Callable[[dict[str, Any]], None]
19
20
21 @dataclass(slots=True)
22 class ToolArgumentSchema:
23 type: type | tuple[type, ...]
24 required: bool = False
25 default: Any = None
26
27
28 @dataclass(slots=True)
29 class ToolSpec:
30 name: str
31 description: str
32 handler: ToolHandler
33 aliases: tuple[str, ...] = ()
34 permission_mode: str = "workspace-write"
35 schema: dict[str, ToolArgumentSchema] | None = None
36 json_schema: dict[str, Any] | None = None
37 concurrency_safe: bool = False
38
39
40 @dataclass(slots=True)
41 class ToolRuntimeContext:
42 requested_name: str
43 canonical_name: str
44 turn_id: str = ""
45 cancel_event: Event | None = None
46 progress_callback: ProgressCallback | None = None
47 metadata: dict[str, Any] = field(default_factory=dict)
48
49 def emit_progress(self, message: str, *, stage: str = "running", metadata: dict[str, Any] | None = None) -> None:
50 if self.progress_callback is None:
51 return
52 payload: dict[str, Any] = {
53 "type": "tool_progress",
54 "tool": {"requested_name": self.requested_name, "name": self.canonical_name},
55 "progress": {"stage": stage, "message": message, "metadata": metadata or {}},
56 }
57 if self.turn_id:
58 payload["turn_id"] = self.turn_id
59 self.progress_callback(payload)
60
61 def emit_terminal(self, line: str, *, stream: str = "stdout") -> None:
62 if self.progress_callback is None:
63 return
64 if not line:
65 return
66 payload: dict[str, Any] = {"type": "terminal", "stream": stream, "line": line}
67 if self.turn_id:
68 payload["turn_id"] = self.turn_id
69 self.progress_callback(payload)
70
71 def is_cancelled(self) -> bool:
72 return self.cancel_event.is_set() if self.cancel_event is not None else False
73
74 def raise_if_cancelled(self, default_reason: str = "Tool execution cancelled") -> None:
75 if self.is_cancelled():
76 raise RuntimeError(str(self.metadata.get("cancel_reason") or default_reason))
77
78
79 class ToolRegistry:
80 def __init__(self, specs: list[ToolSpec] | None = None) -> None:
81 self._specs: dict[str, ToolSpec] = {}
82 self._aliases: dict[str, str] = {}
83 for spec in specs or []:
84 self.register(spec)
85
86 def register(self, spec: ToolSpec) -> None:
87 self._specs[spec.name] = spec
88 for alias in spec.aliases:
89 self._aliases[alias] = spec.name
90
91 def list_tools(self) -> list[dict[str, str]]:
92 return sorted([{"name": spec.name, "description": spec.description, "permission_mode": spec.permission_mode} for spec in self._specs.values()], key=lambda item: item["name"])
93
94 def list_function_tools(self) -> list[dict[str, Any]]:
95 tools = []
96 for spec in sorted(self._specs.values(), key=lambda item: item.name):
97 parameters = spec.json_schema or _argument_schema_to_json_schema(spec.schema or {})
98 tools.append({"type": "function", "function": {"name": spec.name, "description": spec.description, "parameters": parameters}})
99 return tools
100
101 def get_spec(self, name: str) -> ToolSpec | None:
102 return self._specs.get(self.resolve_name(name))
103
104 def resolve_name(self, name: str) -> str:
105 normalized = name.strip()
106 return self._aliases.get(normalized, normalized)
107
108 def validate_arguments(self, name: str, arguments: dict[str, Any]) -> tuple[dict[str, Any] | None, str | None]:
109 spec = self.get_spec(name)
110 if spec is None:
111 return None, f"Unknown tool: {name}"
112 schema = spec.schema or {}
113 normalized = dict(arguments or {})
114 for field_name, field_spec in schema.items():
115 if field_name not in normalized:
116 if field_spec.required and field_spec.default is None:
117 return None, f"Missing required argument '{field_name}' for {spec.name}"
118 if field_spec.default is not None:
119 normalized[field_name] = field_spec.default
120 continue
121 value = normalized[field_name]
122 expected = field_spec.type
123 if expected is bool and isinstance(value, str) and value.lower() in {"true", "false"}:
124 normalized[field_name] = value.lower() == "true"
125 continue
126 if expected is int and isinstance(value, str):
127 try:
128 normalized[field_name] = int(value)
129 continue
130 except ValueError:
131 return None, f"Argument '{field_name}' for {spec.name} must be an integer"
132 if not isinstance(normalized[field_name], expected):
133 expected_name = ", ".join(t.__name__ for t in expected) if isinstance(expected, tuple) else expected.__name__
134 return None, f"Argument '{field_name}' for {spec.name} must be {expected_name}"
135 return normalized, None
136
137 def is_concurrency_safe(self, name: str) -> bool:
138 spec = self.get_spec(name)
139 return bool(spec and spec.concurrency_safe)
140
141 def partition_calls(self, calls: list[ToolCall]) -> list[list[ToolCall]]:
142 batches: list[list[ToolCall]] = []
143 for call in calls:
144 if self.is_concurrency_safe(call.name) and batches and all(self.is_concurrency_safe(item.name) for item in batches[-1]):
145 batches[-1].append(call)
146 else:
147 batches.append([call])
148 return batches
149
150 async def execute(self, name: str, arguments: dict[str, Any], runtime: ToolRuntimeContext | None = None) -> ToolResult:
151 canonical = self.resolve_name(name)
152 spec = self._specs.get(canonical)
153 if spec is None:
154 return ToolResult(name=name, ok=False, content=f"Unknown tool: {name}", metadata={"error_type": "unknown_tool"})
155 handler = spec.handler
156 try:
157 params = inspect.signature(handler).parameters
158 result = handler(arguments, runtime) if runtime is not None and len(params) >= 2 else handler(arguments)
159 if inspect.isawaitable(result):
160 return await result
161 return result
162 except Exception as exc:
163 return ToolResult(name=canonical, ok=False, content=str(exc), metadata={"error_type": "exception"})
164
165
166 def _argument_schema_to_json_schema(schema: dict[str, ToolArgumentSchema]) -> dict[str, Any]:
167 properties: dict[str, Any] = {}
168 required: list[str] = []
169 for field_name, field_spec in schema.items():
170 field_schema = _type_to_json_schema(field_spec.type)
171 if field_spec.default is not None:
172 field_schema["default"] = field_spec.default
173 properties[field_name] = field_schema
174 if field_spec.required and field_spec.default is None:
175 required.append(field_name)
176 payload: dict[str, Any] = {"type": "object", "properties": properties, "additionalProperties": False}
177 if required:
178 payload["required"] = required
179 return payload
180
181
182 def _type_to_json_schema(tp: type | tuple[type, ...]) -> dict[str, Any]:
183 if isinstance(tp, tuple):
184 return {"anyOf": [_type_to_json_schema(item) for item in tp]}
185 return {str: {"type": "string"}, int: {"type": "integer"}, bool: {"type": "boolean"}, dict: {"type": "object", "additionalProperties": True}, list: {"type": "array", "items": {}}}.get(tp, {"type": "string"})
186
187
188 def build_builtin_registry(workspace_root: str | Path, session_index: Any, adapter_specs: list[ToolSpec] | None = None) -> ToolRegistry:
189 root = Path(workspace_root).resolve()
190 view_image = ViewImageHandler(root, session_index)
191
192 def safe_path(raw: Any) -> Path:
193 path = (root / str(raw)).resolve()
194 if root not in path.parents and path != root:
195 raise ValueError(f"Path escapes workspace: {raw}")
196 return path
197
198 def _legacy_virtual_read(raw_path: Any, *, as_json: bool) -> ToolResult | None:
199 """Compatibility for paths older prompts/models may hallucinate.
200
201 The authoritative session state is .vimax/sessions.json and logs are
202 .vimax/logs/*.jsonl, but some model turns ask for per-session files like
203 .working_dir/<session>/session.json or .vimax/logs/<session>.log.
204 """
205 path = safe_path(raw_path)
206 try:
207 rel = path.relative_to(root)
208 except ValueError:
209 return None
210 parts = rel.parts
211 if len(parts) == 3 and parts[0] == ".working_dir" and parts[2] == "session.json":
212 session_id = parts[1]
213 record = session_index.get(session_id)
214 if record is None:
215 return None
216 payload = {
217 "session": record,
218 "artifact_checklist": session_index.artifact_checklist(session_id),
219 "source": ".vimax/sessions.json",
220 "virtual_path": rel.as_posix(),
221 }
222 content = json.dumps(payload, ensure_ascii=False, indent=2)
223 return ToolResult("read_json" if as_json else "read_file", True, content, {"virtual_path": True, "source": ".vimax/sessions.json"})
224 if len(parts) == 3 and parts[0] == ".vimax" and parts[1] == "logs" and parts[2].endswith(".log"):
225 session_id = parts[2][:-4]
226 rows: list[dict[str, Any]] = []
227 for log_name in ("loop_history", "tool_calls", "revisions"):
228 log_path = session_index.logs_dir / f"{log_name}.jsonl"
229 if not log_path.exists():
230 continue
231 for line in log_path.read_text(encoding="utf-8", errors="replace").splitlines():
232 if session_id not in line:
233 continue
234 try:
235 item = json.loads(line)
236 except json.JSONDecodeError:
237 item = {"raw": line}
238 item["_log"] = log_name
239 rows.append(item)
240 payload = {
241 "session_id": session_id,
242 "source": ".vimax/logs/*.jsonl",
243 "virtual_path": rel.as_posix(),
244 "records": rows,
245 }
246 content = json.dumps(payload, ensure_ascii=False, indent=2)
247 return ToolResult("read_json" if as_json else "read_file", True, content, {"virtual_path": True, "source": ".vimax/logs/*.jsonl", "record_count": len(rows)})
248 return None
249
250 def read_file(args: dict[str, Any]) -> ToolResult:
251 path = safe_path(args["path"])
252 if not path.exists():
253 virtual = _legacy_virtual_read(args["path"], as_json=False)
254 if virtual is not None:
255 return virtual
256 return ToolResult("read_file", False, f"File not found: {path}")
257 return ToolResult("read_file", True, path.read_text(encoding="utf-8"))
258
259 def read_json(args: dict[str, Any]) -> ToolResult:
260 path = safe_path(args["path"])
261 if not path.exists():
262 virtual = _legacy_virtual_read(args["path"], as_json=True)
263 if virtual is not None:
264 return virtual
265 return ToolResult("read_json", False, f"File not found: {path}")
266 try:
267 payload = json.loads(path.read_text(encoding="utf-8"))
268 except json.JSONDecodeError as exc:
269 return ToolResult("read_json", False, f"Invalid JSON: {exc}", {"error_type": "invalid_json"})
270 return ToolResult("read_json", True, json.dumps(payload, ensure_ascii=False, indent=2))
271
272 def write_json(args: dict[str, Any]) -> ToolResult:
273 path = safe_path(args["path"])
274 path.parent.mkdir(parents=True, exist_ok=True)
275 path.write_text(json.dumps(args["data"], ensure_ascii=False, indent=2), encoding="utf-8")
276 return ToolResult("write_json", True, f"Wrote JSON {path.relative_to(root)}")
277
278 def list_files(args: dict[str, Any]) -> ToolResult:
279 path = safe_path(args.get("path", "."))
280 if not path.exists():
281 return ToolResult("list_files", False, f"Path not found: {path}")
282 rows = [str(item.relative_to(root)) for item in sorted(path.iterdir())]
283 return ToolResult("list_files", True, "\n".join(rows) or "No entries")
284
285 def glob_files(args: dict[str, Any]) -> ToolResult:
286 pattern = str(args["pattern"])
287 matches = [str(Path(item).resolve().relative_to(root)) for item in glob.glob(str(root / pattern), recursive=True)]
288 return ToolResult("glob_files", True, "\n".join(matches[:200]) or "No matches")
289
290 def search_text(args: dict[str, Any]) -> ToolResult:
291 needle = str(args["query"])
292 base = safe_path(args.get("path", "."))
293 rows: list[str] = []
294 paths = base.rglob("*") if base.is_dir() else [base]
295 for path in paths:
296 if not path.is_file():
297 continue
298 try:
299 text = path.read_text(encoding="utf-8")
300 except UnicodeDecodeError:
301 continue
302 for idx, line in enumerate(text.splitlines(), start=1):
303 if needle in line:
304 rows.append(f"{path.relative_to(root)}:{idx}: {line}")
305 if len(rows) >= int(args.get("max_results", 100)):
306 return ToolResult("search_text", True, "\n".join(rows))
307 return ToolResult("search_text", True, "\n".join(rows) or "No matches")
308
309 def memory_read(args: dict[str, Any]) -> ToolResult:
310 return ToolResult("memory_read", True, session_index.memory_text())
311
312 def memory_write(args: dict[str, Any]) -> ToolResult:
313 session_index.write_memory(str(args["content"]))
314 return ToolResult("memory_write", True, "Updated .vimax/memory.md")
315
316 def todo_path() -> Path:
317 return root / ".vimax" / "todo.json"
318
319 def todo_read(args: dict[str, Any]) -> ToolResult:
320 path = todo_path()
321 if not path.exists():
322 return ToolResult("todo_read", True, json.dumps({"items": []}, ensure_ascii=False, indent=2), {"items": []})
323 try:
324 payload = json.loads(path.read_text(encoding="utf-8"))
325 except json.JSONDecodeError as exc:
326 return ToolResult("todo_read", False, f"Invalid todo JSON: {exc}", {"error_type": "invalid_json"})
327 items = payload.get("items")
328 if not isinstance(items, list):
329 return ToolResult("todo_read", False, "Invalid todo JSON: expected an items array", {"error_type": "invalid_todo"})
330 return ToolResult("todo_read", True, json.dumps({"items": items}, ensure_ascii=False, indent=2), {"items": items})
331
332 def todo_write(args: dict[str, Any]) -> ToolResult:
333 items = args.get("items")
334 if not isinstance(items, list):
335 return ToolResult("todo_write", False, "items must be an array", {"error_type": "invalid_arguments"})
336 normalized: list[dict[str, Any]] = []
337 for index, item in enumerate(items):
338 if not isinstance(item, dict):
339 return ToolResult("todo_write", False, f"items[{index}] must be an object", {"error_type": "invalid_arguments", "index": index})
340 content = str(item.get("content", "")).strip()
341 if not content:
342 return ToolResult("todo_write", False, f"items[{index}].content is required", {"error_type": "invalid_arguments", "index": index})
343 status = str(item.get("status", "pending")).strip() or "pending"
344 if status not in {"pending", "in_progress", "completed"}:
345 return ToolResult("todo_write", False, f"items[{index}].status must be pending, in_progress, or completed", {"error_type": "invalid_arguments", "index": index})
346 normalized.append({"content": content, "status": status})
347 path = todo_path()
348 path.parent.mkdir(parents=True, exist_ok=True)
349 path.write_text(json.dumps({"items": normalized}, ensure_ascii=False, indent=2), encoding="utf-8")
350 return ToolResult("todo_write", True, f"Updated .vimax/todo.json with {len(normalized)} item(s)", {"items": normalized, "item_count": len(normalized)})
351
352 async def sleep_tool(args: dict[str, Any], runtime: ToolRuntimeContext | None = None) -> ToolResult:
353 seconds = float(args.get("seconds", 0))
354 if seconds < 0 or seconds > 300:
355 return ToolResult("sleep", False, "seconds must be between 0 and 300")
356 if runtime:
357 runtime.emit_progress(f"Sleeping for {seconds:g}s", stage="running")
358 await asyncio.sleep(seconds)
359 return ToolResult("sleep", True, f"Slept for {seconds:g}s")
360
361 async def run_shell(args: dict[str, Any], runtime: ToolRuntimeContext | None = None) -> ToolResult:
362 if os.environ.get("VIMAX_ENABLE_RUN_SHELL") != "1":
363 return ToolResult("run_shell", False, "run_shell is disabled by default. Set VIMAX_ENABLE_RUN_SHELL=1 to enable bounded shell commands.", {"error_type": "disabled"})
364 command = str(args["command"]).strip()
365 timeout_seconds = min(max(int(args.get("timeout_seconds", 30)), 1), 120)
366 output_limit = min(max(int(args.get("output_limit", 20000)), 1000), 50000)
367 denied_tokens = ["rm ", "rm -", "sudo", "chmod", "chown", "mkfs", "dd ", ":(){", "curl ", "wget ", "ssh ", "printenv", "env", "export"]
368 lowered = command.lower()
369 if any(token in lowered for token in denied_tokens):
370 return ToolResult("run_shell", False, "Command rejected by run_shell policy.", {"error_type": "command_rejected"})
371 if runtime:
372 runtime.emit_progress("Starting shell command", stage="starting", metadata={"command": command, "timeout_seconds": timeout_seconds})
373 proc = await asyncio.create_subprocess_shell(command, cwd=root, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
374 try:
375 stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout_seconds)
376 except asyncio.TimeoutError:
377 proc.kill()
378 await proc.communicate()
379 return ToolResult("run_shell", False, f"Command timed out after {timeout_seconds}s", {"error_type": "timeout", "timeout_seconds": timeout_seconds})
380 content = ""
381 if stdout:
382 content += stdout.decode(errors="replace")
383 if stderr:
384 content += stderr.decode(errors="replace")
385 truncated = len(content) > output_limit
386 if truncated:
387 content = content[:output_limit] + "\n...[truncated]"
388 return ToolResult("run_shell", proc.returncode == 0, content, {"returncode": proc.returncode, "truncated": truncated})
389
390 specs = [
391 ToolSpec("read_file", "Read a UTF-8 text file inside the workspace. Also resolves virtual legacy session paths like .vimax/logs/<session>.log.", read_file, schema={"path": ToolArgumentSchema(str, True)}, concurrency_safe=True),
392 ToolSpec("read_json", "Read and parse a JSON file inside the workspace. Also resolves virtual legacy session paths like .working_dir/<session>/session.json.", read_json, schema={"path": ToolArgumentSchema(str, True)}, concurrency_safe=True),
393 ToolSpec("write_json", "Write formatted JSON inside the workspace.", write_json, schema={"path": ToolArgumentSchema(str, True), "data": ToolArgumentSchema((dict, list), True)}),
394 ToolSpec("list_files", "List direct children of a workspace path.", list_files, schema={"path": ToolArgumentSchema(str, False, ".")}, concurrency_safe=True),
395 ToolSpec("glob_files", "Find workspace files with a glob pattern.", glob_files, schema={"pattern": ToolArgumentSchema(str, True)}, concurrency_safe=True),
396 ToolSpec("search_text", "Search text in workspace files.", search_text, schema={"query": ToolArgumentSchema(str, True), "path": ToolArgumentSchema(str, False, "."), "max_results": ToolArgumentSchema(int, False, 100)}, concurrency_safe=True),
397 ToolSpec("view_image", "Load a PNG, JPEG, WebP, or GIF from the active session and present its pixels to the multimodal model. Accepts a session-relative path or a path prefixed by the active .working_dir session.", view_image, permission_mode="read-only", schema={"path": ToolArgumentSchema(str, True)}, concurrency_safe=True),
398 ToolSpec("memory_read", "Read .vimax/memory.md user preferences.", memory_read, schema={}, concurrency_safe=True),
399 ToolSpec("memory_write", "Replace .vimax/memory.md with user preference notes only.", memory_write, schema={"content": ToolArgumentSchema(str, True)}),
400 ToolSpec("todo_read", "Read short-term todo items from .vimax/todo.json. This is not a task or team system.", todo_read, schema={}, concurrency_safe=True),
401 ToolSpec("todo_write", "Replace short-term todo items in .vimax/todo.json. Items require content and may use pending, in_progress, or completed status.", todo_write, schema={"items": ToolArgumentSchema(list, True)}),
402 ToolSpec("sleep", "Wait for a bounded number of seconds.", sleep_tool, schema={"seconds": ToolArgumentSchema(int, False, 0)}, concurrency_safe=True),
403 ToolSpec("run_shell", "Run a bounded shell command in the workspace. Disabled unless VIMAX_ENABLE_RUN_SHELL=1; rejects dangerous commands, enforces timeout, and truncates output.", run_shell, schema={"command": ToolArgumentSchema(str, True), "timeout_seconds": ToolArgumentSchema(int, False, 30), "output_limit": ToolArgumentSchema(int, False, 20000)}),
404 ]
405 for spec in adapter_specs or []:
406 specs.append(spec)
407 return ToolRegistry(specs)
408
408 lines PYTHON