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

Implemented the next parity slice: prompt-budget preflight and context collapse.

Browse Source 
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
This commit is contained in:
Abdelrahman Abdallah
2026-04-11 01:38:19 +02:00
parent aacf0a212a
commit b5e5824a56
23 changed files with 2641 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_agent_context.py
+11
View File
@@ -184,6 +184,17 @@ class AgentContextTests(unittest.TestCase):
self.assertIn('Config sources: 1', snapshot.user_context['configRuntime'])
self.assertIn('Effective keys: 2', snapshot.user_context['configRuntime'])
def test_user_context_loads_lsp_runtime_summary(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir) / 'repo'
workspace.mkdir(parents=True)
(workspace / 'sample.py').write_text('def helper(value):\n return value * 2\n', encoding='utf-8')
snapshot = build_context_snapshot(AgentRuntimeConfig(cwd=workspace))
self.assertIn('lspRuntime', snapshot.user_context)
self.assertIn('Indexed candidate files: 1', snapshot.user_context['lspRuntime'])
def test_user_context_loads_task_runtime_summary(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir) / 'repo'
tests/test_agent_context_usage.py
+1 -1
View File
@@ -28,7 +28,7 @@ class AgentContextUsageTests(unittest.TestCase):
report = collect_context_usage(
session=session,
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
model='test-model',
strategy='test session',
)
rendered = format_context_usage(report)
tests/test_agent_prompting.py
+17
View File
@@ -226,6 +226,23 @@ class AgentPromptingTests(unittest.TestCase):
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)
tests/test_agent_runtime.py
+220 -1
View File
@@ -6,6 +6,7 @@ import unittest
from pathlib import Path
from unittest.mock import patch
from src.agent_session import AgentMessage
from src.agent_runtime import LocalCodingAgent
from src.agent_tools import build_tool_context, default_tool_registry, execute_tool
from src.agent_types import (
@@ -14,9 +15,17 @@ from src.agent_types import (
BudgetConfig,
ModelConfig,
OutputSchemaConfig,
UsageStats,
)
from src.compact import CompactionResult
from src.openai_compat import OpenAICompatClient
from src.session_store import load_agent_session
from src.session_store import (
StoredAgentSession,
load_agent_session,
serialize_model_config,
serialize_runtime_config,
)
from src.token_budget import TokenBudgetSnapshot
class FakeHTTPResponse:
@@ -656,6 +665,216 @@ class AgentRuntimeTests(unittest.TestCase):
self.assertIn('token budget', result.final_output)
self.assertEqual(result.usage.total_tokens, 42)
def test_agent_rejects_prompt_before_backend_when_preflight_input_budget_is_exceeded(self) -> None:
snapshot = TokenBudgetSnapshot(
model='test-model',
context_window_tokens=1000,
projected_input_tokens=240,
message_tokens=220,
chat_overhead_tokens=20,
reserved_output_tokens=128,
reserved_compaction_buffer_tokens=64,
reserved_schema_tokens=0,
hard_input_limit_tokens=40,
soft_input_limit_tokens=0,
overflow_tokens=200,
soft_overflow_tokens=240,
exceeds_hard_limit=True,
exceeds_soft_limit=True,
token_counter_backend='heuristic',
token_counter_source='test',
token_counter_accurate=False,
)
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
with patch(
'src.agent_runtime.calculate_token_budget',
return_value=snapshot,
), patch(
'src.agent_runtime.LocalCodingAgent._reduce_context_pressure',
return_value=False,
), patch(
'src.openai_compat.request.urlopen',
side_effect=AssertionError('backend should not be called'),
):
agent = LocalCodingAgent(
model_config=ModelConfig(
model='test-model',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=AgentRuntimeConfig(
cwd=workspace,
budget_config=BudgetConfig(max_input_tokens=20),
),
)
result = agent.run(
'This prompt is intentionally much longer than the tiny configured input budget. '
* 4
)
self.assertEqual(result.stop_reason, 'prompt_too_long')
self.assertIn('Stopped before the next model call', result.final_output)
def test_agent_auto_compacts_context_before_next_model_call(self) -> None:
responses = [
{
'choices': [
{
'message': {
'role': 'assistant',
'content': 'Recovered after compaction.',
},
'finish_reason': 'stop',
}
],
'usage': {'prompt_tokens': 11, 'completion_tokens': 4},
}
]
def fake_budget(
*,
session,
model,
budget_config,
output_schema=None,
):
compacted = any(
message.metadata.get('kind') == 'compact_summary'
for message in session.messages
)
projected = 120 if compacted else 540
return TokenBudgetSnapshot(
model=model,
context_window_tokens=1000,
projected_input_tokens=projected,
message_tokens=max(projected - 18, 0),
chat_overhead_tokens=18,
reserved_output_tokens=128,
reserved_compaction_buffer_tokens=64,
reserved_schema_tokens=0,
hard_input_limit_tokens=420,
soft_input_limit_tokens=180,
overflow_tokens=max(projected - 420, 0),
soft_overflow_tokens=max(projected - 180, 0),
exceeds_hard_limit=projected > 420,
exceeds_soft_limit=projected > 180,
token_counter_backend='heuristic',
token_counter_source='test',
token_counter_accurate=False,
)
def fake_compact(agent, custom_instructions=None):
session = agent.last_session
assert session is not None
preserved_tail = [session.messages[-1]]
boundary = AgentMessage(
role='user',
content='<system-reminder>Earlier conversation was compacted.</system-reminder>',
message_id='compact_boundary',
metadata={'kind': 'compact_boundary'},
)
summary = AgentMessage(
role='user',
content='Compacted summary.',
message_id='compact_summary',
metadata={'kind': 'compact_summary', 'is_compact_summary': True},
)
session.messages = [session.messages[0], boundary, summary] + preserved_tail
return CompactionResult(
boundary_message=boundary,
summary_messages=[summary],
messages_to_keep=preserved_tail,
pre_compact_token_count=540,
post_compact_token_count=120,
summary_text='Summary',
usage=UsageStats(input_tokens=7, output_tokens=3),
)
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
runtime_config = AgentRuntimeConfig(
cwd=workspace,
compact_preserve_messages=1,
)
stored = StoredAgentSession(
session_id='resume_auto_compact',
model_config=serialize_model_config(
ModelConfig(
model='test-model',
base_url='http://127.0.0.1:8000/v1',
)
),
runtime_config=serialize_runtime_config(runtime_config),
system_prompt_parts=('# System\nYou are helpful.',),
user_context={},
system_context={},
messages=(
{
'role': 'system',
'content': '# System\nYou are helpful.',
'message_id': 'system_0',
'metadata': {'kind': 'system'},
},
{
'role': 'user',
'content': 'First request.',
'message_id': 'user_1',
'metadata': {'kind': 'user'},
},
{
'role': 'assistant',
'content': 'First response.',
'message_id': 'assistant_1',
'metadata': {'kind': 'assistant'},
},
{
'role': 'user',
'content': 'Second request.',
'message_id': 'user_2',
'metadata': {'kind': 'user'},
},
{
'role': 'assistant',
'content': 'Second response.',
'message_id': 'assistant_2',
'metadata': {'kind': 'assistant'},
},
),
turns=2,
tool_calls=0,
usage=UsageStats().to_dict(),
total_cost_usd=0.0,
file_history=(),
budget_state={},
plugin_state={},
scratchpad_directory=None,
)
with patch(
'src.agent_runtime.calculate_token_budget',
side_effect=fake_budget,
), patch(
'src.agent_runtime.LocalCodingAgent._reduce_context_pressure',
return_value=False,
), patch(
'src.agent_runtime.compact_conversation',
side_effect=fake_compact,
), patch(
'src.openai_compat.request.urlopen',
side_effect=make_urlopen_side_effect(responses),
):
agent = LocalCodingAgent(
model_config=ModelConfig(
model='test-model',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=runtime_config,
)
result = agent.resume('Continue after compaction', stored)
self.assertEqual(result.final_output, 'Recovered after compaction.')
self.assertTrue(
any(event.get('type') == 'auto_compact_summary' for event in result.events)
)
self.assertGreaterEqual(result.usage.total_tokens, 25)
def test_agent_continues_when_model_response_is_truncated(self) -> None:
responses = [
{
tests/test_agent_slash_commands.py
+42 -2
View File
@@ -99,7 +99,7 @@ class AgentSlashCommandTests(unittest.TestCase):
workspace = Path(tmp_dir)
(workspace / 'CLAUDE.md').write_text('repo instructions\n', encoding='utf-8')
agent = LocalCodingAgent(
model_config=ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct'),
model_config=ModelConfig(model='test-model'),
runtime_config=AgentRuntimeConfig(cwd=workspace),
)
result = agent.run('/context')
@@ -107,6 +107,18 @@ class AgentSlashCommandTests(unittest.TestCase):
self.assertIn('### Estimated usage by category', result.final_output)
self.assertIn('### Memory Files', result.final_output)
def test_token_budget_command_renders_local_report(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
agent = LocalCodingAgent(
model_config=ModelConfig(model='test-model'),
runtime_config=AgentRuntimeConfig(cwd=workspace),
)
result = agent.run('/token-budget')
self.assertIn('# Token Budget', result.final_output)
self.assertIn('Hard input limit', result.final_output)
self.assertIn('Auto-compact buffer', result.final_output)
def test_mcp_and_resource_commands_render_local_reports(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
@@ -198,6 +210,34 @@ class AgentSlashCommandTests(unittest.TestCase):
self.assertIn('# Search Provider', provider_result.final_output)
self.assertIn('backup-search', provider_result.final_output)
def test_lsp_commands_render_local_reports(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / 'sample.py').write_text(
'def helper(value):\n'
' """Double a value."""\n'
' return value * 2\n'
'\n'
'def run(item):\n'
' return helper(item)\n',
encoding='utf-8',
)
agent = LocalCodingAgent(
model_config=ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct'),
runtime_config=AgentRuntimeConfig(cwd=workspace),
)
summary_result = agent.run('/lsp')
symbols_result = agent.run('/lsp symbols sample.py')
definition_result = agent.run('/lsp definition sample.py 6 12')
diagnostics_result = agent.run('/lsp diagnostics sample.py')
self.assertIn('# LSP', summary_result.final_output)
self.assertIn('Indexed candidate files: 1', summary_result.final_output)
self.assertIn('# LSP Document Symbols', symbols_result.final_output)
self.assertIn('function helper', symbols_result.final_output)
self.assertIn('# LSP Definition', definition_result.final_output)
self.assertIn('function helper', definition_result.final_output)
self.assertIn('# LSP Diagnostics', diagnostics_result.final_output)
def test_remote_commands_render_and_update_local_remote_runtime(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
@@ -387,7 +427,7 @@ class AgentSlashCommandTests(unittest.TestCase):
def test_tools_and_status_commands_render_local_reports(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
agent = LocalCodingAgent(
model_config=ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct'),
model_config=ModelConfig(model='test-model'),
runtime_config=AgentRuntimeConfig(cwd=Path(tmp_dir)),
)
tools_result = agent.run('/tools')
tests/test_extended_tools.py
+36
View File
@@ -6,6 +6,7 @@ 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):
@@ -101,3 +102,38 @@ class ExtendedToolTests(unittest.TestCase):
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')
tests/test_lsp_runtime.py
+74
View File
@@ -0,0 +1,74 @@
from __future__ import annotations
import tempfile
import unittest
from pathlib import Path
from src.lsp_runtime import LSPRuntime
SAMPLE_SOURCE = '''def helper(value):
"""Double a numeric value."""
return value * 2
def orchestrate(item):
return helper(item)
class Greeter:
def greet(self, name):
return helper(len(name))
'''
class LSPRuntimeTests(unittest.TestCase):
def test_runtime_renders_symbols_definitions_references_and_hover(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / 'sample.py').write_text(SAMPLE_SOURCE, encoding='utf-8')
runtime = LSPRuntime.from_workspace(workspace)
summary = runtime.render_summary()
symbols = runtime.render_document_symbols('sample.py')
workspace_symbols = runtime.render_workspace_symbols('helper')
definition = runtime.render_definition('sample.py', 7, 12)
references = runtime.render_references('sample.py', 7, 12)
hover = runtime.render_hover('sample.py', 1, 5)
self.assertIn('Indexed candidate files: 1', summary)
self.assertIn('# LSP Document Symbols', symbols)
self.assertIn('function helper', symbols)
self.assertIn('function orchestrate', symbols)
self.assertIn('class Greeter', symbols)
self.assertIn('# LSP Workspace Symbols', workspace_symbols)
self.assertIn('helper', workspace_symbols)
self.assertIn('# LSP Definition', definition)
self.assertIn('function helper', definition)
self.assertIn('# LSP References', references)
self.assertIn('helper(item)', references)
self.assertIn('# LSP Hover', hover)
self.assertIn('signature=helper(value)', hover)
self.assertIn('Double a numeric value.', hover)
def test_runtime_renders_call_hierarchy_and_diagnostics(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
(workspace / 'sample.py').write_text(SAMPLE_SOURCE, encoding='utf-8')
(workspace / 'broken.py').write_text('def broken(:\n pass\n', encoding='utf-8')
runtime = LSPRuntime.from_workspace(workspace)
hierarchy = runtime.render_prepare_call_hierarchy('sample.py', 6, 12)
incoming = runtime.render_incoming_calls('sample.py', 1, 5)
outgoing = runtime.render_outgoing_calls('sample.py', 6, 12)
diagnostics = runtime.render_diagnostics('broken.py')
self.assertIn('# LSP Call Hierarchy', hierarchy)
self.assertIn('symbol=orchestrate', hierarchy)
self.assertIn('# LSP Incoming Calls', incoming)
self.assertIn('orchestrate', incoming)
self.assertIn('greet', incoming)
self.assertIn('# LSP Outgoing Calls', outgoing)
self.assertIn('helper', outgoing)
self.assertIn('# LSP Diagnostics', diagnostics)
self.assertIn('syntax-error', diagnostics)
tests/test_main.py
+17
View File
@@ -104,6 +104,8 @@ class MainCliTests(unittest.TestCase):
[
'agent-chat',
'First prompt',
'--model',
'test-model',
'--cwd',
str(workspace),
]
@@ -196,6 +198,21 @@ class MainCliTests(unittest.TestCase):
self.assertEqual(args.key_path, 'review.mode')
self.assertEqual(args.cwd, '.')
def test_parser_accepts_lsp_runtime_commands(self) -> None:
parser = build_parser()
args = parser.parse_args(['lsp-definition', 'sample.py', '4', '12', '--cwd', '.'])
self.assertEqual(args.command, 'lsp-definition')
self.assertEqual(args.file_path, 'sample.py')
self.assertEqual(args.line, 4)
self.assertEqual(args.character, 12)
self.assertEqual(args.cwd, '.')
def test_parser_accepts_token_budget_command(self) -> None:
parser = build_parser()
args = parser.parse_args(['token-budget', '--cwd', '.'])
self.assertEqual(args.command, 'token-budget')
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', '.'])
tests/test_token_budget.py
+100
View File
@@ -0,0 +1,100 @@
from __future__ import annotations
import unittest
from unittest.mock import patch
from src.agent_session import AgentSessionState
from src.agent_context_usage import ContextUsageReport, MessageBreakdown
from src.agent_types import BudgetConfig
from src.token_budget import calculate_token_budget, format_token_budget
class TokenBudgetTests(unittest.TestCase):
def test_calculate_token_budget_reports_soft_and_hard_limits(self) -> None:
session = AgentSessionState.create(
['# System\nYou are helpful.'],
'Inspect the repository and summarize the current implementation status.',
user_context={'currentDate': "Today's date is 2026-04-11."},
system_context={'gitStatus': 'Current branch: main'},
)
session.append_assistant('Reading files and checking runtime state.')
fake_usage = ContextUsageReport(
model='test-model',
total_tokens=200,
raw_max_tokens=128_000,
percentage=0.15,
strategy='token_budget',
message_count=len(session.messages),
categories=(),
system_prompt_sections=(),
user_context_entries=(),
system_context_entries=(),
memory_files=(),
message_breakdown=MessageBreakdown(
user_message_tokens=50,
assistant_message_tokens=50,
tool_call_tokens=0,
tool_result_tokens=0,
user_context_tokens=10,
tool_calls_by_type=(),
),
token_counter_backend='heuristic',
token_counter_source='test',
token_counter_accurate=False,
)
with patch('src.token_budget.collect_context_usage', return_value=fake_usage):
snapshot = calculate_token_budget(
session=session,
model='test-model',
budget_config=BudgetConfig(),
)
rendered = format_token_budget(snapshot)
self.assertGreater(snapshot.projected_input_tokens, 0)
self.assertGreater(snapshot.hard_input_limit_tokens, snapshot.soft_input_limit_tokens)
self.assertGreater(snapshot.chat_overhead_tokens, 0)
self.assertIn('# Token Budget', rendered)
self.assertIn('Hard input limit', rendered)
self.assertIn('Auto-compact buffer', rendered)
def test_calculate_token_budget_honors_explicit_max_input_tokens(self) -> None:
session = AgentSessionState.create(
['# System\nYou are helpful.'],
'This prompt is deliberately longer than the tiny configured input budget. ' * 4,
)
fake_usage = ContextUsageReport(
model='test-model',
total_tokens=120,
raw_max_tokens=128_000,
percentage=0.09,
strategy='token_budget',
message_count=len(session.messages),
categories=(),
system_prompt_sections=(),
user_context_entries=(),
system_context_entries=(),
memory_files=(),
message_breakdown=MessageBreakdown(
user_message_tokens=80,
assistant_message_tokens=0,
tool_call_tokens=0,
tool_result_tokens=0,
user_context_tokens=0,
tool_calls_by_type=(),
),
token_counter_backend='heuristic',
token_counter_source='test',
token_counter_accurate=False,
)
with patch('src.token_budget.collect_context_usage', return_value=fake_usage):
snapshot = calculate_token_budget(
session=session,
model='test-model',
budget_config=BudgetConfig(max_input_tokens=20),
)
self.assertTrue(snapshot.exceeds_hard_limit)
self.assertGreater(snapshot.overflow_tokens, 0)
self.assertLessEqual(snapshot.hard_input_limit_tokens, 20)