wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enforce session workspace boundaries

Browse Source
This commit is contained in:
武阳
2026-05-08 17:09:07 +08:00
parent 5b14f55736
commit 0bbba2b936
13 changed files with 509 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_agent_prompting.py
+14 -11
View File
@@ -34,6 +34,9 @@ class AgentPromptingTests(unittest.TestCase):
prompt = render_system_prompt(parts)
self.assertIn('# 系统规则', prompt)
self.assertIn('# 处理任务', prompt)
self.assertIn('# 工作空间边界', prompt)
self.assertIn('当前 session 目录是本轮任务的默认工作区', prompt)
self.assertIn('平台服务代码目录始终只读', prompt)
self.assertIn('# 使用工具', prompt)
self.assertIn('# Skills', prompt)
self.assertIn('product-data', prompt)
@@ -100,7 +103,7 @@ class AgentPromptingTests(unittest.TestCase):
)
prompt = render_system_prompt(parts)
self.assertIn('# Plugins', prompt)
self.assertIn('# 插件', prompt)
def test_prompt_builder_mentions_hook_policy_when_manifest_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -119,7 +122,7 @@ class AgentPromptingTests(unittest.TestCase):
)
prompt = render_system_prompt(parts)
self.assertIn('# Hook Policy', prompt)
self.assertIn('# Hook 策略', prompt)
def test_prompt_builder_mentions_mcp_when_manifest_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -162,7 +165,7 @@ class AgentPromptingTests(unittest.TestCase):
)
prompt = render_system_prompt(parts)
self.assertIn('# Search', prompt)
self.assertIn('# 搜索', prompt)
self.assertIn('web_search', prompt)
def test_prompt_builder_mentions_remote_when_manifest_is_loaded(self) -> None:
@@ -185,7 +188,7 @@ class AgentPromptingTests(unittest.TestCase):
)
prompt = render_system_prompt(parts)
self.assertIn('# Remote', prompt)
self.assertIn('# 远程环境', prompt)
def test_prompt_builder_mentions_account_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -204,7 +207,7 @@ class AgentPromptingTests(unittest.TestCase):
)
prompt = render_system_prompt(parts)
self.assertIn('# Account', prompt)
self.assertIn('# 账号', prompt)
def test_prompt_builder_mentions_ask_user_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -223,7 +226,7 @@ class AgentPromptingTests(unittest.TestCase):
)
prompt = render_system_prompt(parts)
self.assertIn('# Ask User', prompt)
self.assertIn('# 询问用户', prompt)
def test_prompt_builder_mentions_config_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -244,7 +247,7 @@ class AgentPromptingTests(unittest.TestCase):
)
prompt = render_system_prompt(parts)
self.assertIn('# Config', prompt)
self.assertIn('# 配置', prompt)
def test_prompt_builder_mentions_lsp_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -261,7 +264,7 @@ class AgentPromptingTests(unittest.TestCase):
prompt = render_system_prompt(parts)
self.assertIn('# LSP', prompt)
self.assertIn('Use the LSP tool', prompt)
self.assertIn('使用 LSP 工具', prompt)
def test_prompt_builder_mentions_tasks_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -278,7 +281,7 @@ class AgentPromptingTests(unittest.TestCase):
)
prompt = render_system_prompt(parts)
self.assertIn('# Tasks', prompt)
self.assertIn('# 任务', prompt)
def test_prompt_builder_mentions_teams_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -297,7 +300,7 @@ class AgentPromptingTests(unittest.TestCase):
)
prompt = render_system_prompt(parts)
self.assertIn('# Teams', prompt)
self.assertIn('# 团队', prompt)
def test_prompt_builder_mentions_planning_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -317,4 +320,4 @@ class AgentPromptingTests(unittest.TestCase):
)
prompt = render_system_prompt(parts)
self.assertIn('# Planning', prompt)
self.assertIn('# 计划', prompt)
tests/test_data_agent_inputs.py
+11
View File
@@ -30,6 +30,17 @@ class DataAgentInputTests(unittest.TestCase):
self.assertEqual(table['title'], 'Sheet')
self.assertEqual(table['rows'][0], ['query', '预期domain', '0106-prev-domain', 'type', '备注'])
def test_load_input_sources_reads_explicit_external_file(self) -> None:
with tempfile.TemporaryDirectory() as root_dir, tempfile.TemporaryDirectory() as external_dir:
path = Path(external_dir) / 'cases.xlsx'
_write_xlsx(path)
payload = load_input_sources(root_dir, [str(path)])
self.assertEqual(payload['source_count'], 1)
self.assertEqual(payload['sources'][0]['path'], path.resolve().as_posix())
self.assertEqual(payload['sources'][0]['tables'][0]['rows'][1][0], '怎么开启查找设备')
def test_extract_case_evidence_profiles_and_extracts_rows(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
path = Path(tmp_dir) / 'cases.xlsx'
tests/test_extended_tools.py
+160
View File
@@ -49,6 +49,166 @@ class ExtendedToolTests(unittest.TestCase):
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:
tests/test_python_exec_tool.py
+20
View File
@@ -51,6 +51,26 @@ class PythonExecToolTests(TestCase):
self.assertIn(str(scratchpad), result.content)
self.assertEqual(result.metadata.get('scratchpad_directory'), str(scratchpad))
def test_python_exec_runs_from_session_scratchpad(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
root = Path(tmp_dir)
scratchpad = (root / 'session' / 'scratchpad').resolve()
scratchpad.mkdir(parents=True)
config = AgentRuntimeConfig(
cwd=root,
permissions=AgentPermissions(allow_shell_commands=True),
)
context = build_tool_context(config, scratchpad_directory=scratchpad)
result = execute_tool(
default_tool_registry(),
'python_exec',
{'code': 'import os\nprint(os.getcwd())'},
context,
)
self.assertTrue(result.ok)
self.assertIn(str(scratchpad), result.content)
def test_python_exec_prefers_user_python_env(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
root = Path(tmp_dir)