diff --git a/PARITY_CHECKLIST.md b/PARITY_CHECKLIST.md index 6422a61..9707b66 100644 --- a/PARITY_CHECKLIST.md +++ b/PARITY_CHECKLIST.md @@ -82,6 +82,7 @@ Done: - [x] Local manifest/env-backed search runtime discovery - [x] Local search-provider activation persistence - [x] Provider-backed web search execution against configured search backends +- [x] Local heuristic LSP runtime for definitions, references, hover, document symbols, workspace symbols, call hierarchy, and diagnostics - [x] Local persistent task runtime discovery - [x] Local task create/get/list/update runtime flows - [x] Local todo-list replacement runtime flow @@ -98,6 +99,10 @@ Done: - [x] Snipped-message metadata with source lineage id and revision - [x] Resume-time compaction / snipping replay reminder - [x] Resume-time compaction replay of source mutation summaries +- [x] Preflight prompt-length validation before each model call +- [x] Hard prompt-length stop before backend calls when the effective input budget is exceeded +- [x] Token-budget calculation with projected prompt size, chat framing overhead, output reserve, and soft/hard input limits +- [x] Preflight auto-compact/context collapse fallback before the next model call - [x] Query-engine facade that can drive the real Python runtime agent - [x] Query-engine runtime event counters and transcript-kind summaries - [x] Query-engine runtime mutation counters @@ -120,9 +125,9 @@ Missing: - [ ] Full executable plugin lifecycle beyond manifest-driven prompt/tool/session hooks, blocking, aliases, virtual tools, and persisted runtime state - [ ] Full session compaction / snipping parity beyond lineage-aware summaries, mutation-serial compaction metadata, and replay reminders - [ ] Full `QueryEngine.ts` parity (session init, message normalization, SDK-compatible message transforms, attachment handling) -- [ ] Auto-compact and context collapse features from `query.ts` -- [ ] Prompt length validation from `query.ts` -- [ ] Token budget calculations from `query/tokenBudget.ts` +- [x] Auto-compact and context collapse features from `query.ts` +- [x] Prompt length validation from `query.ts` +- [x] Token budget calculations from `query/tokenBudget.ts` ## 2. CLI Entrypoints And Runtime Modes @@ -135,6 +140,7 @@ Done: - [x] `agent-prompt` command - [x] `agent-context` command - [x] `agent-context-raw` command +- [x] `token-budget` command - [x] Local background session mode - [x] Local background session listing (`agent-ps`) - [x] Local background session logs (`agent-logs`) @@ -146,6 +152,7 @@ Done: - [x] Local remote runtime inspection commands (`remote-status`, `remote-profiles`, `remote-disconnect`) - [x] Local account runtime inspection commands (`account-status`, `account-profiles`, `account-login`, `account-logout`) - [x] Local search runtime inspection commands (`search-status`, `search-providers`, `search-activate`, `search`) +- [x] Local LSP runtime inspection commands (`lsp-status`, `lsp-symbols`, `lsp-workspace-symbols`, `lsp-definition`, `lsp-references`, `lsp-hover`, `lsp-diagnostics`, `lsp-call-hierarchy`, `lsp-incoming-calls`, `lsp-outgoing-calls`) - [x] Local MCP runtime inspection commands (`mcp-status`, `mcp-resources`, `mcp-resource`, `mcp-tools`, `mcp-call-tool`) - [x] Inventory/helper commands such as `summary`, `manifest`, `commands`, and `tools` @@ -187,6 +194,7 @@ Done: - [x] Local account-runtime guidance section in the Python system prompt - [x] Local planning guidance section in the Python system prompt - [x] Local task guidance section in the Python system prompt +- [x] Local LSP guidance section in the Python system prompt - [x] Product metadata/branding from `constants/product.ts` — ported to `src/prompt_constants.py` - [x] API limits constants from `constants/apiLimits.ts` — ported to `src/prompt_constants.py` @@ -238,6 +246,7 @@ Done: - [x] Manifest-based remote runtime summary injection - [x] Manifest/env-based search runtime summary injection - [x] Manifest-based account runtime summary injection +- [x] Local LSP runtime summary injection - [x] Manifest-based plan runtime summary injection - [x] Manifest-based task runtime summary injection @@ -265,6 +274,7 @@ Done (37 slash command names in 29 specs): - [x] `/help`, `/commands` - [x] `/context`, `/usage` - [x] `/context-raw`, `/env` +- [x] `/token-budget`, `/budget` - [x] `/mcp` (with subcommands: `tools`, `tool `) - [x] `/search` (with subcommands: `providers`, `provider`, `use`) - [x] `/remote` (with `enter`, `exit`) @@ -368,6 +378,7 @@ Missing npm slash commands (from `src/commands/` — 80+ commands total): - [x] `search_list_providers` - [x] `search_activate_provider` - [x] `web_search` +- [x] `LSP` - [x] `tool_search` - [x] `sleep` - [x] `ask_user_question` @@ -420,7 +431,7 @@ Core tools needing full port: - [ ] `AgentTool` — Sub-agent spawning with built-in agents (explore, general-purpose, verification, plan, claudeCodeGuide, statusline), fork support, agent memory/snapshots, resume agent, color management - [ ] `SkillTool` — Skill execution with bundled skills - [ ] `BriefTool` — Brief mode with attachments and file upload -- [ ] `LSPTool` — Language Server Protocol (diagnostics, go-to-definition, references, hover, symbol search, formatting) +- [ ] `LSPTool` — Full upstream LSP fidelity beyond the current local heuristic runtime (server-backed diagnostics, go-to-definition, references, hover, symbol search, formatting) - [ ] `PowerShellTool` — Full PowerShell execution with security, path validation, CLM types, git safety - [ ] `REPLTool` — Interactive REPL with primitive tools (ant-only) - [ ] `MCPTool` — Full MCP tool execution with collapse classification @@ -773,7 +784,7 @@ Mirrored / scaffold areas needing real implementation: ### Tier 1 — Core Feature Gaps (highest user impact) - [x] Full BashTool security parity (sed validation, path validation, sandbox, destructive command warnings, command semantics) → `src/bash_security.py` -- [ ] LSP tool integration for code intelligence +- [x] LSP tool integration for code intelligence - [ ] Full AgentTool with built-in agent types (explore, general-purpose, verification, plan) - [ ] Auto-compact and context collapse from `query.ts` - [ ] Full compact service (autoCompact, microCompact, sessionMemoryCompact) diff --git a/README.md b/README.md index 0e23e21..d85b56d 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,8 @@ | 🆕 | **Remote Trigger Runtime** | Local remote triggers with create/update/run flows similar to the npm remote trigger surface | | 🆕 | **Worktree Runtime** | Managed git worktrees with mid-session cwd switching, slash commands, and CLI flows | | 🆕 | **Tokenizer-Aware Context** | Cached tokenizer backends with heuristic fallback for `/context`, `/status`, and compaction | +| 🆕 | **Prompt Budget Preflight** | Preflight prompt-length validation, token-budget reporting, and auto-compact/context collapse before backend failures | +| 🆕 | **LSP Runtime** | Local LSP-style code intelligence for definitions, references, hover, symbols, call hierarchy, and diagnostics | | 🆕 | **Daemon Commands** | Local `daemon start/ps/logs/attach/kill` wrapper over background agent sessions | | 🆕 | **Background Sessions** | Local `agent-bg`, `agent-ps`, `agent-logs`, `agent-attach`, and `agent-kill` flows | | 🆕 | **Testing Guide** | Comprehensive [TESTING_GUIDE.md](TESTING_GUIDE.md) with commands for every feature | @@ -99,8 +101,10 @@ Built on the public porting workspace from [instructkr/claw-code](https://github | 🧭 **Workflow Runtime** | Manifest-backed workflows with slash commands, CLI inspection, and recorded runs | | ⏰ **Remote Triggers** | Local remote triggers with create/update/run flows and npm-style trigger actions | | 🪝 **Hook & Policy Runtime** | Trust reporting, safe env, managed settings, tool blocking, and budget overrides | +| 🧠 **LSP Code Intelligence** | Local LSP-style definitions, references, hover, symbols, diagnostics, and call hierarchy | | 🧠 **Context Engine** | Automatic context building with CLAUDE.md discovery, compaction, and snipping | | 🔢 **Tokenizer-Aware Accounting** | Model-aware token counting with cached tokenizer backends and fallback heuristics | +| 📏 **Prompt Budgeting** | Soft/hard prompt-window checks, token-budget reports, and preflight context collapse | | 🔄 **Session Persistence** | Save and resume agent sessions with file-history replay | | 🗂️ **Background Sessions** | `agent-bg` and local daemon wrappers for background runs, logs, attach, and kill | | 💰 **Cost & Budget Control** | Token budgets, cost limits, tool-call caps, model-call caps | @@ -137,6 +141,8 @@ Built on the public porting workspace from [instructkr/claw-code](https://github - [x] Truncated-response continuation flow - [x] Auto-snip and auto-compact context reduction - [x] Reactive compaction retry on prompt-too-long errors +- [x] Preflight prompt-length validation and token-budget reporting +- [x] Preflight auto-compact/context collapse before backend prompt-too-long failures - [x] Cost tracking and usage budget enforcement - [x] Token, tool-call, model-call, and session-turn budgets - [x] Structured output / JSON schema response mode @@ -148,6 +154,7 @@ Built on the public porting workspace from [instructkr/claw-code](https://github - [x] Local remote runtime: manifest discovery, profile listing, connect/disconnect persistence, and CLI/slash flows - [x] Local hook and policy runtime with trust reporting, safe env, tool blocking, and budget overrides - [x] Local config runtime: config discovery, effective settings, source inspection, and config mutation +- [x] Local LSP runtime: definitions, references, hover, symbols, diagnostics, and call hierarchy - [x] Local account runtime: profile discovery, login/logout state, and account CLI/slash flows - [x] Local ask-user runtime: queued answers, history, and ask-user CLI/slash flows - [x] Local team runtime: persisted teams, team messages, and team CLI/slash flows @@ -217,6 +224,8 @@ claw-code/ │ ├── account_runtime.py # Local account profiles, login/logout state, account CLI support │ ├── ask_user_runtime.py # Local ask-user queued answers and interaction history │ ├── config_runtime.py # Local workspace config/settings discovery and mutation +│ ├── lsp_runtime.py # Local LSP-style code intelligence and diagnostics +│ ├── token_budget.py # Prompt-window budgeting and preflight prompt-length validation │ ├── plan_runtime.py # Persistent plan runtime and plan sync │ ├── task_runtime.py # Persistent task runtime and task execution │ ├── task.py # Task state model and task dataclasses @@ -436,6 +445,7 @@ python3 -m src.main agent \ | `agent-prompt` | Show the assembled system prompt | | `agent-context` | Show estimated context usage | | `agent-context-raw` | Show the raw context snapshot | +| `token-budget` | Show prompt-window budget, reserves, and soft/hard input limits | | `agent-resume ` | Resume a saved session | ### Runtime Utility Commands @@ -509,6 +519,7 @@ These are handled **locally** before the model loop: | `/help` | `/commands` | Show built-in slash commands | | `/context` | `/usage` | Show estimated session context usage | | `/context-raw` | `/env` | Show raw environment & context snapshot | +| `/token-budget` | `/budget` | Show prompt-window budget, reserves, and soft/hard input limits | | `/mcp` | — | Show MCP runtime status, tools, or a single MCP tool | | `/resources` | — | List MCP resources | | `/resource` | — | Read an MCP resource by URI | @@ -541,6 +552,7 @@ These are handled **locally** before the model loop: ```bash python3 -m src.main agent "/help" python3 -m src.main agent "/context" --cwd . +python3 -m src.main agent "/token-budget" --cwd . python3 -m src.main agent "/tools" --cwd . python3 -m src.main agent "/status" --cwd . ``` diff --git a/TESTING_GUIDE.md b/TESTING_GUIDE.md index 419e738..a47a38d 100644 --- a/TESTING_GUIDE.md +++ b/TESTING_GUIDE.md @@ -91,6 +91,7 @@ python3 -m unittest tests.test_plan_runtime -v python3 -m unittest tests.test_background_runtime -v python3 -m unittest tests.test_remote_runtime -v python3 -m unittest tests.test_config_runtime -v +python3 -m unittest tests.test_lsp_runtime -v python3 -m unittest tests.test_account_runtime -v python3 -m unittest tests.test_ask_user_runtime -v python3 -m unittest tests.test_team_runtime -v @@ -309,6 +310,30 @@ cat > ./test_cases_notebooks/demo.ipynb <<'EOF' EOF ``` +### 4.2e LSP fixture + +```bash +cat > ./test_cases/sample.py <<'EOF' +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)) +EOF + +cat > ./test_cases/broken.py <<'EOF' +def broken(: + pass +EOF +``` + ### 4.3 Remote fixtures ```bash @@ -568,6 +593,8 @@ python3 -m src.main agent "/context" --cwd ./test_cases python3 -m src.main agent "/usage summarize current session" --cwd ./test_cases python3 -m src.main agent "/context-raw" --cwd ./test_cases python3 -m src.main agent "/env" --cwd ./test_cases +python3 -m src.main agent "/token-budget" --cwd ./test_cases +python3 -m src.main agent "/budget" --cwd ./test_cases python3 -m src.main agent "/prompt" --cwd ./test_cases python3 -m src.main agent "/system-prompt" --cwd ./test_cases python3 -m src.main agent "/permissions" --cwd ./test_cases @@ -650,6 +677,7 @@ python3 -m src.main agent "/mcp (MCP)" --cwd ./test_cases_mcp python3 -m src.main agent-prompt --cwd ./test_cases python3 -m src.main agent-context --cwd ./test_cases python3 -m src.main agent-context-raw --cwd ./test_cases +python3 -m src.main token-budget --cwd ./test_cases ``` ### 6.2 Extra working directories and `CLAUDE.md` toggle @@ -680,6 +708,7 @@ python3 -m src.main agent-prompt \ ```bash python3 -m src.main agent "/status" --cwd ./test_cases python3 -m src.main agent-context --cwd ./test_cases +python3 -m src.main token-budget --cwd ./test_cases ``` Override the tokenizer backend: @@ -691,6 +720,7 @@ export CLAW_CODE_TOKENIZER_MODEL=Qwen/Qwen3-Coder-30B-A3B-Instruct python3 -m src.main agent "/status" --cwd ./test_cases python3 -m src.main agent-context --cwd ./test_cases +python3 -m src.main token-budget --cwd ./test_cases ``` If no tokenizer backend is available, the runtime will fall back to the heuristic counter and `/status` will report that. @@ -1228,6 +1258,46 @@ python3 -m src.main agent \ --show-transcript ``` +## 14B. LSP Runtime + +### 14B.1 CLI status and code intelligence reports + +```bash +python3 -m src.main lsp-status --cwd ./test_cases +python3 -m src.main lsp-symbols sample.py --cwd ./test_cases +python3 -m src.main lsp-workspace-symbols helper --cwd ./test_cases +python3 -m src.main lsp-definition sample.py 6 12 --cwd ./test_cases +python3 -m src.main lsp-references sample.py 6 12 --cwd ./test_cases +python3 -m src.main lsp-hover sample.py 1 5 --cwd ./test_cases +python3 -m src.main lsp-diagnostics --cwd ./test_cases +python3 -m src.main lsp-diagnostics --cwd ./test_cases --file-path broken.py +python3 -m src.main lsp-call-hierarchy sample.py 6 12 --cwd ./test_cases +python3 -m src.main lsp-incoming-calls sample.py 1 5 --cwd ./test_cases +python3 -m src.main lsp-outgoing-calls sample.py 6 12 --cwd ./test_cases +``` + +### 14B.2 Slash commands + +```bash +python3 -m src.main agent "/lsp" --cwd ./test_cases +python3 -m src.main agent "/lsp symbols sample.py" --cwd ./test_cases +python3 -m src.main agent "/lsp workspace helper" --cwd ./test_cases +python3 -m src.main agent "/lsp definition sample.py 6 12" --cwd ./test_cases +python3 -m src.main agent "/lsp references sample.py 6 12" --cwd ./test_cases +python3 -m src.main agent "/lsp hover sample.py 1 5" --cwd ./test_cases +python3 -m src.main agent "/lsp diagnostics broken.py" --cwd ./test_cases +python3 -m src.main agent "/lsp hierarchy sample.py 6 12" --cwd ./test_cases +``` + +### 14B.3 Real tool loop + +```bash +python3 -m src.main agent \ + "Use the LSP tool to find the definition of helper in sample.py, then summarize the result." \ + --cwd ./test_cases \ + --show-transcript +``` + ## 14A. Ask-user Runtime ### 14A.1 CLI status and history diff --git a/src/__init__.py b/src/__init__.py index a068263..e540ac7 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -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', diff --git a/src/agent_context.py b/src/agent_context.py index 7f4cedb..64892fe 100644 --- a/src/agent_context.py +++ b/src/agent_context.py @@ -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() diff --git a/src/agent_prompting.py b/src/agent_prompting.py index 8ee0831..42b5224 100644 --- a/src/agent_prompting.py +++ b/src/agent_prompting.py @@ -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: diff --git a/src/agent_runtime.py b/src/agent_runtime.py index c42180d..8a5a383 100644 --- a/src/agent_runtime.py +++ b/src/agent_runtime.py @@ -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, diff --git a/src/agent_slash_commands.py b/src/agent_slash_commands.py index e61e250..3590962 100644 --- a/src/agent_slash_commands.py +++ b/src/agent_slash_commands.py @@ -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 ') + 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 ') + 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} ', + ) + 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 |workspace |definition |references |hover |hierarchy |incoming |outgoing |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)) diff --git a/src/agent_tools.py b/src/agent_tools.py index b3a3eef..317edd5 100644 --- a/src/agent_tools.py +++ b/src/agent_tools.py @@ -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 diff --git a/src/compact.py b/src/compact.py index 3e4a786..4a322a1 100644 --- a/src/compact.py +++ b/src/compact.py @@ -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, ) diff --git a/src/lsp_runtime.py b/src/lsp_runtime.py new file mode 100644 index 0000000..455d6f4 --- /dev/null +++ b/src/lsp_runtime.py @@ -0,0 +1,1108 @@ +from __future__ import annotations + +import ast +import json +import os +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + + +LSP_MANIFEST_FILES = ( + Path('.claw-lsp.json'), + Path('.claude/lsp.json'), +) +DEFAULT_SUPPORTED_EXTENSIONS = ( + '.py', + '.pyi', + '.js', + '.jsx', + '.ts', + '.tsx', + '.json', +) +DEFAULT_IGNORED_DIRECTORIES = ( + '.git', + '.hg', + '.svn', + '.port_sessions', + '__pycache__', + 'node_modules', + 'dist', + 'build', + '.venv', + 'venv', +) +DEFAULT_MAX_INDEXED_FILES = 500 +DEFAULT_MAX_FILE_BYTES = 1_500_000 +WORD_PATTERN = re.compile(r'[A-Za-z_][A-Za-z0-9_]*') + + +@dataclass(frozen=True) +class LSPSymbol: + name: str + kind: str + path: str + line: int + character: int + end_line: int + end_character: int + container_name: str | None = None + signature: str | None = None + documentation: str | None = None + + @property + def symbol_id(self) -> str: + return f'{self.path}:{self.line}:{self.character}:{self.kind}:{self.name}' + + def contains(self, line: int, character: int) -> bool: + if line < self.line or line > self.end_line: + return False + if line == self.line and character < self.character: + return False + if line == self.end_line and character > self.end_character: + return False + return True + + +@dataclass(frozen=True) +class LSPReference: + name: str + path: str + line: int + character: int + line_text: str + + +@dataclass(frozen=True) +class LSPDiagnostic: + path: str + severity: str + message: str + line: int + character: int + code: str | None = None + source: str | None = None + + +@dataclass(frozen=True) +class LSPCallEdge: + caller_symbol_id: str + caller_name: str + callee_name: str + path: str + line: int + character: int + + +@dataclass(frozen=True) +class IndexedFile: + path: Path + language: str + text: str + symbols: tuple[LSPSymbol, ...] = () + diagnostics: tuple[LSPDiagnostic, ...] = () + call_edges: tuple[LSPCallEdge, ...] = () + + +@dataclass(frozen=True) +class LSPQueryResult: + operation: str + content: str + result_count: int = 0 + file_count: int = 0 + symbol_name: str | None = None + + +@dataclass +class LSPRuntime: + cwd: Path + additional_working_directories: tuple[Path, ...] = () + manifests: tuple[str, ...] = () + supported_extensions: tuple[str, ...] = DEFAULT_SUPPORTED_EXTENSIONS + ignored_directories: tuple[str, ...] = DEFAULT_IGNORED_DIRECTORIES + max_indexed_files: int = DEFAULT_MAX_INDEXED_FILES + max_file_bytes: int = DEFAULT_MAX_FILE_BYTES + _cache: dict[str, tuple[int, int, IndexedFile]] = field(default_factory=dict, init=False, repr=False) + + @classmethod + def from_workspace( + cls, + cwd: Path, + additional_working_directories: tuple[str, ...] = (), + ) -> 'LSPRuntime': + resolved_cwd = cwd.resolve() + resolved_dirs = tuple(Path(path).resolve() for path in additional_working_directories) + manifest_paths = _discover_manifest_paths(resolved_cwd, resolved_dirs) + supported_extensions = set(DEFAULT_SUPPORTED_EXTENSIONS) + ignored_directories = set(DEFAULT_IGNORED_DIRECTORIES) + max_indexed_files = DEFAULT_MAX_INDEXED_FILES + max_file_bytes = DEFAULT_MAX_FILE_BYTES + for manifest_path in manifest_paths: + payload = _load_manifest_payload(manifest_path) + raw_extensions = payload.get('includeExtensions') + if isinstance(raw_extensions, list): + supported_extensions.update( + extension + for extension in raw_extensions + if isinstance(extension, str) and extension.strip() + ) + raw_ignored = payload.get('excludeDirs') + if isinstance(raw_ignored, list): + ignored_directories.update( + item + for item in raw_ignored + if isinstance(item, str) and item.strip() + ) + raw_max_files = payload.get('maxIndexedFiles') + if isinstance(raw_max_files, int) and raw_max_files > 0: + max_indexed_files = raw_max_files + raw_max_bytes = payload.get('maxFileBytes') + if isinstance(raw_max_bytes, int) and raw_max_bytes > 0: + max_file_bytes = raw_max_bytes + return cls( + cwd=resolved_cwd, + additional_working_directories=resolved_dirs, + manifests=tuple(str(path) for path in manifest_paths), + supported_extensions=tuple(sorted(supported_extensions)), + ignored_directories=tuple(sorted(ignored_directories)), + max_indexed_files=max_indexed_files, + max_file_bytes=max_file_bytes, + ) + + def has_lsp_support(self) -> bool: + return bool(self._workspace_files(limit=1)) + + def render_summary(self) -> str: + files = self._workspace_files(limit=self.max_indexed_files) + language_counts: dict[str, int] = {} + for path in files: + language = _language_for_extension(path.suffix) + language_counts[language] = language_counts.get(language, 0) + 1 + lines = [ + f'Heuristic LSP manifests: {len(self.manifests)}', + f'LSP roots: {len(self._roots())}', + f'Indexed candidate files: {len(files)}', + f'Supported extensions: {", ".join(self.supported_extensions)}', + ] + if language_counts: + lines.append('- Languages:') + for language, count in sorted(language_counts.items()): + lines.append(f' - {language}: {count}') + return '\n'.join(lines) + + def render_document_symbols(self, file_path: str) -> str: + indexed = self._indexed_from_user_path(file_path) + lines = ['# LSP Document Symbols', '', f'path={self._display_path(indexed.path)}'] + if not indexed.symbols: + lines.extend(['', 'No document symbols found.']) + return '\n'.join(lines) + lines.append(f'symbol_count={len(indexed.symbols)}') + lines.append('') + for symbol in indexed.symbols: + detail = f'{symbol.kind} {symbol.name} @ {self._display_path(Path(symbol.path))}:{symbol.line}:{symbol.character}' + if symbol.container_name: + detail += f' ; container={symbol.container_name}' + lines.append(f'- {detail}') + return '\n'.join(lines) + + def render_workspace_symbols( + self, + query: str, + *, + max_results: int = 50, + ) -> str: + results = self.workspace_symbols(query, max_results=max_results) + lines = ['# LSP Workspace Symbols', '', f'query={query}', f'result_count={len(results)}', ''] + if not results: + lines.append('No workspace symbols found.') + return '\n'.join(lines) + for symbol in results: + lines.append( + f'- {symbol.kind} {symbol.name} @ {self._display_path(Path(symbol.path))}:{symbol.line}:{symbol.character}' + ) + return '\n'.join(lines) + + def render_definition( + self, + file_path: str, + line: int, + character: int, + *, + max_results: int = 20, + ) -> str: + indexed = self._indexed_from_user_path(file_path) + symbol_name = self._symbol_name_at_position(indexed, line, character) + definitions = self.go_to_definition(file_path, line, character, max_results=max_results) + lines = ['# LSP Definition', ''] + if symbol_name: + lines.append(f'symbol={symbol_name}') + lines.append(f'result_count={len(definitions)}') + lines.append('') + if not definitions: + lines.append('No definitions found.') + return '\n'.join(lines) + for symbol in definitions: + lines.append( + f'- {symbol.kind} {symbol.name} @ {self._display_path(Path(symbol.path))}:{symbol.line}:{symbol.character}' + ) + return '\n'.join(lines) + + def render_references( + self, + file_path: str, + line: int, + character: int, + *, + max_results: int = 50, + ) -> str: + indexed = self._indexed_from_user_path(file_path) + symbol_name = self._symbol_name_at_position(indexed, line, character) + references = self.find_references(file_path, line, character, max_results=max_results) + lines = ['# LSP References', ''] + if symbol_name: + lines.append(f'symbol={symbol_name}') + lines.append(f'result_count={len(references)}') + lines.append(f'file_count={len({reference.path for reference in references})}') + lines.append('') + if not references: + lines.append('No references found.') + return '\n'.join(lines) + for reference in references: + lines.append( + f'- {self._display_path(Path(reference.path))}:{reference.line}:{reference.character} :: {reference.line_text.strip()}' + ) + return '\n'.join(lines) + + def render_hover(self, file_path: str, line: int, character: int) -> str: + hover = self.hover(file_path, line, character) + lines = ['# LSP Hover', ''] + if hover is None: + lines.append('No hover information found.') + return '\n'.join(lines) + lines.append(f'symbol={hover.name}') + lines.append(f'kind={hover.kind}') + lines.append( + f'location={self._display_path(Path(hover.path))}:{hover.line}:{hover.character}' + ) + if hover.signature: + lines.append(f'signature={hover.signature}') + if hover.documentation: + lines.extend(['', hover.documentation]) + return '\n'.join(lines) + + def render_prepare_call_hierarchy(self, file_path: str, line: int, character: int) -> str: + symbol = self.prepare_call_hierarchy(file_path, line, character) + lines = ['# LSP Call Hierarchy', ''] + if symbol is None: + lines.append('No callable symbol found at that position.') + return '\n'.join(lines) + incoming = self.incoming_calls(file_path, line, character, max_results=100) + outgoing = self.outgoing_calls(file_path, line, character, max_results=100) + lines.append(f'symbol={symbol.name}') + lines.append(f'kind={symbol.kind}') + lines.append( + f'location={self._display_path(Path(symbol.path))}:{symbol.line}:{symbol.character}' + ) + lines.append(f'incoming_calls={len(incoming)}') + lines.append(f'outgoing_calls={len(outgoing)}') + return '\n'.join(lines) + + def render_incoming_calls( + self, + file_path: str, + line: int, + character: int, + *, + max_results: int = 50, + ) -> str: + symbol = self.prepare_call_hierarchy(file_path, line, character) + calls = self.incoming_calls(file_path, line, character, max_results=max_results) + lines = ['# LSP Incoming Calls', ''] + if symbol is not None: + lines.append(f'symbol={symbol.name}') + lines.append(f'result_count={len(calls)}') + lines.append('') + if not calls: + lines.append('No incoming calls found.') + return '\n'.join(lines) + for call in calls: + lines.append( + f'- {call.name} @ {self._display_path(Path(call.path))}:{call.line}:{call.character}' + ) + return '\n'.join(lines) + + def render_outgoing_calls( + self, + file_path: str, + line: int, + character: int, + *, + max_results: int = 50, + ) -> str: + symbol = self.prepare_call_hierarchy(file_path, line, character) + calls = self.outgoing_calls(file_path, line, character, max_results=max_results) + lines = ['# LSP Outgoing Calls', ''] + if symbol is not None: + lines.append(f'symbol={symbol.name}') + lines.append(f'result_count={len(calls)}') + lines.append('') + if not calls: + lines.append('No outgoing calls found.') + return '\n'.join(lines) + for call in calls: + lines.append( + f'- {call.name} @ {self._display_path(Path(call.path))}:{call.line}:{call.character}' + ) + return '\n'.join(lines) + + def render_diagnostics(self, file_path: str | None = None) -> str: + diagnostics = self.diagnostics(file_path=file_path) + lines = ['# LSP Diagnostics', ''] + if file_path is not None: + lines.append(f'path={self._display_path(self.resolve_path(file_path))}') + lines.append(f'diagnostic_count={len(diagnostics)}') + lines.append('') + if not diagnostics: + lines.append('No diagnostics found.') + return '\n'.join(lines) + for diagnostic in diagnostics: + code = f' [{diagnostic.code}]' if diagnostic.code else '' + source = f' ({diagnostic.source})' if diagnostic.source else '' + lines.append( + f'- {diagnostic.severity.upper()} {self._display_path(Path(diagnostic.path))}:{diagnostic.line}:{diagnostic.character} {diagnostic.message}{code}{source}' + ) + return '\n'.join(lines) + + def query( + self, + operation: str, + *, + file_path: str, + line: int, + character: int, + query: str | None = None, + max_results: int = 50, + ) -> LSPQueryResult: + if operation == 'documentSymbol': + indexed = self._indexed_from_user_path(file_path) + content = self.render_document_symbols(file_path) + return LSPQueryResult( + operation=operation, + content=content, + result_count=len(indexed.symbols), + file_count=1 if indexed.symbols else 0, + ) + if operation == 'workspaceSymbol': + indexed = self._indexed_from_user_path(file_path) + symbol_name = query or self._symbol_name_at_position(indexed, line, character) + if not symbol_name: + raise KeyError('No symbol name available for workspaceSymbol query') + results = self.workspace_symbols(symbol_name, max_results=max_results) + return LSPQueryResult( + operation=operation, + content=self.render_workspace_symbols(symbol_name, max_results=max_results), + result_count=len(results), + file_count=len({result.path for result in results}), + symbol_name=symbol_name, + ) + if operation == 'goToDefinition': + definitions = self.go_to_definition(file_path, line, character, max_results=max_results) + return LSPQueryResult( + operation=operation, + content=self.render_definition(file_path, line, character, max_results=max_results), + result_count=len(definitions), + file_count=len({definition.path for definition in definitions}), + symbol_name=self._symbol_name_at_position(self._indexed_from_user_path(file_path), line, character), + ) + if operation == 'findReferences': + references = self.find_references(file_path, line, character, max_results=max_results) + return LSPQueryResult( + operation=operation, + content=self.render_references(file_path, line, character, max_results=max_results), + result_count=len(references), + file_count=len({reference.path for reference in references}), + symbol_name=self._symbol_name_at_position(self._indexed_from_user_path(file_path), line, character), + ) + if operation == 'hover': + hover = self.hover(file_path, line, character) + return LSPQueryResult( + operation=operation, + content=self.render_hover(file_path, line, character), + result_count=1 if hover is not None else 0, + file_count=1 if hover is not None else 0, + symbol_name=hover.name if hover is not None else None, + ) + if operation == 'goToImplementation': + definitions = self.go_to_definition(file_path, line, character, max_results=max_results) + content = self.render_definition(file_path, line, character, max_results=max_results).replace( + '# LSP Definition', + '# LSP Implementation', + ) + return LSPQueryResult( + operation=operation, + content=content, + result_count=len(definitions), + file_count=len({definition.path for definition in definitions}), + symbol_name=self._symbol_name_at_position(self._indexed_from_user_path(file_path), line, character), + ) + if operation == 'prepareCallHierarchy': + symbol = self.prepare_call_hierarchy(file_path, line, character) + return LSPQueryResult( + operation=operation, + content=self.render_prepare_call_hierarchy(file_path, line, character), + result_count=1 if symbol is not None else 0, + file_count=1 if symbol is not None else 0, + symbol_name=symbol.name if symbol is not None else None, + ) + if operation == 'incomingCalls': + calls = self.incoming_calls(file_path, line, character, max_results=max_results) + return LSPQueryResult( + operation=operation, + content=self.render_incoming_calls(file_path, line, character, max_results=max_results), + result_count=len(calls), + file_count=len({call.path for call in calls}), + symbol_name=self._symbol_name_at_position(self._indexed_from_user_path(file_path), line, character), + ) + if operation == 'outgoingCalls': + calls = self.outgoing_calls(file_path, line, character, max_results=max_results) + return LSPQueryResult( + operation=operation, + content=self.render_outgoing_calls(file_path, line, character, max_results=max_results), + result_count=len(calls), + file_count=len({call.path for call in calls}), + symbol_name=self._symbol_name_at_position(self._indexed_from_user_path(file_path), line, character), + ) + raise KeyError(operation) + + def workspace_symbols(self, query: str, *, max_results: int = 50) -> tuple[LSPSymbol, ...]: + needle = query.strip().lower() + if not needle: + return () + results: list[LSPSymbol] = [] + for indexed in self._workspace_indexes(): + for symbol in indexed.symbols: + if needle in symbol.name.lower(): + results.append(symbol) + results.sort(key=lambda item: (item.name.lower(), item.path, item.line, item.character)) + return tuple(results[:max_results]) + + def go_to_definition( + self, + file_path: str, + line: int, + character: int, + *, + max_results: int = 20, + ) -> tuple[LSPSymbol, ...]: + indexed = self._indexed_from_user_path(file_path) + symbol_name = self._symbol_name_at_position(indexed, line, character) + if not symbol_name: + return () + results: list[LSPSymbol] = [] + for candidate in self._workspace_indexes(): + for symbol in candidate.symbols: + if symbol.name == symbol_name: + results.append(symbol) + results.sort( + key=lambda item: ( + 0 if Path(item.path) == indexed.path else 1, + abs(item.line - line), + item.path, + item.character, + ) + ) + deduped: list[LSPSymbol] = [] + seen: set[str] = set() + for symbol in results: + if symbol.symbol_id in seen: + continue + seen.add(symbol.symbol_id) + deduped.append(symbol) + return tuple(deduped[:max_results]) + + def find_references( + self, + file_path: str, + line: int, + character: int, + *, + max_results: int = 50, + ) -> tuple[LSPReference, ...]: + indexed = self._indexed_from_user_path(file_path) + symbol_name = self._symbol_name_at_position(indexed, line, character) + if not symbol_name: + return () + references: list[LSPReference] = [] + pattern = re.compile(rf'\b{re.escape(symbol_name)}\b') + for candidate in self._workspace_indexes(): + for line_number, line_text in enumerate(candidate.text.splitlines(), start=1): + for match in pattern.finditer(line_text): + references.append( + LSPReference( + name=symbol_name, + path=str(candidate.path), + line=line_number, + character=match.start() + 1, + line_text=line_text, + ) + ) + if len(references) >= max_results: + return tuple(references) + return tuple(references) + + def hover(self, file_path: str, line: int, character: int) -> LSPSymbol | None: + indexed = self._indexed_from_user_path(file_path) + for symbol in indexed.symbols: + if symbol.contains(line, character): + return symbol + definitions = self.go_to_definition(file_path, line, character, max_results=1) + return definitions[0] if definitions else None + + def prepare_call_hierarchy( + self, + file_path: str, + line: int, + character: int, + ) -> LSPSymbol | None: + indexed = self._indexed_from_user_path(file_path) + callable_symbols = [ + symbol + for symbol in indexed.symbols + if symbol.kind in {'function', 'async_function', 'method'} + ] + containing = [ + symbol + for symbol in callable_symbols + if symbol.contains(line, character) + ] + if containing: + containing.sort(key=lambda item: (item.end_line - item.line, item.end_character - item.character)) + return containing[0] + definitions = self.go_to_definition(file_path, line, character, max_results=1) + if definitions and definitions[0].kind in {'function', 'async_function', 'method'}: + return definitions[0] + return None + + def incoming_calls( + self, + file_path: str, + line: int, + character: int, + *, + max_results: int = 50, + ) -> tuple[LSPSymbol, ...]: + target = self.prepare_call_hierarchy(file_path, line, character) + if target is None: + return () + workspace_indexes = self._workspace_indexes() + by_symbol_id = { + symbol.symbol_id: symbol + for indexed in workspace_indexes + for symbol in indexed.symbols + } + results: list[LSPSymbol] = [] + seen: set[str] = set() + for indexed in workspace_indexes: + for edge in indexed.call_edges: + if edge.callee_name != target.name: + continue + caller = by_symbol_id.get(edge.caller_symbol_id) + if caller is None or caller.symbol_id in seen: + continue + seen.add(caller.symbol_id) + results.append(caller) + if len(results) >= max_results: + return tuple(results) + return tuple(results) + + def outgoing_calls( + self, + file_path: str, + line: int, + character: int, + *, + max_results: int = 50, + ) -> tuple[LSPSymbol, ...]: + target = self.prepare_call_hierarchy(file_path, line, character) + if target is None: + return () + definitions_by_name: dict[str, list[LSPSymbol]] = {} + for indexed in self._workspace_indexes(): + for symbol in indexed.symbols: + definitions_by_name.setdefault(symbol.name, []).append(symbol) + matches: list[LSPSymbol] = [] + seen: set[str] = set() + for indexed in self._workspace_indexes(): + for edge in indexed.call_edges: + if edge.caller_symbol_id != target.symbol_id: + continue + for symbol in definitions_by_name.get(edge.callee_name, []): + if symbol.symbol_id in seen: + continue + seen.add(symbol.symbol_id) + matches.append(symbol) + if len(matches) >= max_results: + return tuple(matches) + return tuple(matches) + + def diagnostics(self, *, file_path: str | None = None) -> tuple[LSPDiagnostic, ...]: + if file_path is not None: + indexed = self._indexed_from_user_path(file_path) + return indexed.diagnostics + diagnostics: list[LSPDiagnostic] = [] + for indexed in self._workspace_indexes(): + diagnostics.extend(indexed.diagnostics) + diagnostics.sort(key=lambda item: (item.path, item.line, item.character)) + return tuple(diagnostics) + + def resolve_path(self, raw_path: str) -> Path: + expanded = Path(raw_path).expanduser() + candidate = expanded if expanded.is_absolute() else self.cwd / expanded + resolved = candidate.resolve(strict=True) + if not self._is_under_roots(resolved): + raise ValueError(f'Path {raw_path!r} escapes the configured LSP roots') + if not resolved.is_file(): + raise FileNotFoundError(raw_path) + return resolved + + def _indexed_from_user_path(self, raw_path: str) -> IndexedFile: + try: + path = self.resolve_path(raw_path) + except FileNotFoundError as exc: + raise KeyError(f'Unknown file: {raw_path}') from exc + except ValueError as exc: + raise KeyError(str(exc)) from exc + return self._index_file(path) + + def _display_path(self, path: Path) -> str: + resolved = path.resolve() + for root in self._roots(): + try: + return str(resolved.relative_to(root)) + except ValueError: + continue + return str(resolved) + + def _roots(self) -> tuple[Path, ...]: + roots = [self.cwd, *self.additional_working_directories] + seen: set[Path] = set() + normalized: list[Path] = [] + for root in roots: + resolved = root.resolve() + if resolved in seen or not resolved.exists(): + continue + seen.add(resolved) + normalized.append(resolved) + return tuple(normalized) + + def _is_under_roots(self, path: Path) -> bool: + for root in self._roots(): + try: + path.relative_to(root) + return True + except ValueError: + continue + return False + + def _workspace_files(self, *, limit: int | None = None) -> tuple[Path, ...]: + max_items = self.max_indexed_files if limit is None else min(limit, self.max_indexed_files) + found: list[Path] = [] + seen: set[Path] = set() + ignored = set(self.ignored_directories) + supported = set(self.supported_extensions) + for root in self._roots(): + for current_root, dir_names, file_names in os.walk(root): + dir_names[:] = [ + name + for name in dir_names + if name not in ignored and not name.startswith('.pytest_cache') + ] + current_path = Path(current_root) + for name in sorted(file_names): + candidate = (current_path / name).resolve() + if candidate in seen or candidate.suffix not in supported: + continue + try: + size = candidate.stat().st_size + except OSError: + continue + if size > self.max_file_bytes: + continue + seen.add(candidate) + found.append(candidate) + if len(found) >= max_items: + return tuple(found) + return tuple(found) + + def _workspace_indexes(self) -> tuple[IndexedFile, ...]: + return tuple(self._index_file(path) for path in self._workspace_files()) + + def _index_file(self, path: Path) -> IndexedFile: + stat = path.stat() + cache_key = str(path) + cached = self._cache.get(cache_key) + if cached is not None: + cached_mtime, cached_size, indexed = cached + if cached_mtime == stat.st_mtime_ns and cached_size == stat.st_size: + return indexed + text = path.read_text(encoding='utf-8', errors='replace') + suffix = path.suffix.lower() + if suffix in {'.py', '.pyi'}: + indexed = _index_python_file(path, text) + elif suffix == '.json': + indexed = _index_json_file(path, text) + else: + indexed = _index_generic_code_file(path, text) + self._cache[cache_key] = (stat.st_mtime_ns, stat.st_size, indexed) + return indexed + + def _symbol_name_at_position(self, indexed: IndexedFile, line: int, character: int) -> str | None: + word = _word_at(indexed.text, line, character) + if word: + return word + for symbol in indexed.symbols: + if symbol.contains(line, character): + return symbol.name + return None + + +def _discover_manifest_paths(cwd: Path, additional_dirs: tuple[Path, ...]) -> tuple[Path, ...]: + directories = [cwd, *additional_dirs] + seen: set[Path] = set() + found: list[Path] = [] + for directory in directories: + for filename in LSP_MANIFEST_FILES: + candidate = (directory / filename).resolve() + if candidate in seen or not candidate.exists(): + continue + seen.add(candidate) + found.append(candidate) + return tuple(found) + + +def _load_manifest_payload(path: Path) -> dict[str, Any]: + try: + payload = json.loads(path.read_text(encoding='utf-8')) + except (OSError, json.JSONDecodeError): + return {} + return payload if isinstance(payload, dict) else {} + + +def _index_python_file(path: Path, text: str) -> IndexedFile: + diagnostics: list[LSPDiagnostic] = [] + symbols: list[LSPSymbol] = [] + call_edges: list[LSPCallEdge] = [] + try: + tree = ast.parse(text, filename=str(path)) + except SyntaxError as exc: + diagnostics.append( + LSPDiagnostic( + path=str(path), + severity='error', + message=exc.msg, + line=exc.lineno or 1, + character=(exc.offset or 1), + code='syntax-error', + source='python', + ) + ) + fallback = _index_generic_code_file(path, text) + return IndexedFile( + path=path, + language='python', + text=text, + symbols=fallback.symbols, + diagnostics=tuple(diagnostics), + call_edges=(), + ) + + class Visitor(ast.NodeVisitor): + def __init__(self) -> None: + self.container_stack: list[LSPSymbol] = [] + self.callable_stack: list[LSPSymbol] = [] + + def visit_ClassDef(self, node: ast.ClassDef) -> Any: + symbol = LSPSymbol( + name=node.name, + kind='class', + path=str(path), + line=node.lineno, + character=node.col_offset + 1, + end_line=getattr(node, 'end_lineno', node.lineno), + end_character=getattr(node, 'end_col_offset', node.col_offset) + 1, + container_name=self.container_stack[-1].name if self.container_stack else None, + documentation=_first_doc_line(ast.get_docstring(node)), + ) + symbols.append(symbol) + self.container_stack.append(symbol) + self.generic_visit(node) + self.container_stack.pop() + + def visit_FunctionDef(self, node: ast.FunctionDef) -> Any: + kind = 'method' if self.container_stack and self.container_stack[-1].kind == 'class' else 'function' + symbol = LSPSymbol( + name=node.name, + kind=kind, + path=str(path), + line=node.lineno, + character=node.col_offset + 1, + end_line=getattr(node, 'end_lineno', node.lineno), + end_character=getattr(node, 'end_col_offset', node.col_offset) + 1, + container_name=self.container_stack[-1].name if self.container_stack else None, + signature=_python_signature(node), + documentation=_first_doc_line(ast.get_docstring(node)), + ) + symbols.append(symbol) + self.container_stack.append(symbol) + self.callable_stack.append(symbol) + self.generic_visit(node) + self.callable_stack.pop() + self.container_stack.pop() + + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any: + kind = 'method' if self.container_stack and self.container_stack[-1].kind == 'class' else 'async_function' + symbol = LSPSymbol( + name=node.name, + kind=kind, + path=str(path), + line=node.lineno, + character=node.col_offset + 1, + end_line=getattr(node, 'end_lineno', node.lineno), + end_character=getattr(node, 'end_col_offset', node.col_offset) + 1, + container_name=self.container_stack[-1].name if self.container_stack else None, + signature='async ' + _python_signature(node), + documentation=_first_doc_line(ast.get_docstring(node)), + ) + symbols.append(symbol) + self.container_stack.append(symbol) + self.callable_stack.append(symbol) + self.generic_visit(node) + self.callable_stack.pop() + self.container_stack.pop() + + def visit_Assign(self, node: ast.Assign) -> Any: + if self.callable_stack: + self.generic_visit(node) + return + for target in node.targets: + if isinstance(target, ast.Name): + symbols.append( + LSPSymbol( + name=target.id, + kind='variable', + path=str(path), + line=target.lineno, + character=target.col_offset + 1, + end_line=getattr(target, 'end_lineno', target.lineno), + end_character=getattr(target, 'end_col_offset', target.col_offset) + 1, + container_name=self.container_stack[-1].name if self.container_stack else None, + ) + ) + self.generic_visit(node) + + def visit_AnnAssign(self, node: ast.AnnAssign) -> Any: + if self.callable_stack: + self.generic_visit(node) + return + target = node.target + if isinstance(target, ast.Name): + symbols.append( + LSPSymbol( + name=target.id, + kind='variable', + path=str(path), + line=target.lineno, + character=target.col_offset + 1, + end_line=getattr(target, 'end_lineno', target.lineno), + end_character=getattr(target, 'end_col_offset', target.col_offset) + 1, + container_name=self.container_stack[-1].name if self.container_stack else None, + ) + ) + self.generic_visit(node) + + def visit_Call(self, node: ast.Call) -> Any: + if self.callable_stack: + callee_name = _python_called_name(node.func) + if callee_name: + caller = self.callable_stack[-1] + call_edges.append( + LSPCallEdge( + caller_symbol_id=caller.symbol_id, + caller_name=caller.name, + callee_name=callee_name, + path=str(path), + line=getattr(node, 'lineno', caller.line), + character=getattr(node, 'col_offset', 0) + 1, + ) + ) + self.generic_visit(node) + + Visitor().visit(tree) + return IndexedFile( + path=path, + language='python', + text=text, + symbols=tuple(symbols), + diagnostics=tuple(diagnostics), + call_edges=tuple(call_edges), + ) + + +def _index_json_file(path: Path, text: str) -> IndexedFile: + diagnostics: list[LSPDiagnostic] = [] + symbols: list[LSPSymbol] = [] + try: + payload = json.loads(text) + except json.JSONDecodeError as exc: + diagnostics.append( + LSPDiagnostic( + path=str(path), + severity='error', + message=exc.msg, + line=exc.lineno, + character=exc.colno, + code='json-error', + source='json', + ) + ) + return IndexedFile( + path=path, + language='json', + text=text, + symbols=(), + diagnostics=tuple(diagnostics), + call_edges=(), + ) + if isinstance(payload, dict): + for key in payload.keys(): + line_number, character = _find_string_literal(text, key) + if line_number is None or character is None: + line_number, character = 1, 1 + symbols.append( + LSPSymbol( + name=str(key), + kind='property', + path=str(path), + line=line_number, + character=character, + end_line=line_number, + end_character=character + len(str(key)), + ) + ) + return IndexedFile( + path=path, + language='json', + text=text, + symbols=tuple(symbols), + diagnostics=tuple(diagnostics), + call_edges=(), + ) + + +def _index_generic_code_file(path: Path, text: str) -> IndexedFile: + symbols: list[LSPSymbol] = [] + patterns = ( + ('class', re.compile(r'^\s*(?:export\s+)?class\s+([A-Za-z_][A-Za-z0-9_]*)')), + ('function', re.compile(r'^\s*(?:export\s+)?function\s+([A-Za-z_][A-Za-z0-9_]*)')), + ('function', re.compile(r'^\s*(?:export\s+)?(?:const|let|var)\s+([A-Za-z_][A-Za-z0-9_]*)\s*=\s*(?:async\s*)?(?:\([^)]*\)|[A-Za-z_][A-Za-z0-9_]*)\s*=>')), + ('variable', re.compile(r'^\s*(?:export\s+)?(?:const|let|var)\s+([A-Za-z_][A-Za-z0-9_]*)\b')), + ('function', re.compile(r'^\s*def\s+([A-Za-z_][A-Za-z0-9_]*)\b')), + ('class', re.compile(r'^\s*class\s+([A-Za-z_][A-Za-z0-9_]*)\b')), + ('interface', re.compile(r'^\s*(?:export\s+)?interface\s+([A-Za-z_][A-Za-z0-9_]*)')), + ('type', re.compile(r'^\s*(?:export\s+)?type\s+([A-Za-z_][A-Za-z0-9_]*)')), + ) + for line_number, line_text in enumerate(text.splitlines(), start=1): + for kind, pattern in patterns: + match = pattern.search(line_text) + if match is None: + continue + name = match.group(1) + symbols.append( + LSPSymbol( + name=name, + kind=kind, + path=str(path), + line=line_number, + character=match.start(1) + 1, + end_line=line_number, + end_character=match.end(1) + 1, + ) + ) + break + return IndexedFile( + path=path, + language=_language_for_extension(path.suffix), + text=text, + symbols=tuple(symbols), + diagnostics=(), + call_edges=(), + ) + + +def _python_signature(node: ast.AST) -> str: + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + return '' + argument_names = [arg.arg for arg in node.args.args] + if node.args.vararg is not None: + argument_names.append('*' + node.args.vararg.arg) + if node.args.kwarg is not None: + argument_names.append('**' + node.args.kwarg.arg) + return f'{node.name}({", ".join(argument_names)})' + + +def _python_called_name(node: ast.AST) -> str | None: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return None + + +def _word_at(text: str, line: int, character: int) -> str | None: + lines = text.splitlines() + if line < 1 or line > len(lines): + return None + line_text = lines[line - 1] + if not line_text: + return None + position = max(0, min(len(line_text) - 1, character - 1)) + for match in WORD_PATTERN.finditer(line_text): + if match.start() <= position < match.end(): + return match.group(0) + if position > 0: + for match in WORD_PATTERN.finditer(line_text): + if match.start() < position <= match.end(): + return match.group(0) + return None + + +def _find_string_literal(text: str, value: str) -> tuple[int | None, int | None]: + quoted = json.dumps(value) + needle = quoted[1:-1] + for line_number, line_text in enumerate(text.splitlines(), start=1): + column = line_text.find(f'"{needle}"') + if column >= 0: + return line_number, column + 2 + return None, None + + +def _first_doc_line(text: str | None) -> str | None: + if not text: + return None + first = text.strip().splitlines()[0].strip() + return first or None + + +def _language_for_extension(extension: str) -> str: + lowered = extension.lower() + if lowered in {'.py', '.pyi'}: + return 'python' + if lowered in {'.js', '.jsx'}: + return 'javascript' + if lowered in {'.ts', '.tsx'}: + return 'typescript' + if lowered == '.json': + return 'json' + return lowered.lstrip('.') or 'text' diff --git a/src/main.py b/src/main.py index da2e943..586c2e5 100644 --- a/src/main.py +++ b/src/main.py @@ -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 diff --git a/src/token_budget.py b/src/token_budget.py new file mode 100644 index 0000000..ea231a5 --- /dev/null +++ b/src/token_budget.py @@ -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 diff --git a/src/tokenizer_runtime.py b/src/tokenizer_runtime.py index b52f8e9..2c791eb 100644 --- a/src/tokenizer_runtime.py +++ b/src/tokenizer_runtime.py @@ -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 diff --git a/tests/test_agent_context.py b/tests/test_agent_context.py index 4b0c3bc..50e36f5 100644 --- a/tests/test_agent_context.py +++ b/tests/test_agent_context.py @@ -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' diff --git a/tests/test_agent_context_usage.py b/tests/test_agent_context_usage.py index 3636830..0e5be48 100644 --- a/tests/test_agent_context_usage.py +++ b/tests/test_agent_context_usage.py @@ -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) diff --git a/tests/test_agent_prompting.py b/tests/test_agent_prompting.py index 4ed2a71..2621763 100644 --- a/tests/test_agent_prompting.py +++ b/tests/test_agent_prompting.py @@ -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) diff --git a/tests/test_agent_runtime.py b/tests/test_agent_runtime.py index eb9cd12..e091ca3 100644 --- a/tests/test_agent_runtime.py +++ b/tests/test_agent_runtime.py @@ -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='Earlier conversation was compacted.', + 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 = [ { diff --git a/tests/test_agent_slash_commands.py b/tests/test_agent_slash_commands.py index 891f77f..3746045 100644 --- a/tests/test_agent_slash_commands.py +++ b/tests/test_agent_slash_commands.py @@ -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') diff --git a/tests/test_extended_tools.py b/tests/test_extended_tools.py index 11493f7..0a00ca0 100644 --- a/tests/test_extended_tools.py +++ b/tests/test_extended_tools.py @@ -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') diff --git a/tests/test_lsp_runtime.py b/tests/test_lsp_runtime.py new file mode 100644 index 0000000..ac11d11 --- /dev/null +++ b/tests/test_lsp_runtime.py @@ -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) diff --git a/tests/test_main.py b/tests/test_main.py index ca9a1a0..d39d8d2 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -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', '.']) diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py new file mode 100644 index 0000000..0060dd1 --- /dev/null +++ b/tests/test_token_budget.py @@ -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)