Files
zk-data-agent/tests/test_agent_context_usage.py
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

42 lines
1.5 KiB
Python

from __future__ import annotations
import unittest
from src.agent_context_usage import collect_context_usage, format_context_usage
from src.agent_session import AgentSessionState
class AgentContextUsageTests(unittest.TestCase):
def test_collect_context_usage_formats_breakdown(self) -> None:
session = AgentSessionState.create(
['# Intro\nhello', '# System\nworld'],
'inspect repo',
user_context={'currentDate': "Today's date is 2026-04-01."},
system_context={'gitStatus': 'Current branch: main'},
)
session.append_assistant(
'Reading files',
(
{
'id': 'call_1',
'type': 'function',
'function': {'name': 'read_file', 'arguments': '{"path":"a.py"}'},
},
),
)
session.append_tool('read_file', 'call_1', '{"ok": true, "content": "print(1)"}')
report = collect_context_usage(
session=session,
model='test-model',
strategy='test session',
)
rendered = format_context_usage(report)
self.assertGreater(report.total_tokens, 0)
self.assertIn('## Context Usage', rendered)
self.assertIn('**Token counter:**', rendered)
self.assertIn('### System Prompt Sections', rendered)
self.assertIn('### Message Breakdown', rendered)
self.assertIn('#### Top Tools', rendered)