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
src/__init__.py
+11
View File
@@ -18,6 +18,7 @@ from .agent_types import AgentPermissions, AgentRunResult, AgentRuntimeConfig, M
from .background_runtime import BackgroundSessionRuntime
from .commands import PORTED_COMMANDS, build_command_backlog
from .config_runtime import ConfigMutation, ConfigRuntime
from .lsp_runtime import LSPCallEdge, LSPDiagnostic, LSPReference, LSPRuntime, LSPSymbol
from .mcp_runtime import MCPRuntime, MCPResource, MCPServerProfile, MCPTool
from .parity_audit import ParityAuditResult, run_parity_audit
from .plan_runtime import PlanRuntime, PlanStep
@@ -32,6 +33,7 @@ from .system_init import build_system_init_message
from .task import PortingTask
from .task_runtime import TaskRuntime
from .team_runtime import TeamDefinition, TeamMessage, TeamRuntime
from .token_budget import TokenBudgetSnapshot, calculate_token_budget, estimate_chat_overhead, format_token_budget
from .tokenizer_runtime import TokenCounterInfo, clear_token_counter_cache, count_tokens, describe_token_counter
from .workflow_runtime import WorkflowDefinition, WorkflowRunRecord, WorkflowRuntime
from .worktree_runtime import WorktreeRuntime, WorktreeSessionState, WorktreeStatusReport
@@ -54,6 +56,11 @@ __all__ = [
'BackgroundSessionRuntime',
'ConfigMutation',
'ConfigRuntime',
'LSPCallEdge',
'LSPDiagnostic',
'LSPReference',
'LSPRuntime',
'LSPSymbol',
'LocalCodingAgent',
'MCPResource',
'MCPRuntime',
@@ -82,6 +89,7 @@ __all__ = [
'TeamDefinition',
'TeamMessage',
'TeamRuntime',
'TokenBudgetSnapshot',
'TokenCounterInfo',
'TurnResult',
'WorkflowDefinition',
@@ -101,9 +109,12 @@ __all__ = [
'clear_context_caches',
'clear_token_counter_cache',
'count_tokens',
'calculate_token_budget',
'default_tool_registry',
'describe_token_counter',
'estimate_chat_overhead',
'execute_tool',
'format_token_budget',
'get_system_context',
'get_user_context',
'load_session',
src/agent_context.py
+4
View File
@@ -13,6 +13,7 @@ from .account_runtime import AccountRuntime
from .ask_user_runtime import AskUserRuntime
from .config_runtime import ConfigRuntime
from .hook_policy import HookPolicyRuntime
from .lsp_runtime import LSPRuntime
from .mcp_runtime import MCPRuntime
from .plan_runtime import PlanRuntime
from .plugin_runtime import PluginRuntime
@@ -254,6 +255,9 @@ def _get_user_context_cached(
config_runtime = ConfigRuntime.from_workspace(Path(cwd))
if config_runtime.has_config():
context['configRuntime'] = config_runtime.render_summary()
lsp_runtime = LSPRuntime.from_workspace(Path(cwd), additional_working_directories)
if lsp_runtime.has_lsp_support():
context['lspRuntime'] = lsp_runtime.render_summary()
plan_runtime = PlanRuntime.from_workspace(Path(cwd))
if plan_runtime.steps:
context['planRuntime'] = plan_runtime.render_summary()
src/agent_prompting.py
+13
View File
@@ -97,6 +97,7 @@ def build_system_prompt_parts(
get_account_guidance_section(prompt_context),
get_ask_user_guidance_section(prompt_context),
get_config_guidance_section(prompt_context),
get_lsp_guidance_section(prompt_context),
get_plan_guidance_section(prompt_context),
get_task_guidance_section(prompt_context),
get_team_guidance_section(prompt_context),
@@ -303,6 +304,18 @@ def get_config_guidance_section(prompt_context: PromptContext) -> str:
return '\n'.join(['# Config', *prepend_bullets(items)])
def get_lsp_guidance_section(prompt_context: PromptContext) -> str:
lsp_runtime = prompt_context.user_context.get('lspRuntime')
if not lsp_runtime:
return ''
items = [
'A local LSP-style code-intelligence runtime may be available for supported source files.',
'Use the LSP tool when you need definitions, references, hover details, document symbols, workspace symbols, or call hierarchy information.',
'Use LSP diagnostics to catch syntax and parse issues before making larger edits when the file type is supported.',
]
return '\n'.join(['# LSP', *prepend_bullets(items)])
def get_task_guidance_section(prompt_context: PromptContext) -> str:
task_runtime = prompt_context.user_context.get('taskRuntime')
if not task_runtime:
src/agent_runtime.py
+402
View File
@@ -12,9 +12,11 @@ from .agent_manager import AgentManager
from .agent_context import clear_context_caches
from .agent_context import render_context_report as render_agent_context_report
from .agent_context_usage import collect_context_usage, estimate_tokens, format_context_usage
from .compact import compact_conversation
from .ask_user_runtime import AskUserRuntime
from .config_runtime import ConfigRuntime
from .hook_policy import HookPolicyRuntime
from .lsp_runtime import LSPRuntime
from .mcp_runtime import MCPRuntime
from .agent_prompting import (
build_prompt_context,
@@ -62,6 +64,7 @@ from .session_store import (
serialize_runtime_config,
usage_from_payload,
)
from .token_budget import calculate_token_budget, format_token_budget
@dataclass(frozen=True)
@@ -70,6 +73,14 @@ class BudgetDecision:
reason: str | None = None
@dataclass(frozen=True)
class PromptPreflightResult:
usage_increment: UsageStats = field(default_factory=UsageStats)
model_calls_increment: int = 0
stop_reason: str | None = None
reason: str | None = None
@dataclass
class LocalCodingAgent:
model_config: ModelConfig
@@ -92,6 +103,7 @@ class LocalCodingAgent:
account_runtime: AccountRuntime | None = None
ask_user_runtime: AskUserRuntime | None = None
config_runtime: ConfigRuntime | None = None
lsp_runtime: LSPRuntime | None = None
plan_runtime: PlanRuntime | None = None
task_runtime: TaskRuntime | None = None
team_runtime: TeamRuntime | None = None
@@ -153,6 +165,11 @@ class LocalCodingAgent:
)
if self.config_runtime is None:
self.config_runtime = ConfigRuntime.from_workspace(self.runtime_config.cwd)
if self.lsp_runtime is None:
self.lsp_runtime = LSPRuntime.from_workspace(
self.runtime_config.cwd,
tuple(str(path) for path in self.runtime_config.additional_working_directories),
)
if self.plan_runtime is None:
self.plan_runtime = PlanRuntime.from_workspace(self.runtime_config.cwd)
if self.task_runtime is None:
@@ -191,6 +208,7 @@ class LocalCodingAgent:
account_runtime=self.account_runtime,
ask_user_runtime=self.ask_user_runtime,
config_runtime=self.config_runtime,
lsp_runtime=self.lsp_runtime,
mcp_runtime=self.mcp_runtime,
remote_runtime=self.remote_runtime,
remote_trigger_runtime=self.remote_trigger_runtime,
@@ -489,6 +507,69 @@ class LocalCodingAgent:
stream_events,
turn_index=turn_index,
)
preflight = self._preflight_prompt_length(
session,
stream_events,
turn_index=turn_index,
)
if preflight.usage_increment.total_tokens or preflight.model_calls_increment:
total_usage = total_usage + preflight.usage_increment
total_cost_usd = self.model_config.pricing.estimate_cost_usd(total_usage)
model_calls += preflight.model_calls_increment
budget_after_preflight = self._check_budget(
total_usage,
total_cost_usd,
tool_calls=tool_calls,
delegated_tasks=delegated_tasks,
model_calls=model_calls,
session_turns=starting_session_turns + turn_index,
)
if budget_after_preflight.exceeded:
result = AgentRunResult(
final_output=(
budget_after_preflight.reason
or 'Stopped because the runtime budget was exceeded.'
),
turns=turn_index,
tool_calls=tool_calls,
transcript=session.transcript(),
events=tuple(stream_events),
usage=total_usage,
total_cost_usd=total_cost_usd,
stop_reason='budget_exceeded',
file_history=tuple(file_history),
session_id=session_id,
scratchpad_directory=(
str(scratchpad_directory) if scratchpad_directory is not None else None
),
)
result = self._persist_session(session, result)
self.last_run_result = result
return result
if preflight.stop_reason is not None:
result = AgentRunResult(
final_output=preflight.reason or 'Stopped before the next model call.',
turns=max(turn_index - 1, 0),
tool_calls=tool_calls,
transcript=session.transcript(),
events=tuple(stream_events),
usage=total_usage,
total_cost_usd=total_cost_usd,
stop_reason=preflight.stop_reason,
file_history=tuple(file_history),
session_id=session_id,
scratchpad_directory=(
str(scratchpad_directory) if scratchpad_directory is not None else None
),
)
result = self._append_runtime_after_turn_events(
result,
prompt=effective_prompt,
turn_index=max(turn_index - 1, 0),
)
result = self._persist_session(session, result)
self.last_run_result = result
return result
try:
turn, turn_events = self._query_model(session, tool_specs)
except OpenAICompatError as exc:
@@ -1220,6 +1301,162 @@ class LocalCodingAgent:
)
return BudgetDecision(exceeded=False)
def _preflight_prompt_length(
self,
session: AgentSessionState,
stream_events: list[dict[str, object]],
*,
turn_index: int,
) -> PromptPreflightResult:
snapshot = calculate_token_budget(
session=session,
model=self.model_config.model,
budget_config=self.runtime_config.budget_config,
output_schema=self.runtime_config.output_schema,
)
if not snapshot.exceeds_soft_limit and not snapshot.exceeds_hard_limit:
return PromptPreflightResult()
stream_events.append(
{
'type': 'prompt_length_check',
'turn_index': turn_index,
'projected_input_tokens': snapshot.projected_input_tokens,
'soft_input_limit_tokens': snapshot.soft_input_limit_tokens,
'hard_input_limit_tokens': snapshot.hard_input_limit_tokens,
'soft_overflow_tokens': snapshot.soft_overflow_tokens,
'overflow_tokens': snapshot.overflow_tokens,
'exceeds_hard_limit': snapshot.exceeds_hard_limit,
}
)
target_tokens = snapshot.soft_input_limit_tokens
if snapshot.exceeds_hard_limit:
target_tokens = snapshot.hard_input_limit_tokens
if target_tokens < 0:
target_tokens = 0
if self._reduce_context_pressure(
session,
stream_events,
turn_index=turn_index,
target_tokens=target_tokens,
allow_compaction=True,
):
recovered = calculate_token_budget(
session=session,
model=self.model_config.model,
budget_config=self.runtime_config.budget_config,
output_schema=self.runtime_config.output_schema,
)
stream_events.append(
{
'type': 'prompt_length_recovery',
'turn_index': turn_index,
'strategy': 'heuristic',
'projected_input_tokens': recovered.projected_input_tokens,
'soft_input_limit_tokens': recovered.soft_input_limit_tokens,
'hard_input_limit_tokens': recovered.hard_input_limit_tokens,
'exceeds_hard_limit': recovered.exceeds_hard_limit,
'exceeds_soft_limit': recovered.exceeds_soft_limit,
}
)
if not recovered.exceeds_soft_limit and not recovered.exceeds_hard_limit:
return PromptPreflightResult()
snapshot = recovered
if self._can_auto_compact_with_summary(session):
compact_result = compact_conversation(
self,
custom_instructions=(
'Automatically collapse earlier conversation context to fit the next model '
'turn. Preserve the active task, recent file changes, failures, pending work, '
'and exact next step.'
),
)
if compact_result.error is None:
recovered = calculate_token_budget(
session=session,
model=self.model_config.model,
budget_config=self.runtime_config.budget_config,
output_schema=self.runtime_config.output_schema,
)
stream_events.append(
{
'type': 'auto_compact_summary',
'turn_index': turn_index,
'pre_compact_token_count': compact_result.pre_compact_token_count,
'post_compact_token_count': compact_result.post_compact_token_count,
'summary_usage_tokens': compact_result.usage.total_tokens,
'projected_input_tokens': recovered.projected_input_tokens,
'soft_input_limit_tokens': recovered.soft_input_limit_tokens,
'hard_input_limit_tokens': recovered.hard_input_limit_tokens,
'exceeds_hard_limit': recovered.exceeds_hard_limit,
'exceeds_soft_limit': recovered.exceeds_soft_limit,
}
)
if not recovered.exceeds_soft_limit and not recovered.exceeds_hard_limit:
return PromptPreflightResult(
usage_increment=compact_result.usage,
model_calls_increment=1,
)
snapshot = recovered
if compact_result.usage.total_tokens:
return PromptPreflightResult(
usage_increment=compact_result.usage,
model_calls_increment=1,
stop_reason=(
'prompt_too_long'
if recovered.exceeds_hard_limit
else None
),
reason=(
self._build_prompt_length_error(recovered)
if recovered.exceeds_hard_limit
else None
),
)
else:
stream_events.append(
{
'type': 'auto_compact_failed',
'turn_index': turn_index,
'reason': compact_result.error,
}
)
if snapshot.exceeds_hard_limit:
return PromptPreflightResult(
stop_reason='prompt_too_long',
reason=self._build_prompt_length_error(snapshot),
)
stream_events.append(
{
'type': 'prompt_length_warning',
'turn_index': turn_index,
'projected_input_tokens': snapshot.projected_input_tokens,
'soft_input_limit_tokens': snapshot.soft_input_limit_tokens,
'hard_input_limit_tokens': snapshot.hard_input_limit_tokens,
'soft_overflow_tokens': snapshot.soft_overflow_tokens,
}
)
return PromptPreflightResult()
def _can_auto_compact_with_summary(self, session: AgentSessionState) -> bool:
prefix_count = self._compact_prefix_count(session)
preserve_count = max(self.runtime_config.compact_preserve_messages, 1)
return len(session.messages) - prefix_count > preserve_count
def _build_prompt_length_error(self, snapshot) -> str:
return (
'Stopped before the next model call because the prompt would exceed the '
'effective input budget. '
f'Projected prompt tokens: {snapshot.projected_input_tokens:,}; '
f'hard input limit: {snapshot.hard_input_limit_tokens:,}; '
f'soft input limit: {snapshot.soft_input_limit_tokens:,}.'
)
def _snip_session_if_needed(
self,
session: AgentSessionState,
@@ -2991,6 +3228,138 @@ class LocalCodingAgent:
return '# Config\n\nNo local config runtime is available.'
return '\n'.join(['# Config', '', self.config_runtime.render_summary()])
def render_lsp_report(self) -> str:
if self.lsp_runtime is None or not self.lsp_runtime.has_lsp_support():
return '# LSP\n\nNo local LSP runtime is available.'
return '\n'.join(['# LSP', '', self.lsp_runtime.render_summary()])
def render_lsp_document_symbols_report(self, file_path: str) -> str:
if self.lsp_runtime is None or not self.lsp_runtime.has_lsp_support():
return '# LSP Document Symbols\n\nNo local LSP runtime is available.'
try:
return self.lsp_runtime.render_document_symbols(file_path)
except KeyError as exc:
return f'# LSP Document Symbols\n\n{exc}'
def render_lsp_workspace_symbols_report(
self,
query: str,
*,
max_results: int = 50,
) -> str:
if self.lsp_runtime is None or not self.lsp_runtime.has_lsp_support():
return '# LSP Workspace Symbols\n\nNo local LSP runtime is available.'
return self.lsp_runtime.render_workspace_symbols(query, max_results=max_results)
def render_lsp_definition_report(
self,
file_path: str,
line: int,
character: int,
*,
max_results: int = 20,
) -> str:
if self.lsp_runtime is None or not self.lsp_runtime.has_lsp_support():
return '# LSP Definition\n\nNo local LSP runtime is available.'
try:
return self.lsp_runtime.render_definition(
file_path,
line,
character,
max_results=max_results,
)
except KeyError as exc:
return f'# LSP Definition\n\n{exc}'
def render_lsp_references_report(
self,
file_path: str,
line: int,
character: int,
*,
max_results: int = 50,
) -> str:
if self.lsp_runtime is None or not self.lsp_runtime.has_lsp_support():
return '# LSP References\n\nNo local LSP runtime is available.'
try:
return self.lsp_runtime.render_references(
file_path,
line,
character,
max_results=max_results,
)
except KeyError as exc:
return f'# LSP References\n\n{exc}'
def render_lsp_hover_report(self, file_path: str, line: int, character: int) -> str:
if self.lsp_runtime is None or not self.lsp_runtime.has_lsp_support():
return '# LSP Hover\n\nNo local LSP runtime is available.'
try:
return self.lsp_runtime.render_hover(file_path, line, character)
except KeyError as exc:
return f'# LSP Hover\n\n{exc}'
def render_lsp_diagnostics_report(self, file_path: str | None = None) -> str:
if self.lsp_runtime is None or not self.lsp_runtime.has_lsp_support():
return '# LSP Diagnostics\n\nNo local LSP runtime is available.'
try:
return self.lsp_runtime.render_diagnostics(file_path)
except KeyError as exc:
return f'# LSP Diagnostics\n\n{exc}'
def render_lsp_prepare_call_hierarchy_report(
self,
file_path: str,
line: int,
character: int,
) -> str:
if self.lsp_runtime is None or not self.lsp_runtime.has_lsp_support():
return '# LSP Call Hierarchy\n\nNo local LSP runtime is available.'
try:
return self.lsp_runtime.render_prepare_call_hierarchy(file_path, line, character)
except KeyError as exc:
return f'# LSP Call Hierarchy\n\n{exc}'
def render_lsp_incoming_calls_report(
self,
file_path: str,
line: int,
character: int,
*,
max_results: int = 50,
) -> str:
if self.lsp_runtime is None or not self.lsp_runtime.has_lsp_support():
return '# LSP Incoming Calls\n\nNo local LSP runtime is available.'
try:
return self.lsp_runtime.render_incoming_calls(
file_path,
line,
character,
max_results=max_results,
)
except KeyError as exc:
return f'# LSP Incoming Calls\n\n{exc}'
def render_lsp_outgoing_calls_report(
self,
file_path: str,
line: int,
character: int,
*,
max_results: int = 50,
) -> str:
if self.lsp_runtime is None or not self.lsp_runtime.has_lsp_support():
return '# LSP Outgoing Calls\n\nNo local LSP runtime is available.'
try:
return self.lsp_runtime.render_outgoing_calls(
file_path,
line,
character,
max_results=max_results,
)
except KeyError as exc:
return f'# LSP Outgoing Calls\n\n{exc}'
def render_config_effective_report(self) -> str:
if self.config_runtime is None:
return '# Config Effective\n\nNo local config runtime is available.'
@@ -3287,6 +3656,19 @@ class LocalCodingAgent:
f'- Session ID: {self.active_session_id or "none"}',
f'- Last session loaded: {"yes" if self.last_session is not None else "no"}',
]
if self.last_session is not None:
token_budget = calculate_token_budget(
session=self.last_session,
model=self.model_config.model,
budget_config=self.runtime_config.budget_config,
output_schema=self.runtime_config.output_schema,
)
lines.append(
f'- Prompt budget: {token_budget.projected_input_tokens:,} / {token_budget.soft_input_limit_tokens:,} soft'
)
lines.append(
f'- Prompt hard limit: {token_budget.hard_input_limit_tokens:,}'
)
if self.hook_policy_runtime is not None and self.hook_policy_runtime.manifests:
lines.append(
f'- Workspace trust mode: {"trusted" if self.hook_policy_runtime.is_trusted() else "untrusted"}'
@@ -3325,6 +3707,10 @@ class LocalCodingAgent:
lines.append(
f'- Effective config keys: {len(self.config_runtime.list_keys())}'
)
if self.lsp_runtime is not None and self.lsp_runtime.has_lsp_support():
lines.append(
f'- LSP indexed files: {len(self.lsp_runtime._workspace_files())}'
)
if self.plan_runtime is not None and self.plan_runtime.steps:
lines.append(f'- Local plan steps: {len(self.plan_runtime.steps)}')
if self.task_runtime is not None and self.task_runtime.tasks:
@@ -3353,6 +3739,16 @@ class LocalCodingAgent:
lines.extend(self.agent_manager.summary_lines())
return '\n'.join(lines)
def render_token_budget_report(self) -> str:
session = self.last_session or self.build_session()
snapshot = calculate_token_budget(
session=session,
model=self.model_config.model,
budget_config=self.runtime_config.budget_config,
output_schema=self.runtime_config.output_schema,
)
return format_token_budget(snapshot)
def _finalize_managed_agent(self, result: AgentRunResult) -> None:
if self.managed_agent_id is None or self.agent_manager is None:
self.resume_source_session_id = None
@@ -3463,6 +3859,7 @@ class LocalCodingAgent:
account_runtime=self.account_runtime,
ask_user_runtime=self.ask_user_runtime,
config_runtime=self.config_runtime,
lsp_runtime=self.lsp_runtime,
remote_runtime=self.remote_runtime,
remote_trigger_runtime=self.remote_trigger_runtime,
plan_runtime=self.plan_runtime,
@@ -3514,6 +3911,10 @@ class LocalCodingAgent:
additional_dirs,
)
self.config_runtime = ConfigRuntime.from_workspace(self.runtime_config.cwd)
self.lsp_runtime = LSPRuntime.from_workspace(
self.runtime_config.cwd,
additional_dirs,
)
self.task_runtime = TaskRuntime.from_workspace(self.runtime_config.cwd)
self.plan_runtime = PlanRuntime.from_workspace(self.runtime_config.cwd)
self.team_runtime = TeamRuntime.from_workspace(
@@ -3547,6 +3948,7 @@ class LocalCodingAgent:
account_runtime=self.account_runtime,
ask_user_runtime=self.ask_user_runtime,
config_runtime=self.config_runtime,
lsp_runtime=self.lsp_runtime,
mcp_runtime=self.mcp_runtime,
remote_runtime=self.remote_runtime,
remote_trigger_runtime=self.remote_trigger_runtime,
src/agent_slash_commands.py
+89
View File
@@ -109,6 +109,11 @@ def get_slash_command_specs() -> tuple[SlashCommandSpec, ...]:
description='Show the raw environment, user context, and system context snapshot.',
handler=_handle_context_raw,
),
SlashCommandSpec(
names=('token-budget', 'budget'),
description='Show the current token-budget window, reserves, and prompt-length limits.',
handler=_handle_token_budget,
),
SlashCommandSpec(
names=('mcp',),
description='Show discovered local MCP manifests and resource counts.',
@@ -154,6 +159,11 @@ def get_slash_command_specs() -> tuple[SlashCommandSpec, ...]:
description='Show local config runtime state, effective config, config sources, or a config value.',
handler=_handle_config,
),
SlashCommandSpec(
names=('lsp',),
description='Show local LSP runtime status or run document symbols, definition, references, hover, call hierarchy, and diagnostics queries.',
handler=_handle_lsp,
),
SlashCommandSpec(
names=('remotes',),
description='List configured local remote profiles.',
@@ -395,6 +405,10 @@ def _handle_context_raw(agent: 'LocalCodingAgent', _args: str, input_text: str)
return _local_result(input_text, agent.render_context_snapshot_report())
def _handle_token_budget(agent: 'LocalCodingAgent', _args: str, input_text: str) -> SlashCommandResult:
return _local_result(input_text, agent.render_token_budget_report())
def _handle_mcp(agent: 'LocalCodingAgent', args: str, input_text: str) -> SlashCommandResult:
command = args.strip()
if not command:
@@ -521,6 +535,81 @@ def _handle_config(agent: 'LocalCodingAgent', args: str, input_text: str) -> Sla
)
def _handle_lsp(agent: 'LocalCodingAgent', args: str, input_text: str) -> SlashCommandResult:
command = args.strip()
if not command:
return _local_result(input_text, agent.render_lsp_report())
if command == 'diagnostics':
return _local_result(input_text, agent.render_lsp_diagnostics_report())
if command.startswith('diagnostics '):
file_path = command.split(' ', 1)[1].strip()
if not file_path:
return _local_result(input_text, 'Usage: /lsp diagnostics [file-path]')
return _local_result(input_text, agent.render_lsp_diagnostics_report(file_path))
if command.startswith('symbols '):
file_path = command.split(' ', 1)[1].strip()
if not file_path:
return _local_result(input_text, 'Usage: /lsp symbols <file-path>')
return _local_result(input_text, agent.render_lsp_document_symbols_report(file_path))
if command.startswith('workspace '):
query = command.split(' ', 1)[1].strip()
if not query:
return _local_result(input_text, 'Usage: /lsp workspace <query>')
return _local_result(input_text, agent.render_lsp_workspace_symbols_report(query))
parts = command.split()
if len(parts) == 4 and parts[0] in {
'definition',
'references',
'hover',
'hierarchy',
'incoming',
'outgoing',
}:
subcommand, file_path, line_text, character_text = parts
try:
line = int(line_text)
character = int(character_text)
except ValueError:
return _local_result(
input_text,
f'Usage: /lsp {subcommand} <file-path> <line> <character>',
)
if subcommand == 'definition':
return _local_result(
input_text,
agent.render_lsp_definition_report(file_path, line, character),
)
if subcommand == 'references':
return _local_result(
input_text,
agent.render_lsp_references_report(file_path, line, character),
)
if subcommand == 'hover':
return _local_result(
input_text,
agent.render_lsp_hover_report(file_path, line, character),
)
if subcommand == 'hierarchy':
return _local_result(
input_text,
agent.render_lsp_prepare_call_hierarchy_report(file_path, line, character),
)
if subcommand == 'incoming':
return _local_result(
input_text,
agent.render_lsp_incoming_calls_report(file_path, line, character),
)
if subcommand == 'outgoing':
return _local_result(
input_text,
agent.render_lsp_outgoing_calls_report(file_path, line, character),
)
return _local_result(
input_text,
'Usage: /lsp [symbols <file>|workspace <query>|definition <file> <line> <character>|references <file> <line> <character>|hover <file> <line> <character>|hierarchy <file> <line> <character>|incoming <file> <line> <character>|outgoing <file> <line> <character>|diagnostics [file]]',
)
def _handle_remotes(agent: 'LocalCodingAgent', args: str, input_text: str) -> SlashCommandResult:
query = args or None
return _local_result(input_text, agent.render_remote_profiles_report(query))
src/agent_tools.py
+78
View File
@@ -20,6 +20,7 @@ if TYPE_CHECKING:
from .account_runtime import AccountRuntime
from .ask_user_runtime import AskUserRuntime
from .config_runtime import ConfigRuntime
from .lsp_runtime import LSPRuntime
from .mcp_runtime import MCPRuntime
from .plan_runtime import PlanRuntime
from .remote_runtime import RemoteRuntime
@@ -51,6 +52,7 @@ class ToolExecutionContext:
account_runtime: 'AccountRuntime | None' = None
ask_user_runtime: 'AskUserRuntime | None' = None
config_runtime: 'ConfigRuntime | None' = None
lsp_runtime: 'LSPRuntime | None' = None
mcp_runtime: 'MCPRuntime | None' = None
remote_runtime: 'RemoteRuntime | None' = None
remote_trigger_runtime: 'RemoteTriggerRuntime | None' = None
@@ -126,6 +128,7 @@ def build_tool_context(
account_runtime: 'AccountRuntime | None' = None,
ask_user_runtime: 'AskUserRuntime | None' = None,
config_runtime: 'ConfigRuntime | None' = None,
lsp_runtime: 'LSPRuntime | None' = None,
mcp_runtime: 'MCPRuntime | None' = None,
remote_runtime: 'RemoteRuntime | None' = None,
remote_trigger_runtime: 'RemoteTriggerRuntime | None' = None,
@@ -146,6 +149,7 @@ def build_tool_context(
account_runtime=account_runtime,
ask_user_runtime=ask_user_runtime,
config_runtime=config_runtime,
lsp_runtime=lsp_runtime,
mcp_runtime=mcp_runtime,
remote_runtime=remote_runtime,
remote_trigger_runtime=remote_trigger_runtime,
@@ -313,6 +317,36 @@ def default_tool_registry() -> dict[str, AgentTool]:
},
handler=_run_bash,
),
AgentTool(
name='LSP',
description='Use local LSP-style code intelligence for definitions, references, hover, symbols, and call hierarchy.',
parameters={
'type': 'object',
'properties': {
'operation': {
'type': 'string',
'enum': [
'goToDefinition',
'findReferences',
'hover',
'documentSymbol',
'workspaceSymbol',
'goToImplementation',
'prepareCallHierarchy',
'incomingCalls',
'outgoingCalls',
],
},
'file_path': {'type': 'string'},
'line': {'type': 'integer', 'minimum': 1},
'character': {'type': 'integer', 'minimum': 1},
'query': {'type': 'string'},
'max_results': {'type': 'integer', 'minimum': 1, 'maximum': 200},
},
'required': ['operation', 'file_path', 'line', 'character'],
},
handler=_lsp_query,
),
AgentTool(
name='web_fetch',
description='Fetch a text resource from http, https, or file URLs and return a truncated text response.',
@@ -2729,6 +2763,42 @@ def _delegate_agent_placeholder(
)
def _lsp_query(arguments: dict[str, Any], context: ToolExecutionContext):
runtime = _require_lsp_runtime(context)
operation = _require_string(arguments, 'operation')
file_path = _require_string(arguments, 'file_path')
line = _coerce_int(arguments, 'line', 1)
character = _coerce_int(arguments, 'character', 1)
query = arguments.get('query')
if query is not None and not isinstance(query, str):
raise ToolExecutionError('query must be a string')
max_results = _coerce_int(arguments, 'max_results', 50)
try:
result = runtime.query(
operation,
file_path=file_path,
line=line,
character=character,
query=query,
max_results=max_results,
)
except KeyError as exc:
raise ToolExecutionError(str(exc)) from exc
return (
result.content,
{
'action': 'lsp_query',
'operation': result.operation,
'file_path': file_path,
'line': line,
'character': character,
'result_count': result.result_count,
'file_count': result.file_count,
'symbol_name': result.symbol_name,
},
)
def _require_account_runtime(context: ToolExecutionContext):
if context.account_runtime is None:
raise ToolExecutionError('No local account runtime is available.')
@@ -2756,6 +2826,14 @@ def _require_config_runtime(context: ToolExecutionContext):
return context.config_runtime
def _require_lsp_runtime(context: ToolExecutionContext):
if context.lsp_runtime is None or not context.lsp_runtime.has_lsp_support():
raise ToolExecutionError(
'No local LSP runtime is available. Add supported source files to the workspace or a .claw-lsp.json manifest.'
)
return context.lsp_runtime
def _require_mcp_runtime(context: ToolExecutionContext):
if (
context.mcp_runtime is None
src/compact.py
+3
View File
@@ -18,6 +18,7 @@ from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from .agent_context_usage import estimate_tokens
from .agent_types import UsageStats
from .agent_session import AgentMessage
if TYPE_CHECKING:
@@ -285,6 +286,7 @@ class CompactionResult:
pre_compact_token_count: int = 0
post_compact_token_count: int = 0
summary_text: str = ''
usage: UsageStats = field(default_factory=UsageStats)
error: str | None = None
@@ -421,6 +423,7 @@ def compact_conversation(
pre_compact_token_count=pre_tokens,
post_compact_token_count=post_tokens,
summary_text=summary_text,
usage=turn.usage,
)
src/lsp_runtime.py
+1108
View File
File diff suppressed because it is too large Load Diff
src/main.py
+163
View File
@@ -24,6 +24,7 @@ from .bootstrap_graph import build_bootstrap_graph
from .command_graph import build_command_graph
from .commands import execute_command, get_command, get_commands, render_command_index
from .config_runtime import ConfigRuntime
from .lsp_runtime import LSPRuntime
from .mcp_runtime import MCPRuntime
from .parity_audit import run_parity_audit
from .permissions import ToolPermissionContext
@@ -680,6 +681,52 @@ def build_parser() -> argparse.ArgumentParser:
config_set_parser.add_argument('value_json')
config_set_parser.add_argument('--source', default='local')
config_set_parser.add_argument('--cwd', default='.')
lsp_status_parser = subparsers.add_parser('lsp-status', help='show local LSP runtime summary')
lsp_status_parser.add_argument('--cwd', default='.')
lsp_symbols_parser = subparsers.add_parser('lsp-symbols', help='show local LSP document symbols for one file')
lsp_symbols_parser.add_argument('file_path')
lsp_symbols_parser.add_argument('--cwd', default='.')
lsp_workspace_parser = subparsers.add_parser('lsp-workspace-symbols', help='search workspace symbols through the local LSP runtime')
lsp_workspace_parser.add_argument('query')
lsp_workspace_parser.add_argument('--max-results', type=int, default=50)
lsp_workspace_parser.add_argument('--cwd', default='.')
lsp_definition_parser = subparsers.add_parser('lsp-definition', help='run a local LSP definition query')
lsp_definition_parser.add_argument('file_path')
lsp_definition_parser.add_argument('line', type=int)
lsp_definition_parser.add_argument('character', type=int)
lsp_definition_parser.add_argument('--max-results', type=int, default=20)
lsp_definition_parser.add_argument('--cwd', default='.')
lsp_references_parser = subparsers.add_parser('lsp-references', help='run a local LSP references query')
lsp_references_parser.add_argument('file_path')
lsp_references_parser.add_argument('line', type=int)
lsp_references_parser.add_argument('character', type=int)
lsp_references_parser.add_argument('--max-results', type=int, default=50)
lsp_references_parser.add_argument('--cwd', default='.')
lsp_hover_parser = subparsers.add_parser('lsp-hover', help='run a local LSP hover query')
lsp_hover_parser.add_argument('file_path')
lsp_hover_parser.add_argument('line', type=int)
lsp_hover_parser.add_argument('character', type=int)
lsp_hover_parser.add_argument('--cwd', default='.')
lsp_diagnostics_parser = subparsers.add_parser('lsp-diagnostics', help='show local LSP diagnostics')
lsp_diagnostics_parser.add_argument('--file-path')
lsp_diagnostics_parser.add_argument('--cwd', default='.')
lsp_hierarchy_parser = subparsers.add_parser('lsp-call-hierarchy', help='show local LSP call hierarchy at a position')
lsp_hierarchy_parser.add_argument('file_path')
lsp_hierarchy_parser.add_argument('line', type=int)
lsp_hierarchy_parser.add_argument('character', type=int)
lsp_hierarchy_parser.add_argument('--cwd', default='.')
lsp_incoming_parser = subparsers.add_parser('lsp-incoming-calls', help='show local LSP incoming calls at a position')
lsp_incoming_parser.add_argument('file_path')
lsp_incoming_parser.add_argument('line', type=int)
lsp_incoming_parser.add_argument('character', type=int)
lsp_incoming_parser.add_argument('--max-results', type=int, default=50)
lsp_incoming_parser.add_argument('--cwd', default='.')
lsp_outgoing_parser = subparsers.add_parser('lsp-outgoing-calls', help='show local LSP outgoing calls at a position')
lsp_outgoing_parser.add_argument('file_path')
lsp_outgoing_parser.add_argument('line', type=int)
lsp_outgoing_parser.add_argument('character', type=int)
lsp_outgoing_parser.add_argument('--max-results', type=int, default=50)
lsp_outgoing_parser.add_argument('--cwd', default='.')
workflow_list_parser = subparsers.add_parser('workflow-list', help='list local workflow definitions')
workflow_list_parser.add_argument('--cwd', default='.')
workflow_list_parser.add_argument('--query')
@@ -824,6 +871,9 @@ def build_parser() -> argparse.ArgumentParser:
context_raw_parser = subparsers.add_parser('agent-context-raw', help='render the raw Python agent context snapshot')
_add_agent_common_args(context_raw_parser, include_backend=False)
token_budget_parser = subparsers.add_parser('token-budget', help='render the current token budget and prompt-length limits')
_add_agent_common_args(token_budget_parser, include_backend=False)
return parser
@@ -1105,6 +1155,114 @@ def main(argv: list[str] | None = None) -> int:
print(f'effective_key_count={mutation.effective_key_count}')
print(runtime.render_value(args.key_path))
return 0
if args.command == 'lsp-status':
runtime = LSPRuntime.from_workspace(Path(args.cwd).resolve())
print('# LSP')
print()
print(runtime.render_summary())
return 0
if args.command == 'lsp-symbols':
runtime = LSPRuntime.from_workspace(Path(args.cwd).resolve())
try:
print(runtime.render_document_symbols(args.file_path))
except KeyError as exc:
print(exc)
return 1
return 0
if args.command == 'lsp-workspace-symbols':
runtime = LSPRuntime.from_workspace(Path(args.cwd).resolve())
print(runtime.render_workspace_symbols(args.query, max_results=args.max_results))
return 0
if args.command == 'lsp-definition':
runtime = LSPRuntime.from_workspace(Path(args.cwd).resolve())
try:
print(
runtime.render_definition(
args.file_path,
args.line,
args.character,
max_results=args.max_results,
)
)
except KeyError as exc:
print(exc)
return 1
return 0
if args.command == 'lsp-references':
runtime = LSPRuntime.from_workspace(Path(args.cwd).resolve())
try:
print(
runtime.render_references(
args.file_path,
args.line,
args.character,
max_results=args.max_results,
)
)
except KeyError as exc:
print(exc)
return 1
return 0
if args.command == 'lsp-hover':
runtime = LSPRuntime.from_workspace(Path(args.cwd).resolve())
try:
print(runtime.render_hover(args.file_path, args.line, args.character))
except KeyError as exc:
print(exc)
return 1
return 0
if args.command == 'lsp-diagnostics':
runtime = LSPRuntime.from_workspace(Path(args.cwd).resolve())
try:
print(runtime.render_diagnostics(args.file_path))
except KeyError as exc:
print(exc)
return 1
return 0
if args.command == 'lsp-call-hierarchy':
runtime = LSPRuntime.from_workspace(Path(args.cwd).resolve())
try:
print(
runtime.render_prepare_call_hierarchy(
args.file_path,
args.line,
args.character,
)
)
except KeyError as exc:
print(exc)
return 1
return 0
if args.command == 'lsp-incoming-calls':
runtime = LSPRuntime.from_workspace(Path(args.cwd).resolve())
try:
print(
runtime.render_incoming_calls(
args.file_path,
args.line,
args.character,
max_results=args.max_results,
)
)
except KeyError as exc:
print(exc)
return 1
return 0
if args.command == 'lsp-outgoing-calls':
runtime = LSPRuntime.from_workspace(Path(args.cwd).resolve())
try:
print(
runtime.render_outgoing_calls(
args.file_path,
args.line,
args.character,
max_results=args.max_results,
)
)
except KeyError as exc:
print(exc)
return 1
return 0
if args.command == 'workflow-list':
runtime = WorkflowRuntime.from_workspace(Path(args.cwd).resolve())
print(runtime.render_workflows_index(query=args.query))
@@ -1344,6 +1502,11 @@ def main(argv: list[str] | None = None) -> int:
agent = _build_agent(args)
print(agent.render_context_snapshot_report())
return 0
if args.command == 'token-budget':
agent = _build_agent(args)
agent.last_session = agent.build_session()
print(agent.render_token_budget_report())
return 0
parser.error(f'unknown command: {args.command}')
return 2
src/token_budget.py
+144
View File
@@ -0,0 +1,144 @@
from __future__ import annotations
from dataclasses import dataclass
from .agent_context_usage import collect_context_usage, infer_context_window
from .agent_session import AgentSessionState
from .agent_types import BudgetConfig, OutputSchemaConfig
from .compact import AUTOCOMPACT_BUFFER_TOKENS
DEFAULT_OUTPUT_RESERVE_TOKENS = 4096
MIN_OUTPUT_RESERVE_TOKENS = 1024
CHAT_MESSAGE_OVERHEAD_TOKENS = 5
CHAT_TOOL_CALL_OVERHEAD_TOKENS = 12
CHAT_NAME_OVERHEAD_TOKENS = 2
OUTPUT_SCHEMA_OVERHEAD_TOKENS = 256
@dataclass(frozen=True)
class TokenBudgetSnapshot:
model: str
context_window_tokens: int
projected_input_tokens: int
message_tokens: int
chat_overhead_tokens: int
reserved_output_tokens: int
reserved_compaction_buffer_tokens: int
reserved_schema_tokens: int
hard_input_limit_tokens: int
soft_input_limit_tokens: int
overflow_tokens: int
soft_overflow_tokens: int
exceeds_hard_limit: bool
exceeds_soft_limit: bool
token_counter_backend: str
token_counter_source: str
token_counter_accurate: bool
def calculate_token_budget(
*,
session: AgentSessionState,
model: str,
budget_config: BudgetConfig,
output_schema: OutputSchemaConfig | None = None,
) -> TokenBudgetSnapshot:
usage = collect_context_usage(
session=session,
model=model,
strategy='token_budget',
)
context_window_tokens = infer_context_window(model)
chat_overhead_tokens = estimate_chat_overhead(session)
reserved_output_tokens = _resolve_output_reserve(
context_window_tokens,
budget_config,
)
reserved_schema_tokens = OUTPUT_SCHEMA_OVERHEAD_TOKENS if output_schema is not None else 0
hard_input_limit_tokens = max(
context_window_tokens - reserved_output_tokens - reserved_schema_tokens,
0,
)
if budget_config.max_input_tokens is not None:
hard_input_limit_tokens = min(
hard_input_limit_tokens,
max(budget_config.max_input_tokens, 0),
)
reserved_compaction_buffer_tokens = min(
AUTOCOMPACT_BUFFER_TOKENS,
max(context_window_tokens // 10, MIN_OUTPUT_RESERVE_TOKENS),
)
soft_input_limit_tokens = max(
hard_input_limit_tokens - reserved_compaction_buffer_tokens,
0,
)
projected_input_tokens = usage.total_tokens + chat_overhead_tokens
overflow_tokens = max(projected_input_tokens - hard_input_limit_tokens, 0)
soft_overflow_tokens = max(projected_input_tokens - soft_input_limit_tokens, 0)
return TokenBudgetSnapshot(
model=model,
context_window_tokens=context_window_tokens,
projected_input_tokens=projected_input_tokens,
message_tokens=usage.total_tokens,
chat_overhead_tokens=chat_overhead_tokens,
reserved_output_tokens=reserved_output_tokens,
reserved_compaction_buffer_tokens=reserved_compaction_buffer_tokens,
reserved_schema_tokens=reserved_schema_tokens,
hard_input_limit_tokens=hard_input_limit_tokens,
soft_input_limit_tokens=soft_input_limit_tokens,
overflow_tokens=overflow_tokens,
soft_overflow_tokens=soft_overflow_tokens,
exceeds_hard_limit=overflow_tokens > 0,
exceeds_soft_limit=soft_overflow_tokens > 0,
token_counter_backend=usage.token_counter_backend,
token_counter_source=usage.token_counter_source,
token_counter_accurate=usage.token_counter_accurate,
)
def format_token_budget(snapshot: TokenBudgetSnapshot) -> str:
lines = [
'# Token Budget',
'',
f'- Model: {snapshot.model}',
f'- Context window: {snapshot.context_window_tokens:,}',
f'- Prompt tokens: {snapshot.projected_input_tokens:,}',
f'- Message/body tokens: {snapshot.message_tokens:,}',
f'- Chat framing overhead: {snapshot.chat_overhead_tokens:,}',
f'- Reserved output tokens: {snapshot.reserved_output_tokens:,}',
f'- Reserved schema tokens: {snapshot.reserved_schema_tokens:,}',
f'- Auto-compact buffer: {snapshot.reserved_compaction_buffer_tokens:,}',
f'- Hard input limit: {snapshot.hard_input_limit_tokens:,}',
f'- Soft input limit: {snapshot.soft_input_limit_tokens:,}',
f'- Token counter: {snapshot.token_counter_backend} ({snapshot.token_counter_source})'
+ (' [accurate]' if snapshot.token_counter_accurate else ' [fallback]'),
]
if snapshot.exceeds_hard_limit:
lines.append(f'- Hard overflow: {snapshot.overflow_tokens:,}')
elif snapshot.exceeds_soft_limit:
lines.append(f'- Soft overflow: {snapshot.soft_overflow_tokens:,}')
else:
remaining = max(snapshot.soft_input_limit_tokens - snapshot.projected_input_tokens, 0)
lines.append(f'- Remaining soft headroom: {remaining:,}')
return '\n'.join(lines)
def estimate_chat_overhead(session: AgentSessionState) -> int:
total = 0
for message in session.messages:
total += CHAT_MESSAGE_OVERHEAD_TOKENS
if message.name:
total += CHAT_NAME_OVERHEAD_TOKENS
total += len(message.tool_calls) * CHAT_TOOL_CALL_OVERHEAD_TOKENS
return total + 3
def _resolve_output_reserve(
context_window_tokens: int,
budget_config: BudgetConfig,
) -> int:
if budget_config.max_output_tokens is not None:
return max(budget_config.max_output_tokens, 0)
suggested = min(DEFAULT_OUTPUT_RESERVE_TOKENS, max(context_window_tokens // 16, MIN_OUTPUT_RESERVE_TOKENS))
return suggested
src/tokenizer_runtime.py
+10 -1
View File
@@ -156,7 +156,7 @@ def _try_build_transformers_counter(
model_ref: str | None,
trust_remote_code: str | None,
) -> ResolvedTokenCounter | None:
if model_ref is None:
if model_ref is None or not _should_try_transformers(model_ref):
return None
try:
from transformers import AutoTokenizer
@@ -196,6 +196,15 @@ def _try_build_transformers_counter(
)
def _should_try_transformers(model_ref: str) -> bool:
if os.path.exists(model_ref):
return True
normalized = model_ref.strip()
if not normalized:
return False
return '/' in normalized or '\\' in normalized
def _heuristic_count(text: str) -> int:
if not text:
return 0