Files
zk-data-agent/tests/test_agent_prompting.py
T
Abdelrahman Abdallah b5e5824a56 Implemented the next parity slice: prompt-budget preflight and context collapse.
Core changes:

  - Added claw-code/src/token_budget.py for projected prompt size, chat-framing overhead, output reserve, and soft/hard input limits.
  - Wired preflight prompt-length validation and auto-compact/context collapse into claw-code/src/agent_runtime.py.
  - Extended claw-code/src/compact.py so compaction reports usage back to the runtime.
  - Added inspection surfaces in claw-code/src/agent_slash_commands.py and claw-code/src/main.py:
      - /token-budget and /budget
      - token-budget
  - Hardened claw-code/src/tokenizer_runtime.py so arbitrary simple model names fall back cleanly instead of trying a slow Transformers
    lookup.
  - Exported the new helpers in claw-code/src/__init__.py.

  Docs and tracking:

  - Updated claw-code/PARITY_CHECKLIST.md to mark prompt-length validation, token-budget calculation, and auto-compact/context collapse as
    done.
  - Updated claw-code/README.md and claw-code/TESTING_GUIDE.md with the new commands and behavior.

  Tests:

  - Added claw-code/tests/test_token_budget.py.
  - Updated claw-code/tests/test_agent_runtime.py, claw-code/tests/test_agent_slash_commands.py, claw-code/tests/test_main.py, and claw-code/
    tests/test_agent_context_usage.py.
  - Verified with:
      - /data/fs201059/aa17626/miniconda3/bin/python3 -m compileall src tests
      - /data/fs201059/aa17626/miniconda3/bin/python3 -m unittest -v tests.test_token_budget
        tests.test_agent_runtime.AgentRuntimeTests.test_agent_rejects_prompt_before_backend_when_preflight_input_budget_is_exceeded
        tests.test_agent_runtime.AgentRuntimeTests.test_agent_auto_compacts_context_before_next_model_call tests.test_agent_slash_commands
        tests.test_main tests.test_compact tests.test_tokenizer_runtime tests.test_agent_context_usage
      - Result: 71 tests, OK
2026-04-11 01:38:19 +02:00

301 lines
14 KiB
Python

from __future__ import annotations
import tempfile
import unittest
from pathlib import Path
from src.agent_prompting import build_prompt_context, build_system_prompt_parts, render_system_prompt
from src.plan_runtime import PlanRuntime
from src.agent_runtime import LocalCodingAgent
from src.agent_session import AgentSessionState
from src.agent_tools import default_tool_registry
from src.agent_types import AgentPermissions, AgentRuntimeConfig, ModelConfig
from src.task_runtime import TaskRuntime
class AgentPromptingTests(unittest.TestCase):
def test_prompt_builder_contains_expected_sections(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
runtime_config = AgentRuntimeConfig(
cwd=Path(tmp_dir),
permissions=AgentPermissions(
allow_file_write=True,
allow_shell_commands=False,
),
)
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('# System', prompt)
self.assertIn('# Doing tasks', prompt)
self.assertIn('# Using your tools', prompt)
self.assertIn('# Environment', prompt)
self.assertIn('__SYSTEM_PROMPT_DYNAMIC_BOUNDARY__', prompt)
self.assertIn('Primary working directory:', prompt)
def test_session_state_exports_messages_in_order(self) -> None:
state = AgentSessionState.create(['sys one', 'sys two'], 'hello')
state.append_assistant('working', ())
state.append_tool('read_file', 'call_1', '{"ok": true}')
messages = state.to_openai_messages()
self.assertEqual(messages[0]['role'], 'system')
self.assertEqual(messages[1]['role'], 'user')
self.assertEqual(messages[2]['role'], 'assistant')
self.assertEqual(messages[3]['role'], 'tool')
self.assertEqual(messages[3]['tool_call_id'], 'call_1')
def test_agent_can_render_prompt_without_contacting_model(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
agent = LocalCodingAgent(
model_config=ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct'),
runtime_config=AgentRuntimeConfig(cwd=Path(tmp_dir)),
)
prompt = agent.render_system_prompt()
self.assertIn('Claw Code Python', prompt)
self.assertIn('# System', prompt)
self.assertIn('# Environment', prompt)
def test_prompt_builder_mentions_plugins_when_cache_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
plugin_cache = workspace / '.port_sessions' / 'plugin_cache.json'
plugin_cache.parent.mkdir(parents=True, exist_ok=True)
plugin_cache.write_text(
'{"plugins":[{"name":"example-plugin","enabled":true}]}',
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('# Plugins', prompt)
def test_prompt_builder_mentions_hook_policy_when_manifest_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / '.claw-policy.json').write_text(
'{"trusted": false, "hooks": {"beforePrompt": ["Follow workspace policy."]}}',
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('# Hook Policy', prompt)
def test_prompt_builder_mentions_mcp_when_manifest_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / 'notes.txt').write_text('mcp notes\n', encoding='utf-8')
(workspace / '.claw-mcp.json').write_text(
(
'{"servers":[{"name":"workspace","resources":['
'{"uri":"mcp://workspace/notes","name":"Notes","path":"notes.txt"}'
']}]}'
),
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('# MCP', prompt)
def test_prompt_builder_mentions_search_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / '.claw-search.json').write_text(
'{"providers":[{"name":"local-search","provider":"searxng","baseUrl":"http://127.0.0.1:8080"}]}',
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('# Search', prompt)
self.assertIn('web_search', prompt)
def test_prompt_builder_mentions_remote_when_manifest_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / '.claw-remote.json').write_text(
(
'{"profiles":[{"name":"staging","mode":"ssh","target":"dev@staging",'
'"workspaceCwd":"/srv/app"}]}'
),
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('# Remote', prompt)
def test_prompt_builder_mentions_account_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / '.claw-account.json').write_text(
'{"profiles":[{"name":"local","provider":"openai","identity":"dev@example.com"}]}',
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('# 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)
claude_dir = workspace / '.claude'
claude_dir.mkdir()
(claude_dir / 'settings.json').write_text(
'{"review":{"mode":"strict"}}',
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('# Config', prompt)
def test_prompt_builder_mentions_lsp_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / 'sample.py').write_text('def helper(value):\n return value * 2\n', 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('# LSP', prompt)
self.assertIn('Use the LSP tool', prompt)
def test_prompt_builder_mentions_tasks_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
runtime = TaskRuntime.from_workspace(workspace)
runtime.create_task(title='Inspect runtime tasks')
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('# 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)
runtime = PlanRuntime.from_workspace(workspace)
runtime.update_plan(
[{'step': 'Inspect runtime planning', 'status': 'pending'}],
explanation='Track the current plan.',
)
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('# Planning', prompt)