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

add ىثص ؤخقث شلثىف قعىفهةث

Browse Source
This commit is contained in:
Abdelrahman Abdallah
2026-04-02 22:08:24 +02:00
parent 2c6763eb08
commit 9d90d6a10a
10 changed files with 1699 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_agent_runtime.py
+171
View File
@@ -362,6 +362,14 @@ class AgentRuntimeTests(unittest.TestCase):
self.assertIsNotNone(result.scratchpad_directory)
assert result.scratchpad_directory is not None
self.assertTrue(Path(result.scratchpad_directory).is_dir())
assistant_messages = [
message for message in result.transcript if message.get('role') == 'assistant'
]
self.assertEqual(len(assistant_messages), 1)
metadata = assistant_messages[0].get('metadata', {})
mutation_totals = metadata.get('mutation_totals', {})
self.assertGreaterEqual(mutation_totals.get('assistant_delta_append', 0), 2)
self.assertEqual(mutation_totals.get('assistant_finalize', 0), 1)
def test_agent_streams_tool_calls_and_reconstructs_arguments(self) -> None:
responses = [
@@ -439,6 +447,15 @@ class AgentRuntimeTests(unittest.TestCase):
self.assertEqual(result.usage.output_tokens, 6)
assistant_messages = [message for message in result.transcript if message.get('role') == 'assistant']
self.assertTrue(any(message.get('tool_calls') for message in assistant_messages))
assistant_with_tool = next(
message
for message in assistant_messages
if message.get('tool_calls')
)
metadata = assistant_with_tool.get('metadata', {})
mutation_totals = metadata.get('mutation_totals', {})
self.assertGreaterEqual(mutation_totals.get('assistant_tool_call_delta', 0), 2)
self.assertEqual(mutation_totals.get('assistant_finalize', 0), 1)
def test_transcript_entries_include_structured_blocks(self) -> None:
responses = [
@@ -578,6 +595,10 @@ class AgentRuntimeTests(unittest.TestCase):
if message.get('metadata', {}).get('kind') == 'compact_boundary'
]
self.assertEqual(len(transcript_compact_messages), 1)
compact_metadata = transcript_compact_messages[0].get('metadata', {})
self.assertEqual(compact_metadata.get('compaction_depth'), 1)
self.assertEqual(compact_metadata.get('nested_compaction_count'), 0)
self.assertIn('preserved_tail_ids', compact_metadata)
def test_agent_enforces_total_token_budget(self) -> None:
responses = [
@@ -1386,6 +1407,9 @@ class AgentRuntimeTests(unittest.TestCase):
self.assertEqual(len(metadata.get('child_session_ids', [])), 2)
self.assertEqual(len(metadata.get('child_results', [])), 2)
self.assertIn('Delegated agent completed 2 sequential subtasks.', tool_messages[0].get('content', ''))
self.assertTrue(any(event.get('type') == 'delegate_group_result' for event in result.events))
child_events = [event for event in result.events if event.get('type') == 'delegate_subtask_result']
self.assertEqual(len(child_events), 2)
second_child_request = recorded_payloads[2]['messages']
assert isinstance(second_child_request, list)
self.assertTrue(
@@ -1497,6 +1521,153 @@ class AgentRuntimeTests(unittest.TestCase):
summary = '\n'.join(manager.summary_lines())
self.assertIn(f'group={group.group_id}', summary)
def test_agent_can_delegate_into_resumed_child_session(self) -> None:
responses = [
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Seed child result.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 5, 'completion_tokens': 2},
},
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Delegating into a resumed child session.',
'tool_calls': [
{
'id': 'call_1',
'type': 'function',
'function': {
'name': 'delegate_agent',
'arguments': '{}',
},
}
],
},
'finish_reason': 'tool_calls',
}
],
'usage': {'prompt_tokens': 8, 'completion_tokens': 3},
},
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Resumed child result.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 6, 'completion_tokens': 2},
},
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Parent completed after resumed child.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 7, 'completion_tokens': 2},
},
]
recorded_payloads: list[dict[str, object]] = []
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
session_dir = workspace / '.port_sessions' / 'agent'
with patch(
'src.openai_compat.request.urlopen',
side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads),
):
seed_agent = LocalCodingAgent(
model_config=ModelConfig(
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=AgentRuntimeConfig(
cwd=workspace,
session_directory=session_dir,
),
)
seeded = seed_agent.run('Seed the delegated child')
resumed_child_id = seeded.session_id or ''
delegate_arguments = json.dumps(
{
'subtasks': [
{
'label': 'resume_child',
'prompt': 'Continue the delegated child.',
'resume_session_id': resumed_child_id,
'max_turns': 2,
}
],
'max_turns': 2,
}
)
responses[1]['choices'][0]['message']['tool_calls'][0]['function']['arguments'] = delegate_arguments
parent_agent = LocalCodingAgent(
model_config=ModelConfig(
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=AgentRuntimeConfig(
cwd=workspace,
session_directory=session_dir,
),
)
result = parent_agent.run('Delegate into the resumed child')
self.assertEqual(result.final_output, 'Parent completed after resumed child.')
tool_messages = [message for message in result.transcript if message.get('role') == 'tool']
self.assertEqual(len(tool_messages), 1)
metadata = tool_messages[0].get('metadata', {})
self.assertEqual(metadata.get('resumed_children'), 1)
child_results = metadata.get('child_results', [])
self.assertEqual(len(child_results), 1)
self.assertEqual(child_results[0].get('resume_used'), True)
self.assertEqual(child_results[0].get('resumed_from_session_id'), resumed_child_id)
self.assertTrue(
any(
event.get('type') == 'delegate_subtask_result'
and event.get('resume_used')
for event in result.events
)
)
resumed_child_messages = recorded_payloads[2]['messages']
assert isinstance(resumed_child_messages, list)
resumed_contents = [
message.get('content')
for message in resumed_child_messages
if isinstance(message, dict)
]
self.assertIn('Seed the delegated child', resumed_contents)
self.assertIn('Seed child result.', resumed_contents)
self.assertIn('Continue the delegated child.', resumed_contents)
manager = parent_agent.agent_manager
assert manager is not None
child_records = [
record
for record in manager.completed_records()
if record.parent_agent_id == parent_agent.managed_agent_id
]
self.assertEqual(len(child_records), 1)
self.assertEqual(child_records[0].resumed_from_session_id, resumed_child_id)
summary = '\n'.join(manager.summary_lines())
self.assertIn(f'resumed_from={resumed_child_id}', summary)
def test_agent_enforces_reasoning_token_budget(self) -> None:
responses = [
{
tests/test_query_engine_runtime.py
+540 -1
View File
@@ -10,7 +10,7 @@ from src.agent_runtime import LocalCodingAgent
from src.agent_types import AgentRuntimeConfig, ModelConfig
from src.openai_compat import OpenAICompatClient
from src.plugin_runtime import PluginRuntime
from src.query_engine import QueryEnginePort
from src.query_engine import QueryEngineConfig, QueryEnginePort
class FakeHTTPResponse:
@@ -27,6 +27,27 @@ class FakeHTTPResponse:
return None
class FakeStreamingHTTPResponse:
def __init__(self, payloads: list[dict[str, object]]) -> None:
self.lines: list[bytes] = []
for payload in payloads:
chunk = f'data: {json.dumps(payload)}\n\n'
self.lines.extend(part.encode('utf-8') for part in chunk.splitlines(keepends=True))
done_chunk = 'data: [DONE]\n\n'
self.lines.extend(part.encode('utf-8') for part in done_chunk.splitlines(keepends=True))
def readline(self) -> bytes:
if not self.lines:
return b''
return self.lines.pop(0)
def __enter__(self) -> 'FakeStreamingHTTPResponse':
return self
def __exit__(self, exc_type, exc, tb) -> None:
return None
def make_recording_urlopen_side_effect(
responses: list[dict[str, object]],
recorded_payloads: list[dict[str, object]],
@@ -50,6 +71,17 @@ def make_urlopen_side_effect(responses: list[dict[str, object]]):
return _fake_urlopen
def make_streaming_urlopen_side_effect(
responses: list[list[dict[str, object]]],
):
queued = [FakeStreamingHTTPResponse(payloads) for payloads in responses]
def _fake_urlopen(request_obj, timeout=None): # noqa: ANN001
return queued.pop(0)
return _fake_urlopen
class QueryEngineRuntimeTests(unittest.TestCase):
def test_plugin_runtime_discovers_local_manifest(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -249,6 +281,112 @@ class QueryEngineRuntimeTests(unittest.TestCase):
)
)
def test_runtime_agent_executes_plugin_virtual_tool(self) -> None:
responses = [
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Calling the plugin virtual tool.',
'tool_calls': [
{
'id': 'call_1',
'type': 'function',
'function': {
'name': 'demo_virtual',
'arguments': '{"topic": "plugins"}',
},
}
],
},
'finish_reason': 'tool_calls',
}
],
'usage': {'prompt_tokens': 8, 'completion_tokens': 3},
},
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Plugin virtual tool completed.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 7, 'completion_tokens': 2},
},
]
recorded_payloads: list[dict[str, object]] = []
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
plugin_dir = workspace / 'plugins' / 'demo'
plugin_dir.mkdir(parents=True)
(plugin_dir / 'plugin.json').write_text(
json.dumps(
{
'name': 'demo-plugin',
'virtualTools': [
{
'name': 'demo_virtual',
'description': 'Return a rendered plugin response.',
'responseTemplate': 'Plugin says {topic}',
'parameters': {
'type': 'object',
'properties': {'topic': {'type': 'string'}},
'required': ['topic'],
},
}
],
'toolHooks': {
'demo_virtual': {
'afterResult': 'Use the virtual tool result in the next reply.',
}
},
}
),
encoding='utf-8',
)
with patch(
'src.openai_compat.request.urlopen',
side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads),
):
agent = LocalCodingAgent(
model_config=ModelConfig(
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=AgentRuntimeConfig(cwd=workspace),
)
result = agent.run('Use the plugin virtual tool')
self.assertEqual(result.final_output, 'Plugin virtual tool completed.')
self.assertTrue(
any(event.get('type') == 'plugin_virtual_tool_result' for event in result.events)
)
tool_names = [
item['function']['name']
for item in recorded_payloads[0]['tools']
if isinstance(item, dict) and isinstance(item.get('function'), dict)
]
self.assertIn('demo_virtual', tool_names)
tool_messages = [message for message in result.transcript if message.get('role') == 'tool']
self.assertEqual(len(tool_messages), 1)
metadata = tool_messages[0].get('metadata', {})
self.assertEqual(metadata.get('action'), 'plugin_virtual_tool')
self.assertEqual(metadata.get('plugin_name'), 'demo-plugin')
self.assertEqual(metadata.get('virtual_tool'), 'demo_virtual')
second_messages = recorded_payloads[1]['messages']
assert isinstance(second_messages, list)
self.assertTrue(
any(
isinstance(message, dict)
and 'Use the virtual tool result in the next reply.' in str(message.get('content', ''))
for message in second_messages
)
)
def test_runtime_agent_injects_plugin_tool_runtime_guidance(self) -> None:
responses = [
{
@@ -337,6 +475,93 @@ class QueryEngineRuntimeTests(unittest.TestCase):
)
)
def test_runtime_agent_supports_plugin_before_tool_guidance(self) -> None:
responses = [
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Reading through plugin guidance.',
'tool_calls': [
{
'id': 'call_1',
'type': 'function',
'function': {
'name': 'read_file',
'arguments': '{"path": "guide.txt"}',
},
}
],
},
'finish_reason': 'tool_calls',
}
],
'usage': {'prompt_tokens': 8, 'completion_tokens': 3},
},
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Plugin before-tool guidance consumed.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 7, 'completion_tokens': 2},
},
]
recorded_payloads: list[dict[str, object]] = []
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / 'guide.txt').write_text('plugin before tool\n', encoding='utf-8')
plugin_dir = workspace / 'plugins' / 'demo'
plugin_dir.mkdir(parents=True)
(plugin_dir / 'plugin.json').write_text(
json.dumps(
{
'name': 'demo-plugin',
'toolHooks': {
'read_file': {
'beforeTool': 'Validate the path before reading.',
}
},
}
),
encoding='utf-8',
)
with patch(
'src.openai_compat.request.urlopen',
side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads),
):
agent = LocalCodingAgent(
model_config=ModelConfig(
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=AgentRuntimeConfig(cwd=workspace),
)
result = agent.run('Read the file and continue')
self.assertEqual(result.final_output, 'Plugin before-tool guidance consumed.')
self.assertTrue(any(event.get('type') == 'plugin_tool_preflight' for event in result.events))
runtime_messages = [
message for message in result.transcript
if message.get('metadata', {}).get('kind') == 'plugin_tool_runtime'
]
self.assertEqual(len(runtime_messages), 1)
self.assertIn('Before tool: demo-plugin: Validate the path before reading.', runtime_messages[0].get('content', ''))
second_messages = recorded_payloads[1]['messages']
assert isinstance(second_messages, list)
self.assertTrue(
any(
isinstance(message, dict)
and 'Before tool: demo-plugin: Validate the path before reading.' in str(message.get('content', ''))
for message in second_messages
)
)
def test_runtime_agent_blocks_tool_via_plugin_manifest(self) -> None:
responses = [
{
@@ -495,3 +720,317 @@ class QueryEngineRuntimeTests(unittest.TestCase):
self.assertIn('## Runtime Message Kinds', summary)
self.assertIn('- plugin_tool_runtime=1', summary)
self.assertIn('- transcript_messages=', summary)
def test_query_engine_runtime_stream_emits_runtime_summary_event(self) -> None:
responses = [
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Streaming summary ready.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 8, 'completion_tokens': 3},
}
]
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
with patch(
'src.openai_compat.request.urlopen',
side_effect=make_urlopen_side_effect(responses),
):
agent = LocalCodingAgent(
model_config=ModelConfig(
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=AgentRuntimeConfig(cwd=workspace),
)
engine = QueryEnginePort.from_runtime_agent(agent)
events = list(engine.stream_submit_message('Summarize the repo'))
runtime_summary_events = [
event for event in events if event.get('type') == 'runtime_summary'
]
self.assertEqual(len(runtime_summary_events), 1)
summary_event = runtime_summary_events[0]
self.assertIn('runtime_event_counts', summary_event)
self.assertIn('transcript_store_entries', summary_event)
self.assertIn('transcript_store_compactions', summary_event)
def test_query_engine_compacts_transcript_store_with_summary_entry(self) -> None:
engine = QueryEnginePort.from_workspace()
engine.config = QueryEngineConfig(max_turns=6, compact_after_turns=2)
engine.submit_message('first prompt')
engine.submit_message('second prompt')
engine.submit_message('third prompt')
replay = engine.replay_user_messages()
self.assertTrue(any('[transcript-compaction ' in entry for entry in replay))
summary = engine.render_summary()
self.assertIn('## Transcript Store', summary)
self.assertIn('Transcript compactions:', summary)
def test_query_engine_runtime_summary_tracks_mutation_counts(self) -> None:
responses = [
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Reading the large file first.',
'tool_calls': [
{
'id': 'call_1',
'type': 'function',
'function': {
'name': 'read_file',
'arguments': '{"path": "large.txt"}',
},
}
],
},
'finish_reason': 'tool_calls',
}
],
'usage': {'prompt_tokens': 8, 'completion_tokens': 3},
},
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Mutation summary completed.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 7, 'completion_tokens': 2},
},
]
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / 'large.txt').write_text(('alpha beta gamma\n' * 400), encoding='utf-8')
with patch(
'src.openai_compat.request.urlopen',
side_effect=make_urlopen_side_effect(responses),
):
agent = LocalCodingAgent(
model_config=ModelConfig(
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=AgentRuntimeConfig(
cwd=workspace,
auto_snip_threshold_tokens=120,
compact_preserve_messages=0,
),
)
engine = QueryEnginePort.from_runtime_agent(agent)
turn = engine.submit_message('Read the large file and summarize it')
summary = engine.render_summary()
self.assertEqual(turn.output, 'Mutation summary completed.')
self.assertIn('## Runtime Mutations', summary)
self.assertIn('- snip_tombstone=1', summary)
self.assertIn('- tool_finalize_replace=1', summary)
self.assertIn('## Runtime Context Reduction', summary)
self.assertIn('- snipped_messages=1', summary)
def test_query_engine_runtime_summary_tracks_compaction_lineage(self) -> None:
responses = [
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Read the large file first.',
'tool_calls': [
{
'id': 'call_1',
'type': 'function',
'function': {
'name': 'read_file',
'arguments': '{"path": "large.txt"}',
},
}
],
},
'finish_reason': 'tool_calls',
}
],
'usage': {'prompt_tokens': 8, 'completion_tokens': 3},
},
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Compaction lineage summary completed.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 7, 'completion_tokens': 2},
},
]
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / 'large.txt').write_text(('alpha beta gamma\n' * 400), encoding='utf-8')
with patch(
'src.openai_compat.request.urlopen',
side_effect=make_urlopen_side_effect(responses),
):
agent = LocalCodingAgent(
model_config=ModelConfig(
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=AgentRuntimeConfig(
cwd=workspace,
auto_compact_threshold_tokens=120,
compact_preserve_messages=0,
),
)
engine = QueryEnginePort.from_runtime_agent(agent)
turn = engine.submit_message('Read the large file and summarize it')
summary = engine.render_summary()
self.assertEqual(turn.output, 'Compaction lineage summary completed.')
self.assertIn('## Runtime Context Reduction', summary)
self.assertIn('- compact_boundaries=1', summary)
self.assertIn('- compacted_lineages=', summary)
self.assertIn('## Runtime Lineage', summary)
self.assertIn('- seen_lineages=', summary)
self.assertIn('- compacted_lineages=', summary)
self.assertIn('- max_source_revision=', summary)
def test_query_engine_runtime_summary_tracks_assistant_stream_mutations(self) -> None:
responses = [
[
{'choices': [{'delta': {'content': 'Streaming '}, 'finish_reason': None}]},
{'choices': [{'delta': {'content': 'works.'}, 'finish_reason': None}]},
{
'choices': [{'delta': {}, 'finish_reason': 'stop'}],
'usage': {'prompt_tokens': 14, 'completion_tokens': 5},
},
]
]
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
with patch(
'src.openai_compat.request.urlopen',
side_effect=make_streaming_urlopen_side_effect(responses),
):
agent = LocalCodingAgent(
model_config=ModelConfig(
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=AgentRuntimeConfig(
cwd=workspace,
stream_model_responses=True,
),
)
engine = QueryEnginePort.from_runtime_agent(agent)
turn = engine.submit_message('Say streaming works')
summary = engine.render_summary()
self.assertEqual(turn.output, 'Streaming works.')
self.assertIn('## Runtime Mutations', summary)
self.assertIn('- assistant_delta_append=2', summary)
self.assertIn('- assistant_finalize=1', summary)
def test_query_engine_runtime_summary_tracks_delegate_orchestration(self) -> None:
responses = [
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Delegating multiple subtasks.',
'tool_calls': [
{
'id': 'call_1',
'type': 'function',
'function': {
'name': 'delegate_agent',
'arguments': json.dumps(
{
'subtasks': [
{'label': 'scan', 'prompt': 'Scan the project.'},
{'label': 'summarize', 'prompt': 'Summarize the project.'},
],
'max_turns': 2,
}
),
},
}
],
},
'finish_reason': 'tool_calls',
}
],
'usage': {'prompt_tokens': 8, 'completion_tokens': 3},
},
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Child scan result.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 5, 'completion_tokens': 2},
},
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Child summary result.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 6, 'completion_tokens': 2},
},
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Parent completed after multi-delegate.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 7, 'completion_tokens': 2},
},
]
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
with patch(
'src.openai_compat.request.urlopen',
side_effect=make_urlopen_side_effect(responses),
):
agent = LocalCodingAgent(
model_config=ModelConfig(
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=AgentRuntimeConfig(cwd=workspace),
)
engine = QueryEnginePort.from_runtime_agent(agent)
turn = engine.submit_message('Use multiple delegated subtasks')
summary = engine.render_summary()
self.assertEqual(turn.output, 'Parent completed after multi-delegate.')
self.assertIn('## Runtime Orchestration', summary)
self.assertIn('- group_status:completed=1', summary)
self.assertIn('- child_stop:stop=2', summary)