diff --git a/PARITY_CHECKLIST.md b/PARITY_CHECKLIST.md index fa7080a..30ace9a 100644 --- a/PARITY_CHECKLIST.md +++ b/PARITY_CHECKLIST.md @@ -2,7 +2,7 @@ This document tracks what is already implemented in Python and what is still missing compared with the upstream npm runtime. -This is a functionality-oriented checklist, not a line-by-line source equivalence claim. Large parts of the mirrored Python workspace still act as inventory or scaffolding, while the working Python runtime currently lives mainly in [`src/agent_runtime.py`](src/agent_runtime.py), [`src/agent_tools.py`](src/agent_tools.py), [`src/agent_prompting.py`](src/agent_prompting.py), [`src/agent_context.py`](src/agent_context.py), [`src/agent_slash_commands.py`](src/agent_slash_commands.py), and [`src/openai_compat.py`](src/openai_compat.py). +This is a functionality-oriented checklist, not a line-by-line source equivalence claim. Large parts of the mirrored Python workspace still act as inventory or scaffolding, while the working Python runtime currently lives mainly in [`src/agent_runtime.py`](src/agent_runtime.py), [`src/query_engine.py`](src/query_engine.py), [`src/agent_tools.py`](src/agent_tools.py), [`src/agent_prompting.py`](src/agent_prompting.py), [`src/agent_context.py`](src/agent_context.py), [`src/agent_manager.py`](src/agent_manager.py), [`src/plugin_runtime.py`](src/plugin_runtime.py), [`src/agent_slash_commands.py`](src/agent_slash_commands.py), and [`src/openai_compat.py`](src/openai_compat.py). ## 1. Core Agent Runtime @@ -10,6 +10,7 @@ Done: - [x] One-shot agent loop with iterative tool calling - [x] OpenAI-compatible `chat/completions` client +- [x] Streaming token-by-token assistant output - [x] Local-model execution against `vLLM` - [x] Local-model execution through `Ollama` - [x] Local-model execution through `LiteLLM Proxy` @@ -17,20 +18,47 @@ Done: - [x] Session save and resume support - [x] Configurable max-turn execution - [x] Permission-aware tool execution +- [x] Structured output / JSON schema request mode +- [x] Cost tracking and usage budget enforcement +- [x] Scratchpad directory integration +- [x] File history journaling for write/edit/shell tool actions +- [x] Incremental `bash` tool-result streaming events +- [x] Incremental tool-result streaming for read-only text tools +- [x] Mutable tool transcript updates during tool execution +- [x] Transcript mutation history for replaced/tombstoned messages +- [x] Structured transcript block export for messages, tool calls, and tool results +- [x] Resume-time file-history replay reminders +- [x] Resume-time file-history snapshot previews for file edits +- [x] Truncated-response continuation flow for `finish_reason=length` +- [x] Basic snipping of older tool/tool-call messages for context control +- [x] Basic automatic compact-boundary insertion with preserved recent tail +- [x] Reactive compaction retry after prompt-too-long backend failures +- [x] Reasoning-token budget enforcement +- [x] Tool-call and delegated-task budget enforcement +- [x] Basic nested-agent delegation tool +- [x] Sequential multi-subtask delegation with parent-context carryover +- [x] Basic agent-manager lineage tracking for nested agents +- [x] Managed agent-group membership tracking with child indices +- [x] Plugin-cache discovery and prompt-context injection +- [x] Manifest-based plugin runtime discovery +- [x] Manifest-defined plugin hooks for before-prompt and after-turn runtime injection +- [x] Manifest-defined plugin tool aliases over base runtime tools +- [x] Manifest-defined plugin tool blocking +- [x] Manifest-defined plugin tool-result guidance injected back into the transcript +- [x] Compaction metadata with compacted message ids +- [x] Resume-time compaction / snipping replay reminder +- [x] Query-engine facade that can drive the real Python runtime agent +- [x] Query-engine runtime event counters and transcript-kind summaries Missing: -- [ ] Streaming token-by-token assistant output -- [ ] Partial tool-result streaming -- [ ] Rich transcript mutation behavior like the npm runtime -- [ ] Structured output / JSON schema response modes -- [ ] Reasoning budgets and task budgets -- [ ] Cost accounting and usage budget enforcement -- [ ] Multi-agent orchestration parity -- [ ] File history snapshots and replay flows -- [ ] Scratchpad integration -- [ ] Plugin cache integration in the query engine -- [ ] Session compaction / snipping behavior +- [ ] Full partial tool-result streaming parity across the complete tool surface +- [ ] Full rich transcript mutation behavior like the npm runtime +- [ ] Full reasoning budgets and task budgets parity +- [ ] Full multi-agent orchestration parity +- [ ] Full file history snapshots and replay flows +- [ ] Full executable plugin lifecycle beyond runtime guidance, blocking, and aliases +- [ ] Full session compaction / snipping parity - [ ] Full `QueryEngine.ts` parity ## 2. CLI Entrypoints And Runtime Modes @@ -98,6 +126,8 @@ Done: - [x] Extra directory injection through `--add-dir` - [x] Session context usage report - [x] Raw context inspection command +- [x] Plugin cache snapshot injection +- [x] Manifest-based plugin runtime summary injection Missing: diff --git a/src/__init__.py b/src/__init__.py index 5898d42..464ba7a 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -8,12 +8,14 @@ from .agent_context import ( get_user_context, set_system_prompt_injection, ) +from .agent_manager import AgentManager from .agent_runtime import LocalCodingAgent from .agent_session import AgentMessage, AgentSessionState from .agent_tools import build_tool_context, default_tool_registry, execute_tool from .agent_types import AgentPermissions, AgentRunResult, AgentRuntimeConfig, ModelConfig from .commands import PORTED_COMMANDS, build_command_backlog from .parity_audit import ParityAuditResult, run_parity_audit +from .plugin_runtime import PluginRuntime from .port_manifest import PortManifest, build_port_manifest from .query_engine import QueryEnginePort, TurnResult from .runtime import PortRuntime, RuntimeSession @@ -23,6 +25,7 @@ from .tools import PORTED_TOOLS, build_tool_backlog __all__ = [ 'AgentContextSnapshot', + 'AgentManager', 'AgentPermissions', 'AgentRunResult', 'AgentRuntimeConfig', @@ -33,6 +36,7 @@ __all__ = [ 'ParityAuditResult', 'PortManifest', 'PortRuntime', + 'PluginRuntime', 'QueryEnginePort', 'RuntimeSession', 'StoredSession', diff --git a/src/agent_context.py b/src/agent_context.py index b9265a4..cdb15b8 100644 --- a/src/agent_context.py +++ b/src/agent_context.py @@ -8,6 +8,8 @@ from datetime import date from functools import lru_cache from pathlib import Path +from .agent_plugin_cache import load_plugin_cache_summary +from .plugin_runtime import PluginRuntime from .agent_types import AgentRuntimeConfig MAX_STATUS_CHARS = 2000 @@ -30,6 +32,7 @@ class AgentContextSnapshot: current_date: str is_git_repo: bool is_git_worktree: bool + scratchpad_directory: str | None additional_working_directories: tuple[str, ...] user_context: dict[str, str] system_context: dict[str, str] @@ -51,7 +54,11 @@ def set_system_prompt_injection(value: str | None) -> None: clear_context_caches() -def build_context_snapshot(runtime_config: AgentRuntimeConfig) -> AgentContextSnapshot: +def build_context_snapshot( + runtime_config: AgentRuntimeConfig, + *, + scratchpad_directory: Path | None = None, +) -> AgentContextSnapshot: cwd = runtime_config.cwd.resolve() additional_dirs = tuple( str(path.resolve()) for path in runtime_config.additional_working_directories @@ -64,13 +71,17 @@ def build_context_snapshot(runtime_config: AgentRuntimeConfig) -> AgentContextSn current_date=date.today().isoformat(), is_git_repo=_is_git_repo(cwd), is_git_worktree=_is_git_worktree(cwd), + scratchpad_directory=( + str(scratchpad_directory.resolve()) if scratchpad_directory is not None else None + ), additional_working_directories=additional_dirs, user_context=get_user_context( cwd, additional_dirs, runtime_config.disable_claude_md_discovery, + scratchpad_directory=scratchpad_directory, ), - system_context=get_system_context(cwd), + system_context=get_system_context(cwd, scratchpad_directory=scratchpad_directory), ) @@ -78,14 +89,20 @@ def get_git_status(cwd: Path) -> str | None: return _get_git_status_cached(str(cwd.resolve())) -def get_system_context(cwd: Path) -> dict[str, str]: - return dict(_get_system_context_cached(str(cwd.resolve()))) +def get_system_context( + cwd: Path, + *, + scratchpad_directory: Path | None = None, +) -> dict[str, str]: + scratchpad = str(scratchpad_directory.resolve()) if scratchpad_directory is not None else '' + return dict(_get_system_context_cached(str(cwd.resolve()), scratchpad)) def get_user_context( cwd: Path, additional_working_directories: tuple[str, ...] = (), disable_claude_md_discovery: bool = False, + scratchpad_directory: Path | None = None, ) -> dict[str, str]: normalized_dirs = tuple( str(Path(path).resolve()) for path in additional_working_directories @@ -95,6 +112,7 @@ def get_user_context( str(cwd.resolve()), normalized_dirs, disable_claude_md_discovery, + str(scratchpad_directory.resolve()) if scratchpad_directory is not None else '', ) ) @@ -113,6 +131,8 @@ def render_context_report(snapshot: AgentContextSnapshot, model: str) -> str: f'- Is a git worktree: {snapshot.is_git_worktree}', f'- Current date: {snapshot.current_date}', ] + if snapshot.scratchpad_directory: + lines.append(f'- Scratchpad directory: {snapshot.scratchpad_directory}') if snapshot.additional_working_directories: lines.extend( [ @@ -137,7 +157,7 @@ def render_context_report(snapshot: AgentContextSnapshot, model: str) -> str: @lru_cache(maxsize=32) -def _get_system_context_cached(cwd: str) -> dict[str, str]: +def _get_system_context_cached(cwd: str, scratchpad_directory: str) -> dict[str, str]: context: dict[str, str] = {} git_status = _get_git_status_cached(cwd) if git_status is not None: @@ -145,6 +165,8 @@ def _get_system_context_cached(cwd: str) -> dict[str, str]: injection = get_system_prompt_injection() if injection: context['cacheBreaker'] = f'[CACHE_BREAKER: {injection}]' + if scratchpad_directory: + context['scratchpadDirectory'] = scratchpad_directory return context @@ -153,16 +175,28 @@ def _get_user_context_cached( cwd: str, additional_working_directories: tuple[str, ...], disable_claude_md_discovery: bool, + scratchpad_directory: str, ) -> dict[str, str]: context: dict[str, str] = { 'currentDate': f"Today's date is {date.today().isoformat()}.", } + if scratchpad_directory: + context['scratchpad'] = ( + 'Use this session-specific scratchpad directory for temporary files instead ' + f'of /tmp when you need throwaway workspace: {scratchpad_directory}' + ) if disable_claude_md_discovery: return context memory_bundle = _load_memory_bundle(Path(cwd), additional_working_directories) if memory_bundle: context['claudeMd'] = memory_bundle + plugin_cache = load_plugin_cache_summary(Path(cwd), additional_working_directories) + if plugin_cache: + context['pluginCache'] = plugin_cache + plugin_runtime = PluginRuntime.from_workspace(Path(cwd), additional_working_directories) + if plugin_runtime.manifests: + context['pluginRuntime'] = plugin_runtime.render_summary() return context diff --git a/src/agent_manager.py b/src/agent_manager.py new file mode 100644 index 0000000..6643742 --- /dev/null +++ b/src/agent_manager.py @@ -0,0 +1,214 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class ManagedAgentRecord: + agent_id: str + prompt: str + parent_agent_id: str | None = None + group_id: str | None = None + child_index: int | None = None + label: str | None = None + session_id: str | None = None + session_path: str | None = None + status: str = 'running' + turns: int = 0 + tool_calls: int = 0 + stop_reason: str | None = None + + +@dataclass(frozen=True) +class ManagedAgentGroup: + group_id: str + label: str | None = None + parent_agent_id: str | None = None + child_agent_ids: tuple[str, ...] = () + status: str = 'running' + completed_children: int = 0 + failed_children: int = 0 + + +@dataclass +class AgentManager: + records: dict[str, ManagedAgentRecord] = field(default_factory=dict) + groups: dict[str, ManagedAgentGroup] = field(default_factory=dict) + _counter: int = 0 + _group_counter: int = 0 + + def start_agent( + self, + *, + prompt: str, + parent_agent_id: str | None = None, + group_id: str | None = None, + child_index: int | None = None, + label: str | None = None, + ) -> str: + self._counter += 1 + agent_id = f'agent_{self._counter}' + self.records[agent_id] = ManagedAgentRecord( + agent_id=agent_id, + prompt=prompt, + parent_agent_id=parent_agent_id, + group_id=group_id, + child_index=child_index, + label=label, + ) + if group_id is not None: + self.register_group_child(group_id, agent_id, child_index=child_index) + return agent_id + + def start_group( + self, + *, + label: str | None = None, + parent_agent_id: str | None = None, + ) -> str: + self._group_counter += 1 + group_id = f'group_{self._group_counter}' + self.groups[group_id] = ManagedAgentGroup( + group_id=group_id, + label=label, + parent_agent_id=parent_agent_id, + ) + return group_id + + def register_group_child( + self, + group_id: str, + agent_id: str, + *, + child_index: int | None = None, + ) -> None: + group = self.groups.get(group_id) + if group is None: + return + if agent_id in group.child_agent_ids: + updated_children = group.child_agent_ids + else: + updated_children = (*group.child_agent_ids, agent_id) + self.groups[group_id] = ManagedAgentGroup( + group_id=group.group_id, + label=group.label, + parent_agent_id=group.parent_agent_id, + child_agent_ids=updated_children, + status=group.status, + completed_children=group.completed_children, + failed_children=group.failed_children, + ) + record = self.records.get(agent_id) + if record is None: + return + if record.group_id == group_id and record.child_index == child_index: + return + self.records[agent_id] = ManagedAgentRecord( + agent_id=record.agent_id, + prompt=record.prompt, + parent_agent_id=record.parent_agent_id, + group_id=group_id, + child_index=child_index, + label=record.label, + session_id=record.session_id, + session_path=record.session_path, + status=record.status, + turns=record.turns, + tool_calls=record.tool_calls, + stop_reason=record.stop_reason, + ) + + def finish_group( + self, + group_id: str, + *, + status: str, + completed_children: int, + failed_children: int, + ) -> None: + group = self.groups.get(group_id) + if group is None: + return + self.groups[group_id] = ManagedAgentGroup( + group_id=group.group_id, + label=group.label, + parent_agent_id=group.parent_agent_id, + child_agent_ids=group.child_agent_ids, + status=status, + completed_children=completed_children, + failed_children=failed_children, + ) + + def finish_agent( + self, + agent_id: str, + *, + session_id: str | None, + session_path: str | None, + turns: int, + tool_calls: int, + stop_reason: str | None, + ) -> None: + record = self.records.get(agent_id) + if record is None: + return + self.records[agent_id] = ManagedAgentRecord( + agent_id=record.agent_id, + prompt=record.prompt, + parent_agent_id=record.parent_agent_id, + group_id=record.group_id, + child_index=record.child_index, + label=record.label, + session_id=session_id, + session_path=session_path, + status='completed', + turns=turns, + tool_calls=tool_calls, + stop_reason=stop_reason, + ) + + def children_of(self, agent_id: str) -> tuple[ManagedAgentRecord, ...]: + return tuple( + record + for record in self.records.values() + if record.parent_agent_id == agent_id + ) + + def completed_records(self) -> tuple[ManagedAgentRecord, ...]: + return tuple( + record for record in self.records.values() if record.status == 'completed' + ) + + def summary_lines(self) -> list[str]: + lines = [ + f'- Managed agents: {len(self.records)}', + f'- Completed agents: {len(self.completed_records())}', + ] + child_count = sum(1 for record in self.records.values() if record.parent_agent_id) + lines.append(f'- Child agents: {child_count}') + lines.append(f'- Agent groups: {len(self.groups)}') + completed_groups = sum(1 for group in self.groups.values() if group.status == 'completed') + lines.append(f'- Completed groups: {completed_groups}') + for record in sorted(self.records.values(), key=lambda item: item.agent_id)[:8]: + label = record.label or record.agent_id + group_bits: list[str] = [] + if record.group_id is not None: + group_bits.append(f'group={record.group_id}') + if record.child_index is not None: + group_bits.append(f'child_index={record.child_index}') + group_suffix = f" {' '.join(group_bits)}" if group_bits else '' + lines.append( + f'- {label}: status={record.status} turns={record.turns} ' + f'tool_calls={record.tool_calls} stop={record.stop_reason or "n/a"}{group_suffix}' + ) + if len(self.records) > 8: + lines.append(f'- ... plus {len(self.records) - 8} more managed agents') + for group in sorted(self.groups.values(), key=lambda item: item.group_id)[:6]: + label = group.label or group.group_id + lines.append( + f'- {label}: group_status={group.status} children={len(group.child_agent_ids)} ' + f'completed={group.completed_children} failed={group.failed_children}' + ) + if len(self.groups) > 6: + lines.append(f'- ... plus {len(self.groups) - 6} more agent groups') + return lines diff --git a/src/agent_plugin_cache.py b/src/agent_plugin_cache.py new file mode 100644 index 0000000..7eab470 --- /dev/null +++ b/src/agent_plugin_cache.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +MAX_PLUGIN_LINES = 12 +MAX_PLUGIN_PREVIEW_CHARS = 4000 + + +@dataclass(frozen=True) +class PluginCacheEntry: + name: str + enabled: bool = True + version: str | None = None + source: str | None = None + + +def load_plugin_cache_summary( + cwd: Path, + additional_working_directories: tuple[str, ...] = (), +) -> str | None: + snapshot = discover_plugin_cache(cwd, additional_working_directories) + if snapshot is None: + return None + return snapshot + + +def discover_plugin_cache( + cwd: Path, + additional_working_directories: tuple[str, ...] = (), +) -> str | None: + for path in _discover_candidate_paths(cwd, additional_working_directories): + try: + payload = json.loads(path.read_text(encoding='utf-8')) + except (OSError, json.JSONDecodeError): + continue + entries = _extract_entries(payload) + if not entries: + continue + lines = [ + f'Plugin cache loaded from: {path}', + f'Plugin entries discovered: {len(entries)}', + ] + enabled = [entry for entry in entries if entry.enabled] + disabled = [entry for entry in entries if not entry.enabled] + lines.append(f'Enabled plugins: {len(enabled)}') + if disabled: + lines.append(f'Disabled plugins: {len(disabled)}') + for entry in entries[:MAX_PLUGIN_LINES]: + details = [entry.name] + if entry.version: + details.append(f'version={entry.version}') + if entry.source: + details.append(f'source={entry.source}') + if not entry.enabled: + details.append('disabled') + lines.append(f"- {'; '.join(details)}") + if len(entries) > MAX_PLUGIN_LINES: + lines.append(f'- ... plus {len(entries) - MAX_PLUGIN_LINES} more plugin entries') + rendered = '\n'.join(lines) + if len(rendered) > MAX_PLUGIN_PREVIEW_CHARS: + rendered = rendered[: MAX_PLUGIN_PREVIEW_CHARS - 3] + '...' + return rendered + return None + + +def _discover_candidate_paths( + cwd: Path, + additional_working_directories: tuple[str, ...], +) -> list[Path]: + candidates: list[Path] = [] + seen: set[Path] = set() + relative_paths = ( + '.port_sessions/plugin_cache.json', + '.port_sessions/plugins.json', + '.claude/plugins/cache.json', + '.claw/plugins/cache.json', + 'plugins/cache.json', + '.plugins/cache.json', + ) + + def remember(path: Path) -> None: + resolved = path.resolve() + if resolved in seen or not resolved.exists() or not resolved.is_file(): + return + seen.add(resolved) + candidates.append(resolved) + + for root in _walk_upwards(cwd.resolve()): + for relative in relative_paths: + remember(root / relative) + + for raw_path in additional_working_directories: + directory = Path(raw_path).resolve() + for relative in relative_paths: + remember(directory / relative) + + return candidates + + +def _walk_upwards(path: Path) -> list[Path]: + current = path + walked: list[Path] = [] + while True: + walked.append(current) + if current.parent == current: + break + current = current.parent + return walked + + +def _extract_entries(payload: Any) -> list[PluginCacheEntry]: + entries: list[PluginCacheEntry] = [] + raw_entries: list[Any] = [] + if isinstance(payload, list): + raw_entries = payload + elif isinstance(payload, dict): + if isinstance(payload.get('plugins'), list): + raw_entries = payload['plugins'] + elif isinstance(payload.get('entries'), list): + raw_entries = payload['entries'] + else: + raw_entries = [ + {'name': key, **value} + for key, value in payload.items() + if isinstance(value, dict) + ] + + for item in raw_entries: + entry = _coerce_entry(item) + if entry is not None: + entries.append(entry) + return entries + + +def _coerce_entry(item: Any) -> PluginCacheEntry | None: + if isinstance(item, str) and item.strip(): + return PluginCacheEntry(name=item.strip()) + if not isinstance(item, dict): + return None + name = item.get('name') or item.get('plugin') or item.get('id') + if not isinstance(name, str) or not name.strip(): + return None + source = item.get('source') or item.get('path') or item.get('module') + version = item.get('version') + enabled = item.get('enabled') + return PluginCacheEntry( + name=name.strip(), + enabled=True if enabled is None else bool(enabled), + version=version if isinstance(version, str) and version else None, + source=source if isinstance(source, str) and source else None, + ) diff --git a/src/agent_prompting.py b/src/agent_prompting.py index 2f1f78e..ad6b6f7 100644 --- a/src/agent_prompting.py +++ b/src/agent_prompting.py @@ -20,6 +20,7 @@ class PromptContext: current_date: str is_git_repo: bool is_git_worktree: bool + scratchpad_directory: str | None = None additional_working_directories: tuple[str, ...] = () user_context: dict[str, str] = field(default_factory=dict) system_context: dict[str, str] = field(default_factory=dict) @@ -29,6 +30,7 @@ def build_prompt_context( runtime_config: AgentRuntimeConfig, model_config: ModelConfig, additional_working_directories: tuple[str, ...] = (), + scratchpad_directory: Path | None = None, ) -> PromptContext: merged_directories = tuple(runtime_config.additional_working_directories) for raw_path in additional_working_directories: @@ -39,7 +41,10 @@ def build_prompt_context( runtime_config, additional_working_directories=merged_directories, ) - snapshot = build_context_snapshot(context_runtime) + snapshot = build_context_snapshot( + context_runtime, + scratchpad_directory=scratchpad_directory, + ) return PromptContext( cwd=snapshot.cwd, model=model_config.model, @@ -49,6 +54,7 @@ def build_prompt_context( current_date=snapshot.current_date, is_git_repo=snapshot.is_git_repo, is_git_worktree=snapshot.is_git_worktree, + scratchpad_directory=snapshot.scratchpad_directory, additional_working_directories=snapshot.additional_working_directories, user_context=snapshot.user_context, system_context=snapshot.system_context, @@ -84,6 +90,7 @@ def build_system_prompt_parts( get_doing_tasks_section(), get_actions_section(), get_using_your_tools_section(enabled_tool_names), + get_plugin_guidance_section(prompt_context), get_tone_and_style_section(), get_output_efficiency_section(), SYSTEM_PROMPT_DYNAMIC_BOUNDARY, @@ -186,6 +193,20 @@ def get_tone_and_style_section() -> str: return '\n'.join(['# Tone and style', *prepend_bullets(items)]) +def get_plugin_guidance_section(prompt_context: PromptContext) -> str: + plugin_cache = prompt_context.user_context.get('pluginCache') + plugin_runtime = prompt_context.user_context.get('pluginRuntime') + if not plugin_cache and not plugin_runtime: + return '' + items = [ + 'Local plugin runtime data may be available in the injected user context.', + 'Use cached plugin information as advisory runtime context, not as proof that a plugin executed successfully.', + 'Manifest-based plugin runtime data can hint at plugin tools and hooks that may exist in the workspace.', + 'When a task depends on plugin behavior, prefer verifying against files or explicit tool results before making strong claims.', + ] + return '\n'.join(['# Plugins', *prepend_bullets(items)]) + + def get_output_efficiency_section() -> str: return """# Communicating with the user @@ -222,6 +243,8 @@ def compute_simple_env_info(prompt_context: PromptContext) -> str: if prompt_context.additional_working_directories: items.append('Additional working directories:') items.append(list(prompt_context.additional_working_directories)) + if prompt_context.scratchpad_directory: + items.append(f'Session scratchpad directory: {prompt_context.scratchpad_directory}') items.extend( [ f'Platform: {prompt_context.platform_name}', diff --git a/src/agent_runtime.py b/src/agent_runtime.py index ccb671a..0485f2e 100644 --- a/src/agent_runtime.py +++ b/src/agent_runtime.py @@ -1,11 +1,14 @@ from __future__ import annotations from dataclasses import dataclass, field, replace +from datetime import datetime, timezone import json +from pathlib import Path from uuid import uuid4 +from .agent_manager import AgentManager from .agent_context import render_context_report as render_agent_context_report -from .agent_context_usage import collect_context_usage, format_context_usage +from .agent_context_usage import collect_context_usage, estimate_tokens, format_context_usage from .agent_prompting import ( build_prompt_context, build_system_prompt_parts, @@ -17,11 +20,24 @@ from .agent_tools import ( AgentTool, build_tool_context, default_tool_registry, - execute_tool, + execute_tool_streaming, serialize_tool_result, ) -from .agent_types import AgentRunResult, AgentRuntimeConfig, ModelConfig -from .openai_compat import OpenAICompatClient +from .agent_types import ( + AgentRunResult, + AgentPermissions, + AgentRuntimeConfig, + AssistantTurn, + BudgetConfig, + ModelConfig, + OutputSchemaConfig, + StreamEvent, + ToolCall, + ToolExecutionResult, + UsageStats, +) +from .openai_compat import OpenAICompatClient, OpenAICompatError +from .plugin_runtime import PluginRuntime from .session_store import ( StoredAgentSession, save_agent_session, @@ -30,6 +46,12 @@ from .session_store import ( ) +@dataclass(frozen=True) +class BudgetDecision: + exceeded: bool + reason: str | None = None + + @dataclass class LocalCodingAgent: model_config: ModelConfig @@ -38,14 +60,30 @@ class LocalCodingAgent: append_system_prompt: str | None = None override_system_prompt: str | None = None tool_registry: dict[str, AgentTool] | None = None + agent_manager: AgentManager | None = None + parent_agent_id: str | None = None + managed_group_id: str | None = None + managed_child_index: int | None = None + plugin_runtime: PluginRuntime | None = None last_session: AgentSessionState | None = field(default=None, init=False, repr=False) last_run_result: AgentRunResult | None = field(default=None, init=False, repr=False) active_session_id: str | None = field(default=None, init=False, repr=False) last_session_path: str | None = field(default=None, init=False, repr=False) + managed_agent_id: str | None = field(default=None, init=False, repr=False) def __post_init__(self) -> None: if self.tool_registry is None: self.tool_registry = default_tool_registry() + if self.agent_manager is None: + self.agent_manager = AgentManager() + if self.plugin_runtime is None: + self.plugin_runtime = PluginRuntime.from_workspace( + self.runtime_config.cwd, + tuple(str(path) for path in self.runtime_config.additional_working_directories), + ) + plugin_tools = self.plugin_runtime.register_tool_aliases(self.tool_registry) + if plugin_tools: + self.tool_registry = {**self.tool_registry, **plugin_tools} self.client = OpenAICompatClient(self.model_config) self.tool_context = build_tool_context(self.runtime_config) @@ -59,8 +97,12 @@ class LocalCodingAgent: self.active_session_id = None self.last_session_path = None - def build_prompt_context(self): - return build_prompt_context(self.runtime_config, self.model_config) + def build_prompt_context(self, scratchpad_directory: Path | None = None): + return build_prompt_context( + self.runtime_config, + self.model_config, + scratchpad_directory=scratchpad_directory, + ) def build_system_prompt_parts(self, prompt_context=None) -> list[str]: if prompt_context is None: @@ -74,8 +116,13 @@ class LocalCodingAgent: override_system_prompt=self.override_system_prompt, ) - def build_session(self, user_prompt: str | None = None) -> AgentSessionState: - prompt_context = self.build_prompt_context() + def build_session( + self, + user_prompt: str | None = None, + *, + scratchpad_directory: Path | None = None, + ) -> AgentSessionState: + prompt_context = self.build_prompt_context(scratchpad_directory) system_prompt_parts = self.build_system_prompt_parts(prompt_context) return AgentSessionState.create( system_prompt_parts, @@ -85,30 +132,60 @@ class LocalCodingAgent: ) def run(self, prompt: str) -> AgentRunResult: - return self._run_prompt(prompt, base_session=None, session_id=None) + self.managed_agent_id = None + session_id = uuid4().hex + scratchpad_directory = self._ensure_scratchpad_directory(session_id) + result = self._run_prompt( + prompt, + base_session=None, + session_id=session_id, + scratchpad_directory=scratchpad_directory, + existing_file_history=(), + ) + self._finalize_managed_agent(result) + return result def resume(self, prompt: str, stored_session: StoredAgentSession) -> AgentRunResult: + self.managed_agent_id = None session = AgentSessionState.from_persisted( system_prompt_parts=stored_session.system_prompt_parts, user_context=stored_session.user_context, system_context=stored_session.system_context, messages=stored_session.messages, ) + self._append_file_history_replay_if_needed( + session, + stored_session.file_history, + ) + self._append_compaction_replay_if_needed(session) self.active_session_id = stored_session.session_id self.last_session = session - self.last_session_path = str(self.runtime_config.session_directory / f'{stored_session.session_id}.json') - return self._run_prompt( + self.last_session_path = str( + self.runtime_config.session_directory / f'{stored_session.session_id}.json' + ) + scratchpad_directory = ( + Path(stored_session.scratchpad_directory) + if stored_session.scratchpad_directory + else self._ensure_scratchpad_directory(stored_session.session_id) + ) + result = self._run_prompt( prompt, base_session=session, session_id=stored_session.session_id, + scratchpad_directory=scratchpad_directory, + existing_file_history=stored_session.file_history, ) + self._finalize_managed_agent(result) + return result def _run_prompt( self, prompt: str, *, base_session: AgentSessionState | None, - session_id: str | None, + session_id: str, + scratchpad_directory: Path | None, + existing_file_history: tuple[dict[str, object], ...], ) -> AgentRunResult: slash_result = preprocess_slash_command(self, prompt) if slash_result.handled and not slash_result.should_query: @@ -119,75 +196,1463 @@ class LocalCodingAgent: transcript=slash_result.transcript, session_id=self.active_session_id, session_path=self.last_session_path, + scratchpad_directory=( + str(scratchpad_directory) if scratchpad_directory is not None else None + ), ) - effective_prompt = slash_result.prompt or prompt - session = base_session if base_session is not None else self.build_session(None) + effective_prompt = self._apply_plugin_before_prompt_hooks(slash_result.prompt or prompt) + self.managed_agent_id = self.agent_manager.start_agent( + prompt=effective_prompt, + parent_agent_id=self.parent_agent_id, + group_id=self.managed_group_id, + child_index=self.managed_child_index, + label='root' if base_session is None else 'resume', + ) + session = ( + base_session + if base_session is not None + else self.build_session( + None, + scratchpad_directory=scratchpad_directory, + ) + ) session.append_user(effective_prompt) - if session_id is None: - session_id = uuid4().hex self.last_session = session self.active_session_id = session_id tool_specs = [tool.to_openai_tool() for tool in self.tool_registry.values()] tool_calls = 0 last_content = '' + total_usage = UsageStats() + total_cost_usd = 0.0 + file_history = list(existing_file_history) + stream_events: list[dict[str, object]] = [] + assistant_response_segments: list[str] = [] + delegated_tasks = sum( + 1 for entry in file_history if entry.get('action') == 'delegate_agent' + ) + + initial_budget = self._check_budget( + total_usage, + total_cost_usd, + tool_calls=tool_calls, + delegated_tasks=delegated_tasks, + ) + if initial_budget.exceeded: + result = AgentRunResult( + final_output=initial_budget.reason or 'Stopped before the first model call.', + turns=0, + tool_calls=0, + transcript=session.transcript(), + session_id=session_id, + usage=total_usage, + total_cost_usd=total_cost_usd, + stop_reason='budget_exceeded', + file_history=tuple(file_history), + 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 for turn_index in range(1, self.runtime_config.max_turns + 1): - turn = self.client.complete(session.to_openai_messages(), tool_specs) - assistant_tool_calls = () - if turn.tool_calls: - assistant_tool_calls = tuple( - { - 'id': tool_call.id, - 'type': 'function', - 'function': { - 'name': tool_call.name, - 'arguments': json.dumps(tool_call.arguments, ensure_ascii=True), - }, - } - for tool_call in turn.tool_calls + self._snip_session_if_needed( + session, + stream_events, + turn_index=turn_index, + ) + self._compact_session_if_needed( + session, + stream_events, + turn_index=turn_index, + ) + try: + turn, turn_events = self._query_model(session, tool_specs) + except OpenAICompatError as exc: + if self._is_prompt_too_long_error(exc) and self._reactive_compact_session( + session, + stream_events, + turn_index=turn_index, + ): + try: + turn, turn_events = self._query_model(session, tool_specs) + except OpenAICompatError as retry_exc: + exc = retry_exc + else: + stream_events.extend( + { + 'type': 'reactive_compact_retry', + 'turn_index': turn_index, + } + for _ in [0] + ) + stream_events.extend(event.to_dict() for event in turn_events) + total_usage = total_usage + turn.usage + total_cost_usd = self.model_config.pricing.estimate_cost_usd(total_usage) + last_content = turn.content + + budget_after_model = self._check_budget( + total_usage, + total_cost_usd, + tool_calls=tool_calls, + delegated_tasks=delegated_tasks, + ) + if budget_after_model.exceeded: + result = AgentRunResult( + final_output=( + budget_after_model.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 not turn.tool_calls: + assistant_response_segments.append(turn.content) + if self._should_continue_response(turn): + session.append_user( + self._build_continuation_prompt(), + metadata={ + 'kind': 'continuation_request', + 'continuation_index': len(assistant_response_segments), + }, + message_id=f'continuation_{turn_index}', + ) + stream_events.append( + { + 'type': 'continuation_request', + 'reason': turn.finish_reason, + 'continuation_index': len(assistant_response_segments), + } + ) + last_content = ''.join(assistant_response_segments) + continue + result = AgentRunResult( + final_output=''.join(assistant_response_segments), + 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=turn.finish_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._persist_session(session, result) + self.last_run_result = result + return result + # fall through to the normal tool-call branch below + # normal error path if not recovered + result = AgentRunResult( + final_output=str(exc), + 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='backend_error', + file_history=tuple(file_history), + session_id=session_id, + scratchpad_directory=( + str(scratchpad_directory) if scratchpad_directory is not None else None + ), ) - session.append_assistant(turn.content, assistant_tool_calls) + result = self._append_plugin_after_turn_events( + result, + prompt=effective_prompt, + turn_index=turn_index, + ) + result = self._persist_session(session, result) + self.last_run_result = result + return result + + stream_events.extend(event.to_dict() for event in turn_events) + total_usage = total_usage + turn.usage + total_cost_usd = self.model_config.pricing.estimate_cost_usd(total_usage) last_content = turn.content - if not turn.tool_calls: + budget_after_model = self._check_budget( + total_usage, + total_cost_usd, + tool_calls=tool_calls, + delegated_tasks=delegated_tasks, + ) + if budget_after_model.exceeded: result = AgentRunResult( - final_output=turn.content, + final_output=( + budget_after_model.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 not turn.tool_calls: + assistant_response_segments.append(turn.content) + if self._should_continue_response(turn): + session.append_user( + self._build_continuation_prompt(), + metadata={ + 'kind': 'continuation_request', + 'continuation_index': len(assistant_response_segments), + }, + message_id=f'continuation_{turn_index}', + ) + stream_events.append( + { + 'type': 'continuation_request', + 'reason': turn.finish_reason, + 'continuation_index': len(assistant_response_segments), + } + ) + last_content = ''.join(assistant_response_segments) + continue + result = AgentRunResult( + final_output=''.join(assistant_response_segments), + 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=turn.finish_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_plugin_after_turn_events( + result, + prompt=effective_prompt, + turn_index=turn_index, ) result = self._persist_session(session, result) self.last_run_result = result return result for tool_call in turn.tool_calls: + assistant_response_segments.clear() tool_calls += 1 - result = execute_tool( - self.tool_registry, - tool_call.name, - tool_call.arguments, - self.tool_context, + if tool_call.name == 'delegate_agent': + delegated_tasks += self._delegated_task_units(tool_call.arguments) + budget_after_tool_request = self._check_budget( + total_usage, + total_cost_usd, + tool_calls=tool_calls, + delegated_tasks=delegated_tasks, ) - session.append_tool( + if budget_after_tool_request.exceeded: + stream_events.append( + { + 'type': 'task_budget_exceeded', + 'turn_index': turn_index, + 'tool_name': tool_call.name, + 'tool_call_id': tool_call.id, + 'reason': budget_after_tool_request.reason, + } + ) + result = AgentRunResult( + final_output=( + budget_after_tool_request.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 + tool_result = None + tool_message_index = session.start_tool( name=tool_call.name, tool_call_id=tool_call.id, - content=serialize_tool_result(result), + message_id=f'tool_{len(session.messages)}', + metadata={'phase': 'starting'}, ) + stream_events.append( + { + 'type': 'tool_start', + 'tool_name': tool_call.name, + 'tool_call_id': tool_call.id, + 'message_id': session.messages[tool_message_index].message_id, + } + ) + plugin_block_message = self._plugin_block_message(tool_call.name) + if plugin_block_message is not None: + tool_result = ToolExecutionResult( + name=tool_call.name, + ok=False, + content=plugin_block_message, + metadata={ + 'action': 'plugin_block', + 'plugin_blocked': True, + 'plugin_block_message': plugin_block_message, + }, + ) + stream_events.append( + { + 'type': 'plugin_tool_block', + 'tool_name': tool_call.name, + 'tool_call_id': tool_call.id, + 'message_id': session.messages[tool_message_index].message_id, + 'message': plugin_block_message, + } + ) + if tool_call.name == 'delegate_agent': + if tool_result is None: + tool_result = self._execute_delegate_agent(tool_call.arguments) + elif tool_result is None: + for update in execute_tool_streaming( + self.tool_registry, + tool_call.name, + tool_call.arguments, + self.tool_context, + ): + if update.kind == 'delta': + session.append_tool_delta( + tool_message_index, + update.content, + metadata={'last_stream': update.stream or 'tool'}, + ) + stream_events.append( + { + 'type': 'tool_delta', + 'tool_name': tool_call.name, + 'tool_call_id': tool_call.id, + 'message_id': session.messages[tool_message_index].message_id, + 'stream': update.stream, + 'delta': update.content, + } + ) + continue + tool_result = update.result + if tool_result is None: + raise RuntimeError(f'Tool executor returned no final result for {tool_call.name}') + plugin_messages = self._plugin_tool_result_messages(tool_call.name) + if plugin_messages: + merged_metadata = dict(tool_result.metadata) + merged_metadata['plugin_messages'] = list(plugin_messages) + tool_result = ToolExecutionResult( + name=tool_result.name, + ok=tool_result.ok, + content=tool_result.content, + metadata=merged_metadata, + ) + for message in plugin_messages: + stream_events.append( + { + 'type': 'plugin_tool_hook', + 'tool_name': tool_call.name, + 'tool_call_id': tool_call.id, + 'message_id': session.messages[tool_message_index].message_id, + 'message': message, + } + ) + plugin_runtime_message = self._build_plugin_tool_runtime_message( + tool_name=tool_call.name, + block_message=plugin_block_message, + plugin_messages=plugin_messages, + ) + if plugin_runtime_message is not None: + session.append_user( + plugin_runtime_message, + metadata={ + 'kind': 'plugin_tool_runtime', + 'tool_name': tool_call.name, + 'tool_call_id': tool_call.id, + 'plugin_blocked': plugin_block_message is not None, + 'plugin_message_count': len(plugin_messages), + }, + message_id=f'plugin_tool_runtime_{tool_call.id}', + ) + stream_events.append( + { + 'type': 'plugin_tool_context', + 'tool_name': tool_call.name, + 'tool_call_id': tool_call.id, + 'message_id': f'plugin_tool_runtime_{tool_call.id}', + 'blocked': plugin_block_message is not None, + 'message_count': len(plugin_messages), + } + ) + session.finalize_tool( + tool_message_index, + content=serialize_tool_result(tool_result), + metadata={ + 'phase': 'completed', + **dict(tool_result.metadata), + }, + stop_reason='tool_completed', + ) + stream_events.append( + { + 'type': 'tool_result', + 'tool_name': tool_call.name, + 'tool_call_id': tool_call.id, + 'message_id': session.messages[tool_message_index].message_id, + 'ok': tool_result.ok, + 'metadata': dict(tool_result.metadata), + } + ) + history_entry = self._build_file_history_entry( + tool_call=tool_call, + tool_result=tool_result, + turn_index=turn_index, + ) + if history_entry is not None: + file_history.append(history_entry) result = AgentRunResult( - final_output=last_content or 'Stopped: max turns reached before the model produced a final answer.', + final_output=( + last_content + or 'Stopped: max turns reached before the model produced a final answer.' + ), turns=self.runtime_config.max_turns, tool_calls=tool_calls, transcript=session.transcript(), + events=tuple(stream_events), + usage=total_usage, + total_cost_usd=total_cost_usd, + stop_reason='max_turns', + 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_plugin_after_turn_events( + result, + prompt=effective_prompt, + turn_index=self.runtime_config.max_turns, ) result = self._persist_session(session, result) self.last_run_result = result return result - def _persist_session(self, session: AgentSessionState, result: AgentRunResult) -> AgentRunResult: + def _query_model( + self, + session: AgentSessionState, + tool_specs: list[dict[str, object]], + ) -> tuple[AssistantTurn, tuple[StreamEvent, ...]]: + if not self.runtime_config.stream_model_responses: + turn = self.client.complete( + session.to_openai_messages(), + tool_specs, + output_schema=self.runtime_config.output_schema, + ) + assistant_tool_calls = tuple( + { + 'id': tool_call.id, + 'type': 'function', + 'function': { + 'name': tool_call.name, + 'arguments': json.dumps( + tool_call.arguments, + ensure_ascii=True, + ), + }, + } + for tool_call in turn.tool_calls + ) + session.append_assistant( + turn.content, + assistant_tool_calls, + message_id=f'assistant_{len(session.messages)}', + stop_reason=turn.finish_reason, + usage=turn.usage, + ) + return turn, () + + assistant_index = session.start_assistant( + message_id=f'assistant_{len(session.messages)}' + ) + usage = UsageStats() + finish_reason: str | None = None + events: list[StreamEvent] = [] + for event in self.client.stream( + session.to_openai_messages(), + tool_specs, + output_schema=self.runtime_config.output_schema, + ): + events.append(event) + if event.type == 'content_delta': + session.append_assistant_delta(assistant_index, event.delta) + elif event.type == 'tool_call_delta': + session.merge_assistant_tool_call_delta( + assistant_index, + tool_call_index=event.tool_call_index or 0, + tool_call_id=event.tool_call_id, + tool_name=event.tool_name, + arguments_delta=event.arguments_delta, + ) + elif event.type == 'usage': + usage = usage + event.usage + elif event.type == 'message_stop': + finish_reason = event.finish_reason + + session.finalize_assistant( + assistant_index, + finish_reason=finish_reason, + usage=usage, + ) + assistant_message = session.messages[assistant_index] + turn = AssistantTurn( + content=assistant_message.content, + tool_calls=self._tool_calls_from_message(assistant_message.tool_calls), + finish_reason=finish_reason, + raw_message=assistant_message.to_openai_message(), + usage=usage, + ) + return turn, tuple(events) + + def _tool_calls_from_message( + self, + tool_calls: tuple[dict[str, object], ...], + ) -> tuple[ToolCall, ...]: + parsed: list[ToolCall] = [] + for index, raw_tool_call in enumerate(tool_calls): + function_block = raw_tool_call.get('function') + if not isinstance(function_block, dict): + continue + name = function_block.get('name') + if not isinstance(name, str) or not name: + continue + raw_arguments = function_block.get('arguments', '') + if isinstance(raw_arguments, str) and raw_arguments.strip(): + arguments = json.loads(raw_arguments) + if not isinstance(arguments, dict): + raise OpenAICompatError( + f'Tool arguments must decode to an object, got {type(arguments).__name__}' + ) + else: + arguments = {} + call_id = raw_tool_call.get('id') + if not isinstance(call_id, str) or not call_id: + call_id = f'call_{index}' + parsed.append( + ToolCall( + id=call_id, + name=name, + arguments=arguments, + ) + ) + return tuple(parsed) + + def _should_continue_response(self, turn: AssistantTurn) -> bool: + return turn.finish_reason in {'length', 'max_tokens'} + + def _build_continuation_prompt(self) -> str: + return ( + '\n' + 'Your previous answer was truncated because the model stopped early. ' + 'Continue exactly where you left off. Do not repeat completed text.\n' + '' + ) + + def _check_budget( + self, + usage: UsageStats, + total_cost_usd: float, + *, + tool_calls: int, + delegated_tasks: int, + ) -> BudgetDecision: + budget = self.runtime_config.budget_config + token_reason = self._check_token_budget(usage, budget) + if token_reason is not None: + return BudgetDecision(exceeded=True, reason=token_reason) + if ( + budget.max_total_cost_usd is not None + and total_cost_usd > budget.max_total_cost_usd + ): + return BudgetDecision( + exceeded=True, + reason=( + 'Stopped because the total estimated cost ' + f'(${total_cost_usd:.6f}) exceeded the configured budget ' + f'(${budget.max_total_cost_usd:.6f}).' + ), + ) + if ( + budget.max_tool_calls is not None + and tool_calls > budget.max_tool_calls + ): + return BudgetDecision( + exceeded=True, + reason=( + 'Stopped because the tool-call budget was exceeded ' + f'({tool_calls} > {budget.max_tool_calls}).' + ), + ) + if ( + budget.max_delegated_tasks is not None + and delegated_tasks > budget.max_delegated_tasks + ): + return BudgetDecision( + exceeded=True, + reason=( + 'Stopped because the delegated-task budget was exceeded ' + f'({delegated_tasks} > {budget.max_delegated_tasks}).' + ), + ) + return BudgetDecision(exceeded=False) + + def _snip_session_if_needed( + self, + session: AgentSessionState, + stream_events: list[dict[str, object]], + *, + turn_index: int, + ) -> None: + threshold = self.runtime_config.auto_snip_threshold_tokens + if threshold is None or threshold <= 0: + return + self._reduce_context_pressure( + session, + stream_events, + turn_index=turn_index, + target_tokens=threshold, + allow_compaction=False, + ) + + def _compact_session_if_needed( + self, + session: AgentSessionState, + stream_events: list[dict[str, object]], + *, + turn_index: int, + ) -> None: + threshold = self.runtime_config.auto_compact_threshold_tokens + if threshold is None or threshold <= 0: + return + self._reduce_context_pressure( + session, + stream_events, + turn_index=turn_index, + target_tokens=threshold, + allow_compaction=True, + ) + + def _reactive_compact_session( + self, + session: AgentSessionState, + stream_events: list[dict[str, object]], + *, + turn_index: int, + ) -> bool: + return self._reduce_context_pressure( + session, + stream_events, + turn_index=turn_index, + target_tokens=0, + allow_compaction=True, + reactive=True, + ) + + def _reduce_context_pressure( + self, + session: AgentSessionState, + stream_events: list[dict[str, object]], + *, + turn_index: int, + target_tokens: int, + allow_compaction: bool, + reactive: bool = False, + ) -> bool: + changed = False + for _ in range(6): + usage_report = collect_context_usage( + session=session, + model=self.model_config.model, + strategy='reactive_compact' if reactive else 'context_pressure', + ) + if usage_report.total_tokens <= target_tokens: + break + if self._snip_session_pass( + session, + stream_events, + turn_index=turn_index, + target_tokens=target_tokens, + current_total=usage_report.total_tokens, + reactive=reactive, + ): + changed = True + continue + if allow_compaction and self._compact_session_pass( + session, + stream_events, + turn_index=turn_index, + usage_total=usage_report.total_tokens, + reactive=reactive, + ): + changed = True + if reactive: + continue + break + break + return changed + + def _snip_session_pass( + self, + session: AgentSessionState, + stream_events: list[dict[str, object]], + *, + turn_index: int, + target_tokens: int, + current_total: int, + reactive: bool, + ) -> bool: + prefix_count = self._compact_prefix_count(session) + tail_count = min( + max(self.runtime_config.compact_preserve_messages, 0), + max(len(session.messages) - prefix_count, 0), + ) + candidate_indexes = [ + index + for index in range(prefix_count, max(len(session.messages) - tail_count, prefix_count)) + if self._message_can_be_snipped(session.messages[index]) + ] + if not candidate_indexes: + return False + snipped_count = 0 + tokens_removed = 0 + snipped_message_ids: list[str] = [] + for index in candidate_indexes: + if current_total <= target_tokens and not reactive: + break + message = session.messages[index] + original_tokens = estimate_tokens(message.content) + replacement = self._build_snipped_message_content(message) + replacement_tokens = estimate_tokens(replacement) + if replacement_tokens >= original_tokens: + continue + session.tombstone_message( + index, + summary=replacement, + stop_reason='snipped_for_context', + mutation_kind='snip_tombstone', + metadata={ + 'kind': 'snipped_message', + 'original_token_estimate': original_tokens, + 'replacement_token_estimate': replacement_tokens, + 'snipped_turn_index': turn_index, + }, + ) + delta = original_tokens - replacement_tokens + current_total -= delta + tokens_removed += delta + snipped_count += 1 + if session.messages[index].message_id: + snipped_message_ids.append(session.messages[index].message_id) + if reactive and snipped_count >= 3: + break + if not snipped_count: + return False + stream_events.append( + { + 'type': 'reactive_snip_boundary' if reactive else 'snip_boundary', + 'turn_index': turn_index, + 'snipped_message_count': snipped_count, + 'estimated_tokens_removed': tokens_removed, + 'snipped_message_ids': snipped_message_ids, + } + ) + return True + + def _compact_session_pass( + self, + session: AgentSessionState, + stream_events: list[dict[str, object]], + *, + turn_index: int, + usage_total: int, + reactive: bool, + ) -> bool: + prefix_count = self._compact_prefix_count(session) + preserve_messages = max(self.runtime_config.compact_preserve_messages, 0) + if reactive: + preserve_messages = max(preserve_messages // 2, 1) + tail_count = min( + preserve_messages, + max(len(session.messages) - prefix_count, 0), + ) + compact_end = len(session.messages) - tail_count + if compact_end <= prefix_count: + return False + candidates = session.messages[prefix_count:compact_end] + if not candidates: + return False + compacted_tokens = sum( + usage.tokens + for usage in ( + collect_context_usage( + session=AgentSessionState( + system_prompt_parts=session.system_prompt_parts, + user_context=session.user_context, + system_context=session.system_context, + messages=list(candidates), + ), + model=self.model_config.model, + strategy='compacted_segment', + ).categories + ) + if usage.name != 'Free space' + ) + compact_message = self._build_compact_boundary_message( + candidates, + turn_index=turn_index, + estimated_tokens_before=usage_total, + estimated_tokens_removed=compacted_tokens, + preserved_tail_count=tail_count, + ) + session.messages = ( + session.messages[:prefix_count] + + [compact_message] + + session.messages[compact_end:] + ) + stream_events.append( + { + 'type': 'reactive_compact_boundary' if reactive else 'compact_boundary', + 'turn_index': turn_index, + 'compacted_message_count': len(candidates), + 'estimated_tokens_before': usage_total, + 'estimated_tokens_removed': compacted_tokens, + 'preserved_tail_count': tail_count, + 'compacted_message_ids': [ + message.message_id for message in candidates if message.message_id + ], + } + ) + return True + + def _check_token_budget( + self, + usage: UsageStats, + budget: BudgetConfig, + ) -> str | None: + if budget.max_total_tokens is not None and usage.total_tokens > budget.max_total_tokens: + return ( + 'Stopped because the total token budget was exceeded ' + f'({usage.total_tokens} > {budget.max_total_tokens}).' + ) + if budget.max_input_tokens is not None and usage.input_tokens > budget.max_input_tokens: + return ( + 'Stopped because the input token budget was exceeded ' + f'({usage.input_tokens} > {budget.max_input_tokens}).' + ) + if budget.max_output_tokens is not None and usage.output_tokens > budget.max_output_tokens: + return ( + 'Stopped because the output token budget was exceeded ' + f'({usage.output_tokens} > {budget.max_output_tokens}).' + ) + if ( + budget.max_reasoning_tokens is not None + and usage.reasoning_tokens > budget.max_reasoning_tokens + ): + return ( + 'Stopped because the reasoning token budget was exceeded ' + f'({usage.reasoning_tokens} > {budget.max_reasoning_tokens}).' + ) + return None + + def _build_file_history_entry( + self, + *, + tool_call: ToolCall, + tool_result, + turn_index: int, + ) -> dict[str, object] | None: + if not tool_result.metadata: + return None + if ( + 'path' not in tool_result.metadata + and 'command' not in tool_result.metadata + and tool_result.metadata.get('action') != 'delegate_agent' + ): + return None + return { + 'timestamp': datetime.now(timezone.utc).isoformat(), + 'turn_index': turn_index, + 'tool_call_id': tool_call.id, + 'tool_name': tool_call.name, + 'ok': tool_result.ok, + **dict(tool_result.metadata), + } + + def _compact_prefix_count(self, session: AgentSessionState) -> int: + prefix_count = 0 + for message in session.messages: + if prefix_count == 0 and message.role == 'system': + prefix_count += 1 + continue + if ( + prefix_count == 1 + and message.role == 'user' + and message.content.startswith('') + ): + prefix_count += 1 + continue + break + return prefix_count + + def _message_can_be_snipped(self, message) -> bool: + if message.metadata.get('kind') in { + 'compact_boundary', + 'snipped_message', + 'file_history_replay', + }: + return False + if message.role == 'tool': + return True + if message.role == 'assistant' and (message.tool_calls or len(message.content) > 600): + return True + if ( + message.role == 'user' + and message.metadata.get('kind') in {'continuation_request', 'file_history_replay'} + ): + return True + return False + + def _build_snipped_message_content(self, message) -> str: + preview = ' '.join(message.content.split()) + if len(preview) > 120: + preview = preview[:117] + '...' + if message.role == 'tool': + label = f'tool result ({message.name or "tool"})' + elif message.role == 'assistant': + label = 'assistant message with tool calls' + else: + label = message.role + return ( + '\n' + f'Older {label} was snipped to save context.\n' + f'Message id: {message.message_id or "(none)"}\n' + f'Preview: {preview or "(empty)"}\n' + '' + ) + + def _build_compact_boundary_message( + self, + messages, + *, + turn_index: int, + estimated_tokens_before: int, + estimated_tokens_removed: int, + preserved_tail_count: int, + ): + summary_lines = [ + '', + 'Earlier conversation history was compacted to keep the session within the context budget.', + '', + 'Compacted summary:', + ] + remaining = 24 + for message in messages: + if remaining <= 0: + break + label = message.role + if message.role == 'tool' and message.name: + label = f'tool:{message.name}' + snippet = ' '.join(message.content.split()) + if len(snippet) > 160: + snippet = snippet[:157] + '...' + if not snippet: + snippet = '(empty)' + summary_lines.append(f'- {label}: {snippet}') + remaining -= 1 + if len(messages) > 24: + summary_lines.append(f'- ... plus {len(messages) - 24} older messages') + summary_lines.extend( + [ + '', + 'Keep using the preserved recent tail as the active working set.', + '', + ] + ) + from .agent_session import AgentMessage + + return AgentMessage( + role='system', + content='\n'.join(summary_lines), + message_id=f'compact_boundary_{turn_index}_{len(messages)}', + metadata={ + 'kind': 'compact_boundary', + 'turn_index': turn_index, + 'compacted_message_count': len(messages), + 'estimated_tokens_before': estimated_tokens_before, + 'estimated_tokens_removed': estimated_tokens_removed, + 'preserved_tail_count': preserved_tail_count, + 'compacted_message_ids': [ + message.message_id for message in messages if message.message_id + ], + }, + ) + + def _is_prompt_too_long_error(self, exc: Exception) -> bool: + text = str(exc).lower() + patterns = ( + 'prompt is too long', + 'maximum context length', + 'context length exceeded', + 'too many tokens', + 'input too long', + 'context window', + ) + return any(pattern in text for pattern in patterns) + + def _execute_delegate_agent( + self, + arguments: dict[str, object], + ) -> ToolExecutionResult: + max_turns = arguments.get('max_turns') + if max_turns is not None and (isinstance(max_turns, bool) or not isinstance(max_turns, int) or max_turns < 1): + return ToolExecutionResult( + name='delegate_agent', + ok=False, + content='max_turns must be an integer >= 1', + ) + subtasks = self._normalize_delegate_subtasks(arguments) + if not subtasks: + return ToolExecutionResult( + name='delegate_agent', + ok=False, + content='prompt must be a non-empty string or subtasks must contain at least one prompt', + ) + child_permissions = AgentPermissions( + allow_file_write=( + self.runtime_config.permissions.allow_file_write + and bool(arguments.get('allow_write', False)) + ), + allow_shell_commands=( + self.runtime_config.permissions.allow_shell_commands + and bool(arguments.get('allow_shell', False)) + ), + allow_destructive_shell_commands=False, + ) + child_runtime_config = replace( + self.runtime_config, + max_turns=max_turns or min(self.runtime_config.max_turns, 6), + permissions=child_permissions, + auto_compact_threshold_tokens=self.runtime_config.auto_compact_threshold_tokens, + ) + child_tools = { + name: tool + for name, tool in self.tool_registry.items() + if name != 'delegate_agent' + } + include_parent_context = bool(arguments.get('include_parent_context', True)) + continue_on_error = bool(arguments.get('continue_on_error', True)) + child_summaries: list[dict[str, object]] = [] + child_session_ids: list[str] = [] + prior_results: list[dict[str, str]] = [] + group_id: str | None = None + if self.agent_manager is not None and len(subtasks) > 1: + group_id = self.agent_manager.start_group( + label=str(arguments.get('label') or 'delegated_group'), + parent_agent_id=self.managed_agent_id, + ) + failed_children = 0 + child_result = None + for index, subtask in enumerate(subtasks, start=1): + child_agent = LocalCodingAgent( + model_config=self.model_config, + runtime_config=replace( + child_runtime_config, + max_turns=subtask.get('max_turns', child_runtime_config.max_turns), + ), + custom_system_prompt=self.custom_system_prompt, + append_system_prompt=self.append_system_prompt, + override_system_prompt=self.override_system_prompt, + tool_registry=child_tools, + agent_manager=self.agent_manager, + parent_agent_id=self.managed_agent_id, + managed_group_id=group_id, + managed_child_index=index, + ) + if group_id is not None and child_agent.managed_agent_id is not None: + self.agent_manager.register_group_child( + group_id, + child_agent.managed_agent_id, + child_index=index, + ) + child_prompt = str(subtask['prompt']) + if include_parent_context and prior_results: + child_prompt = self._prepend_delegate_context(child_prompt, prior_results) + child_result = child_agent.run(child_prompt) + if group_id is not None and child_agent.managed_agent_id is not None: + self.agent_manager.register_group_child( + group_id, + child_agent.managed_agent_id, + child_index=index, + ) + summary = { + 'index': index, + 'label': str(subtask.get('label') or f'subtask_{index}'), + 'session_id': child_result.session_id or '', + 'turns': child_result.turns, + 'tool_calls': child_result.tool_calls, + 'stop_reason': child_result.stop_reason or 'stop', + 'output_preview': self._preview_text(child_result.final_output, 220), + } + child_summaries.append(summary) + if child_result.session_id: + child_session_ids.append(child_result.session_id) + prior_results.append( + { + 'label': summary['label'], + 'output_preview': str(summary['output_preview']), + } + ) + if child_result.stop_reason in {'backend_error', 'budget_exceeded'}: + failed_children += 1 + if not continue_on_error: + break + assert child_result is not None + completed_children = len(child_summaries) - failed_children + group_status = 'completed' + if failed_children and completed_children: + group_status = 'partial' + elif failed_children: + group_status = 'failed' + if group_id is not None and self.agent_manager is not None: + self.agent_manager.finish_group( + group_id, + status=group_status, + completed_children=completed_children, + failed_children=failed_children, + ) + summary_lines = [ + ( + 'Delegated agent completed the subtask.' + if len(child_summaries) == 1 + else f'Delegated agent completed {len(child_summaries)} sequential subtasks.' + ), + ] + if group_id is not None: + summary_lines.append(f'group_id={group_id}') + summary_lines.append(f'group_status={group_status}') + summary_lines.append('') + for summary in child_summaries: + summary_lines.extend( + [ + f"[{summary['label']}]", + f"session_id={summary['session_id']}", + f"turns={summary['turns']}", + f"tool_calls={summary['tool_calls']}", + f"stop_reason={summary['stop_reason']}", + f"output_preview={summary['output_preview']}", + '', + ] + ) + summary_lines.append('Final delegated output:') + summary_lines.append(child_result.final_output) + return ToolExecutionResult( + name='delegate_agent', + ok=True, + content='\n'.join(summary_lines).strip(), + metadata={ + 'action': 'delegate_agent', + 'child_session_id': child_result.session_id, + 'child_session_ids': child_session_ids, + 'child_turns': child_result.turns, + 'child_tool_calls': child_result.tool_calls, + 'child_stop_reason': child_result.stop_reason, + 'child_results': child_summaries, + 'subtask_count': len(child_summaries), + 'group_id': group_id, + 'group_status': group_status, + 'failed_children': failed_children, + 'completed_children': completed_children, + }, + ) + + def _normalize_delegate_subtasks( + self, + arguments: dict[str, object], + ) -> list[dict[str, object]]: + subtasks: list[dict[str, object]] = [] + raw_subtasks = arguments.get('subtasks') + if isinstance(raw_subtasks, list): + for index, item in enumerate(raw_subtasks, start=1): + if isinstance(item, str) and item.strip(): + subtasks.append({'prompt': item.strip(), 'label': f'subtask_{index}'}) + continue + if isinstance(item, dict): + prompt = item.get('prompt') + if not isinstance(prompt, str) or not prompt.strip(): + continue + label = item.get('label') + max_turns = item.get('max_turns') + task: dict[str, object] = { + 'prompt': prompt.strip(), + 'label': label if isinstance(label, str) and label.strip() else f'subtask_{index}', + } + if isinstance(max_turns, int) and not isinstance(max_turns, bool) and max_turns > 0: + task['max_turns'] = max_turns + subtasks.append(task) + prompt = arguments.get('prompt') + if isinstance(prompt, str) and prompt.strip(): + if not subtasks: + subtasks.append({'prompt': prompt.strip(), 'label': 'subtask_1'}) + return subtasks[:8] + + def _delegated_task_units( + self, + arguments: dict[str, object], + ) -> int: + subtasks = arguments.get('subtasks') + if isinstance(subtasks, list): + count = sum( + 1 + for item in subtasks + if ( + isinstance(item, str) + and item.strip() + ) or ( + isinstance(item, dict) + and isinstance(item.get('prompt'), str) + and item.get('prompt', '').strip() + ) + ) + if count: + return count + return 1 + + def _prepend_delegate_context( + self, + prompt: str, + prior_results: list[dict[str, str]], + ) -> str: + lines = [ + '', + 'Prior delegated subtask summaries:', + ] + for result in prior_results[-4:]: + lines.append(f"- {result['label']}: {result['output_preview']}") + lines.extend(['', '', prompt]) + return '\n'.join(lines) + + def _preview_text(self, text: str, limit: int) -> str: + normalized = ' '.join(text.split()) + if len(normalized) <= limit: + return normalized + return normalized[: limit - 3] + '...' + + def _ensure_scratchpad_directory(self, session_id: str) -> Path: + scratchpad_directory = (self.runtime_config.scratchpad_root / session_id).resolve() + scratchpad_directory.mkdir(parents=True, exist_ok=True) + return scratchpad_directory + + def _append_file_history_replay_if_needed( + self, + session: AgentSessionState, + file_history: tuple[dict[str, object], ...], + ) -> None: + if not file_history: + return + replay_count = len(file_history) + for message in reversed(session.messages): + if message.metadata.get('kind') != 'file_history_replay': + continue + if message.metadata.get('file_history_count') == replay_count: + return + break + session.append_user( + self._render_file_history_replay(file_history), + metadata={ + 'kind': 'file_history_replay', + 'file_history_count': replay_count, + }, + message_id=f'file_history_replay_{replay_count}', + ) + + def _render_file_history_replay( + self, + file_history: tuple[dict[str, object], ...], + ) -> str: + lines = [ + '', + 'Recent file history from this saved session:', + ] + for entry in file_history[-10:]: + action = str(entry.get('action', entry.get('tool_name', 'tool'))) + turn = entry.get('turn_index') + path = entry.get('path') + command = entry.get('command') + details = [f'action={action}'] + if turn is not None: + details.append(f'turn={turn}') + if path: + details.append(f'path={path}') + if command: + details.append(f'command={command}') + child_session_ids = entry.get('child_session_ids') + if isinstance(child_session_ids, list) and child_session_ids: + details.append(f'child_sessions={len(child_session_ids)}') + lines.append(f"- {'; '.join(details)}") + before_preview = entry.get('before_preview') + if isinstance(before_preview, str) and before_preview: + lines.append(f' before: {before_preview}') + after_preview = entry.get('after_preview') + if isinstance(after_preview, str) and after_preview: + lines.append(f' after: {after_preview}') + if len(file_history) > 10: + lines.append(f'- ... plus {len(file_history) - 10} older file-history entries') + lines.extend( + [ + '', + 'Use this replayed history when continuing the task so you avoid repeating prior edits or commands.', + '', + ] + ) + return '\n'.join(lines) + + def _append_compaction_replay_if_needed( + self, + session: AgentSessionState, + ) -> None: + compact_messages = [ + message for message in session.messages + if message.metadata.get('kind') == 'compact_boundary' + ] + snipped_messages = [ + message for message in session.messages + if message.metadata.get('kind') == 'snipped_message' + ] + if not compact_messages and not snipped_messages: + return + for message in reversed(session.messages): + if message.metadata.get('kind') != 'compaction_replay': + continue + return + session.append_user( + self._render_compaction_replay(compact_messages, snipped_messages), + metadata={ + 'kind': 'compaction_replay', + 'compact_boundary_count': len(compact_messages), + 'snipped_message_count': len(snipped_messages), + }, + message_id=( + f'compaction_replay_{len(compact_messages)}_{len(snipped_messages)}' + ), + ) + + def _render_compaction_replay( + self, + compact_messages, + snipped_messages, + ) -> str: + lines = [ + '', + 'This resumed session already contains compacted or snipped history.', + f'- Compact boundaries: {len(compact_messages)}', + f'- Snipped/tombstoned messages: {len(snipped_messages)}', + ] + latest_boundary = compact_messages[-1] if compact_messages else None + if latest_boundary is not None: + lines.append( + f"- Latest compact boundary id: {latest_boundary.message_id or '(none)'}" + ) + if snipped_messages: + last_ids = [ + message.message_id or '(none)' + for message in snipped_messages[-3:] + ] + lines.append(f"- Recent snipped ids: {', '.join(last_ids)}") + lines.extend( + [ + '', + 'Use the surviving transcript plus the compacted summaries as the authoritative context when continuing.', + '', + ] + ) + return '\n'.join(lines) + + def _build_plugin_tool_runtime_message( + self, + *, + tool_name: str, + block_message: str | None, + plugin_messages: tuple[str, ...], + ) -> str | None: + if block_message is None and not plugin_messages: + return None + lines = [ + '', + f'Plugin tool runtime guidance for `{tool_name}`:', + ] + if block_message is not None: + lines.append(f'- Blocked: {block_message}') + for message in plugin_messages: + lines.append(f'- {message}') + lines.extend( + [ + '', + 'Use this plugin guidance when deciding the next tool call or assistant response.', + '', + ] + ) + return '\n'.join(lines) + + def _plugin_block_message(self, tool_name: str) -> str | None: + if self.plugin_runtime is None: + return None + return self.plugin_runtime.blocked_tool_message(tool_name) + + def _plugin_tool_result_messages(self, tool_name: str) -> tuple[str, ...]: + if self.plugin_runtime is None: + return () + return self.plugin_runtime.tool_result_injections(tool_name) + + def _persist_session( + self, + session: AgentSessionState, + result: AgentRunResult, + ) -> AgentRunResult: if result.session_id is None: return result stored = StoredAgentSession( @@ -200,6 +1665,10 @@ class LocalCodingAgent: messages=session.transcript(), turns=result.turns, tool_calls=result.tool_calls, + usage=result.usage.to_dict(), + total_cost_usd=result.total_cost_usd, + file_history=result.file_history, + scratchpad_directory=result.scratchpad_directory, ) path = save_agent_session( stored, @@ -267,6 +1736,7 @@ class LocalCodingAgent: '', f'- Model: {self.model_config.model}', f'- Registered tools: {len(self.tool_registry)}', + f'- Streaming model responses: {self.runtime_config.stream_model_responses}', f'- Session ID: {self.active_session_id or "none"}', f'- Last session loaded: {"yes" if self.last_session is not None else "no"}', ] @@ -277,8 +1747,64 @@ class LocalCodingAgent: [ f'- Last run turns: {self.last_run_result.turns}', f'- Last run tool calls: {self.last_run_result.tool_calls}', + f'- Last run total tokens: {self.last_run_result.usage.total_tokens}', + f'- Last run total cost: ${self.last_run_result.total_cost_usd:.6f}', ] ) + if self.last_run_result.scratchpad_directory is not None: + lines.append( + f'- Scratchpad directory: {self.last_run_result.scratchpad_directory}' + ) else: lines.append('- Last run: none') + if self.agent_manager is not None: + lines.extend(self.agent_manager.summary_lines()) return '\n'.join(lines) + + def _finalize_managed_agent(self, result: AgentRunResult) -> None: + if self.managed_agent_id is None or self.agent_manager is None: + return + self.agent_manager.finish_agent( + self.managed_agent_id, + session_id=result.session_id, + session_path=result.session_path, + turns=result.turns, + tool_calls=result.tool_calls, + stop_reason=result.stop_reason, + ) + + def _apply_plugin_before_prompt_hooks(self, prompt: str) -> str: + if self.plugin_runtime is None: + return prompt + injections = self.plugin_runtime.before_prompt_injections() + if not injections: + return prompt + lines = ['', 'Plugin before-prompt hooks:'] + lines.extend(f'- {entry}' for entry in injections) + lines.extend(['', '', prompt]) + return '\n'.join(lines) + + def _append_plugin_after_turn_events( + self, + result: AgentRunResult, + *, + prompt: str, + turn_index: int, + ) -> AgentRunResult: + if self.plugin_runtime is None: + return result + injections = self.plugin_runtime.after_turn_injections() + if not injections: + return result + appended = list(result.events) + for entry in injections: + appended.append( + { + 'type': 'plugin_after_turn', + 'turn_index': turn_index, + 'message': entry, + 'prompt_preview': self._preview_text(prompt, 120), + 'stop_reason': result.stop_reason, + } + ) + return replace(result, events=tuple(appended)) diff --git a/src/agent_session.py b/src/agent_session.py index 5063f06..02a204a 100644 --- a/src/agent_session.py +++ b/src/agent_session.py @@ -1,10 +1,12 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any +from .agent_types import UsageStats JSONDict = dict[str, Any] +MAX_MUTATION_HISTORY = 8 @dataclass(frozen=True) @@ -14,6 +16,12 @@ class AgentMessage: name: str | None = None tool_call_id: str | None = None tool_calls: tuple[JSONDict, ...] = () + blocks: tuple[JSONDict, ...] = () + message_id: str | None = None + state: str = 'final' + stop_reason: str | None = None + usage: UsageStats = field(default_factory=UsageStats) + metadata: JSONDict = field(default_factory=dict) def to_openai_message(self) -> JSONDict: payload: JSONDict = { @@ -28,6 +36,23 @@ class AgentMessage: payload['tool_calls'] = list(self.tool_calls) return payload + def to_transcript_entry(self) -> JSONDict: + payload = self.to_openai_message() + blocks = self.blocks or _derive_blocks(self) + if blocks: + payload['blocks'] = [dict(block) for block in blocks] + if self.message_id is not None: + payload['message_id'] = self.message_id + if self.state != 'final': + payload['state'] = self.state + if self.stop_reason is not None: + payload['stop_reason'] = self.stop_reason + if self.usage.total_tokens: + payload['usage'] = self.usage.to_dict() + if self.metadata: + payload['metadata'] = dict(self.metadata) + return payload + @classmethod def from_openai_message(cls, payload: JSONDict) -> 'AgentMessage': tool_calls = payload.get('tool_calls') @@ -36,12 +61,28 @@ class AgentMessage: normalized_tool_calls = tuple( item for item in tool_calls if isinstance(item, dict) ) + blocks = payload.get('blocks') + normalized_blocks: tuple[JSONDict, ...] = () + if isinstance(blocks, list): + normalized_blocks = tuple( + item for item in blocks if isinstance(item, dict) + ) return cls( role=str(payload.get('role', 'user')), content='' if payload.get('content') is None else str(payload.get('content', '')), name=str(payload['name']) if isinstance(payload.get('name'), str) else None, tool_call_id=str(payload['tool_call_id']) if isinstance(payload.get('tool_call_id'), str) else None, tool_calls=normalized_tool_calls, + blocks=normalized_blocks, + message_id=str(payload['message_id']) if isinstance(payload.get('message_id'), str) else None, + state=str(payload.get('state', 'final')), + stop_reason=str(payload['stop_reason']) if isinstance(payload.get('stop_reason'), str) else None, + usage=_usage_from_payload(payload.get('usage')), + metadata=( + dict(payload['metadata']) + if isinstance(payload.get('metadata'), dict) + else {} + ), ) @@ -72,6 +113,7 @@ class AgentSessionState: content='\n\n'.join( _append_system_context(system_prompt_parts, state.system_context) ), + blocks=_text_blocks('\n\n'.join(_append_system_context(system_prompt_parts, state.system_context))), ) ) if state.user_context: @@ -79,6 +121,7 @@ class AgentSessionState: AgentMessage( role='user', content=_render_user_context_reminder(state.user_context), + blocks=_text_blocks(_render_user_context_reminder(state.user_context)), ) ) if user_prompt is not None: @@ -86,6 +129,7 @@ class AgentSessionState: AgentMessage( role='user', content=user_prompt, + blocks=_text_blocks(user_prompt), ) ) return state @@ -94,20 +138,115 @@ class AgentSessionState: self, content: str, tool_calls: tuple[JSONDict, ...] = (), + *, + message_id: str | None = None, + stop_reason: str | None = None, + usage: UsageStats | None = None, ) -> None: self.messages.append( AgentMessage( role='assistant', content=content, tool_calls=tool_calls, + blocks=_assistant_blocks(content, tool_calls), + message_id=message_id, + stop_reason=stop_reason, + usage=usage or UsageStats(), ) ) - def append_user(self, content: str) -> None: + def start_assistant( + self, + *, + message_id: str | None = None, + ) -> int: + self.messages.append( + AgentMessage( + role='assistant', + content='', + tool_calls=(), + blocks=(), + message_id=message_id, + state='streaming', + ) + ) + return len(self.messages) - 1 + + def append_assistant_delta(self, index: int, delta: str) -> None: + message = self.messages[index] + self.messages[index] = replace( + message, + content=message.content + delta, + blocks=_assistant_blocks(message.content + delta, message.tool_calls), + ) + + def merge_assistant_tool_call_delta( + self, + index: int, + *, + tool_call_index: int, + tool_call_id: str | None = None, + tool_name: str | None = None, + arguments_delta: str = '', + ) -> None: + message = self.messages[index] + tool_calls = [dict(item) for item in message.tool_calls] + while len(tool_calls) <= tool_call_index: + tool_calls.append( + { + 'id': None, + 'type': 'function', + 'function': { + 'name': '', + 'arguments': '', + }, + } + ) + tool_call = tool_calls[tool_call_index] + function_block = tool_call.setdefault('function', {}) + if tool_call_id: + tool_call['id'] = tool_call_id + if tool_name: + function_block['name'] = tool_name + if arguments_delta: + current_arguments = function_block.get('arguments', '') + function_block['arguments'] = f'{current_arguments}{arguments_delta}' + self.messages[index] = replace( + message, + tool_calls=tuple(tool_calls), + blocks=_assistant_blocks(message.content, tuple(tool_calls)), + ) + + def finalize_assistant( + self, + index: int, + *, + finish_reason: str | None, + usage: UsageStats | None = None, + ) -> None: + message = self.messages[index] + self.messages[index] = replace( + message, + state='final', + stop_reason=finish_reason, + usage=usage or message.usage, + blocks=_assistant_blocks(message.content, message.tool_calls), + ) + + def append_user( + self, + content: str, + *, + metadata: dict[str, Any] | None = None, + message_id: str | None = None, + ) -> None: self.messages.append( AgentMessage( role='user', content=content, + blocks=_text_blocks(content), + metadata=dict(metadata or {}), + message_id=message_id, ) ) @@ -118,14 +257,148 @@ class AgentSessionState: content=content, name=name, tool_call_id=tool_call_id, + blocks=_tool_blocks(name, tool_call_id, content), ) ) + def start_tool( + self, + *, + name: str, + tool_call_id: str, + message_id: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> int: + self.messages.append( + AgentMessage( + role='tool', + content='', + name=name, + tool_call_id=tool_call_id, + blocks=(), + message_id=message_id, + state='streaming', + metadata=dict(metadata or {}), + ) + ) + return len(self.messages) - 1 + + def append_tool_delta( + self, + index: int, + delta: str, + *, + metadata: dict[str, Any] | None = None, + ) -> None: + message = self.messages[index] + merged_metadata = dict(message.metadata) + if metadata: + merged_metadata.update(metadata) + self.messages[index] = replace( + message, + content=message.content + delta, + blocks=_tool_blocks(message.name, message.tool_call_id, message.content + delta), + metadata=merged_metadata, + ) + + def finalize_tool( + self, + index: int, + *, + content: str, + metadata: dict[str, Any] | None = None, + stop_reason: str | None = None, + ) -> None: + message = self.messages[index] + merged_metadata = dict(message.metadata) + if message.content and message.content != content: + merged_metadata.setdefault('stream_preview', message.content) + merged_metadata = _record_mutation( + merged_metadata, + mutation_kind='tool_finalize_replace', + previous_content=message.content, + previous_state=message.state, + previous_stop_reason=message.stop_reason, + ) + if metadata: + merged_metadata.update(metadata) + self.messages[index] = replace( + message, + content=content, + blocks=_tool_blocks(message.name, message.tool_call_id, content), + state='final', + stop_reason=stop_reason, + metadata=merged_metadata, + ) + + def update_message( + self, + index: int, + *, + content: str | None = None, + state: str | None = None, + stop_reason: str | None = None, + metadata: dict[str, Any] | None = None, + mutation_kind: str | None = None, + ) -> None: + message = self.messages[index] + merged_metadata = dict(message.metadata) + new_content = message.content if content is None else content + new_state = message.state if state is None else state + new_stop_reason = message.stop_reason if stop_reason is None else stop_reason + if mutation_kind and ( + new_content != message.content + or new_state != message.state + or new_stop_reason != message.stop_reason + ): + merged_metadata = _record_mutation( + merged_metadata, + mutation_kind=mutation_kind, + previous_content=message.content, + previous_state=message.state, + previous_stop_reason=message.stop_reason, + ) + if metadata: + merged_metadata.update(metadata) + self.messages[index] = replace( + message, + content=new_content, + blocks=_derive_blocks( + replace( + message, + content=new_content, + state=new_state, + stop_reason=new_stop_reason, + ) + ), + state=new_state, + stop_reason=new_stop_reason, + metadata=merged_metadata, + ) + + def tombstone_message( + self, + index: int, + *, + summary: str, + metadata: dict[str, Any] | None = None, + mutation_kind: str = 'tombstone', + stop_reason: str | None = None, + ) -> None: + self.update_message( + index, + content=summary, + state='tombstoned', + stop_reason=stop_reason, + metadata=metadata, + mutation_kind=mutation_kind, + ) + def to_openai_messages(self) -> list[JSONDict]: return [message.to_openai_message() for message in self.messages] def transcript(self) -> tuple[JSONDict, ...]: - return tuple(message.to_openai_message() for message in self.messages) + return tuple(message.to_transcript_entry() for message in self.messages) @classmethod def from_persisted( @@ -144,6 +417,120 @@ class AgentSessionState: ) +def _usage_from_payload(payload: Any) -> UsageStats: + if not isinstance(payload, dict): + return UsageStats() + + def _as_int(name: str) -> int: + value = payload.get(name, 0) + if isinstance(value, bool): + return 0 + if isinstance(value, int): + return value + try: + return int(value) + except (TypeError, ValueError): + return 0 + + return UsageStats( + input_tokens=_as_int('input_tokens'), + output_tokens=_as_int('output_tokens'), + cache_creation_input_tokens=_as_int('cache_creation_input_tokens'), + cache_read_input_tokens=_as_int('cache_read_input_tokens'), + reasoning_tokens=_as_int('reasoning_tokens'), + ) + + +def _record_mutation( + metadata: JSONDict, + *, + mutation_kind: str, + previous_content: str, + previous_state: str, + previous_stop_reason: str | None, +) -> JSONDict: + mutations = metadata.get('mutations') + if not isinstance(mutations, list): + mutations = [] + else: + mutations = [entry for entry in mutations if isinstance(entry, dict)] + preview = ' '.join(previous_content.split()) + if len(preview) > 120: + preview = preview[:117] + '...' + mutations.append( + { + 'kind': mutation_kind, + 'previous_state': previous_state, + 'previous_stop_reason': previous_stop_reason, + 'previous_content_length': len(previous_content), + 'previous_content_preview': preview or '(empty)', + } + ) + if len(mutations) > MAX_MUTATION_HISTORY: + mutations = mutations[-MAX_MUTATION_HISTORY:] + metadata['mutations'] = mutations + metadata['mutation_count'] = len(mutations) + metadata['last_mutation_kind'] = mutation_kind + return metadata + + +def _text_blocks(text: str) -> tuple[JSONDict, ...]: + if not text: + return () + return ({'type': 'text', 'text': text},) + + +def _assistant_blocks( + content: str, + tool_calls: tuple[JSONDict, ...], +) -> tuple[JSONDict, ...]: + blocks: list[JSONDict] = [] + if content: + blocks.append({'type': 'text', 'text': content}) + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + continue + function_block = tool_call.get('function') + if not isinstance(function_block, dict): + continue + blocks.append( + { + 'type': 'tool_call', + 'id': tool_call.get('id'), + 'name': function_block.get('name'), + 'arguments': function_block.get('arguments', ''), + } + ) + return tuple(blocks) + + +def _tool_blocks( + name: str | None, + tool_call_id: str | None, + content: str, +) -> tuple[JSONDict, ...]: + if not content: + return () + return ( + { + 'type': 'tool_result', + 'name': name, + 'tool_call_id': tool_call_id, + 'text': content, + }, + ) + + +def _derive_blocks(message: AgentMessage) -> tuple[JSONDict, ...]: + if message.blocks: + return message.blocks + if message.role == 'assistant': + return _assistant_blocks(message.content, message.tool_calls) + if message.role == 'tool': + return _tool_blocks(message.name, message.tool_call_id, message.content) + return _text_blocks(message.content) + + def _append_system_context( system_prompt_parts: list[str], system_context: dict[str, str], diff --git a/src/agent_tools.py b/src/agent_tools.py index 8115b62..6e810d6 100644 --- a/src/agent_tools.py +++ b/src/agent_tools.py @@ -1,11 +1,14 @@ from __future__ import annotations +import hashlib import json import re +import selectors import subprocess -from dataclasses import dataclass +import time +from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Callable +from typing import Any, Callable, Iterator, Union from .agent_types import AgentPermissions, AgentRuntimeConfig, ToolExecutionResult @@ -26,7 +29,10 @@ class ToolExecutionContext: permissions: AgentPermissions -ToolHandler = Callable[[dict[str, Any], ToolExecutionContext], str] +ToolHandler = Callable[ + [dict[str, Any], ToolExecutionContext], + Union[str, tuple[str, dict[str, Any]]], +] @dataclass(frozen=True) @@ -48,12 +54,25 @@ class AgentTool: def execute(self, arguments: dict[str, Any], context: ToolExecutionContext) -> ToolExecutionResult: try: - content = self.handler(arguments, context) - return ToolExecutionResult(name=self.name, ok=True, content=content) + result = self.handler(arguments, context) + if isinstance(result, tuple): + content, metadata = result + else: + content, metadata = result, {} + return ToolExecutionResult(name=self.name, ok=True, content=content, metadata=metadata) except (ToolPermissionError, ToolExecutionError, OSError, subprocess.SubprocessError) as exc: return ToolExecutionResult(name=self.name, ok=False, content=str(exc)) +@dataclass(frozen=True) +class ToolStreamUpdate: + kind: str + content: str = '' + stream: str | None = None + result: ToolExecutionResult | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + def build_tool_context(config: AgentRuntimeConfig) -> ToolExecutionContext: return ToolExecutionContext( root=config.cwd.resolve(), @@ -79,6 +98,35 @@ def execute_tool( return tool.execute(arguments, context) +def execute_tool_streaming( + tool_registry: dict[str, AgentTool], + name: str, + arguments: dict[str, Any], + context: ToolExecutionContext, +) -> Iterator[ToolStreamUpdate]: + tool = tool_registry.get(name) + if tool is None: + yield ToolStreamUpdate( + kind='result', + result=ToolExecutionResult( + name=name, + ok=False, + content=f'Unknown tool: {name}', + ), + ) + return + + if name == 'bash': + yield from _stream_bash(arguments, context) + return + + result = tool.execute(arguments, context) + if name in {'list_dir', 'read_file', 'glob_search', 'grep_search'} and result.ok: + yield from _stream_static_text_result(result) + return + yield ToolStreamUpdate(kind='result', result=result) + + def default_tool_registry() -> dict[str, AgentTool]: tools = [ AgentTool( @@ -174,6 +222,38 @@ def default_tool_registry() -> dict[str, AgentTool]: }, handler=_run_bash, ), + AgentTool( + name='delegate_agent', + description='Delegate a subtask to a nested Python coding agent and return its summary.', + parameters={ + 'type': 'object', + 'properties': { + 'prompt': {'type': 'string'}, + 'subtasks': { + 'type': 'array', + 'items': { + 'oneOf': [ + {'type': 'string'}, + { + 'type': 'object', + 'properties': { + 'prompt': {'type': 'string'}, + 'label': {'type': 'string'}, + 'max_turns': {'type': 'integer', 'minimum': 1, 'maximum': 20}, + }, + 'required': ['prompt'], + }, + ] + }, + }, + 'max_turns': {'type': 'integer', 'minimum': 1, 'maximum': 20}, + 'allow_write': {'type': 'boolean'}, + 'allow_shell': {'type': 'boolean'}, + 'include_parent_context': {'type': 'boolean'}, + }, + }, + handler=_delegate_agent_placeholder, + ), ] return {tool.name: tool for tool in tools} @@ -184,6 +264,8 @@ def serialize_tool_result(result: ToolExecutionResult) -> str: 'ok': result.ok, 'content': result.content, } + if result.metadata: + payload['metadata'] = result.metadata return json.dumps(payload, ensure_ascii=True, indent=2) @@ -195,6 +277,13 @@ def _truncate_output(text: str, limit: int) -> str: return f'{head}\n...[truncated]...\n{tail}' +def _snapshot_text(text: str, limit: int = 240) -> str: + normalized = ' '.join(text.split()) + if len(normalized) <= limit: + return normalized + return normalized[: limit - 3] + '...' + + def _require_string(arguments: dict[str, Any], key: str) -> str: value = arguments.get(key) if not isinstance(value, str) or not value: @@ -304,10 +393,34 @@ def _write_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str content = arguments.get('content') if not isinstance(content, str): raise ToolExecutionError('content must be a string') + previous_text: str | None = None + previous_sha256: str | None = None + if target.exists() and target.is_file(): + previous_text = target.read_text(encoding='utf-8', errors='replace') + previous_sha256 = hashlib.sha256(previous_text.encode('utf-8')).hexdigest() target.parent.mkdir(parents=True, exist_ok=True) target.write_text(content, encoding='utf-8') rel = target.relative_to(context.root) - return f'wrote {rel} ({len(content)} chars)' + new_sha256 = hashlib.sha256(content.encode('utf-8')).hexdigest() + return ( + f'wrote {rel} ({len(content)} chars)', + { + 'action': 'write_file', + 'path': str(rel), + 'before_exists': previous_text is not None, + 'before_sha256': previous_sha256, + 'before_size': len(previous_text) if previous_text is not None else 0, + 'before_preview': ( + _snapshot_text(previous_text) + if previous_text is not None + else None + ), + 'after_sha256': new_sha256, + 'after_size': len(content), + 'after_preview': _snapshot_text(content), + 'content_length': len(content), + }, + ) def _edit_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str: @@ -332,11 +445,28 @@ def _edit_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str: raise ToolExecutionError( f'old_text matched {occurrences} times; pass replace_all=true to replace every match' ) + before_sha256 = hashlib.sha256(current.encode('utf-8')).hexdigest() updated = current.replace(old_text, new_text) if replace_all else current.replace(old_text, new_text, 1) target.write_text(updated, encoding='utf-8') rel = target.relative_to(context.root) replaced = occurrences if replace_all else 1 - return f'edited {rel}; replaced {replaced} occurrence(s)' + after_sha256 = hashlib.sha256(updated.encode('utf-8')).hexdigest() + return ( + f'edited {rel}; replaced {replaced} occurrence(s)', + { + 'action': 'edit_file', + 'path': str(rel), + 'before_sha256': before_sha256, + 'after_sha256': after_sha256, + 'before_size': len(current), + 'after_size': len(updated), + 'before_preview': _snapshot_text(current), + 'after_preview': _snapshot_text(updated), + 'old_text_preview': _snapshot_text(old_text), + 'new_text_preview': _snapshot_text(new_text), + 'replaced_occurrences': replaced, + }, + ) def _glob_search(arguments: dict[str, Any], context: ToolExecutionContext) -> str: @@ -400,4 +530,199 @@ def _run_bash(arguments: dict[str, Any], context: ToolExecutionContext) -> str: '[stderr]', stderr.rstrip(), ] - return _truncate_output('\n'.join(payload).strip(), context.max_output_chars) + return ( + _truncate_output('\n'.join(payload).strip(), context.max_output_chars), + { + 'action': 'bash', + 'command': command, + 'exit_code': completed.returncode, + }, + ) + + +def _stream_bash( + arguments: dict[str, Any], + context: ToolExecutionContext, +) -> Iterator[ToolStreamUpdate]: + try: + command = _require_string(arguments, 'command') + _ensure_shell_allowed(command, context) + process = subprocess.Popen( + command, + shell=True, + executable='/bin/bash', + cwd=context.root, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + ) + except (ToolPermissionError, ToolExecutionError, OSError, subprocess.SubprocessError) as exc: + yield ToolStreamUpdate( + kind='result', + result=ToolExecutionResult(name='bash', ok=False, content=str(exc)), + ) + return + + selector = selectors.DefaultSelector() + stdout_chunks: list[str] = [] + stderr_chunks: list[str] = [] + if process.stdout is not None: + selector.register(process.stdout, selectors.EVENT_READ, data='stdout') + if process.stderr is not None: + selector.register(process.stderr, selectors.EVENT_READ, data='stderr') + + deadline = time.monotonic() + context.command_timeout_seconds + timeout_error: str | None = None + + try: + while selector.get_map(): + remaining = deadline - time.monotonic() + if remaining <= 0: + timeout_error = ( + f'Command timed out after {context.command_timeout_seconds:.1f}s: {command}' + ) + process.kill() + break + events = selector.select(timeout=min(remaining, 0.1)) + if not events and process.poll() is not None: + _drain_registered_streams(selector, stdout_chunks, stderr_chunks) + break + for key, _ in events: + stream_name = str(key.data) + line = key.fileobj.readline() + if line == '': + try: + selector.unregister(key.fileobj) + except Exception: + pass + try: + key.fileobj.close() + except Exception: + pass + continue + if stream_name == 'stdout': + stdout_chunks.append(line) + else: + stderr_chunks.append(line) + yield ToolStreamUpdate( + kind='delta', + content=line, + stream=stream_name, + ) + finally: + try: + selector.close() + except Exception: + pass + + exit_code = process.wait() + if timeout_error is not None: + yield ToolStreamUpdate( + kind='result', + result=ToolExecutionResult( + name='bash', + ok=False, + content=timeout_error, + metadata={ + 'action': 'bash', + 'command': command, + 'exit_code': exit_code, + 'timed_out': True, + }, + ), + ) + return + + stdout = ''.join(stdout_chunks) + stderr = ''.join(stderr_chunks) + payload = [ + f'exit_code={exit_code}', + '[stdout]', + stdout.rstrip(), + '[stderr]', + stderr.rstrip(), + ] + yield ToolStreamUpdate( + kind='result', + result=ToolExecutionResult( + name='bash', + ok=True, + content=_truncate_output('\n'.join(payload).strip(), context.max_output_chars), + metadata={ + 'action': 'bash', + 'command': command, + 'exit_code': exit_code, + 'streamed': True, + }, + ), + ) + + +def _delegate_agent_placeholder( + arguments: dict[str, Any], + context: ToolExecutionContext, +) -> str: + raise ToolExecutionError( + 'delegate_agent must be handled by the runtime and is not available as a standalone tool handler' + ) + + +def _drain_registered_streams( + selector: selectors.BaseSelector, + stdout_chunks: list[str], + stderr_chunks: list[str], +) -> None: + for key in list(selector.get_map().values()): + try: + remainder = key.fileobj.read() + except Exception: + remainder = '' + if not remainder: + try: + selector.unregister(key.fileobj) + except Exception: + pass + try: + key.fileobj.close() + except Exception: + pass + continue + if key.data == 'stdout': + stdout_chunks.append(remainder) + else: + stderr_chunks.append(remainder) + try: + selector.unregister(key.fileobj) + except Exception: + pass + try: + key.fileobj.close() + except Exception: + pass + + +def _stream_static_text_result( + result: ToolExecutionResult, + *, + chunk_size: int = 400, +) -> Iterator[ToolStreamUpdate]: + content = result.content + if content: + for start in range(0, len(content), chunk_size): + yield ToolStreamUpdate( + kind='delta', + content=content[start:start + chunk_size], + stream='tool', + ) + metadata = dict(result.metadata) + metadata.setdefault('streamed', True) + yield ToolStreamUpdate( + kind='result', + result=ToolExecutionResult( + name=result.name, + ok=result.ok, + content=result.content, + metadata=metadata, + ), + ) diff --git a/src/agent_types.py b/src/agent_types.py index 24f4682..a496e70 100644 --- a/src/agent_types.py +++ b/src/agent_types.py @@ -8,6 +8,87 @@ from typing import Any JSONDict = dict[str, Any] +@dataclass(frozen=True) +class UsageStats: + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + reasoning_tokens: int = 0 + + @property + def total_tokens(self) -> int: + return ( + self.input_tokens + + self.output_tokens + + self.cache_creation_input_tokens + + self.cache_read_input_tokens + ) + + def __add__(self, other: 'UsageStats') -> 'UsageStats': + return UsageStats( + input_tokens=self.input_tokens + other.input_tokens, + output_tokens=self.output_tokens + other.output_tokens, + cache_creation_input_tokens=( + self.cache_creation_input_tokens + other.cache_creation_input_tokens + ), + cache_read_input_tokens=( + self.cache_read_input_tokens + other.cache_read_input_tokens + ), + reasoning_tokens=self.reasoning_tokens + other.reasoning_tokens, + ) + + def to_dict(self) -> JSONDict: + return { + 'input_tokens': self.input_tokens, + 'output_tokens': self.output_tokens, + 'cache_creation_input_tokens': self.cache_creation_input_tokens, + 'cache_read_input_tokens': self.cache_read_input_tokens, + 'reasoning_tokens': self.reasoning_tokens, + 'total_tokens': self.total_tokens, + } + + +@dataclass(frozen=True) +class ModelPricing: + input_cost_per_million_tokens_usd: float = 0.0 + output_cost_per_million_tokens_usd: float = 0.0 + cache_creation_input_cost_per_million_tokens_usd: float = 0.0 + cache_read_input_cost_per_million_tokens_usd: float = 0.0 + + def estimate_cost_usd(self, usage: UsageStats) -> float: + return ( + (usage.input_tokens / 1_000_000.0) * self.input_cost_per_million_tokens_usd + + (usage.output_tokens / 1_000_000.0) * self.output_cost_per_million_tokens_usd + + ( + usage.cache_creation_input_tokens / 1_000_000.0 + ) + * self.cache_creation_input_cost_per_million_tokens_usd + + ( + usage.cache_read_input_tokens / 1_000_000.0 + ) + * self.cache_read_input_cost_per_million_tokens_usd + ) + + +@dataclass(frozen=True) +class BudgetConfig: + max_total_tokens: int | None = None + max_input_tokens: int | None = None + max_output_tokens: int | None = None + max_reasoning_tokens: int | None = None + max_total_cost_usd: float | None = None + max_tool_calls: int | None = None + max_delegated_tasks: int | None = None + + +@dataclass(frozen=True) +class OutputSchemaConfig: + name: str + schema: JSONDict + strict: bool = False + + @dataclass(frozen=True) class ModelConfig: model: str @@ -15,6 +96,7 @@ class ModelConfig: api_key: str = 'local-token' temperature: float = 0.0 timeout_seconds: float = 120.0 + pricing: ModelPricing = field(default_factory=ModelPricing) @dataclass(frozen=True) @@ -30,6 +112,33 @@ class AssistantTurn: tool_calls: tuple[ToolCall, ...] = () finish_reason: str | None = None raw_message: JSONDict = field(default_factory=dict) + usage: UsageStats = field(default_factory=UsageStats) + + +@dataclass(frozen=True) +class StreamEvent: + type: str + delta: str = '' + tool_call_index: int | None = None + tool_call_id: str | None = None + tool_name: str | None = None + arguments_delta: str = '' + finish_reason: str | None = None + usage: UsageStats = field(default_factory=UsageStats) + raw_event: JSONDict = field(default_factory=dict) + + def to_dict(self) -> JSONDict: + return { + 'type': self.type, + 'delta': self.delta, + 'tool_call_index': self.tool_call_index, + 'tool_call_id': self.tool_call_id, + 'tool_name': self.tool_name, + 'arguments_delta': self.arguments_delta, + 'finish_reason': self.finish_reason, + 'usage': self.usage.to_dict(), + 'raw_event': dict(self.raw_event), + } @dataclass(frozen=True) @@ -45,10 +154,17 @@ class AgentRuntimeConfig: max_turns: int = 12 command_timeout_seconds: float = 30.0 max_output_chars: int = 12000 + stream_model_responses: bool = False + auto_snip_threshold_tokens: int | None = None + auto_compact_threshold_tokens: int | None = None + compact_preserve_messages: int = 4 permissions: AgentPermissions = field(default_factory=AgentPermissions) additional_working_directories: tuple[Path, ...] = () disable_claude_md_discovery: bool = False + budget_config: BudgetConfig = field(default_factory=BudgetConfig) + output_schema: OutputSchemaConfig | None = None session_directory: Path = field(default_factory=lambda: (Path('.port_sessions') / 'agent').resolve()) + scratchpad_root: Path = field(default_factory=lambda: (Path('.port_sessions') / 'scratchpad').resolve()) @dataclass(frozen=True) @@ -56,6 +172,7 @@ class ToolExecutionResult: name: str ok: bool content: str + metadata: JSONDict = field(default_factory=dict) @dataclass(frozen=True) @@ -64,5 +181,11 @@ class AgentRunResult: turns: int tool_calls: int transcript: tuple[JSONDict, ...] + events: tuple[JSONDict, ...] = () + usage: UsageStats = field(default_factory=UsageStats) + total_cost_usd: float = 0.0 + stop_reason: str | None = None + file_history: tuple[JSONDict, ...] = () session_id: str | None = None session_path: str | None = None + scratchpad_directory: str | None = None diff --git a/src/main.py b/src/main.py index 8a24375..738d2be 100644 --- a/src/main.py +++ b/src/main.py @@ -4,9 +4,17 @@ import argparse import os from pathlib import Path from dataclasses import replace +import json from .agent_runtime import LocalCodingAgent -from .agent_types import AgentPermissions, AgentRuntimeConfig, ModelConfig +from .agent_types import ( + AgentPermissions, + AgentRuntimeConfig, + BudgetConfig, + ModelConfig, + ModelPricing, + OutputSchemaConfig, +) 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 @@ -36,12 +44,29 @@ def _add_agent_common_args(parser: argparse.ArgumentParser, *, include_backend: parser.add_argument('--api-key', default=os.environ.get('OPENAI_API_KEY', 'local-token')) parser.add_argument('--temperature', type=float, default=0.0) parser.add_argument('--timeout-seconds', type=float, default=120.0) + parser.add_argument('--input-cost-per-million', type=float, default=0.0) + parser.add_argument('--output-cost-per-million', type=float, default=0.0) parser.add_argument('--cwd', default='.') parser.add_argument('--add-dir', action='append', default=[]) parser.add_argument('--disable-claude-md', action='store_true') parser.add_argument('--allow-write', action='store_true') parser.add_argument('--allow-shell', action='store_true') parser.add_argument('--unsafe', action='store_true') + parser.add_argument('--stream', action='store_true') + parser.add_argument('--auto-snip-threshold', type=int) + parser.add_argument('--auto-compact-threshold', type=int) + parser.add_argument('--compact-preserve-messages', type=int, default=4) + parser.add_argument('--max-total-tokens', type=int) + parser.add_argument('--max-input-tokens', type=int) + parser.add_argument('--max-output-tokens', type=int) + parser.add_argument('--max-reasoning-tokens', type=int) + parser.add_argument('--max-budget-usd', type=float) + parser.add_argument('--max-tool-calls', type=int) + parser.add_argument('--max-delegated-tasks', type=int) + parser.add_argument('--response-schema-file') + parser.add_argument('--response-schema-name') + parser.add_argument('--response-schema-strict', action='store_true') + parser.add_argument('--scratchpad-root') parser.add_argument('--system-prompt') parser.add_argument('--append-system-prompt') parser.add_argument('--override-system-prompt') @@ -56,9 +81,28 @@ def _build_runtime_config(args: argparse.Namespace) -> AgentRuntimeConfig: allow_shell_commands=args.allow_shell, allow_destructive_shell_commands=args.unsafe, ), + stream_model_responses=bool(getattr(args, 'stream', False)), + auto_snip_threshold_tokens=getattr(args, 'auto_snip_threshold', None), + auto_compact_threshold_tokens=getattr(args, 'auto_compact_threshold', None), + compact_preserve_messages=max(0, int(getattr(args, 'compact_preserve_messages', 4))), additional_working_directories=tuple(Path(path).resolve() for path in args.add_dir), disable_claude_md_discovery=args.disable_claude_md, + budget_config=BudgetConfig( + max_total_tokens=getattr(args, 'max_total_tokens', None), + max_input_tokens=getattr(args, 'max_input_tokens', None), + max_output_tokens=getattr(args, 'max_output_tokens', None), + max_reasoning_tokens=getattr(args, 'max_reasoning_tokens', None), + max_total_cost_usd=getattr(args, 'max_budget_usd', None), + max_tool_calls=getattr(args, 'max_tool_calls', None), + max_delegated_tasks=getattr(args, 'max_delegated_tasks', None), + ), + output_schema=_load_output_schema_config(args), session_directory=(Path('.port_sessions') / 'agent').resolve(), + scratchpad_root=( + Path(getattr(args, 'scratchpad_root')).resolve() + if getattr(args, 'scratchpad_root', None) + else (Path('.port_sessions') / 'scratchpad').resolve() + ), ) @@ -69,6 +113,29 @@ def _build_model_config(args: argparse.Namespace) -> ModelConfig: api_key=getattr(args, 'api_key', os.environ.get('OPENAI_API_KEY', 'local-token')), temperature=getattr(args, 'temperature', 0.0), timeout_seconds=getattr(args, 'timeout_seconds', 120.0), + pricing=ModelPricing( + input_cost_per_million_tokens_usd=float( + getattr(args, 'input_cost_per_million', 0.0) or 0.0 + ), + output_cost_per_million_tokens_usd=float( + getattr(args, 'output_cost_per_million', 0.0) or 0.0 + ), + ), + ) + + +def _load_output_schema_config(args: argparse.Namespace) -> OutputSchemaConfig | None: + schema_file = getattr(args, 'response_schema_file', None) + if not schema_file: + return None + payload = json.loads(Path(schema_file).read_text(encoding='utf-8')) + if not isinstance(payload, dict): + raise ValueError('response schema file must contain a top-level JSON object') + name = getattr(args, 'response_schema_name', None) or Path(schema_file).stem + return OutputSchemaConfig( + name=name, + schema=payload, + strict=bool(getattr(args, 'response_schema_strict', False)), ) @@ -92,9 +159,26 @@ def _add_agent_resume_args(parser: argparse.ArgumentParser) -> None: parser.add_argument('--api-key') parser.add_argument('--temperature', type=float) parser.add_argument('--timeout-seconds', type=float) + parser.add_argument('--input-cost-per-million', type=float) + parser.add_argument('--output-cost-per-million', type=float) parser.add_argument('--allow-write', action='store_true') parser.add_argument('--allow-shell', action='store_true') parser.add_argument('--unsafe', action='store_true') + parser.add_argument('--stream', action='store_true') + parser.add_argument('--auto-snip-threshold', type=int) + parser.add_argument('--auto-compact-threshold', type=int) + parser.add_argument('--compact-preserve-messages', type=int) + parser.add_argument('--max-total-tokens', type=int) + parser.add_argument('--max-input-tokens', type=int) + parser.add_argument('--max-output-tokens', type=int) + parser.add_argument('--max-reasoning-tokens', type=int) + parser.add_argument('--max-budget-usd', type=float) + parser.add_argument('--max-tool-calls', type=int) + parser.add_argument('--max-delegated-tasks', type=int) + parser.add_argument('--response-schema-file') + parser.add_argument('--response-schema-name') + parser.add_argument('--response-schema-strict', action='store_true') + parser.add_argument('--scratchpad-root') def _build_resumed_agent(args: argparse.Namespace) -> tuple[LocalCodingAgent, StoredAgentSession]: @@ -112,6 +196,23 @@ def _build_resumed_agent(args: argparse.Namespace) -> tuple[LocalCodingAgent, St model_config = replace(model_config, temperature=args.temperature) if args.timeout_seconds is not None: model_config = replace(model_config, timeout_seconds=args.timeout_seconds) + if args.input_cost_per_million is not None or args.output_cost_per_million is not None: + model_config = replace( + model_config, + pricing=replace( + model_config.pricing, + input_cost_per_million_tokens_usd=( + args.input_cost_per_million + if args.input_cost_per_million is not None + else model_config.pricing.input_cost_per_million_tokens_usd + ), + output_cost_per_million_tokens_usd=( + args.output_cost_per_million + if args.output_cost_per_million is not None + else model_config.pricing.output_cost_per_million_tokens_usd + ), + ), + ) if args.max_turns is not None: runtime_config = replace(runtime_config, max_turns=args.max_turns) @@ -124,6 +225,88 @@ def _build_resumed_agent(args: argparse.Namespace) -> tuple[LocalCodingAgent, St allow_destructive_shell_commands=runtime_config.permissions.allow_destructive_shell_commands or args.unsafe, ), ) + if args.stream: + runtime_config = replace(runtime_config, stream_model_responses=True) + if ( + args.auto_snip_threshold is not None + or args.auto_compact_threshold is not None + or args.compact_preserve_messages is not None + ): + runtime_config = replace( + runtime_config, + auto_snip_threshold_tokens=( + args.auto_snip_threshold + if args.auto_snip_threshold is not None + else runtime_config.auto_snip_threshold_tokens + ), + auto_compact_threshold_tokens=( + args.auto_compact_threshold + if args.auto_compact_threshold is not None + else runtime_config.auto_compact_threshold_tokens + ), + compact_preserve_messages=( + max(0, args.compact_preserve_messages) + if args.compact_preserve_messages is not None + else runtime_config.compact_preserve_messages + ), + ) + if ( + args.max_total_tokens is not None + or args.max_input_tokens is not None + or args.max_output_tokens is not None + or args.max_reasoning_tokens is not None + or args.max_budget_usd is not None + or args.max_tool_calls is not None + or args.max_delegated_tasks is not None + ): + runtime_config = replace( + runtime_config, + budget_config=BudgetConfig( + max_total_tokens=( + args.max_total_tokens + if args.max_total_tokens is not None + else runtime_config.budget_config.max_total_tokens + ), + max_input_tokens=( + args.max_input_tokens + if args.max_input_tokens is not None + else runtime_config.budget_config.max_input_tokens + ), + max_output_tokens=( + args.max_output_tokens + if args.max_output_tokens is not None + else runtime_config.budget_config.max_output_tokens + ), + max_reasoning_tokens=( + args.max_reasoning_tokens + if args.max_reasoning_tokens is not None + else runtime_config.budget_config.max_reasoning_tokens + ), + max_total_cost_usd=( + args.max_budget_usd + if args.max_budget_usd is not None + else runtime_config.budget_config.max_total_cost_usd + ), + max_tool_calls=( + args.max_tool_calls + if args.max_tool_calls is not None + else runtime_config.budget_config.max_tool_calls + ), + max_delegated_tasks=( + args.max_delegated_tasks + if args.max_delegated_tasks is not None + else runtime_config.budget_config.max_delegated_tasks + ), + ), + ) + output_schema = _load_output_schema_config(args) + if output_schema is not None: + runtime_config = replace(runtime_config, output_schema=output_schema) + if args.scratchpad_root: + runtime_config = replace( + runtime_config, + scratchpad_root=Path(args.scratchpad_root).resolve(), + ) agent = LocalCodingAgent( model_config=model_config, @@ -134,11 +317,20 @@ def _build_resumed_agent(args: argparse.Namespace) -> tuple[LocalCodingAgent, St def _print_agent_result(result, *, show_transcript: bool) -> None: print(result.final_output) + print('\n# Usage') + print(f'total_tokens={result.usage.total_tokens}') + print(f'input_tokens={result.usage.input_tokens}') + print(f'output_tokens={result.usage.output_tokens}') + print(f'total_cost_usd={result.total_cost_usd:.6f}') + if result.stop_reason: + print(f'stop_reason={result.stop_reason}') if result.session_id: print('\n# Session') print(f'session_id={result.session_id}') if result.session_path: print(f'session_path={result.session_path}') + if result.scratchpad_directory: + print(f'scratchpad_directory={result.scratchpad_directory}') if show_transcript: print('\n# Transcript') for message in result.transcript: diff --git a/src/openai_compat.py b/src/openai_compat.py index ae59004..c30981f 100644 --- a/src/openai_compat.py +++ b/src/openai_compat.py @@ -1,10 +1,17 @@ from __future__ import annotations import json -from typing import Any +from typing import Any, Iterator from urllib import error, request -from .agent_types import AssistantTurn, ModelConfig, ToolCall +from .agent_types import ( + AssistantTurn, + ModelConfig, + OutputSchemaConfig, + StreamEvent, + ToolCall, + UsageStats, +) class OpenAICompatError(RuntimeError): @@ -66,6 +73,64 @@ def _parse_tool_arguments(raw_arguments: Any) -> dict[str, Any]: ) +def _optional_int(value: Any) -> int: + if isinstance(value, bool): + return 0 + if isinstance(value, int): + return value + if isinstance(value, float): + return int(value) + if isinstance(value, str): + try: + return int(value) + except ValueError: + return 0 + return 0 + + +def _parse_usage(payload: Any) -> UsageStats: + if not isinstance(payload, dict): + return UsageStats() + completion_details = payload.get('completion_tokens_details') + if not isinstance(completion_details, dict): + completion_details = {} + return UsageStats( + input_tokens=( + _optional_int(payload.get('input_tokens')) + or _optional_int(payload.get('prompt_tokens')) + or _optional_int(payload.get('prompt_eval_count')) + ), + output_tokens=( + _optional_int(payload.get('output_tokens')) + or _optional_int(payload.get('completion_tokens')) + or _optional_int(payload.get('eval_count')) + ), + cache_creation_input_tokens=_optional_int( + payload.get('cache_creation_input_tokens') + ), + cache_read_input_tokens=_optional_int(payload.get('cache_read_input_tokens')), + reasoning_tokens=( + _optional_int(payload.get('reasoning_tokens')) + or _optional_int(completion_details.get('reasoning_tokens')) + ), + ) + + +def _build_response_format( + schema: OutputSchemaConfig | None, +) -> dict[str, Any] | None: + if schema is None: + return None + return { + 'type': 'json_schema', + 'json_schema': { + 'name': schema.name, + 'schema': schema.schema, + 'strict': schema.strict, + }, + } + + class OpenAICompatClient: """Minimal OpenAI-compatible chat client for local model servers.""" @@ -76,15 +141,81 @@ class OpenAICompatClient: self, messages: list[dict[str, Any]], tools: list[dict[str, Any]], + *, + output_schema: OutputSchemaConfig | None = None, ) -> AssistantTurn: - payload = { - 'model': self.config.model, - 'messages': messages, - 'tools': tools, - 'tool_choice': 'auto', - 'temperature': self.config.temperature, - 'stream': False, - } + payload = self._request_json( + self._build_payload( + messages=messages, + tools=tools, + stream=False, + output_schema=output_schema, + ) + ) + choices = payload.get('choices') + if not isinstance(choices, list) or not choices: + raise OpenAICompatError('Local model backend returned no choices') + first_choice = choices[0] + if not isinstance(first_choice, dict): + raise OpenAICompatError('Local model backend returned malformed choice data') + + message = first_choice.get('message') + if not isinstance(message, dict): + raise OpenAICompatError('Local model backend returned no assistant message') + + content = _normalize_content(message.get('content')) + tool_calls = self._parse_tool_calls_from_message(message) + + finish_reason = first_choice.get('finish_reason') + if finish_reason is not None and not isinstance(finish_reason, str): + finish_reason = str(finish_reason) + + return AssistantTurn( + content=content, + tool_calls=tuple(tool_calls), + finish_reason=finish_reason, + raw_message=message, + usage=_parse_usage(payload.get('usage')), + ) + + def stream( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]], + *, + output_schema: OutputSchemaConfig | None = None, + ) -> Iterator[StreamEvent]: + payload = self._build_payload( + messages=messages, + tools=tools, + stream=True, + output_schema=output_schema, + ) + req = request.Request( + _join_url(self.config.base_url, '/chat/completions'), + data=json.dumps(payload).encode('utf-8'), + headers={ + 'Authorization': f'Bearer {self.config.api_key}', + 'Content-Type': 'application/json', + }, + method='POST', + ) + try: + with request.urlopen(req, timeout=self.config.timeout_seconds) as response: + yield StreamEvent(type='message_start') + for event_payload in self._iter_sse_payloads(response): + yield from self._parse_stream_payload(event_payload) + except error.HTTPError as exc: + detail = exc.read().decode('utf-8', errors='replace') + raise OpenAICompatError( + f'HTTP {exc.code} from local model backend: {detail}' + ) from exc + except error.URLError as exc: + raise OpenAICompatError( + f'Unable to reach local model backend at {self.config.base_url}: {exc.reason}' + ) from exc + + def _request_json(self, payload: dict[str, Any]) -> dict[str, Any]: body = json.dumps(payload).encode('utf-8') req = request.Request( _join_url(self.config.base_url, '/chat/completions'), @@ -112,19 +243,34 @@ class OpenAICompatClient: payload = json.loads(raw.decode('utf-8')) except json.JSONDecodeError as exc: raise OpenAICompatError('Local model backend returned invalid JSON') from exc + if not isinstance(payload, dict): + raise OpenAICompatError('Local model backend returned malformed JSON payload') + return payload - choices = payload.get('choices') - if not isinstance(choices, list) or not choices: - raise OpenAICompatError('Local model backend returned no choices') - first_choice = choices[0] - if not isinstance(first_choice, dict): - raise OpenAICompatError('Local model backend returned malformed choice data') + def _build_payload( + self, + *, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]], + stream: bool, + output_schema: OutputSchemaConfig | None, + ) -> dict[str, Any]: + payload: dict[str, Any] = { + 'model': self.config.model, + 'messages': messages, + 'tools': tools, + 'tool_choice': 'auto', + 'temperature': self.config.temperature, + 'stream': stream, + } + if stream: + payload['stream_options'] = {'include_usage': True} + response_format = _build_response_format(output_schema) + if response_format is not None: + payload['response_format'] = response_format + return payload - message = first_choice.get('message') - if not isinstance(message, dict): - raise OpenAICompatError('Local model backend returned no assistant message') - - content = _normalize_content(message.get('content')) + def _parse_tool_calls_from_message(self, message: dict[str, Any]) -> list[ToolCall]: tool_calls: list[ToolCall] = [] raw_tool_calls = message.get('tool_calls') if isinstance(raw_tool_calls, list): @@ -149,14 +295,119 @@ class OpenAICompatClient: raise OpenAICompatError('Function call missing name') arguments = _parse_tool_arguments(function_call.get('arguments')) tool_calls.append(ToolCall(id='call_0', name=name, arguments=arguments)) + return tool_calls - finish_reason = first_choice.get('finish_reason') - if finish_reason is not None and not isinstance(finish_reason, str): - finish_reason = str(finish_reason) + def _iter_sse_payloads(self, response: Any) -> Iterator[dict[str, Any]]: + buffer: list[str] = [] + while True: + line = response.readline() + if not line: + break + if isinstance(line, bytes): + text = line.decode('utf-8', errors='replace') + else: + text = str(line) + stripped = text.strip() + if not stripped: + if not buffer: + continue + joined = '\n'.join(buffer) + buffer.clear() + if joined == '[DONE]': + break + try: + payload = json.loads(joined) + except json.JSONDecodeError as exc: + raise OpenAICompatError( + f'Invalid JSON in streaming response: {joined!r}' + ) from exc + if not isinstance(payload, dict): + raise OpenAICompatError('Malformed SSE payload from model backend') + yield payload + continue + if stripped.startswith('data:'): + buffer.append(stripped[5:].strip()) - return AssistantTurn( - content=content, - tool_calls=tuple(tool_calls), - finish_reason=finish_reason, - raw_message=message, - ) + if buffer: + joined = '\n'.join(buffer) + if joined != '[DONE]': + try: + payload = json.loads(joined) + except json.JSONDecodeError as exc: + raise OpenAICompatError( + f'Invalid trailing JSON in streaming response: {joined!r}' + ) from exc + if not isinstance(payload, dict): + raise OpenAICompatError('Malformed trailing SSE payload from model backend') + yield payload + + def _parse_stream_payload( + self, + payload: dict[str, Any], + ) -> Iterator[StreamEvent]: + usage = _parse_usage(payload.get('usage')) + if usage.total_tokens: + yield StreamEvent( + type='usage', + usage=usage, + raw_event=payload, + ) + + choices = payload.get('choices') + if not isinstance(choices, list): + return + + for choice in choices: + if not isinstance(choice, dict): + continue + delta = choice.get('delta') + if not isinstance(delta, dict): + delta = {} + content = delta.get('content') + if isinstance(content, str) and content: + yield StreamEvent( + type='content_delta', + delta=content, + raw_event=choice, + ) + tool_calls = delta.get('tool_calls') + if isinstance(tool_calls, list): + for raw_tool_call in tool_calls: + if not isinstance(raw_tool_call, dict): + continue + function_block = raw_tool_call.get('function') + if not isinstance(function_block, dict): + function_block = {} + yield StreamEvent( + type='tool_call_delta', + tool_call_index=( + raw_tool_call.get('index') + if isinstance(raw_tool_call.get('index'), int) + else 0 + ), + tool_call_id=( + raw_tool_call.get('id') + if isinstance(raw_tool_call.get('id'), str) + else None + ), + tool_name=( + function_block.get('name') + if isinstance(function_block.get('name'), str) + else None + ), + arguments_delta=( + function_block.get('arguments') + if isinstance(function_block.get('arguments'), str) + else '' + ), + raw_event=raw_tool_call, + ) + finish_reason = choice.get('finish_reason') + if finish_reason is not None: + if not isinstance(finish_reason, str): + finish_reason = str(finish_reason) + yield StreamEvent( + type='message_stop', + finish_reason=finish_reason, + raw_event=choice, + ) diff --git a/src/plugin_runtime.py b/src/plugin_runtime.py new file mode 100644 index 0000000..512929a --- /dev/null +++ b/src/plugin_runtime.py @@ -0,0 +1,319 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class PluginToolAlias: + name: str + base_tool: str + description: str | None = None + + +@dataclass(frozen=True) +class PluginToolHook: + tool_name: str + after_result: str | None = None + block_message: str | None = None + + +@dataclass(frozen=True) +class PluginManifest: + name: str + path: str + version: str | None = None + description: str | None = None + tool_names: tuple[str, ...] = () + hook_names: tuple[str, ...] = () + tool_aliases: tuple[PluginToolAlias, ...] = () + tool_hooks: tuple[PluginToolHook, ...] = () + blocked_tools: tuple[str, ...] = () + before_prompt: str | None = None + after_turn: str | None = None + + +@dataclass +class PluginRuntime: + manifests: tuple[PluginManifest, ...] = field(default_factory=tuple) + + @classmethod + def from_workspace( + cls, + cwd: Path, + additional_working_directories: tuple[str, ...] = (), + ) -> 'PluginRuntime': + manifests: list[PluginManifest] = [] + for path in _discover_plugin_manifest_paths(cwd, additional_working_directories): + manifest = _load_manifest(path) + if manifest is not None: + manifests.append(manifest) + return cls(manifests=tuple(manifests)) + + def instruction_blocks(self) -> tuple[str, ...]: + blocks: list[str] = [] + for manifest in self.manifests: + lines = [ + f'Plugin: {manifest.name}', + ] + if manifest.description: + lines.append(f'Description: {manifest.description}') + if manifest.tool_names: + lines.append(f'Tools: {", ".join(manifest.tool_names)}') + if manifest.hook_names: + lines.append(f'Hooks: {", ".join(manifest.hook_names)}') + if manifest.tool_aliases: + lines.append( + 'Tool aliases: ' + + ', '.join(alias.name for alias in manifest.tool_aliases) + ) + if manifest.blocked_tools: + lines.append( + 'Blocked tools: ' + + ', '.join(manifest.blocked_tools) + ) + blocks.append('\n'.join(lines)) + return tuple(blocks) + + def before_prompt_injections(self) -> tuple[str, ...]: + return tuple( + manifest.before_prompt + for manifest in self.manifests + if manifest.before_prompt + ) + + def after_turn_injections(self) -> tuple[str, ...]: + return tuple( + manifest.after_turn + for manifest in self.manifests + if manifest.after_turn + ) + + def register_tool_aliases( + self, + base_registry: dict[str, Any], + ) -> dict[str, Any]: + from .agent_tools import AgentTool + + aliases: dict[str, AgentTool] = {} + for manifest in self.manifests: + for alias in manifest.tool_aliases: + base_tool = base_registry.get(alias.base_tool) + if base_tool is None or alias.name in base_registry or alias.name in aliases: + continue + aliases[alias.name] = AgentTool( + name=alias.name, + description=( + alias.description + or f'Plugin alias from {manifest.name} for base tool {alias.base_tool}.' + ), + parameters=base_tool.parameters, + handler=base_tool.handler, + ) + return aliases + + def blocked_tool_message(self, tool_name: str) -> str | None: + for manifest in self.manifests: + if tool_name in manifest.blocked_tools: + return f'Plugin {manifest.name} blocked tool {tool_name}.' + for hook in manifest.tool_hooks: + if hook.tool_name == tool_name and hook.block_message: + return hook.block_message + return None + + def tool_result_injections(self, tool_name: str) -> tuple[str, ...]: + messages: list[str] = [] + for manifest in self.manifests: + for hook in manifest.tool_hooks: + if hook.tool_name == tool_name and hook.after_result: + messages.append(f'{manifest.name}: {hook.after_result}') + return tuple(messages) + + def render_summary(self) -> str: + if not self.manifests: + return 'No local plugin manifests discovered.' + lines = [f'Local plugin manifests: {len(self.manifests)}'] + for manifest in self.manifests[:10]: + details = [manifest.name] + if manifest.version: + details.append(f'version={manifest.version}') + if manifest.tool_names: + details.append(f'tools={len(manifest.tool_names)}') + if manifest.hook_names: + details.append(f'hooks={len(manifest.hook_names)}') + if manifest.tool_aliases: + details.append(f'aliases={len(manifest.tool_aliases)}') + if manifest.blocked_tools: + details.append(f'blocked={len(manifest.blocked_tools)}') + if manifest.tool_hooks: + details.append(f'tool_hooks={len(manifest.tool_hooks)}') + lines.append(f"- {'; '.join(details)}") + if len(self.manifests) > 10: + lines.append(f'- ... plus {len(self.manifests) - 10} more plugin manifests') + return '\n'.join(lines) + + +def _discover_plugin_manifest_paths( + cwd: Path, + additional_working_directories: tuple[str, ...], +) -> tuple[Path, ...]: + candidates: list[Path] = [] + seen: set[Path] = set() + + def remember(path: Path) -> None: + resolved = path.resolve() + if resolved in seen or not resolved.exists() or not resolved.is_file(): + return + seen.add(resolved) + candidates.append(resolved) + + roots = _walk_upwards(cwd.resolve()) + roots.extend(Path(path).resolve() for path in additional_working_directories) + for root in roots: + remember(root / '.codex-plugin' / 'plugin.json') + remember(root / '.claw-plugin' / 'plugin.json') + plugins_dir = root / 'plugins' + if plugins_dir.is_dir(): + for candidate in sorted(plugins_dir.glob('*/plugin.json')): + remember(candidate) + return tuple(candidates) + + +def _walk_upwards(path: Path) -> list[Path]: + walked: list[Path] = [] + current = path + while True: + walked.append(current) + if current.parent == current: + break + current = current.parent + return walked + + +def _load_manifest(path: Path) -> PluginManifest | None: + try: + payload = json.loads(path.read_text(encoding='utf-8')) + except (OSError, json.JSONDecodeError): + return None + if not isinstance(payload, dict): + return None + name = payload.get('name') + if not isinstance(name, str) or not name.strip(): + return None + before_prompt, after_turn, hook_names = _parse_hooks(payload.get('hooks')) + return PluginManifest( + name=name.strip(), + path=str(path), + version=_optional_string(payload.get('version')), + description=_optional_string(payload.get('description')), + tool_names=_extract_string_tuple(payload.get('tools')), + hook_names=hook_names, + tool_aliases=_extract_tool_aliases(payload), + tool_hooks=_extract_tool_hooks(payload), + blocked_tools=_extract_string_tuple( + payload.get('blocked_tools') + if payload.get('blocked_tools') is not None + else payload.get('blockedTools') + ), + before_prompt=before_prompt, + after_turn=after_turn, + ) + + +def _optional_string(value: Any) -> str | None: + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _extract_string_tuple(value: Any) -> tuple[str, ...]: + if isinstance(value, list): + return tuple(item for item in value if isinstance(item, str) and item.strip()) + if isinstance(value, dict): + names = [key for key in value if isinstance(key, str) and key.strip()] + return tuple(names) + return () + + +def _extract_tool_aliases(payload: dict[str, Any]) -> tuple[PluginToolAlias, ...]: + raw_aliases = payload.get('tool_aliases') + if raw_aliases is None: + raw_aliases = payload.get('toolAliases') + aliases: list[PluginToolAlias] = [] + if isinstance(raw_aliases, list): + for item in raw_aliases: + if not isinstance(item, dict): + continue + name = item.get('name') + base_tool = item.get('base_tool') + if base_tool is None: + base_tool = item.get('baseTool') + if not isinstance(name, str) or not name.strip(): + continue + if not isinstance(base_tool, str) or not base_tool.strip(): + continue + aliases.append( + PluginToolAlias( + name=name.strip(), + base_tool=base_tool.strip(), + description=_optional_string(item.get('description')), + ) + ) + return tuple(aliases) + + +def _parse_hooks(value: Any) -> tuple[str | None, str | None, tuple[str, ...]]: + if isinstance(value, list): + names = tuple(item for item in value if isinstance(item, str) and item.strip()) + return None, None, names + if isinstance(value, dict): + names = tuple(key for key in value if isinstance(key, str) and key.strip()) + before_prompt = value.get('beforePrompt') + if before_prompt is None: + before_prompt = value.get('before_prompt') + after_turn = value.get('afterTurn') + if after_turn is None: + after_turn = value.get('after_turn') + return ( + _optional_string(before_prompt), + _optional_string(after_turn), + names, + ) + return None, None, () + + +def _extract_tool_hooks(payload: dict[str, Any]) -> tuple[PluginToolHook, ...]: + raw_hooks = payload.get('tool_hooks') + if raw_hooks is None: + raw_hooks = payload.get('toolHooks') + hooks: list[PluginToolHook] = [] + if isinstance(raw_hooks, dict): + for tool_name, value in raw_hooks.items(): + if not isinstance(tool_name, str) or not tool_name.strip(): + continue + if isinstance(value, str): + hooks.append( + PluginToolHook( + tool_name=tool_name.strip(), + after_result=value.strip() or None, + ) + ) + continue + if not isinstance(value, dict): + continue + after_result = value.get('afterResult') + if after_result is None: + after_result = value.get('after_result') + block_message = value.get('blockMessage') + if block_message is None: + block_message = value.get('block_message') + hooks.append( + PluginToolHook( + tool_name=tool_name.strip(), + after_result=_optional_string(after_result), + block_message=_optional_string(block_message), + ) + ) + return tuple(hooks) diff --git a/src/query_engine.py b/src/query_engine.py index dd30cda..d018a07 100644 --- a/src/query_engine.py +++ b/src/query_engine.py @@ -1,193 +1,378 @@ -from __future__ import annotations - -import json -from dataclasses import dataclass, field -from uuid import uuid4 - -from .commands import build_command_backlog -from .models import PermissionDenial, UsageSummary -from .port_manifest import PortManifest, build_port_manifest -from .session_store import StoredSession, load_session, save_session -from .tools import build_tool_backlog -from .transcript import TranscriptStore - - -@dataclass(frozen=True) -class QueryEngineConfig: - max_turns: int = 8 - max_budget_tokens: int = 2000 - compact_after_turns: int = 12 - structured_output: bool = False - structured_retry_limit: int = 2 - - -@dataclass(frozen=True) -class TurnResult: - prompt: str - output: str - matched_commands: tuple[str, ...] - matched_tools: tuple[str, ...] - permission_denials: tuple[PermissionDenial, ...] - usage: UsageSummary - stop_reason: str - - -@dataclass -class QueryEnginePort: - manifest: PortManifest - config: QueryEngineConfig = field(default_factory=QueryEngineConfig) - session_id: str = field(default_factory=lambda: uuid4().hex) - mutable_messages: list[str] = field(default_factory=list) - permission_denials: list[PermissionDenial] = field(default_factory=list) - total_usage: UsageSummary = field(default_factory=UsageSummary) - transcript_store: TranscriptStore = field(default_factory=TranscriptStore) - - @classmethod - def from_workspace(cls) -> 'QueryEnginePort': - return cls(manifest=build_port_manifest()) - - @classmethod - def from_saved_session(cls, session_id: str) -> 'QueryEnginePort': - stored = load_session(session_id) - transcript = TranscriptStore(entries=list(stored.messages), flushed=True) - return cls( - manifest=build_port_manifest(), - session_id=stored.session_id, - mutable_messages=list(stored.messages), - total_usage=UsageSummary(stored.input_tokens, stored.output_tokens), - transcript_store=transcript, - ) - - def submit_message( - self, - prompt: str, - matched_commands: tuple[str, ...] = (), - matched_tools: tuple[str, ...] = (), - denied_tools: tuple[PermissionDenial, ...] = (), - ) -> TurnResult: - if len(self.mutable_messages) >= self.config.max_turns: - output = f'Max turns reached before processing prompt: {prompt}' - return TurnResult( - prompt=prompt, - output=output, - matched_commands=matched_commands, - matched_tools=matched_tools, - permission_denials=denied_tools, - usage=self.total_usage, - stop_reason='max_turns_reached', - ) - - summary_lines = [ - f'Prompt: {prompt}', - f'Matched commands: {", ".join(matched_commands) if matched_commands else "none"}', - f'Matched tools: {", ".join(matched_tools) if matched_tools else "none"}', - f'Permission denials: {len(denied_tools)}', - ] - output = self._format_output(summary_lines) - projected_usage = self.total_usage.add_turn(prompt, output) - stop_reason = 'completed' - if projected_usage.input_tokens + projected_usage.output_tokens > self.config.max_budget_tokens: - stop_reason = 'max_budget_reached' - self.mutable_messages.append(prompt) - self.transcript_store.append(prompt) - self.permission_denials.extend(denied_tools) - self.total_usage = projected_usage - self.compact_messages_if_needed() - return TurnResult( - prompt=prompt, - output=output, - matched_commands=matched_commands, - matched_tools=matched_tools, - permission_denials=denied_tools, - usage=self.total_usage, - stop_reason=stop_reason, - ) - - def stream_submit_message( - self, - prompt: str, - matched_commands: tuple[str, ...] = (), - matched_tools: tuple[str, ...] = (), - denied_tools: tuple[PermissionDenial, ...] = (), - ): - yield {'type': 'message_start', 'session_id': self.session_id, 'prompt': prompt} - if matched_commands: - yield {'type': 'command_match', 'commands': matched_commands} - if matched_tools: - yield {'type': 'tool_match', 'tools': matched_tools} - if denied_tools: - yield {'type': 'permission_denial', 'denials': [denial.tool_name for denial in denied_tools]} - result = self.submit_message(prompt, matched_commands, matched_tools, denied_tools) - yield {'type': 'message_delta', 'text': result.output} - yield { - 'type': 'message_stop', - 'usage': {'input_tokens': result.usage.input_tokens, 'output_tokens': result.usage.output_tokens}, - 'stop_reason': result.stop_reason, - 'transcript_size': len(self.transcript_store.entries), - } - - def compact_messages_if_needed(self) -> None: - if len(self.mutable_messages) > self.config.compact_after_turns: - self.mutable_messages[:] = self.mutable_messages[-self.config.compact_after_turns :] - self.transcript_store.compact(self.config.compact_after_turns) - - def replay_user_messages(self) -> tuple[str, ...]: - return self.transcript_store.replay() - - def flush_transcript(self) -> None: - self.transcript_store.flush() - - def persist_session(self) -> str: - self.flush_transcript() - path = save_session( - StoredSession( - session_id=self.session_id, - messages=tuple(self.mutable_messages), - input_tokens=self.total_usage.input_tokens, - output_tokens=self.total_usage.output_tokens, - ) - ) - return str(path) - - def _format_output(self, summary_lines: list[str]) -> str: - if self.config.structured_output: - payload = { - 'summary': summary_lines, - 'session_id': self.session_id, - } - return self._render_structured_output(payload) - return '\n'.join(summary_lines) - - def _render_structured_output(self, payload: dict[str, object]) -> str: - last_error: Exception | None = None - for _ in range(self.config.structured_retry_limit): - try: - return json.dumps(payload, indent=2) - except (TypeError, ValueError) as exc: # pragma: no cover - defensive branch - last_error = exc - payload = {'summary': ['structured output retry'], 'session_id': self.session_id} - raise RuntimeError('structured output rendering failed') from last_error - - def render_summary(self) -> str: - command_backlog = build_command_backlog() - tool_backlog = build_tool_backlog() - sections = [ - '# Python Porting Workspace Summary', - '', - self.manifest.to_markdown(), - '', - f'Command surface: {len(command_backlog.modules)} mirrored entries', - *command_backlog.summary_lines()[:10], - '', - f'Tool surface: {len(tool_backlog.modules)} mirrored entries', - *tool_backlog.summary_lines()[:10], - '', - f'Session id: {self.session_id}', - f'Conversation turns stored: {len(self.mutable_messages)}', - f'Permission denials tracked: {len(self.permission_denials)}', - f'Usage totals: in={self.total_usage.input_tokens} out={self.total_usage.output_tokens}', - f'Max turns: {self.config.max_turns}', - f'Max budget tokens: {self.config.max_budget_tokens}', - f'Transcript flushed: {self.transcript_store.flushed}', - ] - return '\n'.join(sections) +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path +from uuid import uuid4 + +from .agent_runtime import LocalCodingAgent +from .commands import build_command_backlog +from .models import PermissionDenial, UsageSummary +from .plugin_runtime import PluginRuntime +from .port_manifest import PortManifest, build_port_manifest +from .session_store import StoredSession, load_agent_session, load_session, save_session +from .tools import build_tool_backlog +from .transcript import TranscriptStore + + +@dataclass(frozen=True) +class QueryEngineConfig: + max_turns: int = 8 + max_budget_tokens: int = 2000 + compact_after_turns: int = 12 + structured_output: bool = False + structured_retry_limit: int = 2 + use_runtime_agent: bool = False + + +@dataclass(frozen=True) +class TurnResult: + prompt: str + output: str + matched_commands: tuple[str, ...] + matched_tools: tuple[str, ...] + permission_denials: tuple[PermissionDenial, ...] + usage: UsageSummary + stop_reason: str + session_id: str | None = None + session_path: str | None = None + tool_calls: int = 0 + total_cost_usd: float = 0.0 + events: tuple[dict[str, object], ...] = () + transcript: tuple[dict[str, object], ...] = () + + +@dataclass +class QueryEnginePort: + manifest: PortManifest + config: QueryEngineConfig = field(default_factory=QueryEngineConfig) + session_id: str = field(default_factory=lambda: uuid4().hex) + mutable_messages: list[str] = field(default_factory=list) + permission_denials: list[PermissionDenial] = field(default_factory=list) + total_usage: UsageSummary = field(default_factory=UsageSummary) + transcript_store: TranscriptStore = field(default_factory=TranscriptStore) + runtime_agent: LocalCodingAgent | None = None + plugin_runtime: PluginRuntime | None = None + runtime_event_counts: dict[str, int] = field(default_factory=dict) + runtime_message_kind_counts: dict[str, int] = field(default_factory=dict) + runtime_transcript_size: int = 0 + last_turn: TurnResult | None = field(default=None, init=False, repr=False) + + @classmethod + def from_workspace(cls) -> 'QueryEnginePort': + return cls( + manifest=build_port_manifest(), + plugin_runtime=PluginRuntime.from_workspace(Path.cwd()), + ) + + @classmethod + def from_saved_session(cls, session_id: str) -> 'QueryEnginePort': + stored = load_session(session_id) + transcript = TranscriptStore(entries=list(stored.messages), flushed=True) + return cls( + manifest=build_port_manifest(), + session_id=stored.session_id, + mutable_messages=list(stored.messages), + total_usage=UsageSummary(stored.input_tokens, stored.output_tokens), + transcript_store=transcript, + plugin_runtime=PluginRuntime.from_workspace(Path.cwd()), + ) + + @classmethod + def from_runtime_agent( + cls, + agent: LocalCodingAgent, + *, + manifest: PortManifest | None = None, + ) -> 'QueryEnginePort': + return cls( + manifest=manifest or build_port_manifest(), + config=QueryEngineConfig(use_runtime_agent=True), + session_id=agent.active_session_id or uuid4().hex, + runtime_agent=agent, + plugin_runtime=PluginRuntime.from_workspace( + agent.runtime_config.cwd, + tuple(str(path) for path in agent.runtime_config.additional_working_directories), + ), + ) + + def submit_message( + self, + prompt: str, + matched_commands: tuple[str, ...] = (), + matched_tools: tuple[str, ...] = (), + denied_tools: tuple[PermissionDenial, ...] = (), + ) -> TurnResult: + if self.config.use_runtime_agent and self.runtime_agent is not None: + result = self._submit_runtime_message(prompt) + turn = TurnResult( + prompt=prompt, + output=result.final_output, + matched_commands=matched_commands, + matched_tools=matched_tools, + permission_denials=denied_tools, + usage=UsageSummary( + input_tokens=result.usage.input_tokens, + output_tokens=result.usage.output_tokens, + ), + stop_reason=result.stop_reason or 'completed', + session_id=result.session_id, + session_path=result.session_path, + tool_calls=result.tool_calls, + total_cost_usd=result.total_cost_usd, + events=result.events, + transcript=result.transcript, + ) + self._record_turn(prompt, turn, denied_tools) + return turn + + if len(self.mutable_messages) >= self.config.max_turns: + output = f'Max turns reached before processing prompt: {prompt}' + return TurnResult( + prompt=prompt, + output=output, + matched_commands=matched_commands, + matched_tools=matched_tools, + permission_denials=denied_tools, + usage=self.total_usage, + stop_reason='max_turns_reached', + ) + + summary_lines = [ + f'Prompt: {prompt}', + f'Matched commands: {", ".join(matched_commands) if matched_commands else "none"}', + f'Matched tools: {", ".join(matched_tools) if matched_tools else "none"}', + f'Permission denials: {len(denied_tools)}', + ] + output = self._format_output(summary_lines) + projected_usage = self.total_usage.add_turn(prompt, output) + stop_reason = 'completed' + if projected_usage.input_tokens + projected_usage.output_tokens > self.config.max_budget_tokens: + stop_reason = 'max_budget_reached' + turn = TurnResult( + prompt=prompt, + output=output, + matched_commands=matched_commands, + matched_tools=matched_tools, + permission_denials=denied_tools, + usage=projected_usage, + stop_reason=stop_reason, + ) + self._record_turn(prompt, turn, denied_tools) + self.compact_messages_if_needed() + return turn + + def stream_submit_message( + self, + prompt: str, + matched_commands: tuple[str, ...] = (), + matched_tools: tuple[str, ...] = (), + denied_tools: tuple[PermissionDenial, ...] = (), + ): + yield {'type': 'message_start', 'session_id': self.session_id, 'prompt': prompt} + if matched_commands: + yield {'type': 'command_match', 'commands': matched_commands} + if matched_tools: + yield {'type': 'tool_match', 'tools': matched_tools} + if denied_tools: + yield {'type': 'permission_denial', 'denials': [denial.tool_name for denial in denied_tools]} + result = self.submit_message(prompt, matched_commands, matched_tools, denied_tools) + if self.config.use_runtime_agent: + for event in result.events: + yield event + yield { + 'type': 'message_stop', + 'usage': { + 'input_tokens': result.usage.input_tokens, + 'output_tokens': result.usage.output_tokens, + }, + 'stop_reason': result.stop_reason, + 'session_id': result.session_id, + 'transcript_size': len(result.transcript), + } + return + yield {'type': 'message_delta', 'text': result.output} + yield { + 'type': 'message_stop', + 'usage': {'input_tokens': result.usage.input_tokens, 'output_tokens': result.usage.output_tokens}, + 'stop_reason': result.stop_reason, + 'transcript_size': len(self.transcript_store.entries), + } + + def compact_messages_if_needed(self) -> None: + if len(self.mutable_messages) > self.config.compact_after_turns: + self.mutable_messages[:] = self.mutable_messages[-self.config.compact_after_turns :] + self.transcript_store.compact(self.config.compact_after_turns) + + def replay_user_messages(self) -> tuple[str, ...]: + return self.transcript_store.replay() + + def flush_transcript(self) -> None: + self.transcript_store.flush() + + def persist_session(self) -> str: + if self.config.use_runtime_agent and self.last_turn is not None and self.last_turn.session_path: + return self.last_turn.session_path + self.flush_transcript() + path = save_session( + StoredSession( + session_id=self.session_id, + messages=tuple(self.mutable_messages), + input_tokens=self.total_usage.input_tokens, + output_tokens=self.total_usage.output_tokens, + ) + ) + return str(path) + + def render_summary(self) -> str: + command_backlog = build_command_backlog() + tool_backlog = build_tool_backlog() + sections = [ + '# Python Porting Workspace Summary', + '', + self.manifest.to_markdown(), + '', + f'Command surface: {len(command_backlog.modules)} mirrored entries', + *command_backlog.summary_lines()[:10], + '', + f'Tool surface: {len(tool_backlog.modules)} mirrored entries', + *tool_backlog.summary_lines()[:10], + '', + f'Session id: {self.session_id}', + f'Conversation turns stored: {len(self.mutable_messages)}', + f'Permission denials tracked: {len(self.permission_denials)}', + f'Usage totals: in={self.total_usage.input_tokens} out={self.total_usage.output_tokens}', + f'Max turns: {self.config.max_turns}', + f'Max budget tokens: {self.config.max_budget_tokens}', + f'Transcript flushed: {self.transcript_store.flushed}', + f'Real runtime agent mode: {self.config.use_runtime_agent}', + ] + if self.plugin_runtime is not None: + sections.extend(['', '## Plugin Runtime', self.plugin_runtime.render_summary()]) + if self.runtime_agent is not None and self.runtime_agent.agent_manager is not None: + sections.extend(['', '## Agent Manager', *self.runtime_agent.agent_manager.summary_lines()]) + if self.runtime_event_counts: + sections.extend(['', '## Runtime Events']) + sections.extend( + f'- {name}={count}' + for name, count in sorted(self.runtime_event_counts.items()) + ) + sections.append(f'- runtime_transcript_size={self.runtime_transcript_size}') + if self.runtime_message_kind_counts: + sections.extend(['', '## Runtime Message Kinds']) + sections.extend( + f'- {name}={count}' + for name, count in sorted(self.runtime_message_kind_counts.items()) + ) + if self.last_turn is not None: + sections.extend( + [ + '', + '## Last Turn', + f'- stop_reason={self.last_turn.stop_reason}', + f'- tool_calls={self.last_turn.tool_calls}', + f'- session_id={self.last_turn.session_id or "none"}', + f'- transcript_messages={len(self.last_turn.transcript)}', + ] + ) + return '\n'.join(sections) + + def _format_output(self, summary_lines: list[str]) -> str: + if self.config.structured_output: + payload = { + 'summary': summary_lines, + 'session_id': self.session_id, + } + return self._render_structured_output(payload) + return '\n'.join(summary_lines) + + def _render_structured_output(self, payload: dict[str, object]) -> str: + last_error: Exception | None = None + for _ in range(self.config.structured_retry_limit): + try: + return json.dumps(payload, indent=2) + except (TypeError, ValueError) as exc: # pragma: no cover - defensive branch + last_error = exc + payload = {'summary': ['structured output retry'], 'session_id': self.session_id} + raise RuntimeError('structured output rendering failed') from last_error + + def _record_turn( + self, + prompt: str, + turn: TurnResult, + denied_tools: tuple[PermissionDenial, ...], + ) -> None: + self.mutable_messages.append(prompt) + self.transcript_store.append(prompt) + self.transcript_store.append(turn.output) + if self.config.use_runtime_agent: + self._record_runtime_turn(turn) + self.permission_denials.extend(denied_tools) + self.total_usage = turn.usage + self.last_turn = turn + if turn.session_id is not None: + self.session_id = turn.session_id + + def _submit_runtime_message(self, prompt: str): + assert self.runtime_agent is not None + if self.last_turn is None or not self.last_turn.session_id: + return self.runtime_agent.run(prompt) + stored = load_agent_session( + self.last_turn.session_id, + directory=self.runtime_agent.runtime_config.session_directory, + ) + return self.runtime_agent.resume(prompt, stored) + + def _record_runtime_turn(self, turn: TurnResult) -> None: + self.runtime_transcript_size = len(turn.transcript) + event_counts: dict[str, int] = {} + for event in turn.events: + event_type = event.get('type') + if not isinstance(event_type, str) or not event_type: + continue + event_counts[event_type] = event_counts.get(event_type, 0) + 1 + self.runtime_event_counts[event_type] = ( + self.runtime_event_counts.get(event_type, 0) + 1 + ) + kind_counts: dict[str, int] = {} + for entry in turn.transcript: + if not isinstance(entry, dict): + continue + metadata = entry.get('metadata') + if not isinstance(metadata, dict): + continue + kind = metadata.get('kind') + if not isinstance(kind, str) or not kind: + continue + kind_counts[kind] = kind_counts.get(kind, 0) + 1 + self.runtime_message_kind_counts[kind] = ( + self.runtime_message_kind_counts.get(kind, 0) + 1 + ) + summary = self._summarize_runtime_turn(event_counts, kind_counts, len(turn.transcript)) + if summary: + self.transcript_store.append(summary) + + def _summarize_runtime_turn( + self, + event_counts: dict[str, int], + kind_counts: dict[str, int], + transcript_size: int, + ) -> str: + parts = [f'runtime_transcript={transcript_size}'] + if event_counts: + parts.append( + 'events=' + + ', '.join( + f'{name}:{count}' + for name, count in sorted(event_counts.items()) + ) + ) + if kind_counts: + parts.append( + 'kinds=' + + ', '.join( + f'{name}:{count}' + for name, count in sorted(kind_counts.items()) + ) + ) + return '[runtime] ' + ' | '.join(parts) diff --git a/src/session_store.py b/src/session_store.py index 374b70e..6c46c72 100644 --- a/src/session_store.py +++ b/src/session_store.py @@ -5,7 +5,15 @@ from dataclasses import asdict, dataclass from pathlib import Path from typing import Any -from .agent_types import AgentPermissions, AgentRuntimeConfig, ModelConfig +from .agent_types import ( + AgentPermissions, + AgentRuntimeConfig, + BudgetConfig, + ModelConfig, + ModelPricing, + OutputSchemaConfig, + UsageStats, +) @dataclass(frozen=True) @@ -53,6 +61,10 @@ class StoredAgentSession: messages: tuple[JSONDict, ...] turns: int tool_calls: int + usage: JSONDict + total_cost_usd: float + file_history: tuple[JSONDict, ...] + scratchpad_directory: str | None = None def save_agent_session(session: StoredAgentSession, directory: Path | None = None) -> Path: @@ -78,6 +90,16 @@ def load_agent_session(session_id: str, directory: Path | None = None) -> Stored ), turns=int(data['turns']), tool_calls=int(data['tool_calls']), + usage=dict(data.get('usage', {})), + total_cost_usd=float(data.get('total_cost_usd', 0.0)), + file_history=tuple( + entry for entry in data.get('file_history', []) if isinstance(entry, dict) + ), + scratchpad_directory=( + str(data['scratchpad_directory']) + if isinstance(data.get('scratchpad_directory'), str) + else None + ), ) @@ -88,6 +110,12 @@ def serialize_model_config(model_config: ModelConfig) -> JSONDict: 'api_key': model_config.api_key, 'temperature': model_config.temperature, 'timeout_seconds': model_config.timeout_seconds, + 'pricing': { + 'input_cost_per_million_tokens_usd': model_config.pricing.input_cost_per_million_tokens_usd, + 'output_cost_per_million_tokens_usd': model_config.pricing.output_cost_per_million_tokens_usd, + 'cache_creation_input_cost_per_million_tokens_usd': model_config.pricing.cache_creation_input_cost_per_million_tokens_usd, + 'cache_read_input_cost_per_million_tokens_usd': model_config.pricing.cache_read_input_cost_per_million_tokens_usd, + }, } @@ -98,6 +126,7 @@ def deserialize_model_config(payload: JSONDict) -> ModelConfig: api_key=str(payload.get('api_key', 'local-token')), temperature=float(payload.get('temperature', 0.0)), timeout_seconds=float(payload.get('timeout_seconds', 120.0)), + pricing=_deserialize_pricing(payload.get('pricing')), ) @@ -107,6 +136,10 @@ def serialize_runtime_config(runtime_config: AgentRuntimeConfig) -> JSONDict: 'max_turns': runtime_config.max_turns, 'command_timeout_seconds': runtime_config.command_timeout_seconds, 'max_output_chars': runtime_config.max_output_chars, + 'stream_model_responses': runtime_config.stream_model_responses, + 'auto_snip_threshold_tokens': runtime_config.auto_snip_threshold_tokens, + 'auto_compact_threshold_tokens': runtime_config.auto_compact_threshold_tokens, + 'compact_preserve_messages': runtime_config.compact_preserve_messages, 'permissions': { 'allow_file_write': runtime_config.permissions.allow_file_write, 'allow_shell_commands': runtime_config.permissions.allow_shell_commands, @@ -114,7 +147,26 @@ def serialize_runtime_config(runtime_config: AgentRuntimeConfig) -> JSONDict: }, 'additional_working_directories': [str(path) for path in runtime_config.additional_working_directories], 'disable_claude_md_discovery': runtime_config.disable_claude_md_discovery, + 'budget_config': { + 'max_total_tokens': runtime_config.budget_config.max_total_tokens, + 'max_input_tokens': runtime_config.budget_config.max_input_tokens, + 'max_output_tokens': runtime_config.budget_config.max_output_tokens, + 'max_reasoning_tokens': runtime_config.budget_config.max_reasoning_tokens, + 'max_total_cost_usd': runtime_config.budget_config.max_total_cost_usd, + 'max_tool_calls': runtime_config.budget_config.max_tool_calls, + 'max_delegated_tasks': runtime_config.budget_config.max_delegated_tasks, + }, + 'output_schema': ( + { + 'name': runtime_config.output_schema.name, + 'schema': runtime_config.output_schema.schema, + 'strict': runtime_config.output_schema.strict, + } + if runtime_config.output_schema is not None + else None + ), 'session_directory': str(runtime_config.session_directory), + 'scratchpad_root': str(runtime_config.scratchpad_root), } @@ -122,11 +174,19 @@ def deserialize_runtime_config(payload: JSONDict) -> AgentRuntimeConfig: permissions_payload = payload.get('permissions') if not isinstance(permissions_payload, dict): permissions_payload = {} + budget_payload = payload.get('budget_config') + if not isinstance(budget_payload, dict): + budget_payload = {} + output_schema_payload = payload.get('output_schema') return AgentRuntimeConfig( cwd=Path(str(payload['cwd'])).resolve(), max_turns=int(payload.get('max_turns', 12)), command_timeout_seconds=float(payload.get('command_timeout_seconds', 30.0)), max_output_chars=int(payload.get('max_output_chars', 12000)), + stream_model_responses=bool(payload.get('stream_model_responses', False)), + auto_snip_threshold_tokens=_optional_int(payload.get('auto_snip_threshold_tokens')), + auto_compact_threshold_tokens=_optional_int(payload.get('auto_compact_threshold_tokens')), + compact_preserve_messages=int(payload.get('compact_preserve_messages', 4)), permissions=AgentPermissions( allow_file_write=bool(permissions_payload.get('allow_file_write', False)), allow_shell_commands=bool(permissions_payload.get('allow_shell_commands', False)), @@ -137,5 +197,83 @@ def deserialize_runtime_config(payload: JSONDict) -> AgentRuntimeConfig: for path in payload.get('additional_working_directories', []) ), disable_claude_md_discovery=bool(payload.get('disable_claude_md_discovery', False)), + budget_config=BudgetConfig( + max_total_tokens=_optional_int(budget_payload.get('max_total_tokens')), + max_input_tokens=_optional_int(budget_payload.get('max_input_tokens')), + max_output_tokens=_optional_int(budget_payload.get('max_output_tokens')), + max_reasoning_tokens=_optional_int(budget_payload.get('max_reasoning_tokens')), + max_total_cost_usd=_optional_float(budget_payload.get('max_total_cost_usd')), + max_tool_calls=_optional_int(budget_payload.get('max_tool_calls')), + max_delegated_tasks=_optional_int(budget_payload.get('max_delegated_tasks')), + ), + output_schema=_deserialize_output_schema(output_schema_payload), session_directory=Path(str(payload.get('session_directory', DEFAULT_AGENT_SESSION_DIR))).resolve(), + scratchpad_root=Path(str(payload.get('scratchpad_root', DEFAULT_SESSION_DIR / 'scratchpad'))).resolve(), ) + + +def usage_from_payload(payload: JSONDict | None) -> UsageStats: + if not isinstance(payload, dict): + return UsageStats() + return UsageStats( + input_tokens=_optional_int(payload.get('input_tokens')) or 0, + output_tokens=_optional_int(payload.get('output_tokens')) or 0, + cache_creation_input_tokens=_optional_int(payload.get('cache_creation_input_tokens')) or 0, + cache_read_input_tokens=_optional_int(payload.get('cache_read_input_tokens')) or 0, + reasoning_tokens=_optional_int(payload.get('reasoning_tokens')) or 0, + ) + + +def _deserialize_pricing(payload: Any) -> ModelPricing: + if not isinstance(payload, dict): + return ModelPricing() + return ModelPricing( + input_cost_per_million_tokens_usd=_optional_float(payload.get('input_cost_per_million_tokens_usd')) or 0.0, + output_cost_per_million_tokens_usd=_optional_float(payload.get('output_cost_per_million_tokens_usd')) or 0.0, + cache_creation_input_cost_per_million_tokens_usd=( + _optional_float(payload.get('cache_creation_input_cost_per_million_tokens_usd')) + or 0.0 + ), + cache_read_input_cost_per_million_tokens_usd=( + _optional_float(payload.get('cache_read_input_cost_per_million_tokens_usd')) + or 0.0 + ), + ) + + +def _deserialize_output_schema(payload: Any) -> OutputSchemaConfig | None: + if not isinstance(payload, dict): + return None + schema = payload.get('schema') + if not isinstance(schema, dict): + return None + name = payload.get('name') + if not isinstance(name, str) or not name: + return None + return OutputSchemaConfig( + name=name, + schema=dict(schema), + strict=bool(payload.get('strict', False)), + ) + + +def _optional_int(value: Any) -> int | None: + if value is None or isinstance(value, bool): + return None + if isinstance(value, int): + return value + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _optional_float(value: Any) -> float | None: + if value is None or isinstance(value, bool): + return None + if isinstance(value, (int, float)): + return float(value) + try: + return float(value) + except (TypeError, ValueError): + return None diff --git a/tests/test_agent_context.py b/tests/test_agent_context.py index 316b2ff..7ed34f7 100644 --- a/tests/test_agent_context.py +++ b/tests/test_agent_context.py @@ -40,6 +40,23 @@ class AgentContextTests(unittest.TestCase): self.assertEqual(snapshot.system_context['cacheBreaker'], '[CACHE_BREAKER: debug-token]') + def test_user_context_loads_plugin_cache_snapshot(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) / 'repo' + workspace.mkdir(parents=True) + plugin_cache = workspace / '.port_sessions' / 'plugin_cache.json' + plugin_cache.parent.mkdir(parents=True, exist_ok=True) + plugin_cache.write_text( + '{"plugins":[{"name":"demo-plugin","version":"1.2.3","enabled":true}]}', + encoding='utf-8', + ) + + snapshot = build_context_snapshot(AgentRuntimeConfig(cwd=workspace)) + + self.assertIn('pluginCache', snapshot.user_context) + self.assertIn('demo-plugin', snapshot.user_context['pluginCache']) + self.assertIn('1.2.3', snapshot.user_context['pluginCache']) + @unittest.skipIf(shutil.which('git') is None, 'git is required for git context tests') def test_git_status_snapshot_contains_branch_and_status(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: diff --git a/tests/test_agent_prompting.py b/tests/test_agent_prompting.py index 139a8b2..327e5f5 100644 --- a/tests/test_agent_prompting.py +++ b/tests/test_agent_prompting.py @@ -58,3 +58,24 @@ class AgentPromptingTests(unittest.TestCase): self.assertIn('Claw Code Python', prompt) self.assertIn('# System', prompt) self.assertIn('# Environment', prompt) + + def test_prompt_builder_mentions_plugins_when_cache_is_loaded(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + plugin_cache = workspace / '.port_sessions' / 'plugin_cache.json' + plugin_cache.parent.mkdir(parents=True, exist_ok=True) + plugin_cache.write_text( + '{"plugins":[{"name":"example-plugin","enabled":true}]}', + encoding='utf-8', + ) + runtime_config = AgentRuntimeConfig(cwd=workspace) + model_config = ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct') + prompt_context = build_prompt_context(runtime_config, model_config) + parts = build_system_prompt_parts( + prompt_context=prompt_context, + runtime_config=runtime_config, + tools=default_tool_registry(), + ) + + prompt = render_system_prompt(parts) + self.assertIn('# Plugins', prompt) diff --git a/tests/test_agent_runtime.py b/tests/test_agent_runtime.py index e70b0b7..65c718b 100644 --- a/tests/test_agent_runtime.py +++ b/tests/test_agent_runtime.py @@ -8,7 +8,13 @@ from unittest.mock import patch from src.agent_runtime import LocalCodingAgent from src.agent_tools import build_tool_context, default_tool_registry, execute_tool -from src.agent_types import AgentRuntimeConfig, ModelConfig +from src.agent_types import ( + AgentPermissions, + AgentRuntimeConfig, + BudgetConfig, + ModelConfig, + OutputSchemaConfig, +) from src.openai_compat import OpenAICompatClient from src.session_store import load_agent_session @@ -27,6 +33,27 @@ class FakeHTTPResponse: return None +class FakeStreamingHTTPResponse: + def __init__(self, payloads: list[dict[str, object]]) -> None: + self.lines: list[bytes] = [] + for payload in payloads: + chunk = f'data: {json.dumps(payload)}\n\n' + self.lines.extend(part.encode('utf-8') for part in chunk.splitlines(keepends=True)) + done_chunk = 'data: [DONE]\n\n' + self.lines.extend(part.encode('utf-8') for part in done_chunk.splitlines(keepends=True)) + + def readline(self) -> bytes: + if not self.lines: + return b'' + return self.lines.pop(0) + + def __enter__(self) -> 'FakeStreamingHTTPResponse': + return self + + def __exit__(self, exc_type, exc, tb) -> None: + return None + + def make_urlopen_side_effect(responses: list[dict[str, object]]): queued = [FakeHTTPResponse(payload) for payload in responses] @@ -50,6 +77,31 @@ def make_recording_urlopen_side_effect( return _fake_urlopen +def make_streaming_urlopen_side_effect( + responses: list[list[dict[str, object]]], +): + queued = [FakeStreamingHTTPResponse(payloads) for payloads in responses] + + def _fake_urlopen(request_obj, timeout=None): # noqa: ANN001 + return queued.pop(0) + + return _fake_urlopen + + +def make_recording_streaming_urlopen_side_effect( + responses: list[list[dict[str, object]]], + recorded_payloads: list[dict[str, object]], +): + queued = [FakeStreamingHTTPResponse(payloads) for payloads in responses] + + def _fake_urlopen(request_obj, timeout=None): # noqa: ANN001 + body = request_obj.data.decode('utf-8') + recorded_payloads.append(json.loads(body)) + return queued.pop(0) + + return _fake_urlopen + + class AgentRuntimeTests(unittest.TestCase): def test_openai_client_parses_tool_calls(self) -> None: responses = [ @@ -91,6 +143,43 @@ class AgentRuntimeTests(unittest.TestCase): self.assertEqual(turn.tool_calls[0].name, 'read_file') self.assertEqual(turn.tool_calls[0].arguments['path'], 'hello.txt') + def test_openai_client_streams_content_and_usage(self) -> None: + responses = [ + [ + {'choices': [{'delta': {'content': 'Hello '}, 'finish_reason': None}]}, + {'choices': [{'delta': {'content': 'world'}, 'finish_reason': None}]}, + { + 'choices': [{'delta': {}, 'finish_reason': 'stop'}], + 'usage': {'prompt_tokens': 10, 'completion_tokens': 3}, + }, + ] + ] + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_streaming_urlopen_side_effect(responses), + ): + client = OpenAICompatClient( + ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ) + ) + events = list( + client.stream( + messages=[{'role': 'user', 'content': 'say hello'}], + tools=[], + ) + ) + self.assertEqual(events[0].type, 'message_start') + self.assertEqual( + ''.join(event.delta for event in events if event.type == 'content_delta'), + 'Hello world', + ) + usage_events = [event for event in events if event.type == 'usage'] + self.assertEqual(len(usage_events), 1) + self.assertEqual(usage_events[0].usage.input_tokens, 10) + self.assertEqual(usage_events[0].usage.output_tokens, 3) + def test_agent_executes_tool_calls_against_fake_backend(self) -> None: responses = [ { @@ -142,6 +231,7 @@ class AgentRuntimeTests(unittest.TestCase): self.assertEqual(result.final_output, 'The file contains hello world.') self.assertEqual(result.tool_calls, 1) self.assertGreaterEqual(len(result.transcript), 5) + self.assertGreaterEqual(len(result.file_history), 0) def test_write_tool_is_blocked_without_permission(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: @@ -236,3 +326,1360 @@ class AgentRuntimeTests(unittest.TestCase): self.assertIn('Start task', contents) self.assertIn('Initial answer.', contents) self.assertIn('Continue the task', contents) + + def test_agent_streams_runtime_output_and_usage(self) -> None: + responses = [ + [ + {'choices': [{'delta': {'content': 'Streaming '}, 'finish_reason': None}]}, + {'choices': [{'delta': {'content': 'works.'}, 'finish_reason': None}]}, + { + 'choices': [{'delta': {}, 'finish_reason': 'stop'}], + 'usage': {'prompt_tokens': 14, 'completion_tokens': 5}, + }, + ] + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_streaming_urlopen_side_effect(responses), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + stream_model_responses=True, + ), + ) + result = agent.run('Say streaming works') + self.assertEqual(result.final_output, 'Streaming works.') + self.assertEqual(result.usage.input_tokens, 14) + self.assertEqual(result.usage.output_tokens, 5) + self.assertTrue(any(event.get('type') == 'content_delta' for event in result.events)) + self.assertIsNotNone(result.scratchpad_directory) + assert result.scratchpad_directory is not None + self.assertTrue(Path(result.scratchpad_directory).is_dir()) + + def test_agent_streams_tool_calls_and_reconstructs_arguments(self) -> None: + responses = [ + [ + { + 'choices': [ + { + 'delta': { + 'tool_calls': [ + { + 'index': 0, + 'id': 'call_1', + 'function': { + 'name': 'read_file', + 'arguments': '{"path": "hello', + }, + } + ] + }, + 'finish_reason': None, + } + ] + }, + { + 'choices': [ + { + 'delta': { + 'tool_calls': [ + { + 'index': 0, + 'function': { + 'arguments': '.txt"}', + }, + } + ] + }, + 'finish_reason': None, + } + ] + }, + { + 'choices': [{'delta': {}, 'finish_reason': 'tool_calls'}], + 'usage': {'prompt_tokens': 9, 'completion_tokens': 4}, + }, + ], + [ + {'choices': [{'delta': {'content': 'Read done.'}, 'finish_reason': None}]}, + { + 'choices': [{'delta': {}, 'finish_reason': 'stop'}], + 'usage': {'prompt_tokens': 11, 'completion_tokens': 2}, + }, + ], + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + (workspace / 'hello.txt').write_text('hello world\n', encoding='utf-8') + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_streaming_urlopen_side_effect(responses), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + stream_model_responses=True, + ), + ) + result = agent.run('Inspect hello.txt') + self.assertEqual(result.final_output, 'Read done.') + self.assertEqual(result.tool_calls, 1) + self.assertEqual(result.usage.input_tokens, 20) + self.assertEqual(result.usage.output_tokens, 6) + assistant_messages = [message for message in result.transcript if message.get('role') == 'assistant'] + self.assertTrue(any(message.get('tool_calls') for message in assistant_messages)) + + def test_transcript_entries_include_structured_blocks(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'I will inspect the file.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'read_file', + 'arguments': '{"path": "hello.txt"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 6, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Done reading.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + (workspace / 'hello.txt').write_text('hello world\n', encoding='utf-8') + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + result = agent.run('Inspect hello.txt') + assistant_with_tool = next( + message + for message in result.transcript + if message.get('role') == 'assistant' and message.get('tool_calls') + ) + self.assertIn('blocks', assistant_with_tool) + block_types = [block.get('type') for block in assistant_with_tool['blocks']] + self.assertIn('text', block_types) + self.assertIn('tool_call', block_types) + tool_message = next(message for message in result.transcript if message.get('role') == 'tool') + self.assertIn('blocks', tool_message) + self.assertEqual(tool_message['blocks'][0]['type'], 'tool_result') + + def test_agent_inserts_compact_boundary_when_threshold_is_exceeded(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'I will inspect the file and then continue.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'read_file', + 'arguments': '{"path": "hello.txt"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 30, 'completion_tokens': 10}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Compaction test completed.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 12, 'completion_tokens': 4}, + }, + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + (workspace / 'hello.txt').write_text('hello world\n', encoding='utf-8') + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + auto_compact_threshold_tokens=80, + compact_preserve_messages=1, + ), + ) + result = agent.run( + 'Read hello.txt and then continue with a detailed explanation that is intentionally long enough to trigger compaction.' + ) + self.assertEqual(result.final_output, 'Compaction test completed.') + compact_events = [event for event in result.events if event.get('type') == 'compact_boundary'] + self.assertEqual(len(compact_events), 1) + second_request_messages = recorded_payloads[1]['messages'] + assert isinstance(second_request_messages, list) + compact_messages = [ + message for message in second_request_messages + if isinstance(message, dict) + and isinstance(message.get('content'), str) + and 'Earlier conversation history was compacted' in message['content'] + ] + self.assertEqual(len(compact_messages), 1) + transcript_compact_messages = [ + message for message in result.transcript + if message.get('metadata', {}).get('kind') == 'compact_boundary' + ] + self.assertEqual(len(transcript_compact_messages), 1) + + def test_agent_enforces_total_token_budget(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'This would be the answer.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': { + 'prompt_tokens': 30, + 'completion_tokens': 12, + }, + } + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + budget_config=BudgetConfig(max_total_tokens=20), + ), + ) + result = agent.run('Use too many tokens') + self.assertEqual(result.stop_reason, 'budget_exceeded') + self.assertIn('token budget', result.final_output) + self.assertEqual(result.usage.total_tokens, 42) + + def test_agent_continues_when_model_response_is_truncated(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Part 1 ', + }, + 'finish_reason': 'length', + } + ], + 'usage': {'prompt_tokens': 10, 'completion_tokens': 4}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Part 2', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 4, 'completion_tokens': 2}, + }, + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + result = agent.run('Give me a long answer') + self.assertEqual(result.final_output, 'Part 1 Part 2') + continuation_events = [ + event for event in result.events if event.get('type') == 'continuation_request' + ] + self.assertEqual(len(continuation_events), 1) + second_request_messages = recorded_payloads[1]['messages'] + assert isinstance(second_request_messages, list) + self.assertTrue( + any( + isinstance(message, dict) + and 'Continue exactly where you left off' in str(message.get('content', '')) + for message in second_request_messages + ) + ) + + def test_agent_records_file_history_for_write_tool(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Creating the file.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'write_file', + 'arguments': '{"path": "out.txt", "content": "hi"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 4, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Done.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + session_dir = workspace / '.port_sessions' / 'agent' + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + session_directory=session_dir, + permissions=AgentPermissions(allow_file_write=True), + ), + ) + result = agent.run('Create out.txt') + stored = load_agent_session(result.session_id or '', directory=session_dir) + self.assertEqual(len(result.file_history), 1) + self.assertEqual(result.file_history[0]['path'], 'out.txt') + self.assertEqual(stored.file_history[0]['action'], 'write_file') + + def test_agent_streams_bash_tool_output_and_mutates_tool_transcript(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Running a shell command.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'bash', + 'arguments': json.dumps( + { + 'command': ( + "printf 'alpha\\n'; " + "sleep 0.05; " + "printf 'beta\\n' >&2" + ) + } + ), + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 6, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Shell command completed.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + permissions=AgentPermissions(allow_shell_commands=True), + ), + ) + result = agent.run('Run a shell command') + self.assertEqual(result.final_output, 'Shell command completed.') + tool_delta_events = [event for event in result.events if event.get('type') == 'tool_delta'] + self.assertGreaterEqual(len(tool_delta_events), 2) + joined_delta = ''.join(event.get('delta', '') for event in tool_delta_events) + self.assertIn('alpha', joined_delta) + self.assertIn('beta', joined_delta) + tool_messages = [message for message in result.transcript if message.get('role') == 'tool'] + self.assertEqual(len(tool_messages), 1) + tool_message = tool_messages[0] + self.assertIn('exit_code=0', tool_message.get('content', '')) + metadata = tool_message.get('metadata', {}) + self.assertIn('stream_preview', metadata) + self.assertIn('alpha', metadata['stream_preview']) + self.assertIn('beta', metadata['stream_preview']) + + def test_agent_streams_read_file_tool_output_in_chunks(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Reading the file.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'read_file', + 'arguments': '{"path": "large.txt"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 6, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Read finished.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + (workspace / 'large.txt').write_text('alpha\n' * 300, encoding='utf-8') + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + result = agent.run('Read the large file') + self.assertEqual(result.final_output, 'Read finished.') + tool_delta_events = [ + event for event in result.events + if event.get('type') == 'tool_delta' and event.get('tool_name') == 'read_file' + ] + self.assertGreaterEqual(len(tool_delta_events), 2) + tool_messages = [message for message in result.transcript if message.get('role') == 'tool'] + self.assertEqual(len(tool_messages), 1) + self.assertEqual(tool_messages[0].get('metadata', {}).get('streamed'), True) + self.assertIn('stream_preview', tool_messages[0].get('metadata', {})) + + def test_agent_records_tombstone_mutation_history_when_snipping(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Reading the large file first.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'read_file', + 'arguments': '{"path": "large.txt"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Snip run completed.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + (workspace / 'large.txt').write_text(('alpha beta gamma\n' * 400), encoding='utf-8') + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + auto_snip_threshold_tokens=120, + compact_preserve_messages=0, + ), + ) + result = agent.run('Read the large file and summarize it') + tool_messages = [message for message in result.transcript if message.get('role') == 'tool'] + self.assertEqual(len(tool_messages), 1) + metadata = tool_messages[0].get('metadata', {}) + self.assertEqual(tool_messages[0].get('state'), 'tombstoned') + self.assertEqual(metadata.get('kind'), 'snipped_message') + self.assertEqual(metadata.get('last_mutation_kind'), 'snip_tombstone') + self.assertGreaterEqual(metadata.get('mutation_count', 0), 2) + self.assertTrue(any(entry.get('kind') == 'tool_finalize_replace' for entry in metadata.get('mutations', []))) + self.assertTrue(any(event.get('type') == 'snip_boundary' for event in result.events)) + + def test_resume_injects_file_history_replay_reminder(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Creating the file first.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'write_file', + 'arguments': '{"path": "replay.txt", "content": "hello"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Initial write done.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 6, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Resume acknowledged.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + session_dir = workspace / '.port_sessions' / 'agent' + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + session_directory=session_dir, + permissions=AgentPermissions(allow_file_write=True), + ), + ) + first_result = agent.run('Create replay.txt') + stored = load_agent_session(first_result.session_id or '', directory=session_dir) + + resumed_agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + session_directory=session_dir, + permissions=AgentPermissions(allow_file_write=True), + ), + ) + resumed_agent.resume('Continue the work', stored) + resumed_messages = recorded_payloads[-1]['messages'] + assert isinstance(resumed_messages, list) + replay_messages = [ + message for message in resumed_messages + if isinstance(message, dict) + and isinstance(message.get('content'), str) + and 'Recent file history from this saved session:' in message['content'] + ] + self.assertEqual(len(replay_messages), 1) + self.assertIn('path=replay.txt', replay_messages[0]['content']) + self.assertIn('action=write_file', replay_messages[0]['content']) + + def test_resume_replays_file_history_snapshots(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Editing the file.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'edit_file', + 'arguments': json.dumps( + { + 'path': 'draft.txt', + 'old_text': 'hello world', + 'new_text': 'hello mars', + } + ), + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Edit completed.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 6, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Resume processed.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + (workspace / 'draft.txt').write_text('hello world\n', encoding='utf-8') + session_dir = workspace / '.port_sessions' / 'agent' + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + session_directory=session_dir, + permissions=AgentPermissions(allow_file_write=True), + ), + ) + first_result = agent.run('Edit draft.txt') + stored = load_agent_session(first_result.session_id or '', directory=session_dir) + resumed_agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + session_directory=session_dir, + permissions=AgentPermissions(allow_file_write=True), + ), + ) + resumed_agent.resume('Continue after the edit', stored) + resumed_messages = recorded_payloads[-1]['messages'] + assert isinstance(resumed_messages, list) + replay_messages = [ + message for message in resumed_messages + if isinstance(message, dict) + and isinstance(message.get('content'), str) + and 'Recent file history from this saved session:' in message['content'] + ] + self.assertEqual(len(replay_messages), 1) + replay_content = replay_messages[0]['content'] + self.assertIn('before: hello world', replay_content) + self.assertIn('after: hello mars', replay_content) + + def test_resume_injects_compaction_replay_reminder(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Reading the large file first.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'read_file', + 'arguments': '{"path": "large.txt"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Initial compaction run completed.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Resume after compaction processed.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + (workspace / 'large.txt').write_text(('alpha beta gamma\n' * 400), encoding='utf-8') + session_dir = workspace / '.port_sessions' / 'agent' + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + session_directory=session_dir, + auto_snip_threshold_tokens=120, + compact_preserve_messages=0, + ), + ) + first_result = agent.run('Read the large file and summarize it') + stored = load_agent_session(first_result.session_id or '', directory=session_dir) + resumed_agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + session_directory=session_dir, + auto_snip_threshold_tokens=120, + compact_preserve_messages=0, + ), + ) + resumed_agent.resume('Continue after compaction', stored) + resumed_messages = recorded_payloads[-1]['messages'] + assert isinstance(resumed_messages, list) + compaction_messages = [ + message for message in resumed_messages + if isinstance(message, dict) + and isinstance(message.get('content'), str) + and 'This resumed session already contains compacted or snipped history.' in message['content'] + ] + self.assertEqual(len(compaction_messages), 1) + self.assertIn('Snipped/tombstoned messages:', compaction_messages[0]['content']) + + def test_agent_can_delegate_to_nested_agent(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Delegating this task.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'delegate_agent', + 'arguments': json.dumps( + { + 'prompt': 'Summarize the delegated task.', + 'max_turns': 2, + } + ), + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Delegated summary complete.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Parent task completed after delegation.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + result = agent.run('Use a delegate agent') + self.assertEqual(result.final_output, 'Parent task completed after delegation.') + tool_messages = [message for message in result.transcript if message.get('role') == 'tool'] + self.assertEqual(len(tool_messages), 1) + self.assertIn('Delegated agent completed the subtask.', tool_messages[0]['content']) + metadata = tool_messages[0].get('metadata', {}) + self.assertEqual(metadata.get('action'), 'delegate_agent') + self.assertIsNotNone(metadata.get('child_session_id')) + + def test_agent_can_delegate_multiple_subtasks_with_parent_context(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Delegating multiple subtasks.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'delegate_agent', + 'arguments': json.dumps( + { + 'subtasks': [ + {'label': 'scan', 'prompt': 'Scan the project.'}, + {'label': 'summarize', 'prompt': 'Summarize the project.'}, + ], + 'max_turns': 2, + 'include_parent_context': True, + } + ), + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Child scan result.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Child summary result.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 6, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Parent completed after multi-delegate.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + result = agent.run('Use multiple delegated subtasks') + self.assertEqual(result.final_output, 'Parent completed after multi-delegate.') + tool_messages = [message for message in result.transcript if message.get('role') == 'tool'] + self.assertEqual(len(tool_messages), 1) + metadata = tool_messages[0].get('metadata', {}) + self.assertEqual(metadata.get('subtask_count'), 2) + self.assertEqual(len(metadata.get('child_session_ids', [])), 2) + self.assertEqual(len(metadata.get('child_results', [])), 2) + self.assertIn('Delegated agent completed 2 sequential subtasks.', tool_messages[0].get('content', '')) + second_child_request = recorded_payloads[2]['messages'] + assert isinstance(second_child_request, list) + self.assertTrue( + any( + isinstance(message, dict) + and 'Prior delegated subtask summaries:' in str(message.get('content', '')) + for message in second_child_request + ) + ) + + def test_agent_manager_tracks_delegate_group_membership(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Delegating multiple subtasks.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'delegate_agent', + 'arguments': json.dumps( + { + 'subtasks': [ + {'label': 'scan', 'prompt': 'Scan the project.'}, + {'label': 'summarize', 'prompt': 'Summarize the project.'}, + ], + 'max_turns': 2, + } + ), + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Child scan result.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Child summary result.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 6, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Parent completed after multi-delegate.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + result = agent.run('Use multiple delegated subtasks') + + self.assertEqual(result.final_output, 'Parent completed after multi-delegate.') + self.assertIsNotNone(agent.agent_manager) + manager = agent.agent_manager + assert manager is not None + self.assertEqual(len(manager.groups), 1) + group = next(iter(manager.groups.values())) + self.assertEqual(group.completed_children, 2) + child_records = sorted( + ( + record for record in manager.completed_records() + if record.parent_agent_id == agent.managed_agent_id + ), + key=lambda record: (record.child_index or 0), + ) + self.assertEqual(len(child_records), 2) + self.assertEqual([record.child_index for record in child_records], [1, 2]) + self.assertTrue(all(record.group_id == group.group_id for record in child_records)) + summary = '\n'.join(manager.summary_lines()) + self.assertIn(f'group={group.group_id}', summary) + + def test_agent_enforces_reasoning_token_budget(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'This uses too much reasoning.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': { + 'prompt_tokens': 8, + 'completion_tokens': 4, + 'completion_tokens_details': {'reasoning_tokens': 9}, + }, + } + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + budget_config=BudgetConfig(max_reasoning_tokens=5), + ), + ) + result = agent.run('Use too much reasoning') + self.assertEqual(result.stop_reason, 'budget_exceeded') + self.assertIn('reasoning token budget', result.final_output) + self.assertEqual(result.usage.reasoning_tokens, 9) + + def test_agent_enforces_tool_call_budget_before_execution(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'I will create the file.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'write_file', + 'arguments': '{"path": "blocked.txt", "content": "nope"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 4, 'completion_tokens': 2}, + } + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + permissions=AgentPermissions(allow_file_write=True), + budget_config=BudgetConfig(max_tool_calls=0), + ), + ) + result = agent.run('Try to create a file') + self.assertFalse((workspace / 'blocked.txt').exists()) + self.assertEqual(result.stop_reason, 'budget_exceeded') + self.assertIn('tool-call budget', result.final_output) + self.assertFalse(any(message.get('role') == 'tool' for message in result.transcript)) + + def test_agent_enforces_delegated_task_budget_before_child_agent_runs(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'I will delegate this.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'delegate_agent', + 'arguments': json.dumps( + { + 'prompt': 'Do the delegated work.', + 'max_turns': 2, + } + ), + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 4, 'completion_tokens': 2}, + } + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + budget_config=BudgetConfig(max_delegated_tasks=0), + ), + ) + result = agent.run('Try to delegate') + self.assertEqual(result.stop_reason, 'budget_exceeded') + self.assertIn('delegated-task budget', result.final_output) + self.assertFalse(any(message.get('role') == 'tool' for message in result.transcript)) + + def test_agent_sends_response_schema_when_configured(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': '{"status":"ok"}', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 4}, + } + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + output_schema=OutputSchemaConfig( + name='status_response', + schema={ + 'type': 'object', + 'properties': {'status': {'type': 'string'}}, + 'required': ['status'], + }, + strict=True, + ), + ), + ) + agent.run('Return a JSON status payload') + self.assertEqual( + recorded_payloads[0]['response_format'], + { + 'type': 'json_schema', + 'json_schema': { + 'name': 'status_response', + 'schema': { + 'type': 'object', + 'properties': {'status': {'type': 'string'}}, + 'required': ['status'], + }, + 'strict': True, + }, + }, + ) diff --git a/tests/test_query_engine_runtime.py b/tests/test_query_engine_runtime.py new file mode 100644 index 0000000..2cf531b --- /dev/null +++ b/tests/test_query_engine_runtime.py @@ -0,0 +1,497 @@ +from __future__ import annotations + +import json +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +from src.agent_runtime import LocalCodingAgent +from src.agent_types import AgentRuntimeConfig, ModelConfig +from src.openai_compat import OpenAICompatClient +from src.plugin_runtime import PluginRuntime +from src.query_engine import QueryEnginePort + + +class FakeHTTPResponse: + def __init__(self, payload: dict[str, object]) -> None: + self.payload = payload + + def read(self) -> bytes: + return json.dumps(self.payload).encode('utf-8') + + def __enter__(self) -> 'FakeHTTPResponse': + return self + + def __exit__(self, exc_type, exc, tb) -> None: + return None + + +def make_recording_urlopen_side_effect( + responses: list[dict[str, object]], + recorded_payloads: list[dict[str, object]], +): + queued = [FakeHTTPResponse(payload) for payload in responses] + + def _fake_urlopen(request_obj, timeout=None): # noqa: ANN001 + body = request_obj.data.decode('utf-8') + recorded_payloads.append(json.loads(body)) + return queued.pop(0) + + return _fake_urlopen + + +def make_urlopen_side_effect(responses: list[dict[str, object]]): + queued = [FakeHTTPResponse(payload) for payload in responses] + + def _fake_urlopen(request_obj, timeout=None): # noqa: ANN001 + return queued.pop(0) + + return _fake_urlopen + + +class QueryEngineRuntimeTests(unittest.TestCase): + def test_plugin_runtime_discovers_local_manifest(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + plugin_dir = workspace / 'plugins' / 'demo' + plugin_dir.mkdir(parents=True) + (plugin_dir / 'plugin.json').write_text( + json.dumps( + { + 'name': 'demo-plugin', + 'version': '0.1.0', + 'description': 'Demo plugin', + 'tools': ['demo_tool'], + 'hooks': { + 'beforePrompt': 'Run plugin hook before prompt.', + 'afterTurn': 'Plugin after-turn hook.', + }, + 'toolAliases': [ + { + 'name': 'plugin_read', + 'baseTool': 'read_file', + 'description': 'Plugin read alias', + } + ], + } + ), + encoding='utf-8', + ) + runtime = PluginRuntime.from_workspace(workspace) + + self.assertEqual(len(runtime.manifests), 1) + self.assertEqual(runtime.manifests[0].name, 'demo-plugin') + self.assertEqual(runtime.manifests[0].tool_names, ('demo_tool',)) + self.assertIn('beforePrompt', runtime.manifests[0].hook_names) + self.assertIn('afterTurn', runtime.manifests[0].hook_names) + self.assertEqual(runtime.manifests[0].tool_aliases[0].name, 'plugin_read') + self.assertEqual(runtime.manifests[0].before_prompt, 'Run plugin hook before prompt.') + self.assertEqual(runtime.manifests[0].after_turn, 'Plugin after-turn hook.') + + def test_query_engine_can_drive_real_runtime_agent(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Initial runtime answer.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Resumed runtime answer.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 6, 'completion_tokens': 2}, + }, + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + plugin_dir = workspace / '.codex-plugin' + plugin_dir.mkdir(parents=True) + (plugin_dir / 'plugin.json').write_text( + json.dumps({'name': 'runtime-plugin', 'tools': ['runtime_tool']}), + encoding='utf-8', + ) + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + engine = QueryEnginePort.from_runtime_agent(agent) + first = engine.submit_message('Start the task') + second = engine.submit_message('Continue the task') + summary = engine.render_summary() + + self.assertEqual(first.output, 'Initial runtime answer.') + self.assertEqual(second.output, 'Resumed runtime answer.') + self.assertEqual(first.session_id, second.session_id) + self.assertEqual(second.usage.input_tokens, 6) + self.assertIn('Real runtime agent mode: True', summary) + self.assertIn('## Agent Manager', summary) + self.assertIn('runtime-plugin', summary) + self.assertEqual(len(recorded_payloads), 2) + resumed_messages = recorded_payloads[1]['messages'] + assert isinstance(resumed_messages, list) + contents = [message.get('content') for message in resumed_messages if isinstance(message, dict)] + self.assertIn('Start the task', contents) + self.assertIn('Initial runtime answer.', contents) + self.assertIn('Continue the task', contents) + + def test_runtime_agent_uses_plugin_aliases_and_hooks(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Using plugin alias.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'plugin_read', + 'arguments': '{"path": "hello.txt"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Plugin alias completed.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + (workspace / 'hello.txt').write_text('hello plugin\n', encoding='utf-8') + plugin_dir = workspace / 'plugins' / 'demo' + plugin_dir.mkdir(parents=True) + (plugin_dir / 'plugin.json').write_text( + json.dumps( + { + 'name': 'demo-plugin', + 'hooks': { + 'beforePrompt': 'Run plugin hook before prompt.', + 'afterTurn': 'Plugin after-turn hook.', + }, + 'toolAliases': [ + { + 'name': 'plugin_read', + 'baseTool': 'read_file', + 'description': 'Plugin read alias', + } + ], + } + ), + encoding='utf-8', + ) + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + result = agent.run('Read the file through the plugin alias') + + self.assertEqual(result.final_output, 'Plugin alias completed.') + self.assertTrue(any(event.get('type') == 'plugin_after_turn' for event in result.events)) + tool_names = [ + item['function']['name'] + for item in recorded_payloads[0]['tools'] + if isinstance(item, dict) and isinstance(item.get('function'), dict) + ] + self.assertIn('plugin_read', tool_names) + messages = recorded_payloads[0]['messages'] + assert isinstance(messages, list) + self.assertTrue( + any( + isinstance(message, dict) + and 'Run plugin hook before prompt.' in str(message.get('content', '')) + for message in messages + ) + ) + + def test_runtime_agent_injects_plugin_tool_runtime_guidance(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Reading through plugin guidance.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'read_file', + 'arguments': '{"path": "guide.txt"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Plugin runtime guidance consumed.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + (workspace / 'guide.txt').write_text('plugin guidance\n', encoding='utf-8') + plugin_dir = workspace / 'plugins' / 'demo' + plugin_dir.mkdir(parents=True) + (plugin_dir / 'plugin.json').write_text( + json.dumps( + { + 'name': 'demo-plugin', + 'toolHooks': { + 'read_file': { + 'afterResult': 'Summarize the file before making edits.', + } + }, + } + ), + encoding='utf-8', + ) + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + result = agent.run('Read the file and continue') + + self.assertEqual(result.final_output, 'Plugin runtime guidance consumed.') + self.assertTrue(any(event.get('type') == 'plugin_tool_context' for event in result.events)) + runtime_messages = [ + message for message in result.transcript + if message.get('metadata', {}).get('kind') == 'plugin_tool_runtime' + ] + self.assertEqual(len(runtime_messages), 1) + self.assertIn('Summarize the file before making edits.', runtime_messages[0].get('content', '')) + second_messages = recorded_payloads[1]['messages'] + assert isinstance(second_messages, list) + self.assertTrue( + any( + isinstance(message, dict) + and 'Plugin tool runtime guidance for `read_file`:' in str(message.get('content', '')) + and 'Summarize the file before making edits.' in str(message.get('content', '')) + for message in second_messages + ) + ) + + def test_runtime_agent_blocks_tool_via_plugin_manifest(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Trying a blocked shell command.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'bash', + 'arguments': '{"command": "pwd"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Blocked tool handled.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + plugin_dir = workspace / 'plugins' / 'demo' + plugin_dir.mkdir(parents=True) + (plugin_dir / 'plugin.json').write_text( + json.dumps( + { + 'name': 'demo-plugin', + 'blockedTools': ['bash'], + } + ), + encoding='utf-8', + ) + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_recording_urlopen_side_effect(responses, recorded_payloads), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + result = agent.run('Try a blocked tool') + + self.assertEqual(result.final_output, 'Blocked tool handled.') + self.assertTrue(any(event.get('type') == 'plugin_tool_block' for event in result.events)) + self.assertTrue(any(event.get('type') == 'plugin_tool_context' for event in result.events)) + tool_messages = [message for message in result.transcript if message.get('role') == 'tool'] + self.assertEqual(len(tool_messages), 1) + metadata = tool_messages[0].get('metadata', {}) + self.assertEqual(metadata.get('action'), 'plugin_block') + self.assertEqual(metadata.get('plugin_blocked'), True) + second_messages = recorded_payloads[1]['messages'] + assert isinstance(second_messages, list) + self.assertTrue( + any( + isinstance(message, dict) + and 'Plugin tool runtime guidance for `bash`:' in str(message.get('content', '')) + and 'blocked tool bash' in str(message.get('content', '')).lower() + for message in second_messages + ) + ) + + def test_query_engine_runtime_summary_tracks_runtime_events(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Reading through plugin guidance.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'read_file', + 'arguments': '{"path": "guide.txt"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Summary ready.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + (workspace / 'guide.txt').write_text('runtime summary\n', encoding='utf-8') + plugin_dir = workspace / 'plugins' / 'demo' + plugin_dir.mkdir(parents=True) + (plugin_dir / 'plugin.json').write_text( + json.dumps( + { + 'name': 'demo-plugin', + 'toolHooks': { + 'read_file': {'afterResult': 'Summarize the file before editing it.'} + }, + } + ), + encoding='utf-8', + ) + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_urlopen_side_effect(responses), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + engine = QueryEnginePort.from_runtime_agent(agent) + turn = engine.submit_message('Read the file and summarize it') + summary = engine.render_summary() + + self.assertEqual(turn.output, 'Summary ready.') + self.assertIn('## Runtime Events', summary) + self.assertIn('- plugin_tool_context=1', summary) + self.assertIn('- tool_result=1', summary) + self.assertIn('## Runtime Message Kinds', summary) + self.assertIn('- plugin_tool_runtime=1', summary) + self.assertIn('- transcript_messages=', summary)