武阳
300 lines
11 KiB
Python
300 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from src.agent_tools import build_tool_context, default_tool_registry, execute_tool
|
|
from src.agent_types import AgentPermissions, AgentRuntimeConfig
|
|
from src.lsp_runtime import LSPRuntime
|
|
|
|
|
|
class ExtendedToolTests(unittest.TestCase):
|
|
def test_web_fetch_rejects_file_url(self) -> None:
|
|
"""Verify that file:// URLs are blocked to prevent SSRF attacks."""
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
workspace = Path(tmp_dir)
|
|
target = workspace / 'page.txt'
|
|
target.write_text('hello from web fetch\n', encoding='utf-8')
|
|
context = build_tool_context(
|
|
AgentRuntimeConfig(cwd=workspace),
|
|
tool_registry=default_tool_registry(),
|
|
)
|
|
result = execute_tool(
|
|
default_tool_registry(),
|
|
'web_fetch',
|
|
{'url': target.resolve().as_uri()},
|
|
context,
|
|
)
|
|
|
|
self.assertFalse(result.ok)
|
|
self.assertIn('http or https', result.content)
|
|
|
|
def test_tool_search_lists_matching_tools(self) -> None:
|
|
registry = default_tool_registry()
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
context = build_tool_context(
|
|
AgentRuntimeConfig(cwd=Path(tmp_dir)),
|
|
tool_registry=registry,
|
|
)
|
|
result = execute_tool(
|
|
registry,
|
|
'tool_search',
|
|
{'query': 'file'},
|
|
context,
|
|
)
|
|
|
|
self.assertTrue(result.ok)
|
|
self.assertIn('# Tool Search', result.content)
|
|
self.assertIn('read_file', result.content)
|
|
self.assertIn('write_file', result.content)
|
|
|
|
def test_grep_search_skips_generated_dirs_by_default(self) -> None:
|
|
registry = default_tool_registry()
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
workspace = Path(tmp_dir)
|
|
(workspace / '.next' / 'static').mkdir(parents=True)
|
|
(workspace / '.next' / 'static' / 'bundle.js').write_text(
|
|
'router_session_parquet should not be searched by default\n',
|
|
encoding='utf-8',
|
|
)
|
|
(workspace / 'README.md').write_text(
|
|
'router_session_parquet is documented here\n',
|
|
encoding='utf-8',
|
|
)
|
|
context = build_tool_context(
|
|
AgentRuntimeConfig(cwd=workspace),
|
|
tool_registry=registry,
|
|
)
|
|
result = execute_tool(
|
|
registry,
|
|
'grep_search',
|
|
{'pattern': 'router_session_parquet'},
|
|
context,
|
|
)
|
|
|
|
self.assertTrue(result.ok)
|
|
self.assertIn('README.md:1:', result.content)
|
|
self.assertNotIn('.next/static/bundle.js', result.content)
|
|
|
|
def test_grep_search_truncates_very_long_lines(self) -> None:
|
|
registry = default_tool_registry()
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
workspace = Path(tmp_dir)
|
|
long_line = 'prefix needle ' + ('x' * 5000)
|
|
(workspace / 'large.txt').write_text(long_line, encoding='utf-8')
|
|
context = build_tool_context(
|
|
AgentRuntimeConfig(cwd=workspace, max_output_chars=1200),
|
|
tool_registry=registry,
|
|
)
|
|
result = execute_tool(
|
|
registry,
|
|
'grep_search',
|
|
{'pattern': 'needle'},
|
|
context,
|
|
)
|
|
|
|
self.assertTrue(result.ok)
|
|
self.assertIn('large.txt:1:', result.content)
|
|
self.assertIn('[line truncated,', result.content)
|
|
self.assertLess(len(result.content), 1200)
|
|
|
|
def test_read_file_can_read_explicit_external_path(self) -> None:
|
|
registry = default_tool_registry()
|
|
with tempfile.TemporaryDirectory() as workspace_dir, tempfile.TemporaryDirectory() as external_dir:
|
|
external = Path(external_dir) / 'reference.txt'
|
|
external.write_text('external reference\n', encoding='utf-8')
|
|
context = build_tool_context(
|
|
AgentRuntimeConfig(cwd=Path(workspace_dir)),
|
|
tool_registry=registry,
|
|
)
|
|
result = execute_tool(
|
|
registry,
|
|
'read_file',
|
|
{'path': str(external)},
|
|
context,
|
|
)
|
|
|
|
self.assertTrue(result.ok)
|
|
self.assertIn('external reference', result.content)
|
|
|
|
def test_write_file_blocks_platform_code_paths(self) -> None:
|
|
registry = default_tool_registry()
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
workspace = Path(tmp_dir)
|
|
(workspace / 'src').mkdir()
|
|
(workspace / 'src' / 'agent_tools.py').write_text('', encoding='utf-8')
|
|
(workspace / 'backend' / 'api').mkdir(parents=True)
|
|
(workspace / 'backend' / 'api' / 'server.py').write_text('', encoding='utf-8')
|
|
(workspace / 'frontend' / 'app').mkdir(parents=True)
|
|
scratchpad = (
|
|
workspace
|
|
/ '.port_sessions'
|
|
/ 'accounts'
|
|
/ 'user'
|
|
/ 'sessions'
|
|
/ 'thread'
|
|
/ 'scratchpad'
|
|
)
|
|
scratchpad.mkdir(parents=True)
|
|
context = build_tool_context(
|
|
AgentRuntimeConfig(
|
|
cwd=workspace,
|
|
permissions=AgentPermissions(allow_file_write=True),
|
|
),
|
|
scratchpad_directory=scratchpad,
|
|
tool_registry=registry,
|
|
)
|
|
|
|
result = execute_tool(
|
|
registry,
|
|
'write_file',
|
|
{'path': 'src/new_file.py', 'content': 'print(1)\n'},
|
|
context,
|
|
)
|
|
output_result = execute_tool(
|
|
registry,
|
|
'write_file',
|
|
{
|
|
'path': '.port_sessions/accounts/user/sessions/thread/output/report.txt',
|
|
'content': 'ok\n',
|
|
},
|
|
context,
|
|
)
|
|
|
|
self.assertFalse(result.ok)
|
|
self.assertEqual(result.metadata.get('error_kind'), 'permission_denied')
|
|
self.assertTrue(output_result.ok)
|
|
|
|
def test_logical_session_paths_route_to_current_session(self) -> None:
|
|
registry = default_tool_registry()
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
workspace = Path(tmp_dir)
|
|
scratchpad = (
|
|
workspace
|
|
/ '.port_sessions'
|
|
/ 'accounts'
|
|
/ 'user'
|
|
/ 'sessions'
|
|
/ 'thread'
|
|
/ 'scratchpad'
|
|
)
|
|
scratchpad.mkdir(parents=True)
|
|
context = build_tool_context(
|
|
AgentRuntimeConfig(
|
|
cwd=workspace,
|
|
permissions=AgentPermissions(allow_file_write=True),
|
|
),
|
|
scratchpad_directory=scratchpad,
|
|
tool_registry=registry,
|
|
)
|
|
write_result = execute_tool(
|
|
registry,
|
|
'write_file',
|
|
{'path': 'output/report.md', 'content': 'session report\n'},
|
|
context,
|
|
)
|
|
read_result = execute_tool(
|
|
registry,
|
|
'read_file',
|
|
{'path': 'output/report.md'},
|
|
context,
|
|
)
|
|
session_file_exists = (scratchpad.parent / 'output' / 'report.md').is_file()
|
|
root_file_exists = (workspace / 'output' / 'report.md').exists()
|
|
|
|
self.assertTrue(write_result.ok, write_result.content)
|
|
self.assertTrue(read_result.ok, read_result.content)
|
|
self.assertIn('session report', read_result.content)
|
|
self.assertTrue(session_file_exists)
|
|
self.assertFalse(root_file_exists)
|
|
|
|
def test_sleep_tool_waits_briefly_and_returns_metadata(self) -> None:
|
|
registry = default_tool_registry()
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
context = build_tool_context(
|
|
AgentRuntimeConfig(cwd=Path(tmp_dir)),
|
|
tool_registry=registry,
|
|
)
|
|
result = execute_tool(
|
|
registry,
|
|
'sleep',
|
|
{'seconds': 0.01},
|
|
context,
|
|
)
|
|
|
|
self.assertTrue(result.ok)
|
|
self.assertIn('slept for', result.content)
|
|
self.assertEqual(result.metadata.get('action'), 'sleep')
|
|
|
|
def test_notebook_edit_updates_ipynb_cell(self) -> None:
|
|
registry = default_tool_registry()
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
workspace = Path(tmp_dir)
|
|
notebook = workspace / 'demo.ipynb'
|
|
notebook.write_text(
|
|
'{\n'
|
|
' "cells": [\n'
|
|
' {"cell_type": "code", "metadata": {}, "source": ["print(1)\\n"], "outputs": [], "execution_count": null}\n'
|
|
' ],\n'
|
|
' "metadata": {},\n'
|
|
' "nbformat": 4,\n'
|
|
' "nbformat_minor": 5\n'
|
|
'}\n',
|
|
encoding='utf-8',
|
|
)
|
|
context = build_tool_context(
|
|
AgentRuntimeConfig(
|
|
cwd=workspace,
|
|
permissions=AgentPermissions(allow_file_write=True),
|
|
),
|
|
tool_registry=registry,
|
|
)
|
|
result = execute_tool(
|
|
registry,
|
|
'notebook_edit',
|
|
{'path': 'demo.ipynb', 'cell_index': 0, 'source': 'print(2)\n'},
|
|
context,
|
|
)
|
|
updated = notebook.read_text(encoding='utf-8')
|
|
|
|
self.assertTrue(result.ok)
|
|
self.assertIn('updated notebook cell 0', result.content)
|
|
self.assertIn('print(2)', updated)
|
|
self.assertEqual(result.metadata.get('action'), 'notebook_edit')
|
|
|
|
def test_lsp_tool_returns_definition_report(self) -> None:
|
|
registry = default_tool_registry()
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
workspace = Path(tmp_dir)
|
|
(workspace / 'sample.py').write_text(
|
|
'def helper(value):\n'
|
|
' return value * 2\n'
|
|
'\n'
|
|
'def run(item):\n'
|
|
' return helper(item)\n',
|
|
encoding='utf-8',
|
|
)
|
|
context = build_tool_context(
|
|
AgentRuntimeConfig(cwd=workspace),
|
|
tool_registry=registry,
|
|
lsp_runtime=LSPRuntime.from_workspace(workspace),
|
|
)
|
|
result = execute_tool(
|
|
registry,
|
|
'LSP',
|
|
{
|
|
'operation': 'goToDefinition',
|
|
'file_path': 'sample.py',
|
|
'line': 5,
|
|
'character': 12,
|
|
},
|
|
context,
|
|
)
|
|
|
|
self.assertTrue(result.ok)
|
|
self.assertIn('# LSP Definition', result.content)
|
|
self.assertIn('function helper', result.content)
|
|
self.assertEqual(result.metadata.get('action'), 'lsp_query')
|
|
self.assertEqual(result.metadata.get('operation'), 'goToDefinition')
|