Merge pull request #2 from HarnessLab/dev

add new agent components
This commit is contained in:
Unknown
2026-04-02 21:13:26 +02:00
committed by GitHub
19 changed files with 6183 additions and 296 deletions
+42 -12
View File
@@ -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:
+4
View File
@@ -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',
+39 -5
View File
@@ -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
+214
View File
@@ -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
+154
View File
@@ -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,
)
+24 -1
View File
@@ -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}',
+1565 -39
View File
File diff suppressed because it is too large Load Diff
+390 -3
View File
@@ -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],
+333 -8
View File
@@ -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,
),
)
+123
View File
@@ -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
+193 -1
View File
@@ -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:
+280 -29
View File
@@ -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,
)
+319
View File
@@ -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)
+213 -28
View File
@@ -2,12 +2,15 @@ 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_session, save_session
from .session_store import StoredSession, load_agent_session, load_session, save_session
from .tools import build_tool_backlog
from .transcript import TranscriptStore
@@ -19,6 +22,7 @@ class QueryEngineConfig:
compact_after_turns: int = 12
structured_output: bool = False
structured_retry_limit: int = 2
use_runtime_agent: bool = False
@dataclass(frozen=True)
@@ -30,6 +34,12 @@ class TurnResult:
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
@@ -41,10 +51,19 @@ class QueryEnginePort:
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())
return cls(
manifest=build_port_manifest(),
plugin_runtime=PluginRuntime.from_workspace(Path.cwd()),
)
@classmethod
def from_saved_session(cls, session_id: str) -> 'QueryEnginePort':
@@ -56,6 +75,25 @@ class QueryEnginePort:
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(
@@ -65,6 +103,29 @@ class QueryEnginePort:
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(
@@ -88,20 +149,18 @@ class QueryEnginePort:
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(
turn = TurnResult(
prompt=prompt,
output=output,
matched_commands=matched_commands,
matched_tools=matched_tools,
permission_denials=denied_tools,
usage=self.total_usage,
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,
@@ -118,6 +177,20 @@ class QueryEnginePort:
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',
@@ -138,6 +211,8 @@ class QueryEnginePort:
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(
@@ -149,25 +224,6 @@ class QueryEnginePort:
)
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()
@@ -189,5 +245,134 @@ class QueryEnginePort:
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)
+139 -1
View File
@@ -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
+17
View File
@@ -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:
+21
View File
@@ -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)
File diff suppressed because it is too large Load Diff
+497
View File
@@ -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)