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

Implemented the next parity slice.

Browse Source 
New runtime/code:

  - src/ask_user_runtime.py
  - src/team_runtime.py

  New real tools in src/agent_tools.py:

  - ask_user_question
  - team_create
  - team_delete
  - team_list
  - team_get
  - send_message
  - team_messages
  - notebook_edit
This commit is contained in:
Abdelrahman Abdallah
2026-04-07 02:51:30 +02:00
parent a54c90b18f
commit a5629295ac
21 changed files with 1886 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_agent_context.py
+28
View File
@@ -12,9 +12,11 @@ from src.agent_context import (
clear_context_caches,
set_system_prompt_injection,
)
from src.ask_user_runtime import AskUserRuntime
from src.plan_runtime import PlanRuntime
from src.agent_types import AgentRuntimeConfig
from src.task_runtime import TaskRuntime
from src.team_runtime import TeamRuntime
class AgentContextTests(unittest.TestCase):
@@ -151,6 +153,20 @@ class AgentContextTests(unittest.TestCase):
self.assertIn('Configured account profiles: 1', snapshot.user_context['accountRuntime'])
self.assertIn('dev@example.com', snapshot.user_context['accountRuntime'])
def test_user_context_loads_ask_user_runtime_summary(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir) / 'repo'
workspace.mkdir(parents=True)
(workspace / '.claw-ask-user.json').write_text(
'{"answers":[{"question":"Approve deploy?","answer":"yes"}]}',
encoding='utf-8',
)
snapshot = build_context_snapshot(AgentRuntimeConfig(cwd=workspace))
self.assertIn('askUserRuntime', snapshot.user_context)
self.assertIn('Queued answers: 1', snapshot.user_context['askUserRuntime'])
def test_user_context_loads_config_runtime_summary(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir) / 'repo'
@@ -195,6 +211,18 @@ class AgentContextTests(unittest.TestCase):
self.assertIn('planRuntime', snapshot.user_context)
self.assertIn('Total plan steps: 1', snapshot.user_context['planRuntime'])
def test_user_context_loads_team_runtime_summary(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir) / 'repo'
workspace.mkdir(parents=True)
runtime = TeamRuntime.from_workspace(workspace)
runtime.create_team('reviewers', members=['alice', 'bob'])
snapshot = build_context_snapshot(AgentRuntimeConfig(cwd=workspace))
self.assertIn('teamRuntime', snapshot.user_context)
self.assertIn('Configured teams: 1', snapshot.user_context['teamRuntime'])
@unittest.skipIf(shutil.which('git') is None, 'git is required for git context tests')
def test_git_status_snapshot_contains_branch_and_status(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
tests/test_agent_prompting.py
+38
View File
@@ -186,6 +186,25 @@ class AgentPromptingTests(unittest.TestCase):
prompt = render_system_prompt(parts)
self.assertIn('# Account', prompt)
def test_prompt_builder_mentions_ask_user_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / '.claw-ask-user.json').write_text(
'{"answers":[{"question":"Approve deploy?","answer":"yes"}]}',
encoding='utf-8',
)
runtime_config = AgentRuntimeConfig(cwd=workspace)
model_config = ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct')
prompt_context = build_prompt_context(runtime_config, model_config)
parts = build_system_prompt_parts(
prompt_context=prompt_context,
runtime_config=runtime_config,
tools=default_tool_registry(),
)
prompt = render_system_prompt(parts)
self.assertIn('# Ask User', prompt)
def test_prompt_builder_mentions_config_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
@@ -224,6 +243,25 @@ class AgentPromptingTests(unittest.TestCase):
prompt = render_system_prompt(parts)
self.assertIn('# Tasks', prompt)
def test_prompt_builder_mentions_teams_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / '.claw-teams.json').write_text(
'{"teams":[{"name":"reviewers","members":["alice","bob"]}]}',
encoding='utf-8',
)
runtime_config = AgentRuntimeConfig(cwd=workspace)
model_config = ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct')
prompt_context = build_prompt_context(runtime_config, model_config)
parts = build_system_prompt_parts(
prompt_context=prompt_context,
runtime_config=runtime_config,
tools=default_tool_registry(),
)
prompt = render_system_prompt(parts)
self.assertIn('# Teams', prompt)
def test_prompt_builder_mentions_planning_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
tests/test_agent_slash_commands.py
+37
View File
@@ -246,6 +246,43 @@ class AgentSlashCommandTests(unittest.TestCase):
self.assertIn('profile=local', login_result.final_output)
self.assertIn('logged_in=False', logout_result.final_output)
def test_ask_commands_render_local_ask_user_runtime(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / '.claw-ask-user.json').write_text(
'{"answers":[{"question":"Approve deploy?","answer":"yes"}]}',
encoding='utf-8',
)
agent = LocalCodingAgent(
model_config=ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct'),
runtime_config=AgentRuntimeConfig(cwd=workspace),
)
status_result = agent.run('/ask')
history_result = agent.run('/ask history')
self.assertIn('# Ask User', status_result.final_output)
self.assertIn('Queued answers: 1', status_result.final_output)
self.assertIn('# Ask User History', history_result.final_output)
def test_team_commands_render_local_team_runtime(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / '.claw-teams.json').write_text(
'{"teams":[{"name":"reviewers","members":["alice","bob"]}]}',
encoding='utf-8',
)
agent = LocalCodingAgent(
model_config=ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct'),
runtime_config=AgentRuntimeConfig(cwd=workspace),
)
teams_result = agent.run('/teams')
team_result = agent.run('/team reviewers')
messages_result = agent.run('/messages')
self.assertIn('# Teams', teams_result.final_output)
self.assertIn('reviewers', teams_result.final_output)
self.assertIn('# Team', team_result.final_output)
self.assertIn('alice', team_result.final_output)
self.assertIn('# Team Messages', messages_result.final_output)
def test_config_commands_render_local_reports(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
tests/test_ask_user_runtime.py
+55
View File
@@ -0,0 +1,55 @@
from __future__ import annotations
import tempfile
import unittest
from pathlib import Path
from src.ask_user_runtime import AskUserRuntime
from src.agent_tools import build_tool_context, default_tool_registry, execute_tool
from src.agent_types import AgentRuntimeConfig
class AskUserRuntimeTests(unittest.TestCase):
def test_ask_user_runtime_consumes_queued_answers(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / '.claw-ask-user.json').write_text(
'{"answers":[{"question":"Approve deploy?","answer":"yes"}]}',
encoding='utf-8',
)
runtime = AskUserRuntime.from_workspace(workspace)
response = runtime.answer(question='Approve deploy?')
restored = AskUserRuntime.from_workspace(workspace)
self.assertEqual(response.answer, 'yes')
self.assertEqual(response.source, 'queued')
self.assertEqual(len(restored.queued_answers), 0)
self.assertEqual(len(restored.history), 1)
def test_ask_user_tool_executes_against_runtime(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / '.claw-ask-user.json').write_text(
'{"answers":[{"question":"Choose mode","answer":"safe"}]}',
encoding='utf-8',
)
runtime = AskUserRuntime.from_workspace(workspace)
context = build_tool_context(
AgentRuntimeConfig(cwd=workspace),
ask_user_runtime=runtime,
)
result = execute_tool(
default_tool_registry(),
'ask_user_question',
{
'question': 'Choose mode',
'choices': ['safe', 'fast'],
'allow_free_text': False,
},
context,
)
self.assertTrue(result.ok)
self.assertIn('# Ask User', result.content)
self.assertIn('safe', result.content)
self.assertEqual(result.metadata.get('action'), 'ask_user_question')
tests/test_extended_tools.py
+37 -1
View File
@@ -5,7 +5,7 @@ import unittest
from pathlib import Path
from src.agent_tools import build_tool_context, default_tool_registry, execute_tool
from src.agent_types import AgentRuntimeConfig
from src.agent_types import AgentPermissions, AgentRuntimeConfig
class ExtendedToolTests(unittest.TestCase):
@@ -65,3 +65,39 @@ class ExtendedToolTests(unittest.TestCase):
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')
tests/test_main.py
+14
View File
@@ -140,6 +140,12 @@ class MainCliTests(unittest.TestCase):
self.assertEqual(args.command, 'account-profiles')
self.assertEqual(args.cwd, '.')
def test_parser_accepts_ask_user_runtime_commands(self) -> None:
parser = build_parser()
args = parser.parse_args(['ask-history', '--cwd', '.'])
self.assertEqual(args.command, 'ask-history')
self.assertEqual(args.cwd, '.')
def test_parser_accepts_search_runtime_commands(self) -> None:
parser = build_parser()
args = parser.parse_args(['search', 'repo query', '--cwd', '.', '--provider', 'local-search'])
@@ -167,3 +173,11 @@ class MainCliTests(unittest.TestCase):
self.assertEqual(args.command, 'config-get')
self.assertEqual(args.key_path, 'review.mode')
self.assertEqual(args.cwd, '.')
def test_parser_accepts_team_runtime_commands(self) -> None:
parser = build_parser()
args = parser.parse_args(['team-create', 'reviewers', '--member', 'alice', '--cwd', '.'])
self.assertEqual(args.command, 'team-create')
self.assertEqual(args.team_name, 'reviewers')
self.assertEqual(args.member, ['alice'])
self.assertEqual(args.cwd, '.')
tests/test_team_runtime.py
+76
View File
@@ -0,0 +1,76 @@
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.team_runtime import TeamRuntime
class TeamRuntimeTests(unittest.TestCase):
def test_team_runtime_persists_team_and_messages(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
runtime = TeamRuntime.from_workspace(workspace)
runtime.create_team('reviewers', members=['alice', 'bob'])
runtime.send_message(
team_name='reviewers',
text='Please review the patch.',
sender='agent',
recipient='alice',
)
restored = TeamRuntime.from_workspace(workspace)
self.assertEqual(len(restored.teams), 1)
self.assertEqual(restored.teams[0].name, 'reviewers')
self.assertEqual(len(restored.messages), 1)
self.assertIn('Please review the patch.', restored.render_messages())
def test_team_tools_execute_against_runtime(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
runtime = TeamRuntime.from_workspace(workspace)
context = build_tool_context(
AgentRuntimeConfig(
cwd=workspace,
permissions=AgentPermissions(allow_file_write=True),
),
team_runtime=runtime,
)
create_result = execute_tool(
default_tool_registry(),
'team_create',
{'team_name': 'reviewers', 'members': ['alice', 'bob']},
context,
)
send_result = execute_tool(
default_tool_registry(),
'send_message',
{
'team_name': 'reviewers',
'message': 'Check notebook changes',
'sender': 'agent',
},
context,
)
list_result = execute_tool(
default_tool_registry(),
'team_list',
{},
context,
)
message_result = execute_tool(
default_tool_registry(),
'team_messages',
{'team_name': 'reviewers'},
context,
)
self.assertTrue(create_result.ok)
self.assertEqual(create_result.metadata.get('action'), 'team_create')
self.assertTrue(send_result.ok)
self.assertEqual(send_result.metadata.get('action'), 'send_message')
self.assertIn('reviewers', list_result.content)
self.assertIn('Check notebook changes', message_result.content)
tests/test_terminal_bench_local.py
+52
View File
@@ -7,6 +7,7 @@ from pathlib import Path
from benchmarks.run_terminal_bench_local import (
TerminalBenchTask,
build_host_agent_command,
build_verifier_exec_command,
discover_tasks,
filter_tasks,
parse_dockerfile_workdir,
@@ -129,6 +130,57 @@ docker_image = "example/other:latest"
self.assertIn("instruction=$(cat", cmd)
self.assertNotIn("claw-code-agent agent", cmd)
def test_build_verifier_command_fakeroot_adds_flags(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
root = Path(tmp_dir)
workspace_dir = root / "workspace"
verifier_logs_dir = root / "verifier"
task_dir = root / "task"
tests_dir = task_dir / "tests"
workspace_dir.mkdir()
verifier_logs_dir.mkdir()
tests_dir.mkdir(parents=True)
image_path = root / "image.sif"
image_path.write_text("fake", encoding="utf-8")
task = TerminalBenchTask(
task_dir=task_dir,
name="terminal-bench/demo",
short_name="demo",
instruction="solve it",
docker_image="example/demo:latest",
agent_timeout_sec=30.0,
verifier_timeout_sec=30.0,
workdir="/workspace",
has_docker_compose=False,
)
cmd_no_fakeroot = build_verifier_exec_command(
task=task,
image_path=image_path,
workspace_dir=workspace_dir,
task_dir=task_dir,
verifier_logs_dir=verifier_logs_dir,
env={},
fakeroot=False,
)
self.assertNotIn("--fakeroot", cmd_no_fakeroot)
self.assertNotIn("--writable-tmpfs", cmd_no_fakeroot)
cmd_fakeroot = build_verifier_exec_command(
task=task,
image_path=image_path,
workspace_dir=workspace_dir,
task_dir=task_dir,
verifier_logs_dir=verifier_logs_dir,
env={},
fakeroot=True,
)
self.assertIn("--fakeroot", cmd_fakeroot)
self.assertIn("--writable-tmpfs", cmd_fakeroot)
self.assertIn("--contain", cmd_fakeroot)
self.assertIn("TMPDIR", cmd_fakeroot)
self.assertIn("CURL_CA_BUNDLE", cmd_fakeroot)
if __name__ == "__main__":
unittest.main()