| 1 | import asyncio |
| 2 | import json |
| 3 | import tempfile |
| 4 | import unittest |
| 5 | from pathlib import Path |
| 6 | |
| 7 | from PIL import Image |
| 8 | |
| 9 | from agent_runtime.models import ToolCall, TurnControl |
| 10 | from agent_runtime.session_index import SessionIndex |
| 11 | from agent_runtime.tool_executor import ToolExecutor |
| 12 | from agent_runtime.tools import build_builtin_registry |
| 13 | |
| 14 | |
| 15 | class ToolRegistryTests(unittest.IsolatedAsyncioTestCase): |
| 16 | async def test_validation_default_write_json_and_logging(self): |
| 17 | with tempfile.TemporaryDirectory() as tmp: |
| 18 | index = SessionIndex(tmp) |
| 19 | index.create() |
| 20 | registry = build_builtin_registry(tmp, index) |
| 21 | executor = ToolExecutor(registry, index) |
| 22 | record = await executor.execute(ToolCall(name="write_json", arguments={"path": "data/a.json", "data": {"x": 1}}), TurnControl()) |
| 23 | self.assertTrue(record.result.ok) |
| 24 | self.assertEqual(json.loads((Path(tmp) / "data" / "a.json").read_text())["x"], 1) |
| 25 | self.assertTrue((Path(tmp) / ".vimax" / "logs" / "tool_calls.jsonl").exists()) |
| 26 | |
| 27 | async def test_unknown_and_missing_argument_return_tool_errors(self): |
| 28 | with tempfile.TemporaryDirectory() as tmp: |
| 29 | index = SessionIndex(tmp) |
| 30 | registry = build_builtin_registry(tmp, index) |
| 31 | executor = ToolExecutor(registry, index) |
| 32 | missing = await executor.execute(ToolCall(name="read_file", arguments={}), TurnControl()) |
| 33 | self.assertFalse(missing.result.ok) |
| 34 | unknown = await executor.execute(ToolCall(name="does_not_exist", arguments={}), TurnControl()) |
| 35 | self.assertFalse(unknown.result.ok) |
| 36 | |
| 37 | |
| 38 | async def test_run_shell_is_disabled_by_default(self): |
| 39 | with tempfile.TemporaryDirectory() as tmp: |
| 40 | index = SessionIndex(tmp) |
| 41 | registry = build_builtin_registry(tmp, index) |
| 42 | executor = ToolExecutor(registry, index) |
| 43 | record = await executor.execute(ToolCall(name="run_shell", arguments={"command": "pwd"}), TurnControl()) |
| 44 | self.assertFalse(record.result.ok) |
| 45 | self.assertEqual(record.result.metadata["error_type"], "disabled") |
| 46 | |
| 47 | |
| 48 | async def test_todo_read_returns_empty_items_by_default(self): |
| 49 | with tempfile.TemporaryDirectory() as tmp: |
| 50 | index = SessionIndex(tmp) |
| 51 | registry = build_builtin_registry(tmp, index) |
| 52 | executor = ToolExecutor(registry, index) |
| 53 | record = await executor.execute(ToolCall(name="todo_read", arguments={}), TurnControl()) |
| 54 | self.assertTrue(record.result.ok) |
| 55 | self.assertEqual(json.loads(record.result.content)["items"], []) |
| 56 | |
| 57 | async def test_todo_write_then_read_persists_items_and_logs(self): |
| 58 | with tempfile.TemporaryDirectory() as tmp: |
| 59 | index = SessionIndex(tmp) |
| 60 | index.create() |
| 61 | registry = build_builtin_registry(tmp, index) |
| 62 | executor = ToolExecutor(registry, index) |
| 63 | items = [{"content": "实现 TUI", "status": "in_progress"}, {"content": "补测试"}] |
| 64 | written = await executor.execute(ToolCall(name="todo_write", arguments={"items": items}), TurnControl()) |
| 65 | self.assertTrue(written.result.ok) |
| 66 | todo_path = Path(tmp) / ".vimax" / "todo.json" |
| 67 | self.assertTrue(todo_path.exists()) |
| 68 | payload = json.loads(todo_path.read_text()) |
| 69 | self.assertEqual(payload["items"][0]["content"], "实现 TUI") |
| 70 | self.assertEqual(payload["items"][1]["status"], "pending") |
| 71 | |
| 72 | read = await executor.execute(ToolCall(name="todo_read", arguments={}), TurnControl()) |
| 73 | self.assertTrue(read.result.ok) |
| 74 | self.assertEqual(json.loads(read.result.content)["items"], payload["items"]) |
| 75 | logs = (Path(tmp) / ".vimax" / "logs" / "tool_calls.jsonl").read_text() |
| 76 | self.assertIn('"tool": "todo_write"', logs) |
| 77 | |
| 78 | async def test_todo_write_rejects_invalid_payload(self): |
| 79 | with tempfile.TemporaryDirectory() as tmp: |
| 80 | index = SessionIndex(tmp) |
| 81 | registry = build_builtin_registry(tmp, index) |
| 82 | executor = ToolExecutor(registry, index) |
| 83 | not_list = await executor.execute(ToolCall(name="todo_write", arguments={"items": {"content": "x"}}), TurnControl()) |
| 84 | self.assertFalse(not_list.result.ok) |
| 85 | missing_content = await executor.execute(ToolCall(name="todo_write", arguments={"items": [{"status": "pending"}]}), TurnControl()) |
| 86 | self.assertFalse(missing_content.result.ok) |
| 87 | bad_status = await executor.execute(ToolCall(name="todo_write", arguments={"items": [{"content": "x", "status": "blocked"}]}), TurnControl()) |
| 88 | self.assertFalse(bad_status.result.ok) |
| 89 | |
| 90 | async def test_read_json_supports_virtual_session_json_path(self): |
| 91 | with tempfile.TemporaryDirectory() as tmp: |
| 92 | index = SessionIndex(tmp) |
| 93 | record = index.create(session_id="20260630-125442-vimax", idea="surfing") |
| 94 | registry = build_builtin_registry(tmp, index) |
| 95 | executor = ToolExecutor(registry, index) |
| 96 | result = await executor.execute( |
| 97 | ToolCall(name="read_json", arguments={"path": f"{record['working_dir']}/session.json"}), |
| 98 | TurnControl(), |
| 99 | ) |
| 100 | self.assertTrue(result.result.ok) |
| 101 | payload = json.loads(result.result.content) |
| 102 | self.assertEqual(payload["session"]["session_id"], "20260630-125442-vimax") |
| 103 | self.assertEqual(payload["source"], ".vimax/sessions.json") |
| 104 | self.assertTrue(result.result.metadata["virtual_path"]) |
| 105 | |
| 106 | async def test_read_file_supports_virtual_session_log_path(self): |
| 107 | with tempfile.TemporaryDirectory() as tmp: |
| 108 | index = SessionIndex(tmp) |
| 109 | record = index.create(session_id="20260630-125442-vimax", idea="surfing") |
| 110 | index.append_turn_record(record["session_id"], {"turn_id": "turn-1", "status": "completed", "final_assistant_text": "done"}) |
| 111 | registry = build_builtin_registry(tmp, index) |
| 112 | executor = ToolExecutor(registry, index) |
| 113 | result = await executor.execute( |
| 114 | ToolCall(name="read_file", arguments={"path": ".vimax/logs/20260630-125442-vimax.log"}), |
| 115 | TurnControl(), |
| 116 | ) |
| 117 | self.assertTrue(result.result.ok) |
| 118 | payload = json.loads(result.result.content) |
| 119 | self.assertEqual(payload["session_id"], "20260630-125442-vimax") |
| 120 | self.assertEqual(payload["source"], ".vimax/logs/*.jsonl") |
| 121 | self.assertEqual(payload["records"][0]["turn_id"], "turn-1") |
| 122 | self.assertTrue(result.result.metadata["virtual_path"]) |
| 123 | |
| 124 | async def test_view_image_returns_transient_multimodal_content_without_logging_pixels(self): |
| 125 | with tempfile.TemporaryDirectory() as tmp: |
| 126 | index = SessionIndex(tmp) |
| 127 | record = index.create(session_id="visual-session") |
| 128 | image_path = index.working_dir(record["session_id"]) / "idea2video" / "reference.png" |
| 129 | exif = Image.Exif() |
| 130 | exif[271] = "Camera Corp" |
| 131 | exif[272] = "Model One" |
| 132 | exif[37386] = 50.0 |
| 133 | Image.new("RGB", (2400, 1200), (20, 40, 60)).save(image_path, exif=exif) |
| 134 | registry = build_builtin_registry(tmp, index) |
| 135 | executor = ToolExecutor(registry, index) |
| 136 | result = await executor.execute( |
| 137 | ToolCall(name="view_image", arguments={"path": "idea2video/reference.png"}), |
| 138 | TurnControl(), |
| 139 | ) |
| 140 | |
| 141 | self.assertTrue(result.result.ok) |
| 142 | self.assertEqual(result.result.metadata["original_width"], 2400) |
| 143 | self.assertLessEqual(result.result.metadata["display_width"], 1568) |
| 144 | self.assertEqual(result.result.metadata["camera_metadata"]["focal_length_mm"], 50.0) |
| 145 | data_url = result.result.model_content[0]["image_url"]["url"] |
| 146 | self.assertTrue(data_url.startswith("data:image/jpeg;base64,")) |
| 147 | self.assertNotIn("model_content", result.result.as_dict()) |
| 148 | self.assertNotIn("data:image", (index.logs_dir / "tool_calls.jsonl").read_text(encoding="utf-8")) |
| 149 | |
| 150 | async def test_view_image_accepts_active_workspace_prefixed_path(self): |
| 151 | with tempfile.TemporaryDirectory() as tmp: |
| 152 | index = SessionIndex(tmp) |
| 153 | record = index.create(session_id="visual-session") |
| 154 | image_path = index.working_dir(record["session_id"]) / "script2video" / "frame.jpg" |
| 155 | Image.new("RGB", (800, 450), "blue").save(image_path) |
| 156 | registry = build_builtin_registry(tmp, index) |
| 157 | result = await registry.execute("view_image", {"path": f"{record['working_dir']}/script2video/frame.jpg"}) |
| 158 | self.assertTrue(result.ok) |
| 159 | self.assertEqual(result.metadata["path"], "script2video/frame.jpg") |
| 160 | |
| 161 | async def test_view_image_rejects_paths_outside_active_session(self): |
| 162 | with tempfile.TemporaryDirectory() as tmp: |
| 163 | index = SessionIndex(tmp) |
| 164 | active = index.create(session_id="active-session") |
| 165 | other = index.create(session_id="other-session") |
| 166 | other_image = index.working_dir(other["session_id"]) / "idea2video" / "other.png" |
| 167 | Image.new("RGB", (16, 9), "red").save(other_image) |
| 168 | index.set_active(active["session_id"]) |
| 169 | registry = build_builtin_registry(tmp, index) |
| 170 | escaped = await registry.execute("view_image", {"path": "../../outside.png"}) |
| 171 | cross_session = await registry.execute("view_image", {"path": f"{other['working_dir']}/idea2video/other.png"}) |
| 172 | self.assertFalse(escaped.ok) |
| 173 | self.assertFalse(cross_session.ok) |
| 174 | self.assertEqual(escaped.metadata["error_type"], "invalid_input") |
| 175 | self.assertEqual(cross_session.metadata["error_type"], "invalid_input") |
| 176 | |
| 177 | def test_concurrency_partition_groups_read_tools(self): |
| 178 | with tempfile.TemporaryDirectory() as tmp: |
| 179 | registry = build_builtin_registry(tmp, SessionIndex(tmp)) |
| 180 | batches = registry.partition_calls([ToolCall("read_file", {"path": "x"}), ToolCall("glob_files", {"pattern": "*"}), ToolCall("write_json", {"path": "x", "data": {}})]) |
| 181 | self.assertEqual(len(batches), 2) |
| 182 | self.assertEqual(len(batches[0]), 2) |
| 183 |