diff --git a/PARITY_CHECKLIST.md b/PARITY_CHECKLIST.md index dcb3372..18357dc 100644 --- a/PARITY_CHECKLIST.md +++ b/PARITY_CHECKLIST.md @@ -28,6 +28,7 @@ Done: - [x] Mutable tool transcript updates during tool execution - [x] Transcript mutation history for replaced/tombstoned messages - [x] Assistant streaming and tool-call transcript mutation history +- [x] Session-wide mutation serial tracking across transcript updates - [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 @@ -39,32 +40,42 @@ Done: - [x] Reactive compaction retry after prompt-too-long backend failures - [x] Reasoning-token budget enforcement - [x] Tool-call and delegated-task budget enforcement +- [x] Resume-aware cumulative model-call budgets +- [x] Resume-aware cumulative session usage/cost persistence - [x] Basic nested-agent delegation tool - [x] Sequential multi-subtask delegation with parent-context carryover +- [x] Dependency-aware delegated subtasks +- [x] Topological dependency-batch delegation planning - [x] Basic agent-manager lineage tracking for nested agents - [x] Managed agent-group membership tracking with child indices +- [x] Agent-manager strategy and batch summary tracking for delegated groups - [x] Delegated child-session resume by saved session id - [x] Agent-manager tracking for resumed child-session lineage - [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 lifecycle hooks for resume, persist, and delegate phases - [x] Manifest-defined plugin tool aliases over base runtime tools - [x] Manifest-defined executable virtual tools - [x] Manifest-defined plugin tool blocking - [x] Manifest-defined plugin `beforeTool` guidance - [x] Manifest-defined plugin tool-result guidance injected back into the transcript +- [x] Plugin runtime session-state persistence and resume restoration - [x] Compaction metadata with compacted message ids - [x] Compaction metadata with preserved-tail ids and compaction depth - [x] Compaction metadata with compacted/preserved lineage ids and revision summaries +- [x] Compaction metadata with source mutation serials and mutation totals - [x] Snipped-message metadata with source role/kind lineage - [x] Snipped-message metadata with source lineage id and revision - [x] Resume-time compaction / snipping replay reminder +- [x] Resume-time compaction replay of source mutation summaries - [x] Query-engine facade that can drive the real Python runtime agent - [x] Query-engine runtime event counters and transcript-kind summaries - [x] Query-engine runtime mutation counters - [x] Query-engine stream-level runtime summary event - [x] Query-engine transcript-store compaction summaries - [x] Delegate-group and delegated-subtask runtime events +- [x] Delegate-batch runtime events and summaries - [x] Query-engine runtime orchestration summaries for group status and child stop reasons - [x] Query-engine runtime context-reduction summaries - [x] Query-engine runtime lineage summaries @@ -73,12 +84,12 @@ Done: Missing: - [ ] Full partial tool-result streaming parity across the complete upstream/npm tool surface -- [ ] Full rich transcript mutation behavior like the npm runtime -- [ ] Full reasoning budgets and task budgets parity -- [ ] Full multi-agent orchestration parity beyond sequential grouped delegation and resumed-child flows -- [ ] Full file history snapshots and replay flows beyond the current preview/id-based implementation -- [ ] Full executable plugin lifecycle beyond manifest-driven guidance, blocking, aliases, and virtual tools -- [ ] Full session compaction / snipping parity beyond lineage-aware summaries and replay reminders +- [ ] Full rich transcript mutation behavior like the npm runtime beyond the current lineage, counters, block export, and mutation-serial tracking +- [ ] Full reasoning budgets and task budgets parity beyond the current cumulative model/tool/delegation/session-call enforcement +- [ ] Full multi-agent orchestration parity beyond dependency-aware batched delegation, resumed-child flows, and current agent-manager summaries +- [ ] Full file history snapshots and replay flows beyond the current preview/id-based implementation and delegated-batch replay metadata +- [ ] Full executable plugin lifecycle beyond manifest-driven prompt/tool/session hooks, blocking, aliases, virtual tools, and persisted runtime state +- [ ] Full session compaction / snipping parity beyond lineage-aware summaries, mutation-serial compaction metadata, and replay reminders - [ ] Full `QueryEngine.ts` parity ## 2. CLI Entrypoints And Runtime Modes @@ -87,6 +98,7 @@ Done: - [x] Python CLI entrypoint - [x] `agent` command +- [x] `agent-chat` command - [x] `agent-resume` command - [x] `agent-prompt` command - [x] `agent-context` command @@ -292,11 +304,12 @@ Missing: Done: - [x] Non-interactive CLI execution +- [x] Basic interactive REPL-style agent chat loop - [x] Transcript printing for debugging Missing: -- [ ] Interactive REPL parity +- [ ] Interactive REPL parity beyond the current basic `agent-chat` loop - [ ] Ink/TUI component parity - [ ] Screen system parity - [ ] Keyboard interaction parity diff --git a/README.md b/README.md index b7d153c..200ed89 100644 --- a/README.md +++ b/README.md @@ -50,6 +50,10 @@ Built on the public porting workspace from [instructkr/claw-code](https://github ## 📋 Roadmap +### Testing + +- See [TESTING_GUIDE.md](TESTING_GUIDE.md) for concrete commands to verify the current implementation feature by feature. + ### Done - [x] Python CLI agent loop diff --git a/TESTING_GUIDE.md b/TESTING_GUIDE.md new file mode 100644 index 0000000..81680f6 --- /dev/null +++ b/TESTING_GUIDE.md @@ -0,0 +1,649 @@ +# Testing Guide + +This guide gives concrete commands you can run to verify the current Python implementation feature by feature. + +All commands below assume you are inside: + +```bash +cd /path/to/claw-code-agent +``` + +## 1. Environment Setup + +### 1.1 Start `vLLM` with Qwen3-Coder + +```bash +python -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen3-Coder-30B-A3B-Instruct \ + --host 127.0.0.1 \ + --port 8000 \ + --enable-auto-tool-choice \ + --tool-call-parser qwen3_xml +``` + +Verify the server: + +```bash +curl http://127.0.0.1:8000/v1/models +``` + +Set the runtime environment: + +```bash +export OPENAI_BASE_URL=http://127.0.0.1:8000/v1 +export OPENAI_API_KEY=local-token +export OPENAI_MODEL=Qwen/Qwen3-Coder-30B-A3B-Instruct +``` + +### 1.2 Run the unit test suite + +```bash +python3 -m unittest discover -s tests -v +``` + +## 2. Installation And Basic Usage + +### 2.1 Editable install + +```bash +pip install -e . +``` + +### 2.2 Show CLI help + +```bash +python3 -m src.main --help +python3 -m src.main agent --help +python3 -m src.main agent-chat --help +python3 -m src.main agent-resume --help +``` + +### 2.3 Run the packaged entrypoint + +```bash +claw-code-agent agent "/help" +``` + +## 3. Slash Commands + +These are handled locally and do not require the model to answer. + +```bash +python3 -m src.main agent "/help" +python3 -m src.main agent "/commands" +python3 -m src.main agent "/context" --cwd .. +python3 -m src.main agent "/context-raw" --cwd .. +python3 -m src.main agent "/prompt" --cwd .. +python3 -m src.main agent "/permissions" --cwd .. +python3 -m src.main agent "/tools" --cwd .. +python3 -m src.main agent "/memory" --cwd .. +python3 -m src.main agent "/status" --cwd .. +python3 -m src.main agent "/clear" --cwd .. +``` + +## 4. Context And Prompt Inspection + +### 4.1 System prompt rendering + +```bash +python3 -m src.main agent-prompt --cwd .. +``` + +### 4.2 Context usage accounting + +```bash +python3 -m src.main agent-context --cwd .. +``` + +### 4.3 Raw context snapshot + +```bash +python3 -m src.main agent-context-raw --cwd .. +``` + +### 4.4 Additional working directories + +```bash +python3 -m src.main agent-context --cwd .. --add-dir /path/to/directory +``` + +### 4.5 Disable `CLAUDE.md` discovery + +```bash +python3 -m src.main agent-context --cwd .. --disable-claude-md +``` + +## 5. Core Agent Loop + +### 5.1 Read-only run + +```bash +python3 -m src.main agent \ + "Read claw-code-agent/src/agent_runtime.py and summarize how the loop works." \ + --cwd .. +``` + +### 5.2 Show transcript output + +```bash +python3 -m src.main agent \ + "Read claw-code-agent/src/agent_session.py and summarize the session model." \ + --cwd .. \ + --show-transcript +``` + +### 5.3 Streaming model responses + +```bash +python3 -m src.main agent \ + "Inspect the current repository and summarize the architecture." \ + --cwd .. \ + --stream \ + --show-transcript +``` + +### 5.4 Interactive chat loop + +```bash +python3 -m src.main agent-chat --cwd .. +``` + +Optional first prompt: + +```bash +python3 -m src.main agent-chat \ + "Inspect the repository and tell me where the runtime loop lives." \ + --cwd .. +``` + +Inside chat mode: + +- type normal prompts to continue the same session +- use `/exit` or `/quit` to leave + +### 5.5 Resume directly into chat mode + +```bash +python3 -m src.main agent-chat \ + --resume-session-id \ + --cwd .. +``` + +## 6. Tool Execution + +### 6.1 Read files + +```bash +python3 -m src.main agent \ + "Read claw-code-agent/src/agent_tools.py and summarize each tool." \ + --cwd .. +``` + +### 6.2 Write files + +```bash +python3 -m src.main agent \ + "Create TEST_WRITE.md in the current directory with one line: write test ok" \ + --cwd ./test_cases \ + --allow-write +``` + +### 6.3 Edit files + +```bash +python3 -m src.main agent \ + "Create demo.txt with 'hello world', then replace 'world' with 'agent'." \ + --cwd ./test_cases \ + --allow-write +``` + +### 6.4 Glob and grep + +```bash +python3 -m src.main agent \ + "Find Python files in the current directory, then search for 'LocalCodingAgent' and summarize the matches." \ + --cwd .. +``` + +### 6.5 Shell commands + +```bash +python3 -m src.main agent \ + "Run pwd and ls in the current working directory, then summarize the result." \ + --cwd .. \ + --allow-shell \ + --show-transcript +``` + +### 6.6 Unsafe shell mode + +Use only if you intentionally want destructive-shell permission enabled. + +```bash +python3 -m src.main agent \ + "Explain whether destructive shell commands are currently allowed." \ + --cwd .. \ + --allow-shell \ + --unsafe +``` + +## 7. Session Persistence And Resume + +### 7.1 Create a saved session + +```bash +python3 -m src.main agent \ + "Create a short TODO file in the current directory and explain what you wrote." \ + --cwd ./test_cases \ + --allow-write +``` + +At the end of the run, note: + +```text +session_id=... +session_path=... +``` + +### 7.2 Resume a saved session + +```bash +python3 -m src.main agent-resume \ + \ + "Continue the previous task and improve the file." \ + --allow-write \ + --show-transcript +``` + +### 7.3 Inspect saved sessions + +```bash +ls -lt .port_sessions/agent +``` + +## 8. Structured Output / JSON Schema + +Create a schema file: + +```bash +cat > /tmp/claw_schema.json <<'EOF' +{ + "type": "object", + "properties": { + "status": { "type": "string" }, + "summary": { "type": "string" } + }, + "required": ["status", "summary"], + "additionalProperties": false +} +EOF +``` + +Run the agent with schema mode: + +```bash +python3 -m src.main agent \ + "Inspect the current repository and respond in the requested JSON format." \ + --cwd .. \ + --response-schema-file /tmp/claw_schema.json \ + --response-schema-name claw_summary \ + --response-schema-strict +``` + +## 9. Budgets And Limits + +### 9.1 Total token budget + +```bash +python3 -m src.main agent \ + "Give a very long answer about the current repository." \ + --cwd .. \ + --max-total-tokens 50 +``` + +### 9.2 Input / output token budgets + +```bash +python3 -m src.main agent \ + "Read several files and explain them in detail." \ + --cwd .. \ + --max-input-tokens 80 \ + --max-output-tokens 80 +``` + +### 9.3 Reasoning-token budget + +```bash +python3 -m src.main agent \ + "Solve a multi-step task and explain the result." \ + --cwd .. \ + --max-reasoning-tokens 10 +``` + +### 9.4 Tool-call budget + +```bash +python3 -m src.main agent \ + "Read multiple files, search for symbols, and summarize the repo." \ + --cwd .. \ + --max-tool-calls 1 +``` + +### 9.5 Delegated-task budget + +```bash +python3 -m src.main agent \ + "Delegate two subtasks to inspect and summarize the repo." \ + --cwd .. \ + --max-delegated-tasks 1 +``` + +### 9.6 Cost budget + +```bash +python3 -m src.main agent \ + "Inspect the current repository and summarize it." \ + --cwd .. \ + --input-cost-per-million 0.15 \ + --output-cost-per-million 0.60 \ + --max-budget-usd 0.000001 +``` + +### 9.7 Model-call budget + +```bash +python3 -m src.main agent \ + "Continue inspecting the repository until you are done." \ + --cwd .. \ + --max-model-calls 1 +``` + +### 9.8 Session-turn budget + +```bash +python3 -m src.main agent \ + "Work through the repository across multiple turns and keep going." \ + --cwd .. \ + --max-session-turns 1 +``` + +## 10. Streaming, Continuation, And Context Reduction + +### 10.1 Streaming assistant output + +```bash +python3 -m src.main agent \ + "Produce a long explanation of the current repository architecture." \ + --cwd .. \ + --stream \ + --show-transcript +``` + +### 10.2 Automatic continuation after truncation + +Use a small output budget so the backend is more likely to stop early: + +```bash +python3 -m src.main agent \ + "Write a long, structured explanation of the current repository." \ + --cwd .. \ + --max-output-tokens 32 \ + --show-transcript +``` + +### 10.3 Snipping older context + +```bash +python3 -m src.main agent \ + "Read claw-code-agent/src/agent_runtime.py, claw-code-agent/src/agent_session.py, claw-code-agent/src/query_engine.py, and claw-code-agent/src/plugin_runtime.py, then summarize all of them in detail." \ + --cwd .. \ + --auto-snip-threshold 120 \ + --compact-preserve-messages 0 \ + --show-transcript +``` + +### 10.4 Compaction boundaries + +```bash +python3 -m src.main agent \ + "Read several large files from claw-code-agent/src and keep explaining the repository until the context gets compacted." \ + --cwd .. \ + --auto-compact-threshold 120 \ + --compact-preserve-messages 1 \ + --show-transcript +``` + +## 11. File History Replay + +### 11.1 Create file history + +```bash +python3 -m src.main agent \ + "Create notes.txt with one line, then update that line to mention file history." \ + --cwd ./test_cases \ + --allow-write +``` + +### 11.2 Resume and inspect replay + +```bash +python3 -m src.main agent-resume \ + \ + "Continue the previous work and tell me what files were changed before this turn." \ + --allow-write \ + --show-transcript +``` + +Look for `file_history_replay` messages in the transcript. + +## 12. Nested Delegation + +### 12.1 Basic delegated subtask + +```bash +python3 -m src.main agent \ + "Delegate a subtask to inspect claw-code-agent/src/agent_runtime.py and return the summary." \ + --cwd .. \ + --show-transcript +``` + +### 12.2 Multiple delegated subtasks + +```bash +python3 -m src.main agent \ + "Delegate one subtask to scan the repository and another to summarize it after the scan." \ + --cwd .. \ + --show-transcript +``` + +### 12.3 Resume a delegated child session + +1. Seed a normal saved session: + +```bash +python3 -m src.main agent \ + "Inspect claw-code-agent/src/agent_tools.py and give a short summary." \ + --cwd .. +``` + +2. Copy that `session_id`, then run: + +```bash +python3 -m src.main agent \ + "Delegate a subtask that resumes session and continues it." \ + --cwd .. \ + --show-transcript +``` + +### 12.4 Topological dependency batches + +```bash +python3 -m src.main agent \ + "Delegate two subtasks: one named scan, and one named summarize that depends on scan. Use topological batching and then return the final summary." \ + --cwd .. \ + --show-transcript +``` + +Look for: + +- `delegate_batch_result` +- `delegate_group_result` +- `batch_index=...` + +## 13. Plugin Runtime + +Create a local plugin manifest: + +```bash +mkdir -p ./test_cases/plugins/demo +cat > ./test_cases/plugins/demo/plugin.json <<'EOF' +{ + "name": "demo-plugin", + "hooks": { + "beforePrompt": "Inject plugin prompt guidance.", + "afterTurn": "Attach plugin after-turn guidance.", + "onResume": "Reapply plugin state on resume.", + "beforePersist": "Persist plugin state before saving.", + "beforeDelegate": "Add plugin guidance before delegated children run.", + "afterDelegate": "Add plugin guidance after delegated children finish." + }, + "toolAliases": [ + { + "name": "plugin_read", + "baseTool": "read_file", + "description": "Plugin alias for reading files." + } + ], + "virtualTools": [ + { + "name": "demo_virtual", + "description": "Return a rendered plugin response.", + "responseTemplate": "plugin topic: {topic}" + } + ], + "toolHooks": { + "read_file": { + "beforeTool": "Validate the path before reading.", + "afterResult": "Summarize the file before the next action." + } + } +} +EOF +``` + +### 13.1 Plugin prompt/context discovery + +```bash +python3 -m src.main agent-prompt --cwd ./test_cases +python3 -m src.main agent-context-raw --cwd ./test_cases +``` + +### 13.2 Plugin alias tool + +```bash +echo "hello plugin" > ./test_cases/hello.txt +python3 -m src.main agent \ + "Use the plugin_read tool to read hello.txt and summarize it." \ + --cwd ./test_cases \ + --show-transcript +``` + +### 13.3 Plugin virtual tool + +```bash +python3 -m src.main agent \ + "Use the demo_virtual tool with topic plugins and return the result." \ + --cwd ./test_cases \ + --show-transcript +``` + +### 13.4 Plugin before/after tool guidance + +```bash +python3 -m src.main agent \ + "Read hello.txt and follow the plugin guidance around the read_file tool." \ + --cwd ./test_cases \ + --show-transcript +``` + +### 13.5 Plugin lifecycle with resume/persist + +1. Start a session: + +```bash +python3 -m src.main agent \ + "Use the plugin system and create a saved session." \ + --cwd ./test_cases +``` + +2. Resume it: + +```bash +python3 -m src.main agent-resume \ + \ + "Continue and mention any plugin lifecycle guidance you received." \ + --show-transcript +``` + +Look for: + +- `plugin_before_persist` +- `Plugin resume hooks:` +- `Plugin runtime state:` + +## 14. Query Engine And Workspace Commands + +### 14.1 Workspace inventory + +```bash +python3 -m src.main summary +python3 -m src.main manifest +python3 -m src.main subsystems --limit 20 +python3 -m src.main commands --limit 20 +python3 -m src.main tools --limit 20 +``` + +### 14.2 Query routing and bootstrap reports + +```bash +python3 -m src.main route "inspect the runtime and tools" --limit 10 +python3 -m src.main bootstrap "inspect the runtime and tools" --limit 10 +python3 -m src.main turn-loop "inspect the runtime and tools" --limit 5 --max-turns 3 +``` + +### 14.3 Session flushing for the mirrored workspace + +```bash +python3 -m src.main flush-transcript "store a temporary transcript" +python3 -m src.main load-session +``` + +## 15. Remote/Direct Mode Simulations + +These are mirrored workspace simulation commands, not the real agent runtime: + +```bash +python3 -m src.main remote-mode demo-target +python3 -m src.main ssh-mode demo-target +python3 -m src.main teleport-mode demo-target +python3 -m src.main direct-connect-mode demo-target +python3 -m src.main deep-link-mode demo-target +``` + +## 16. Parity Tracking Workflow + +Use this every time a new feature lands: + +```bash +python3 -m unittest discover -s tests -v +``` + +Then update: + +- `PARITY_CHECKLIST.md` +- `TESTING_GUIDE.md` + +Rule for future work: + +- every new implemented feature should add a checked item in `PARITY_CHECKLIST.md` +- every user-testable feature should add at least one concrete command example in `TESTING_GUIDE.md` diff --git a/src/agent_manager.py b/src/agent_manager.py index a6e4910..0aa03e3 100644 --- a/src/agent_manager.py +++ b/src/agent_manager.py @@ -26,9 +26,13 @@ class ManagedAgentGroup: label: str | None = None parent_agent_id: str | None = None child_agent_ids: tuple[str, ...] = () + strategy: str = 'serial' status: str = 'running' completed_children: int = 0 failed_children: int = 0 + batch_count: int = 0 + max_batch_size: int = 0 + dependency_skips: int = 0 @dataclass @@ -68,6 +72,7 @@ class AgentManager: *, label: str | None = None, parent_agent_id: str | None = None, + strategy: str = 'serial', ) -> str: self._group_counter += 1 group_id = f'group_{self._group_counter}' @@ -75,6 +80,7 @@ class AgentManager: group_id=group_id, label=label, parent_agent_id=parent_agent_id, + strategy=strategy, ) return group_id @@ -97,9 +103,13 @@ class AgentManager: label=group.label, parent_agent_id=group.parent_agent_id, child_agent_ids=updated_children, + strategy=group.strategy, status=group.status, completed_children=group.completed_children, failed_children=group.failed_children, + batch_count=group.batch_count, + max_batch_size=group.max_batch_size, + dependency_skips=group.dependency_skips, ) record = self.records.get(agent_id) if record is None: @@ -129,6 +139,9 @@ class AgentManager: status: str, completed_children: int, failed_children: int, + batch_count: int = 0, + max_batch_size: int = 0, + dependency_skips: int = 0, ) -> None: group = self.groups.get(group_id) if group is None: @@ -138,9 +151,13 @@ class AgentManager: label=group.label, parent_agent_id=group.parent_agent_id, child_agent_ids=group.child_agent_ids, + strategy=group.strategy, status=status, completed_children=completed_children, failed_children=failed_children, + batch_count=batch_count, + max_batch_size=max_batch_size, + dependency_skips=dependency_skips, ) def finish_agent( @@ -209,11 +226,15 @@ class AgentManager: return { 'group_id': group.group_id, 'label': group.label, + 'strategy': group.strategy, 'status': group.status, 'child_count': len(children), 'completed_children': group.completed_children, 'failed_children': group.failed_children, 'resumed_children': resumed_children, + 'batch_count': group.batch_count, + 'max_batch_size': group.max_batch_size, + 'dependency_skips': group.dependency_skips, 'stop_reason_counts': stop_reason_counts, } @@ -266,7 +287,9 @@ class AgentManager: lines.append( f'- {label}: group_status={group.status} children={len(group.child_agent_ids)} ' f'completed={group.completed_children} failed={group.failed_children} ' - f"resumed={summary['resumed_children']}{stop_suffix}" + f"resumed={summary['resumed_children']} strategy={group.strategy} " + f"batches={group.batch_count} max_batch_size={group.max_batch_size} " + f"dependency_skips={group.dependency_skips}{stop_suffix}" ) if len(self.groups) > 6: lines.append(f'- ... plus {len(self.groups) - 6} more agent groups') diff --git a/src/agent_runtime.py b/src/agent_runtime.py index 08366b3..623681e 100644 --- a/src/agent_runtime.py +++ b/src/agent_runtime.py @@ -45,6 +45,7 @@ from .session_store import ( save_agent_session, serialize_model_config, serialize_runtime_config, + usage_from_payload, ) @@ -106,6 +107,8 @@ class LocalCodingAgent: self.active_session_id = None self.last_session_path = None self.resume_source_session_id = None + if self.plugin_runtime is not None: + self.plugin_runtime.restore_session_state({}) def build_prompt_context(self, scratchpad_directory: Path | None = None): return build_prompt_context( @@ -144,6 +147,8 @@ class LocalCodingAgent: def run(self, prompt: str) -> AgentRunResult: self.managed_agent_id = None self.resume_source_session_id = None + if self.plugin_runtime is not None: + self.plugin_runtime.restore_session_state({}) session_id = uuid4().hex scratchpad_directory = self._ensure_scratchpad_directory(session_id) result = self._run_prompt( @@ -175,6 +180,8 @@ class LocalCodingAgent: self.last_session_path = str( self.runtime_config.session_directory / f'{stored_session.session_id}.json' ) + if self.plugin_runtime is not None: + self.plugin_runtime.restore_session_state(stored_session.plugin_state) scratchpad_directory = ( Path(stored_session.scratchpad_directory) if stored_session.scratchpad_directory @@ -214,6 +221,10 @@ class LocalCodingAgent: ) effective_prompt = self._apply_plugin_before_prompt_hooks(slash_result.prompt or prompt) + effective_prompt = self._apply_plugin_resume_hooks( + effective_prompt, + resumed=base_session is not None, + ) self.managed_agent_id = self.agent_manager.start_agent( prompt=effective_prompt, parent_agent_id=self.parent_agent_id, @@ -234,22 +245,49 @@ class LocalCodingAgent: self.last_session = session self.active_session_id = session_id tool_specs = [tool.to_openai_tool() for tool in self.tool_registry.values()] - tool_calls = 0 + starting_usage = UsageStats() + starting_cost_usd = 0.0 + starting_tool_calls = 0 + starting_session_turns = 0 + starting_model_calls = 0 + if base_session is not None and self.resume_source_session_id: + try: + stored_resume_state = load_agent_session( + self.resume_source_session_id, + directory=self.runtime_config.session_directory, + ) + except OSError: + stored_resume_state = None + if stored_resume_state is not None: + starting_usage = usage_from_payload(stored_resume_state.usage) + starting_cost_usd = stored_resume_state.total_cost_usd + starting_tool_calls = stored_resume_state.tool_calls + starting_session_turns = stored_resume_state.turns + budget_state = ( + stored_resume_state.budget_state + if isinstance(stored_resume_state.budget_state, dict) + else {} + ) + starting_model_calls = int(budget_state.get('model_calls', 0)) if isinstance(budget_state.get('model_calls', 0), int) else 0 + tool_calls = starting_tool_calls last_content = '' - total_usage = UsageStats() - total_cost_usd = 0.0 + total_usage = starting_usage + total_cost_usd = starting_cost_usd file_history = list(existing_file_history) stream_events: list[dict[str, object]] = [] assistant_response_segments: list[str] = [] delegated_tasks = sum( 1 for entry in file_history if entry.get('action') == 'delegate_agent' ) + model_calls = starting_model_calls initial_budget = self._check_budget( total_usage, total_cost_usd, tool_calls=tool_calls, delegated_tasks=delegated_tasks, + model_calls=model_calls, + session_turns=starting_session_turns, ) if initial_budget.exceeded: result = AgentRunResult( @@ -302,6 +340,7 @@ class LocalCodingAgent: for _ in [0] ) stream_events.extend(event.to_dict() for event in turn_events) + model_calls += 1 total_usage = total_usage + turn.usage total_cost_usd = self.model_config.pricing.estimate_cost_usd(total_usage) last_content = turn.content @@ -311,6 +350,8 @@ class LocalCodingAgent: total_cost_usd, tool_calls=tool_calls, delegated_tasks=delegated_tasks, + model_calls=model_calls, + session_turns=starting_session_turns + turn_index, ) if budget_after_model.exceeded: result = AgentRunResult( @@ -400,6 +441,7 @@ class LocalCodingAgent: return result stream_events.extend(event.to_dict() for event in turn_events) + model_calls += 1 total_usage = total_usage + turn.usage total_cost_usd = self.model_config.pricing.estimate_cost_usd(total_usage) last_content = turn.content @@ -409,6 +451,8 @@ class LocalCodingAgent: total_cost_usd, tool_calls=tool_calls, delegated_tasks=delegated_tasks, + model_calls=model_calls, + session_turns=starting_session_turns + turn_index, ) if budget_after_model.exceeded: result = AgentRunResult( @@ -487,6 +531,8 @@ class LocalCodingAgent: total_cost_usd, tool_calls=tool_calls, delegated_tasks=delegated_tasks, + model_calls=model_calls, + session_turns=starting_session_turns + turn_index, ) if budget_after_tool_request.exceeded: stream_events.append( @@ -534,6 +580,8 @@ class LocalCodingAgent: 'message_id': session.messages[tool_message_index].message_id, } ) + if self.plugin_runtime is not None: + self.plugin_runtime.record_tool_attempt(tool_call.name, blocked=False) plugin_preflight_messages = self._plugin_tool_preflight_messages(tool_call.name) if plugin_preflight_messages: stream_events.append( @@ -547,6 +595,13 @@ class LocalCodingAgent: ) plugin_block_message = self._plugin_block_message(tool_call.name) if plugin_block_message is not None: + if self.plugin_runtime is not None: + blocked_attempts = int( + self.plugin_runtime.session_state.get('blocked_tool_attempts', 0) + ) + self.plugin_runtime.session_state['blocked_tool_attempts'] = ( + blocked_attempts + 1 + ) tool_result = ToolExecutionResult( name=tool_call.name, ok=False, @@ -596,6 +651,12 @@ class LocalCodingAgent: tool_result = update.result if tool_result is None: raise RuntimeError(f'Tool executor returned no final result for {tool_call.name}') + if self.plugin_runtime is not None: + self.plugin_runtime.record_tool_result( + tool_call.name, + ok=tool_result.ok, + metadata=tool_result.metadata, + ) plugin_messages = self._plugin_tool_result_messages(tool_call.name) if plugin_messages: merged_metadata = dict(tool_result.metadata) @@ -646,6 +707,22 @@ class LocalCodingAgent: preflight_messages=plugin_preflight_messages, block_message=plugin_block_message, plugin_messages=plugin_messages, + delegate_preflight_messages=tuple( + message + for message in tool_result.metadata.get( + 'plugin_delegate_preflight_messages', + [], + ) + if isinstance(message, str) and message + ), + delegate_after_messages=tuple( + message + for message in tool_result.metadata.get( + 'plugin_delegate_after_messages', + [], + ) + if isinstance(message, str) and message + ), ) if plugin_runtime_message is not None: session.append_user( @@ -833,6 +910,8 @@ class LocalCodingAgent: *, tool_calls: int, delegated_tasks: int, + model_calls: int, + session_turns: int, ) -> BudgetDecision: budget = self.runtime_config.budget_config token_reason = self._check_token_budget(usage, budget) @@ -872,6 +951,28 @@ class LocalCodingAgent: f'({delegated_tasks} > {budget.max_delegated_tasks}).' ), ) + if ( + budget.max_model_calls is not None + and model_calls > budget.max_model_calls + ): + return BudgetDecision( + exceeded=True, + reason=( + 'Stopped because the model-call budget was exceeded ' + f'({model_calls} > {budget.max_model_calls}).' + ), + ) + if ( + budget.max_session_turns is not None + and session_turns > budget.max_session_turns + ): + return BudgetDecision( + exceeded=True, + reason=( + 'Stopped because the session-turn budget was exceeded ' + f'({session_turns} > {budget.max_session_turns}).' + ), + ) return BudgetDecision(exceeded=False) def _snip_session_if_needed( @@ -1185,6 +1286,12 @@ class LocalCodingAgent: entry['history_kind'] = 'shell' elif action == 'delegate_agent': entry['history_kind'] = 'delegation' + delegate_batches = metadata.get('delegate_batches') + if isinstance(delegate_batches, list): + entry['delegate_batch_count'] = len(delegate_batches) + dependency_skips = metadata.get('dependency_skips') + if isinstance(dependency_skips, int) and not isinstance(dependency_skips, bool): + entry['dependency_skips'] = dependency_skips else: entry['history_kind'] = 'tool' return entry @@ -1292,6 +1399,7 @@ class LocalCodingAgent: ] compaction_depth = (max(prior_depths) if prior_depths else 0) + 1 compacted_kinds: dict[str, int] = {} + source_mutation_totals: dict[str, int] = {} compacted_lineage_ids: list[str] = [] preserved_tail_lineage_ids = [ lineage_id @@ -1301,6 +1409,7 @@ class LocalCodingAgent: if isinstance(lineage_id, str) and lineage_id ] max_source_revision = 0 + max_source_mutation_serial = 0 compacted_revision_total = 0 for message in messages: kind = message.metadata.get('kind') @@ -1313,6 +1422,26 @@ class LocalCodingAgent: if isinstance(revision, int) and not isinstance(revision, bool): max_source_revision = max(max_source_revision, revision) compacted_revision_total += revision + max_mutation_serial = message.metadata.get('max_mutation_serial') + if isinstance(max_mutation_serial, int) and not isinstance(max_mutation_serial, bool): + max_source_mutation_serial = max( + max_source_mutation_serial, + max_mutation_serial, + ) + mutation_totals = message.metadata.get('mutation_totals') + if isinstance(mutation_totals, dict): + for mutation_kind, count in mutation_totals.items(): + if ( + not isinstance(mutation_kind, str) + or not mutation_kind + or isinstance(count, bool) + or not isinstance(count, int) + or count <= 0 + ): + continue + source_mutation_totals[mutation_kind] = ( + source_mutation_totals.get(mutation_kind, 0) + count + ) compact_boundary_id = f'compact_boundary_{turn_index}_{len(messages)}' @@ -1340,6 +1469,8 @@ class LocalCodingAgent: 'compacted_lineage_ids': compacted_lineage_ids, 'preserved_tail_lineage_ids': preserved_tail_lineage_ids, 'max_source_revision': max_source_revision, + 'max_source_mutation_serial': max_source_mutation_serial, + 'source_mutation_totals': source_mutation_totals, 'compacted_revision_total': compacted_revision_total, 'compacted_message_ids': [ message.message_id for message in messages if message.message_id @@ -1401,121 +1532,245 @@ class LocalCodingAgent: } include_parent_context = bool(arguments.get('include_parent_context', True)) continue_on_error = bool(arguments.get('continue_on_error', True)) + max_failures = arguments.get('max_failures') + if isinstance(max_failures, bool) or (max_failures is not None and not isinstance(max_failures, int)): + max_failures = None + if isinstance(max_failures, int) and max_failures < 0: + max_failures = None + strategy = self._normalize_delegate_strategy(arguments.get('strategy')) child_summaries: list[dict[str, object]] = [] child_session_ids: list[str] = [] prior_results: list[dict[str, str]] = [] + completed_labels: set[str] = set() + failed_labels: set[str] = set() + delegate_preflight_messages = ( + self.plugin_runtime.before_delegate_injections() + if self.plugin_runtime is not None + else () + ) + delegate_after_messages: tuple[str, ...] = () group_id: str | None = None if self.agent_manager is not None and len(subtasks) > 1: group_id = self.agent_manager.start_group( label=str(arguments.get('label') or 'delegated_group'), parent_agent_id=self.managed_agent_id, + strategy=strategy, ) + planned_batches = self._plan_delegate_batches(subtasks, strategy) + batch_summaries: list[dict[str, object]] = [] failed_children = 0 + dependency_skips = 0 child_result = None - for index, subtask in enumerate(subtasks, start=1): - subtask_label = str(subtask.get('label') or f'subtask_{index}') - child_agent = LocalCodingAgent( - model_config=self.model_config, - runtime_config=replace( - child_runtime_config, - max_turns=subtask.get('max_turns', child_runtime_config.max_turns), - ), - custom_system_prompt=self.custom_system_prompt, - append_system_prompt=self.append_system_prompt, - override_system_prompt=self.override_system_prompt, - tool_registry=child_tools, - agent_manager=self.agent_manager, - parent_agent_id=self.managed_agent_id, - managed_group_id=group_id, - managed_child_index=index, - managed_label=subtask_label, - ) - if group_id is not None and child_agent.managed_agent_id is not None: - self.agent_manager.register_group_child( - group_id, - child_agent.managed_agent_id, - child_index=index, + stop_processing = False + for batch_index, batch in enumerate(planned_batches, start=1): + if stop_processing: + break + batch_completed = 0 + batch_failed = 0 + batch_skipped = 0 + batch_labels: list[str] = [] + for subtask in batch: + index = int(subtask.get('_delegate_index', len(child_summaries) + 1)) + subtask_label = str(subtask.get('label') or f'subtask_{index}') + batch_labels.append(subtask_label) + dependencies = tuple( + item + for item in subtask.get('depends_on', ()) + if isinstance(item, str) and item ) - child_prompt = str(subtask['prompt']) - if include_parent_context and prior_results: - child_prompt = self._prepend_delegate_context(child_prompt, prior_results) - resume_session_id = subtask.get('resume_session_id') - resume_used = False - if isinstance(resume_session_id, str) and resume_session_id: - try: - stored_child_session = load_agent_session( - resume_session_id, - directory=child_runtime_config.session_directory, + unmet_dependencies = [ + dependency + for dependency in dependencies + if dependency not in completed_labels + ] + blocked_dependencies = [ + dependency + for dependency in dependencies + if dependency in failed_labels + ] + if unmet_dependencies: + skip_reason = ( + 'skipped_dependency' + if blocked_dependencies + else 'pending_dependency' ) - except OSError: child_result = AgentRunResult( - final_output=f'Unable to load delegated session {resume_session_id}.', + final_output=( + 'Skipped delegated subtask because dependencies were not satisfied: ' + + ', '.join(unmet_dependencies) + ), turns=0, tool_calls=0, transcript=(), - stop_reason='resume_load_error', - session_id=resume_session_id, + stop_reason=skip_reason, ) - failed_children += 1 summary = { 'index': index, 'label': subtask_label, - 'session_id': resume_session_id, + 'session_id': '', 'turns': child_result.turns, 'tool_calls': child_result.tool_calls, - 'stop_reason': child_result.stop_reason or 'resume_load_error', + 'stop_reason': skip_reason, 'output_preview': self._preview_text(child_result.final_output, 220), - 'resume_used': True, - 'resumed_from_session_id': resume_session_id, + 'resume_used': False, + 'resumed_from_session_id': '', + 'depends_on': list(dependencies), + 'batch_index': batch_index, } child_summaries.append(summary) - prior_results.append( - { - 'label': summary['label'], - 'output_preview': str(summary['output_preview']), - } - ) + failed_children += 1 + batch_failed += 1 + batch_skipped += 1 + dependency_skips += 1 + failed_labels.add(subtask_label) + if isinstance(max_failures, int) and failed_children > max_failures: + stop_processing = True + break if not continue_on_error: + stop_processing = True break continue - child_result = child_agent.resume(child_prompt, stored_child_session) - resume_used = True - else: - child_result = child_agent.run(child_prompt) - if group_id is not None and child_agent.managed_agent_id is not None: - self.agent_manager.register_group_child( - group_id, - child_agent.managed_agent_id, - child_index=index, + child_agent = LocalCodingAgent( + model_config=self.model_config, + runtime_config=replace( + child_runtime_config, + max_turns=subtask.get('max_turns', child_runtime_config.max_turns), + ), + custom_system_prompt=self.custom_system_prompt, + append_system_prompt=self.append_system_prompt, + override_system_prompt=self.override_system_prompt, + tool_registry=child_tools, + agent_manager=self.agent_manager, + parent_agent_id=self.managed_agent_id, + managed_group_id=group_id, + managed_child_index=index, + managed_label=subtask_label, ) - summary = { - 'index': index, - 'label': subtask_label, - 'session_id': child_result.session_id or '', - 'turns': child_result.turns, - 'tool_calls': child_result.tool_calls, - 'stop_reason': child_result.stop_reason or 'stop', - 'output_preview': self._preview_text(child_result.final_output, 220), - 'resume_used': resume_used, - 'resumed_from_session_id': ( - str(resume_session_id) - if isinstance(resume_session_id, str) and resume_session_id - else '' - ), - } - child_summaries.append(summary) - if child_result.session_id: - child_session_ids.append(child_result.session_id) - prior_results.append( + if group_id is not None and child_agent.managed_agent_id is not None: + self.agent_manager.register_group_child( + group_id, + child_agent.managed_agent_id, + child_index=index, + ) + child_prompt = str(subtask['prompt']) + if delegate_preflight_messages: + child_prompt = self._prepend_plugin_delegate_context( + child_prompt, + delegate_preflight_messages, + ) + if include_parent_context and prior_results: + child_prompt = self._prepend_delegate_context(child_prompt, prior_results) + resume_session_id = subtask.get('resume_session_id') + resume_used = False + if isinstance(resume_session_id, str) and resume_session_id: + try: + stored_child_session = load_agent_session( + resume_session_id, + directory=child_runtime_config.session_directory, + ) + except OSError: + child_result = AgentRunResult( + final_output=f'Unable to load delegated session {resume_session_id}.', + turns=0, + tool_calls=0, + transcript=(), + stop_reason='resume_load_error', + session_id=resume_session_id, + ) + failed_children += 1 + batch_failed += 1 + summary = { + 'index': index, + 'label': subtask_label, + 'session_id': resume_session_id, + 'turns': child_result.turns, + 'tool_calls': child_result.tool_calls, + 'stop_reason': child_result.stop_reason or 'resume_load_error', + 'output_preview': self._preview_text(child_result.final_output, 220), + 'resume_used': True, + 'resumed_from_session_id': resume_session_id, + 'depends_on': list(dependencies), + 'batch_index': batch_index, + } + child_summaries.append(summary) + prior_results.append( + { + 'label': summary['label'], + 'output_preview': str(summary['output_preview']), + } + ) + failed_labels.add(subtask_label) + if isinstance(max_failures, int) and failed_children > max_failures: + stop_processing = True + break + if not continue_on_error: + stop_processing = True + break + continue + child_result = child_agent.resume(child_prompt, stored_child_session) + resume_used = True + else: + child_result = child_agent.run(child_prompt) + if group_id is not None and child_agent.managed_agent_id is not None: + self.agent_manager.register_group_child( + group_id, + child_agent.managed_agent_id, + child_index=index, + ) + summary = { + 'index': index, + 'label': subtask_label, + 'session_id': child_result.session_id or '', + 'turns': child_result.turns, + 'tool_calls': child_result.tool_calls, + 'stop_reason': child_result.stop_reason or 'stop', + 'output_preview': self._preview_text(child_result.final_output, 220), + 'resume_used': resume_used, + 'resumed_from_session_id': ( + str(resume_session_id) + if isinstance(resume_session_id, str) and resume_session_id + else '' + ), + 'depends_on': list(dependencies), + 'batch_index': batch_index, + } + child_summaries.append(summary) + if child_result.session_id: + child_session_ids.append(child_result.session_id) + prior_results.append( + { + 'label': summary['label'], + 'output_preview': str(summary['output_preview']), + } + ) + if child_result.stop_reason in {'backend_error', 'budget_exceeded'}: + failed_children += 1 + batch_failed += 1 + failed_labels.add(subtask_label) + if isinstance(max_failures, int) and failed_children > max_failures: + stop_processing = True + break + if not continue_on_error: + stop_processing = True + break + else: + batch_completed += 1 + completed_labels.add(subtask_label) + batch_status = 'completed' + if batch_failed and batch_completed: + batch_status = 'partial' + elif batch_failed: + batch_status = 'failed' + batch_summaries.append( { - 'label': summary['label'], - 'output_preview': str(summary['output_preview']), + 'batch_index': batch_index, + 'labels': batch_labels, + 'completed_children': batch_completed, + 'failed_children': batch_failed, + 'skipped_children': batch_skipped, + 'status': batch_status, } ) - if child_result.stop_reason in {'backend_error', 'budget_exceeded'}: - failed_children += 1 - if not continue_on_error: - break assert child_result is not None completed_children = len(child_summaries) - failed_children resumed_children = sum( @@ -1526,12 +1781,20 @@ class LocalCodingAgent: group_status = 'partial' elif failed_children: group_status = 'failed' + delegate_after_messages = ( + self.plugin_runtime.after_delegate_injections() + if self.plugin_runtime is not None + else () + ) if group_id is not None and self.agent_manager is not None: self.agent_manager.finish_group( group_id, status=group_status, completed_children=completed_children, failed_children=failed_children, + batch_count=len(batch_summaries), + max_batch_size=max((len(batch['labels']) for batch in batch_summaries), default=0), + dependency_skips=dependency_skips, ) summary_lines = [ ( @@ -1544,21 +1807,43 @@ class LocalCodingAgent: summary_lines.append(f'group_id={group_id}') summary_lines.append(f'group_status={group_status}') summary_lines.append(f'resumed_children={resumed_children}') + summary_lines.append(f'strategy={strategy}') + summary_lines.append(f'batch_count={len(batch_summaries)}') + summary_lines.append(f'dependency_skips={dependency_skips}') + summary_lines.append('') + if delegate_preflight_messages: + summary_lines.append('Plugin delegate preflight:') + summary_lines.extend(f'- {message}' for message in delegate_preflight_messages) + summary_lines.append('') + for batch in batch_summaries: + summary_lines.append( + f"[batch {batch['batch_index']}] status={batch['status']} " + f"labels={','.join(batch['labels']) or '(none)'} " + f"completed={batch['completed_children']} failed={batch['failed_children']} " + f"skipped={batch['skipped_children']}" + ) + if batch_summaries: summary_lines.append('') for summary in child_summaries: summary_lines.extend( [ f"[{summary['label']}]", + f"batch_index={summary['batch_index']}", f"session_id={summary['session_id']}", f"turns={summary['turns']}", f"tool_calls={summary['tool_calls']}", f"stop_reason={summary['stop_reason']}", f"resume_used={summary['resume_used']}", f"resumed_from_session_id={summary['resumed_from_session_id']}", + f"depends_on={','.join(summary.get('depends_on', [])) or '(none)'}", f"output_preview={summary['output_preview']}", '', ] ) + if delegate_after_messages: + summary_lines.append('Plugin delegate completion:') + summary_lines.extend(f'- {message}' for message in delegate_after_messages) + summary_lines.append('') summary_lines.append('Final delegated output:') summary_lines.append(child_result.final_output) return ToolExecutionResult( @@ -1579,6 +1864,12 @@ class LocalCodingAgent: 'failed_children': failed_children, 'completed_children': completed_children, 'resumed_children': resumed_children, + 'strategy': strategy, + 'max_failures': max_failures, + 'delegate_batches': batch_summaries, + 'dependency_skips': dependency_skips, + 'plugin_delegate_preflight_messages': list(delegate_preflight_messages), + 'plugin_delegate_after_messages': list(delegate_after_messages), }, ) @@ -1591,7 +1882,13 @@ class LocalCodingAgent: if isinstance(raw_subtasks, list): for index, item in enumerate(raw_subtasks, start=1): if isinstance(item, str) and item.strip(): - subtasks.append({'prompt': item.strip(), 'label': f'subtask_{index}'}) + subtasks.append( + { + 'prompt': item.strip(), + 'label': f'subtask_{index}', + '_delegate_index': index, + } + ) continue if isinstance(item, dict): prompt = item.get('prompt') @@ -1608,8 +1905,16 @@ class LocalCodingAgent: resume_session_id = item.get('session_id') if isinstance(resume_session_id, str) and resume_session_id.strip(): task['resume_session_id'] = resume_session_id.strip() + depends_on = item.get('depends_on') + if isinstance(depends_on, list): + task['depends_on'] = tuple( + dependency.strip() + for dependency in depends_on + if isinstance(dependency, str) and dependency.strip() + ) if isinstance(max_turns, int) and not isinstance(max_turns, bool) and max_turns > 0: task['max_turns'] = max_turns + task['_delegate_index'] = index subtasks.append(task) prompt = arguments.get('prompt') if isinstance(prompt, str) and prompt.strip(): @@ -1620,8 +1925,71 @@ class LocalCodingAgent: resume_session_id = arguments.get('session_id') if isinstance(resume_session_id, str) and resume_session_id.strip(): task['resume_session_id'] = resume_session_id.strip() + task['_delegate_index'] = 1 subtasks.append(task) - return subtasks[:8] + return [ + { + **task, + '_delegate_index': int(task.get('_delegate_index', index)), + } + for index, task in enumerate(subtasks[:8], start=1) + ] + + def _normalize_delegate_strategy(self, strategy: object) -> str: + if not isinstance(strategy, str) or not strategy.strip(): + return 'serial' + normalized = strategy.strip().lower().replace('-', '_') + if normalized in {'graph', 'topological', 'dependency_graph', 'parallel', 'parallel_batches'}: + return 'topological' + return 'serial' + + def _plan_delegate_batches( + self, + subtasks: list[dict[str, object]], + strategy: str, + ) -> list[list[dict[str, object]]]: + if strategy != 'topological': + return [subtasks] + remaining = list(subtasks) + scheduled_labels: set[str] = set() + known_labels = { + str(task.get('label')) + for task in subtasks + if isinstance(task.get('label'), str) and str(task.get('label')).strip() + } + batches: list[list[dict[str, object]]] = [] + while remaining: + ready: list[dict[str, object]] = [] + blocked: list[dict[str, object]] = [] + for task in remaining: + dependencies = tuple( + item + for item in task.get('depends_on', ()) + if isinstance(item, str) and item + ) + if any(dependency not in known_labels for dependency in dependencies): + blocked.append(task) + continue + if all(dependency in scheduled_labels for dependency in dependencies): + ready.append(task) + else: + blocked.append(task) + if not ready: + batches.append(blocked) + break + batches.append( + sorted( + ready, + key=lambda task: int(task.get('_delegate_index', 0)), + ) + ) + scheduled_labels.update( + str(task.get('label')) + for task in ready + if isinstance(task.get('label'), str) and str(task.get('label')).strip() + ) + remaining = blocked + return batches def _delegated_task_units( self, @@ -1659,6 +2027,21 @@ class LocalCodingAgent: lines.extend(['', '', prompt]) return '\n'.join(lines) + def _prepend_plugin_delegate_context( + self, + prompt: str, + messages: tuple[str, ...], + ) -> str: + if not messages: + return prompt + lines = [ + '', + 'Plugin delegate guidance:', + ] + lines.extend(f'- {message}' for message in messages) + lines.extend(['', '', prompt]) + return '\n'.join(lines) + def _append_runtime_tool_followup_events( self, stream_events: list[dict[str, object]], @@ -1677,8 +2060,46 @@ class LocalCodingAgent: 'virtual_tool': metadata.get('virtual_tool'), } ) + plugin_delegate_preflight = metadata.get('plugin_delegate_preflight_messages') + if isinstance(plugin_delegate_preflight, list) and plugin_delegate_preflight: + stream_events.append( + { + 'type': 'plugin_delegate_preflight', + 'tool_call_id': tool_call.id, + 'tool_name': tool_call.name, + 'message_count': len(plugin_delegate_preflight), + } + ) + plugin_delegate_after = metadata.get('plugin_delegate_after_messages') + if isinstance(plugin_delegate_after, list) and plugin_delegate_after: + stream_events.append( + { + 'type': 'plugin_delegate_after', + 'tool_call_id': tool_call.id, + 'tool_name': tool_call.name, + 'message_count': len(plugin_delegate_after), + } + ) if tool_call.name != 'delegate_agent': return + delegate_batches = metadata.get('delegate_batches') + if isinstance(delegate_batches, list): + for batch in delegate_batches: + if not isinstance(batch, dict): + continue + stream_events.append( + { + 'type': 'delegate_batch_result', + 'tool_call_id': tool_call.id, + 'group_id': metadata.get('group_id'), + 'batch_index': batch.get('batch_index'), + 'status': batch.get('status'), + 'labels': batch.get('labels'), + 'completed_children': batch.get('completed_children'), + 'failed_children': batch.get('failed_children'), + 'skipped_children': batch.get('skipped_children'), + } + ) child_results = metadata.get('child_results') if isinstance(child_results, list): for child in child_results: @@ -1697,6 +2118,8 @@ class LocalCodingAgent: 'turns': child.get('turns'), 'resume_used': child.get('resume_used'), 'resumed_from_session_id': child.get('resumed_from_session_id'), + 'depends_on': child.get('depends_on'), + 'batch_index': child.get('batch_index'), } ) if metadata.get('group_id') is not None: @@ -1710,6 +2133,10 @@ class LocalCodingAgent: 'completed_children': metadata.get('completed_children'), 'failed_children': metadata.get('failed_children'), 'resumed_children': metadata.get('resumed_children'), + 'strategy': metadata.get('strategy'), + 'max_failures': metadata.get('max_failures'), + 'batch_count': len(delegate_batches) if isinstance(delegate_batches, list) else 0, + 'dependency_skips': metadata.get('dependency_skips'), } ) @@ -1819,6 +2246,12 @@ class LocalCodingAgent: child_session_ids = entry.get('child_session_ids') if isinstance(child_session_ids, list) and child_session_ids: details.append(f'child_sessions={len(child_session_ids)}') + delegate_batch_count = entry.get('delegate_batch_count') + if isinstance(delegate_batch_count, int) and not isinstance(delegate_batch_count, bool): + details.append(f'batches={delegate_batch_count}') + dependency_skips = entry.get('dependency_skips') + if isinstance(dependency_skips, int) and not isinstance(dependency_skips, bool): + details.append(f'dependency_skips={dependency_skips}') lines.append(f"- {'; '.join(details)}") before_snapshot_id = entry.get('before_snapshot_id') if isinstance(before_snapshot_id, str) and before_snapshot_id: @@ -1898,6 +2331,28 @@ class LocalCodingAgent: compacted_lineages = latest_boundary.metadata.get('compacted_lineage_ids') if isinstance(compacted_lineages, list) and compacted_lineages: lines.append(f'- Latest compacted lineages: {len(compacted_lineages)}') + max_source_mutation_serial = latest_boundary.metadata.get('max_source_mutation_serial') + if ( + isinstance(max_source_mutation_serial, int) + and not isinstance(max_source_mutation_serial, bool) + and max_source_mutation_serial > 0 + ): + lines.append( + f'- Latest source mutation serial: {max_source_mutation_serial}' + ) + source_mutation_totals = latest_boundary.metadata.get('source_mutation_totals') + if isinstance(source_mutation_totals, dict) and source_mutation_totals: + rendered = ', '.join( + f'{name}:{count}' + for name, count in sorted(source_mutation_totals.items()) + if isinstance(name, str) + and name + and isinstance(count, int) + and not isinstance(count, bool) + and count > 0 + ) + if rendered: + lines.append(f'- Latest compacted mutations: {rendered}') preserved_tail = latest_boundary.metadata.get('preserved_tail_ids') if isinstance(preserved_tail, list) and preserved_tail: lines.append( @@ -1933,8 +2388,16 @@ class LocalCodingAgent: preflight_messages: tuple[str, ...], block_message: str | None, plugin_messages: tuple[str, ...], + delegate_preflight_messages: tuple[str, ...] = (), + delegate_after_messages: tuple[str, ...] = (), ) -> str | None: - if block_message is None and not plugin_messages and not preflight_messages: + if ( + block_message is None + and not plugin_messages + and not preflight_messages + and not delegate_preflight_messages + and not delegate_after_messages + ): return None lines = [ '', @@ -1942,10 +2405,14 @@ class LocalCodingAgent: ] for message in preflight_messages: lines.append(f'- Before tool: {message}') + for message in delegate_preflight_messages: + lines.append(f'- Before delegate: {message}') if block_message is not None: lines.append(f'- Blocked: {block_message}') for message in plugin_messages: lines.append(f'- After result: {message}') + for message in delegate_after_messages: + lines.append(f'- After delegate: {message}') lines.extend( [ '', @@ -1977,6 +2444,51 @@ class LocalCodingAgent: ) -> AgentRunResult: if result.session_id is None: return result + persist_events = list(result.events) + if self.plugin_runtime is not None: + persist_messages = self.plugin_runtime.before_persist_injections() + if persist_messages: + session.append_user( + self._render_plugin_persist_message(persist_messages), + metadata={ + 'kind': 'plugin_persist', + 'message_count': len(persist_messages), + }, + message_id=f'plugin_persist_{result.session_id}', + ) + persist_events.append( + { + 'type': 'plugin_before_persist', + 'session_id': result.session_id, + 'message_count': len(persist_messages), + } + ) + previous_turns = 0 + previous_tool_calls = 0 + previous_budget_state: dict[str, object] = {} + existing_path = self.runtime_config.session_directory / f'{result.session_id}.json' + if existing_path.exists(): + try: + previous = load_agent_session( + result.session_id, + directory=self.runtime_config.session_directory, + ) + except OSError: + previous = None + if previous is not None: + previous_turns = previous.turns + previous_tool_calls = previous.tool_calls + if isinstance(previous.budget_state, dict): + previous_budget_state = dict(previous.budget_state) + budget_state = { + 'model_calls': int(previous_budget_state.get('model_calls', 0)) + + max(result.turns, 0), + 'session_turns': previous_turns + result.turns, + 'tool_calls': previous_tool_calls + result.tool_calls, + 'delegated_tasks': sum( + 1 for entry in result.file_history if entry.get('action') == 'delegate_agent' + ), + } stored = StoredAgentSession( session_id=result.session_id, model_config=serialize_model_config(self.model_config), @@ -1985,11 +2497,17 @@ class LocalCodingAgent: user_context=dict(session.user_context), system_context=dict(session.system_context), messages=session.transcript(), - turns=result.turns, - tool_calls=result.tool_calls, + turns=previous_turns + result.turns, + tool_calls=previous_tool_calls + result.tool_calls, usage=result.usage.to_dict(), total_cost_usd=result.total_cost_usd, file_history=result.file_history, + budget_state=budget_state, + plugin_state=( + self.plugin_runtime.export_session_state() + if self.plugin_runtime is not None + else {} + ), scratchpad_directory=result.scratchpad_directory, ) path = save_agent_session( @@ -1997,7 +2515,12 @@ class LocalCodingAgent: directory=self.runtime_config.session_directory, ) self.last_session_path = str(path) - return replace(result, session_path=self.last_session_path) + return replace( + result, + session_path=self.last_session_path, + events=tuple(persist_events), + transcript=session.transcript(), + ) def render_system_prompt(self) -> str: prompt_context = self.build_prompt_context() @@ -2101,13 +2624,47 @@ class LocalCodingAgent: if self.plugin_runtime is None: return prompt injections = self.plugin_runtime.before_prompt_injections() - if not injections: + state_reminder = self.plugin_runtime.runtime_state_reminder() + if not injections and not state_reminder: return prompt lines = ['', 'Plugin before-prompt hooks:'] lines.extend(f'- {entry}' for entry in injections) + if state_reminder: + lines.extend(['', state_reminder]) lines.extend(['', '', prompt]) return '\n'.join(lines) + def _apply_plugin_resume_hooks( + self, + prompt: str, + *, + resumed: bool, + ) -> str: + if not resumed or self.plugin_runtime is None: + return prompt + injections = self.plugin_runtime.on_resume_injections() + if not injections: + return prompt + lines = ['', 'Plugin resume hooks:'] + lines.extend(f'- {entry}' for entry in injections) + lines.extend(['', '', prompt]) + return '\n'.join(lines) + + def _render_plugin_persist_message( + self, + messages: tuple[str, ...], + ) -> str: + lines = ['', 'Plugin persist hooks:'] + lines.extend(f'- {entry}' for entry in messages) + lines.extend( + [ + '', + 'This session state was persisted with plugin lifecycle guidance.', + '', + ] + ) + return '\n'.join(lines) + def _append_plugin_after_turn_events( self, result: AgentRunResult, diff --git a/src/agent_session.py b/src/agent_session.py index 3c7e290..6504169 100644 --- a/src/agent_session.py +++ b/src/agent_session.py @@ -92,6 +92,7 @@ class AgentSessionState: user_context: dict[str, str] = field(default_factory=dict) system_context: dict[str, str] = field(default_factory=dict) messages: list[AgentMessage] = field(default_factory=list) + mutation_serial: int = 0 @classmethod def create( @@ -200,6 +201,7 @@ class AgentSessionState: previous_content=message.content, previous_state=message.state, previous_stop_reason=message.stop_reason, + mutation_serial=self._next_mutation_serial(), ) merged_metadata = _advance_lineage_revision(merged_metadata) self.messages[index] = replace( @@ -246,6 +248,7 @@ class AgentSessionState: previous_content=message.content, previous_state=message.state, previous_stop_reason=message.stop_reason, + mutation_serial=self._next_mutation_serial(), ) merged_metadata = _advance_lineage_revision(merged_metadata) self.messages[index] = replace( @@ -269,6 +272,7 @@ class AgentSessionState: previous_content=message.content, previous_state=message.state, previous_stop_reason=message.stop_reason, + mutation_serial=self._next_mutation_serial(), ) merged_metadata = _advance_lineage_revision(merged_metadata) self.messages[index] = replace( @@ -362,6 +366,7 @@ class AgentSessionState: previous_content=message.content, previous_state=message.state, previous_stop_reason=message.stop_reason, + mutation_serial=self._next_mutation_serial(), ) merged_metadata = _advance_lineage_revision(merged_metadata) if metadata: @@ -391,6 +396,7 @@ class AgentSessionState: previous_content=message.content, previous_state=message.state, previous_stop_reason=message.stop_reason, + mutation_serial=self._next_mutation_serial(), ) merged_metadata = _advance_lineage_revision(merged_metadata) if metadata: @@ -430,6 +436,7 @@ class AgentSessionState: previous_content=message.content, previous_state=message.state, previous_stop_reason=message.stop_reason, + mutation_serial=self._next_mutation_serial(), ) merged_metadata = _advance_lineage_revision(merged_metadata) if metadata: @@ -474,6 +481,10 @@ class AgentSessionState: def transcript(self) -> tuple[JSONDict, ...]: return tuple(message.to_transcript_entry() for message in self.messages) + def _next_mutation_serial(self) -> int: + self.mutation_serial += 1 + return self.mutation_serial + @classmethod def from_persisted( cls, @@ -488,6 +499,17 @@ class AgentSessionState: user_context=dict(user_context or {}), system_context=dict(system_context or {}), messages=[AgentMessage.from_openai_message(message) for message in messages], + mutation_serial=max( + ( + int(message.get('metadata', {}).get('last_mutation_serial', 0)) + for message in messages + if isinstance(message, dict) + and isinstance(message.get('metadata'), dict) + and isinstance(message.get('metadata', {}).get('last_mutation_serial', 0), int) + and not isinstance(message.get('metadata', {}).get('last_mutation_serial', 0), bool) + ), + default=0, + ), ) @@ -522,6 +544,7 @@ def _record_mutation( previous_content: str, previous_state: str, previous_stop_reason: str | None, + mutation_serial: int, ) -> JSONDict: mutations = metadata.get('mutations') if not isinstance(mutations, list): @@ -538,6 +561,7 @@ def _record_mutation( 'previous_stop_reason': previous_stop_reason, 'previous_content_length': len(previous_content), 'previous_content_preview': preview or '(empty)', + 'serial': mutation_serial, } ) if len(mutations) > MAX_MUTATION_HISTORY: @@ -545,6 +569,11 @@ def _record_mutation( metadata['mutations'] = mutations metadata['mutation_count'] = len(mutations) metadata['last_mutation_kind'] = mutation_kind + metadata['last_mutation_serial'] = mutation_serial + max_mutation_serial = metadata.get('max_mutation_serial') + if isinstance(max_mutation_serial, bool) or not isinstance(max_mutation_serial, int): + max_mutation_serial = 0 + metadata['max_mutation_serial'] = max(max_mutation_serial, mutation_serial) totals = metadata.get('mutation_totals') if not isinstance(totals, dict): totals = {} diff --git a/src/agent_tools.py b/src/agent_tools.py index 731be3f..fc5c1eb 100644 --- a/src/agent_tools.py +++ b/src/agent_tools.py @@ -242,6 +242,10 @@ def default_tool_registry() -> dict[str, AgentTool]: 'max_turns': {'type': 'integer', 'minimum': 1, 'maximum': 20}, 'resume_session_id': {'type': 'string'}, 'session_id': {'type': 'string'}, + 'depends_on': { + 'type': 'array', + 'items': {'type': 'string'}, + }, }, 'required': ['prompt'], }, @@ -255,6 +259,8 @@ def default_tool_registry() -> dict[str, AgentTool]: 'allow_shell': {'type': 'boolean'}, 'include_parent_context': {'type': 'boolean'}, 'continue_on_error': {'type': 'boolean'}, + 'max_failures': {'type': 'integer', 'minimum': 0, 'maximum': 20}, + 'strategy': {'type': 'string'}, }, }, handler=_delegate_agent_placeholder, diff --git a/src/agent_types.py b/src/agent_types.py index a496e70..a540f90 100644 --- a/src/agent_types.py +++ b/src/agent_types.py @@ -80,6 +80,8 @@ class BudgetConfig: max_total_cost_usd: float | None = None max_tool_calls: int | None = None max_delegated_tasks: int | None = None + max_model_calls: int | None = None + max_session_turns: int | None = None @dataclass(frozen=True) diff --git a/src/main.py b/src/main.py index 738d2be..7462f02 100644 --- a/src/main.py +++ b/src/main.py @@ -5,6 +5,7 @@ import os from pathlib import Path from dataclasses import replace import json +from typing import Callable from .agent_runtime import LocalCodingAgent from .agent_types import ( @@ -63,6 +64,8 @@ def _add_agent_common_args(parser: argparse.ArgumentParser, *, include_backend: 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('--max-model-calls', type=int) + parser.add_argument('--max-session-turns', type=int) parser.add_argument('--response-schema-file') parser.add_argument('--response-schema-name') parser.add_argument('--response-schema-strict', action='store_true') @@ -95,6 +98,8 @@ def _build_runtime_config(args: argparse.Namespace) -> AgentRuntimeConfig: 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), + max_model_calls=getattr(args, 'max_model_calls', None), + max_session_turns=getattr(args, 'max_session_turns', None), ), output_schema=_load_output_schema_config(args), session_directory=(Path('.port_sessions') / 'agent').resolve(), @@ -175,6 +180,8 @@ def _add_agent_resume_args(parser: argparse.ArgumentParser) -> None: 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('--max-model-calls', type=int) + parser.add_argument('--max-session-turns', type=int) parser.add_argument('--response-schema-file') parser.add_argument('--response-schema-name') parser.add_argument('--response-schema-strict', action='store_true') @@ -258,6 +265,8 @@ def _build_resumed_agent(args: argparse.Namespace) -> tuple[LocalCodingAgent, St or args.max_budget_usd is not None or args.max_tool_calls is not None or args.max_delegated_tasks is not None + or args.max_model_calls is not None + or args.max_session_turns is not None ): runtime_config = replace( runtime_config, @@ -297,6 +306,16 @@ def _build_resumed_agent(args: argparse.Namespace) -> tuple[LocalCodingAgent, St if args.max_delegated_tasks is not None else runtime_config.budget_config.max_delegated_tasks ), + max_model_calls=( + args.max_model_calls + if args.max_model_calls is not None + else runtime_config.budget_config.max_model_calls + ), + max_session_turns=( + args.max_session_turns + if args.max_session_turns is not None + else runtime_config.budget_config.max_session_turns + ), ), ) output_schema = _load_output_schema_config(args) @@ -339,6 +358,57 @@ def _print_agent_result(result, *, show_transcript: bool) -> None: print(message.get('content', '')) +def _run_agent_chat_loop( + agent: LocalCodingAgent, + *, + initial_prompt: str | None, + resume_session_id: str | None, + show_transcript: bool, + input_func: Callable[[str], str] = input, + output_func: Callable[[str], None] = print, + result_printer: Callable[..., None] = _print_agent_result, +) -> int: + active_session_id = resume_session_id + first_prompt = initial_prompt + + output_func('# Agent Chat') + output_func("Enter a prompt. Use '/exit' or '/quit' to stop.") + if active_session_id: + output_func(f'resuming_session_id={active_session_id}') + + while True: + if first_prompt is not None: + prompt = first_prompt + first_prompt = None + else: + try: + prompt = input_func('user> ') + except EOFError: + output_func('chat_ended=eof') + return 0 + except KeyboardInterrupt: + output_func('\nchat_ended=interrupt') + return 130 + + normalized = prompt.strip() + if not normalized: + continue + if normalized in {'/exit', '/quit'}: + output_func('chat_ended=user_exit') + return 0 + + if active_session_id: + stored_session = load_agent_session( + active_session_id, + directory=agent.runtime_config.session_directory, + ) + result = agent.resume(prompt, stored_session) + else: + result = agent.run(prompt) + result_printer(result, show_transcript=show_transcript) + active_session_id = result.session_id + + def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description='Python porting workspace for the Claude Code rewrite effort') subparsers = parser.add_subparsers(dest='command', required=True) @@ -417,6 +487,13 @@ def build_parser() -> argparse.ArgumentParser: agent_parser.add_argument('--show-transcript', action='store_true') _add_agent_common_args(agent_parser, include_backend=True) + chat_parser = subparsers.add_parser('agent-chat', help='run an interactive Python local-model chat loop') + chat_parser.add_argument('prompt', nargs='?') + chat_parser.add_argument('--resume-session-id') + chat_parser.add_argument('--max-turns', type=int, default=12) + chat_parser.add_argument('--show-transcript', action='store_true') + _add_agent_common_args(chat_parser, include_backend=True) + resume_parser = subparsers.add_parser('agent-resume', help='resume a saved Python local-model agent session') _add_agent_resume_args(resume_parser) @@ -563,6 +640,14 @@ def main(argv: list[str] | None = None) -> int: result = agent.run(args.prompt) _print_agent_result(result, show_transcript=args.show_transcript) return 0 + if args.command == 'agent-chat': + agent = _build_agent(args) + return _run_agent_chat_loop( + agent, + initial_prompt=args.prompt, + resume_session_id=args.resume_session_id, + show_transcript=args.show_transcript, + ) if args.command == 'agent-resume': agent, stored_session = _build_resumed_agent(args) result = agent.resume(args.prompt, stored_session) diff --git a/src/plugin_runtime.py b/src/plugin_runtime.py index a61f71e..3a92c8d 100644 --- a/src/plugin_runtime.py +++ b/src/plugin_runtime.py @@ -44,11 +44,16 @@ class PluginManifest: blocked_tools: tuple[str, ...] = () before_prompt: str | None = None after_turn: str | None = None + on_resume: str | None = None + before_persist: str | None = None + before_delegate: str | None = None + after_delegate: str | None = None @dataclass class PluginRuntime: manifests: tuple[PluginManifest, ...] = field(default_factory=tuple) + session_state: dict[str, Any] = field(default_factory=dict) @classmethod def from_workspace( @@ -94,18 +99,83 @@ class PluginRuntime: return tuple(blocks) def before_prompt_injections(self) -> tuple[str, ...]: - return tuple( + if not self.manifests: + return () + injections = tuple( manifest.before_prompt for manifest in self.manifests if manifest.before_prompt ) + if injections: + self.session_state['before_prompt_calls'] = int( + self.session_state.get('before_prompt_calls', 0) + ) + 1 + return injections def after_turn_injections(self) -> tuple[str, ...]: - return tuple( + if not self.manifests: + return () + injections = tuple( manifest.after_turn for manifest in self.manifests if manifest.after_turn ) + if injections: + self.session_state['after_turn_calls'] = int( + self.session_state.get('after_turn_calls', 0) + ) + 1 + return injections + + def on_resume_injections(self) -> tuple[str, ...]: + if not self.manifests: + return () + injections = tuple( + manifest.on_resume + for manifest in self.manifests + if manifest.on_resume + ) + if injections: + self.session_state['resume_calls'] = int( + self.session_state.get('resume_calls', 0) + ) + 1 + return injections + + def before_persist_injections(self) -> tuple[str, ...]: + if not self.manifests: + return () + injections = tuple( + manifest.before_persist + for manifest in self.manifests + if manifest.before_persist + ) + if injections: + self.session_state['persist_calls'] = int( + self.session_state.get('persist_calls', 0) + ) + 1 + return injections + + def before_delegate_injections(self) -> tuple[str, ...]: + if not self.manifests: + return () + injections = tuple( + manifest.before_delegate + for manifest in self.manifests + if manifest.before_delegate + ) + if injections: + self.session_state['delegate_calls'] = int( + self.session_state.get('delegate_calls', 0) + ) + 1 + return injections + + def after_delegate_injections(self) -> tuple[str, ...]: + if not self.manifests: + return () + return tuple( + manifest.after_delegate + for manifest in self.manifests + if manifest.after_delegate + ) def register_tool_aliases( self, @@ -198,6 +268,111 @@ class PluginRuntime: lines.append(f"- {'; '.join(details)}") if len(self.manifests) > 10: lines.append(f'- ... plus {len(self.manifests) - 10} more plugin manifests') + if self.session_state: + lines.append( + '- runtime_state=' + + ', '.join( + f'{name}={value}' + for name, value in sorted(self.session_state.items()) + if isinstance(value, (int, float, str, bool)) + ) + ) + return '\n'.join(lines) + + def record_tool_attempt(self, tool_name: str, *, blocked: bool) -> None: + attempts = int(self.session_state.get('tool_attempts', 0)) + self.session_state['tool_attempts'] = attempts + 1 + if blocked: + blocked_count = int(self.session_state.get('blocked_tool_attempts', 0)) + self.session_state['blocked_tool_attempts'] = blocked_count + 1 + counts = self.session_state.get('tool_attempt_counts') + if not isinstance(counts, dict): + counts = {} + counts[tool_name] = int(counts.get(tool_name, 0)) + 1 + self.session_state['tool_attempt_counts'] = counts + + def record_tool_result( + self, + tool_name: str, + *, + ok: bool, + metadata: dict[str, Any] | None = None, + ) -> None: + counts = self.session_state.get('tool_result_counts') + if not isinstance(counts, dict): + counts = {} + counts[tool_name] = int(counts.get(tool_name, 0)) + 1 + self.session_state['tool_result_counts'] = counts + key = 'successful_tool_results' if ok else 'failed_tool_results' + self.session_state[key] = int(self.session_state.get(key, 0)) + 1 + if isinstance(metadata, dict) and metadata.get('action') == 'plugin_virtual_tool': + self.session_state['virtual_tool_results'] = int( + self.session_state.get('virtual_tool_results', 0) + ) + 1 + + def export_session_state(self) -> dict[str, Any]: + exported: dict[str, Any] = {} + for key, value in self.session_state.items(): + if isinstance(value, dict): + exported[key] = { + str(name): count + for name, count in value.items() + if isinstance(name, str) + and isinstance(count, int) + and not isinstance(count, bool) + } + elif isinstance(value, (int, float, str, bool)): + exported[key] = value + return exported + + def restore_session_state(self, payload: dict[str, Any] | None) -> None: + if not isinstance(payload, dict): + self.session_state = {} + return + restored: dict[str, Any] = {} + for key, value in payload.items(): + if isinstance(value, dict): + restored[key] = { + str(name): int(count) + for name, count in value.items() + if isinstance(name, str) + and isinstance(count, int) + and not isinstance(count, bool) + } + elif isinstance(value, (int, float, str, bool)): + restored[key] = value + self.session_state = restored + + def runtime_state_reminder(self) -> str | None: + if not self.manifests or not self.session_state: + return None + lines = ['Plugin runtime state:'] + before_prompt_calls = self.session_state.get('before_prompt_calls') + if isinstance(before_prompt_calls, int): + lines.append(f'- before_prompt_calls={before_prompt_calls}') + after_turn_calls = self.session_state.get('after_turn_calls') + if isinstance(after_turn_calls, int): + lines.append(f'- after_turn_calls={after_turn_calls}') + tool_attempts = self.session_state.get('tool_attempts') + if isinstance(tool_attempts, int): + lines.append(f'- tool_attempts={tool_attempts}') + blocked_attempts = self.session_state.get('blocked_tool_attempts') + if isinstance(blocked_attempts, int): + lines.append(f'- blocked_tool_attempts={blocked_attempts}') + resume_calls = self.session_state.get('resume_calls') + if isinstance(resume_calls, int): + lines.append(f'- resume_calls={resume_calls}') + persist_calls = self.session_state.get('persist_calls') + if isinstance(persist_calls, int): + lines.append(f'- persist_calls={persist_calls}') + delegate_calls = self.session_state.get('delegate_calls') + if isinstance(delegate_calls, int): + lines.append(f'- delegate_calls={delegate_calls}') + virtual_results = self.session_state.get('virtual_tool_results') + if isinstance(virtual_results, int): + lines.append(f'- virtual_tool_results={virtual_results}') + if len(lines) == 1: + return None return '\n'.join(lines) @@ -248,7 +423,15 @@ def _load_manifest(path: Path) -> PluginManifest | 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')) + ( + before_prompt, + after_turn, + on_resume, + before_persist, + before_delegate, + after_delegate, + hook_names, + ) = _parse_hooks(payload.get('hooks')) return PluginManifest( name=name.strip(), path=str(path), @@ -266,6 +449,10 @@ def _load_manifest(path: Path) -> PluginManifest | None: ), before_prompt=before_prompt, after_turn=after_turn, + on_resume=on_resume, + before_persist=before_persist, + before_delegate=before_delegate, + after_delegate=after_delegate, ) @@ -311,10 +498,20 @@ def _extract_tool_aliases(payload: dict[str, Any]) -> tuple[PluginToolAlias, ... return tuple(aliases) -def _parse_hooks(value: Any) -> tuple[str | None, str | None, tuple[str, ...]]: +def _parse_hooks( + value: Any, +) -> tuple[ + str | None, + str | None, + str | None, + str | None, + 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 + return None, None, None, None, 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') @@ -323,12 +520,28 @@ def _parse_hooks(value: Any) -> tuple[str | None, str | None, tuple[str, ...]]: after_turn = value.get('afterTurn') if after_turn is None: after_turn = value.get('after_turn') + on_resume = value.get('onResume') + if on_resume is None: + on_resume = value.get('on_resume') + before_persist = value.get('beforePersist') + if before_persist is None: + before_persist = value.get('before_persist') + before_delegate = value.get('beforeDelegate') + if before_delegate is None: + before_delegate = value.get('before_delegate') + after_delegate = value.get('afterDelegate') + if after_delegate is None: + after_delegate = value.get('after_delegate') return ( _optional_string(before_prompt), _optional_string(after_turn), + _optional_string(on_resume), + _optional_string(before_persist), + _optional_string(before_delegate), + _optional_string(after_delegate), names, ) - return None, None, () + return None, None, None, None, None, None, () def _extract_tool_hooks(payload: dict[str, Any]) -> tuple[PluginToolHook, ...]: diff --git a/src/query_engine.py b/src/query_engine.py index 06e1013..ac8c1c9 100644 --- a/src/query_engine.py +++ b/src/query_engine.py @@ -53,6 +53,7 @@ class QueryEnginePort: transcript_store: TranscriptStore = field(default_factory=TranscriptStore) runtime_agent: LocalCodingAgent | None = None plugin_runtime: PluginRuntime | None = None + runtime_cumulative_usage: UsageSummary = field(default_factory=UsageSummary) runtime_event_counts: dict[str, int] = field(default_factory=dict) runtime_message_kind_counts: dict[str, int] = field(default_factory=dict) runtime_mutation_counts: dict[str, int] = field(default_factory=dict) @@ -84,6 +85,7 @@ class QueryEnginePort: session_id=stored.session_id, mutable_messages=list(stored.messages), total_usage=UsageSummary(stored.input_tokens, stored.output_tokens), + runtime_cumulative_usage=UsageSummary(stored.input_tokens, stored.output_tokens), transcript_store=transcript, plugin_runtime=PluginRuntime.from_workspace(Path.cwd()), ) @@ -115,16 +117,34 @@ class QueryEnginePort: ) -> TurnResult: if self.config.use_runtime_agent and self.runtime_agent is not None: result = self._submit_runtime_message(prompt) + cumulative_usage = UsageSummary( + input_tokens=result.usage.input_tokens, + output_tokens=result.usage.output_tokens, + ) + usage = cumulative_usage + if self.runtime_cumulative_usage.input_tokens or self.runtime_cumulative_usage.output_tokens: + usage = UsageSummary( + input_tokens=max( + cumulative_usage.input_tokens - self.runtime_cumulative_usage.input_tokens, + 0, + ), + output_tokens=max( + cumulative_usage.output_tokens - self.runtime_cumulative_usage.output_tokens, + 0, + ), + ) + else: + usage = UsageSummary( + input_tokens=cumulative_usage.input_tokens, + output_tokens=cumulative_usage.output_tokens, + ) 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, - ), + usage=usage, stop_reason=result.stop_reason or 'completed', session_id=result.session_id, session_path=result.session_path, @@ -133,7 +153,12 @@ class QueryEnginePort: events=result.events, transcript=result.transcript, ) - self._record_turn(prompt, turn, denied_tools) + self._record_turn( + prompt, + turn, + denied_tools, + runtime_cumulative_usage=cumulative_usage, + ) return turn if len(self.mutable_messages) >= self.config.max_turns: @@ -345,6 +370,7 @@ class QueryEnginePort: prompt: str, turn: TurnResult, denied_tools: tuple[PermissionDenial, ...], + runtime_cumulative_usage: UsageSummary | None = None, ) -> None: self.mutable_messages.append(prompt) self.transcript_store.append(prompt, kind='prompt') @@ -352,7 +378,11 @@ class QueryEnginePort: if self.config.use_runtime_agent: self._record_runtime_turn(turn) self.permission_denials.extend(denied_tools) - self.total_usage = turn.usage + if runtime_cumulative_usage is not None: + self.runtime_cumulative_usage = runtime_cumulative_usage + self.total_usage = runtime_cumulative_usage + else: + self.total_usage = turn.usage self.last_turn = turn if turn.session_id is not None: self.session_id = turn.session_id @@ -537,6 +567,16 @@ class QueryEnginePort: self.runtime_context_reduction.get('preserved_tail_messages', 0) + preserved_tail_count ) + max_source_mutation_serial = metadata.get('max_source_mutation_serial') + if ( + isinstance(max_source_mutation_serial, int) + and not isinstance(max_source_mutation_serial, bool) + ): + current = self.runtime_context_reduction.get('max_source_mutation_serial', 0) + self.runtime_context_reduction['max_source_mutation_serial'] = max( + current, + max_source_mutation_serial, + ) compacted_lineage_ids = metadata.get('compacted_lineage_ids') if isinstance(compacted_lineage_ids, list): self.runtime_context_reduction['compacted_lineages'] = ( @@ -572,6 +612,16 @@ class QueryEnginePort: if isinstance(revision_count, int) and not isinstance(revision_count, bool): current = self.runtime_lineage_stats.get('max_revision_count', 0) self.runtime_lineage_stats['max_revision_count'] = max(current, revision_count) + max_mutation_serial = metadata.get('max_mutation_serial') + if ( + isinstance(max_mutation_serial, int) + and not isinstance(max_mutation_serial, bool) + ): + current = self.runtime_lineage_stats.get('max_mutation_serial', 0) + self.runtime_lineage_stats['max_mutation_serial'] = max( + current, + max_mutation_serial, + ) kind = metadata.get('kind') if kind == 'snipped_message': diff --git a/src/session_store.py b/src/session_store.py index 6c46c72..437e04e 100644 --- a/src/session_store.py +++ b/src/session_store.py @@ -64,6 +64,8 @@ class StoredAgentSession: usage: JSONDict total_cost_usd: float file_history: tuple[JSONDict, ...] + budget_state: JSONDict + plugin_state: JSONDict scratchpad_directory: str | None = None @@ -95,6 +97,16 @@ def load_agent_session(session_id: str, directory: Path | None = None) -> Stored file_history=tuple( entry for entry in data.get('file_history', []) if isinstance(entry, dict) ), + budget_state=( + dict(data.get('budget_state', {})) + if isinstance(data.get('budget_state'), dict) + else {} + ), + plugin_state=( + dict(data.get('plugin_state', {})) + if isinstance(data.get('plugin_state'), dict) + else {} + ), scratchpad_directory=( str(data['scratchpad_directory']) if isinstance(data.get('scratchpad_directory'), str) @@ -155,6 +167,8 @@ def serialize_runtime_config(runtime_config: AgentRuntimeConfig) -> JSONDict: '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, + 'max_model_calls': runtime_config.budget_config.max_model_calls, + 'max_session_turns': runtime_config.budget_config.max_session_turns, }, 'output_schema': ( { @@ -205,6 +219,8 @@ def deserialize_runtime_config(payload: JSONDict) -> AgentRuntimeConfig: 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')), + max_model_calls=_optional_int(budget_payload.get('max_model_calls')), + max_session_turns=_optional_int(budget_payload.get('max_session_turns')), ), output_schema=_deserialize_output_schema(output_schema_payload), session_directory=Path(str(payload.get('session_directory', DEFAULT_AGENT_SESSION_DIR))).resolve(), diff --git a/tests/test_agent_runtime.py b/tests/test_agent_runtime.py index 87ddf7b..245ce98 100644 --- a/tests/test_agent_runtime.py +++ b/tests/test_agent_runtime.py @@ -1872,6 +1872,457 @@ class AgentRuntimeTests(unittest.TestCase): self.assertIn('delegated-task budget', result.final_output) self.assertFalse(any(message.get('role') == 'tool' for message in result.transcript)) + def test_agent_enforces_cumulative_model_call_budget_across_resume(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'First answer.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 4, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Second answer should exceed the model-call budget.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 4, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + session_dir = workspace / '.port_sessions' / 'agent' + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + session_directory=session_dir, + budget_config=BudgetConfig(max_model_calls=1), + ), + ) + first = agent.run('First prompt') + stored = load_agent_session(first.session_id or '', directory=session_dir) + resumed = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + session_directory=session_dir, + budget_config=BudgetConfig(max_model_calls=1), + ), + ) + second = resumed.resume('Second prompt', stored) + self.assertEqual(first.stop_reason, 'stop') + self.assertEqual(second.stop_reason, 'budget_exceeded') + self.assertIn('model-call budget', second.final_output) + + def test_plugin_runtime_state_is_restored_on_resume(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Calling the virtual plugin tool.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'demo_virtual', + 'arguments': '{"topic": "state"}', + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Initial plugin state stored.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Resume consumed plugin runtime state.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + ] + recorded_payloads: list[dict[str, object]] = [] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + session_dir = workspace / '.port_sessions' / 'agent' + plugin_dir = workspace / 'plugins' / 'demo' + plugin_dir.mkdir(parents=True) + (plugin_dir / 'plugin.json').write_text( + json.dumps( + { + 'name': 'demo-plugin', + 'hooks': {'beforePrompt': 'Plugin hook before prompt.'}, + 'virtualTools': [ + { + 'name': 'demo_virtual', + 'description': 'Emit plugin state.', + 'responseTemplate': 'plugin state {topic}', + } + ], + } + ), + 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, + session_directory=session_dir, + ), + ) + first = agent.run('Use the plugin virtual tool') + stored = load_agent_session(first.session_id or '', directory=session_dir) + resumed = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + session_directory=session_dir, + ), + ) + second = resumed.resume('Continue after plugin state', stored) + self.assertEqual(second.final_output, 'Resume consumed plugin runtime state.') + resumed_messages = recorded_payloads[-1]['messages'] + assert isinstance(resumed_messages, list) + self.assertTrue( + any( + isinstance(message, dict) + and 'Plugin runtime state:' in str(message.get('content', '')) + and 'virtual_tool_results=1' in str(message.get('content', '')) + for message in resumed_messages + ) + ) + + def test_plugin_lifecycle_hooks_are_persisted_and_reapplied_on_resume(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Initial lifecycle run stored.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Resume observed plugin lifecycle hooks.', + }, + '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) + session_dir = workspace / '.port_sessions' / 'agent' + plugin_dir = workspace / 'plugins' / 'demo' + plugin_dir.mkdir(parents=True) + (plugin_dir / 'plugin.json').write_text( + json.dumps( + { + 'name': 'demo-plugin', + 'hooks': { + 'onResume': 'Reapply the saved plugin lifecycle state on resume.', + 'beforePersist': 'Persist plugin lifecycle markers into the saved session.', + }, + } + ), + 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, + session_directory=session_dir, + ), + ) + first = agent.run('Store the plugin lifecycle session') + stored = load_agent_session(first.session_id or '', directory=session_dir) + resumed = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig( + cwd=workspace, + session_directory=session_dir, + ), + ) + second = resumed.resume('Continue after lifecycle persistence', stored) + self.assertEqual(second.final_output, 'Resume observed plugin lifecycle hooks.') + self.assertTrue( + any( + message.get('metadata', {}).get('kind') == 'plugin_persist' + and 'Persist plugin lifecycle markers into the saved session.' in str(message.get('content', '')) + for message in stored.messages + ) + ) + resumed_messages = recorded_payloads[-1]['messages'] + assert isinstance(resumed_messages, list) + self.assertTrue( + any( + isinstance(message, dict) + and 'Plugin resume hooks:' in str(message.get('content', '')) + and 'Reapply the saved plugin lifecycle state on resume.' in str(message.get('content', '')) + for message in resumed_messages + ) + ) + self.assertTrue(any(event.get('type') == 'plugin_before_persist' for event in first.events)) + + def test_agent_can_delegate_with_dependency_aware_subtasks(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Delegating dependency-aware subtasks.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'delegate_agent', + 'arguments': json.dumps( + { + 'subtasks': [ + { + 'label': 'summarize', + 'prompt': 'Summarize the project.', + 'depends_on': ['scan'], + }, + { + 'label': 'scan', + 'prompt': 'Scan the project.', + }, + ], + 'max_turns': 2, + } + ), + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Child scan result.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Parent completed after dependency-aware delegation.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + result = agent.run('Use dependency-aware delegated subtasks') + self.assertEqual(result.final_output, 'Parent completed after dependency-aware delegation.') + child_events = [event for event in result.events if event.get('type') == 'delegate_subtask_result'] + self.assertEqual(len(child_events), 2) + self.assertEqual(child_events[0].get('stop_reason'), 'pending_dependency') + self.assertEqual(child_events[1].get('stop_reason'), 'stop') + + def test_agent_can_delegate_with_topological_batches(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Delegating batched subtasks.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'delegate_agent', + 'arguments': json.dumps( + { + 'subtasks': [ + { + 'label': 'summarize', + 'prompt': 'Summarize the project.', + 'depends_on': ['scan'], + }, + { + 'label': 'scan', + 'prompt': 'Scan the project.', + }, + ], + 'strategy': 'topological', + 'max_turns': 2, + } + ), + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Child scan result.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Child summary result.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Parent completed after topological delegation.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch('src.openai_compat.request.urlopen', side_effect=make_urlopen_side_effect(responses)): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + result = agent.run('Use topological delegated subtasks') + self.assertEqual(result.final_output, 'Parent completed after topological delegation.') + child_events = [event for event in result.events if event.get('type') == 'delegate_subtask_result'] + self.assertEqual([event.get('label') for event in child_events], ['scan', 'summarize']) + self.assertEqual([event.get('batch_index') for event in child_events], [1, 2]) + batch_events = [event for event in result.events if event.get('type') == 'delegate_batch_result'] + self.assertEqual(len(batch_events), 2) + self.assertEqual(batch_events[0].get('status'), 'completed') + delegate_entry = next( + entry for entry in result.file_history if entry.get('action') == 'delegate_agent' + ) + self.assertEqual(delegate_entry.get('delegate_batch_count'), 2) + self.assertEqual(delegate_entry.get('dependency_skips'), 0) + tool_message = next( + message + for message in result.transcript + if message.get('role') == 'tool' + and message.get('metadata', {}).get('action') == 'delegate_agent' + ) + self.assertEqual(tool_message['metadata'].get('strategy'), 'topological') + def test_agent_sends_response_schema_when_configured(self) -> None: responses = [ { diff --git a/tests/test_main.py b/tests/test_main.py new file mode 100644 index 0000000..6f840e4 --- /dev/null +++ b/tests/test_main.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import json +import tempfile +import unittest +from dataclasses import replace +from pathlib import Path +from unittest.mock import patch + +from src.main import _build_runtime_config, _build_agent, _run_agent_chat_loop, build_parser + + +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_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 MainCliTests(unittest.TestCase): + def test_build_runtime_config_parses_model_and_session_budget_flags(self) -> None: + parser = build_parser() + args = parser.parse_args( + [ + 'agent', + 'Summarize the repo', + '--cwd', + '.', + '--max-model-calls', + '3', + '--max-session-turns', + '5', + ] + ) + runtime_config = _build_runtime_config(args) + self.assertEqual(runtime_config.budget_config.max_model_calls, 3) + self.assertEqual(runtime_config.budget_config.max_session_turns, 5) + + def test_agent_chat_loop_runs_multiple_turns_and_reuses_session(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'First chat reply.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Second chat reply.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 6, 'completion_tokens': 2}, + }, + ] + recorded_results: list[str] = [] + recorded_lines: list[str] = [] + prompts = iter(['Second prompt', '/exit']) + + def _input(prompt: str) -> str: + return next(prompts) + + def _output(line: str) -> None: + recorded_lines.append(line) + + def _result_printer(result, *, show_transcript: bool) -> None: # noqa: ANN001 + recorded_results.append(result.final_output) + + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + session_dir = workspace / '.port_sessions' / 'agent' + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_urlopen_side_effect(responses), + ): + parser = build_parser() + args = parser.parse_args( + [ + 'agent-chat', + 'First prompt', + '--cwd', + str(workspace), + ] + ) + agent = _build_agent(args) + agent.runtime_config = replace( + agent.runtime_config, + session_directory=session_dir, + ) + exit_code = _run_agent_chat_loop( + agent, + initial_prompt=args.prompt, + resume_session_id=None, + show_transcript=False, + input_func=_input, + output_func=_output, + result_printer=_result_printer, + ) + + self.assertEqual(exit_code, 0) + self.assertEqual(recorded_results, ['First chat reply.', 'Second chat reply.']) + self.assertIn('# Agent Chat', recorded_lines) + self.assertIn('chat_ended=user_exit', recorded_lines) diff --git a/tests/test_query_engine_runtime.py b/tests/test_query_engine_runtime.py index e535037..2eccd62 100644 --- a/tests/test_query_engine_runtime.py +++ b/tests/test_query_engine_runtime.py @@ -903,6 +903,7 @@ class QueryEngineRuntimeTests(unittest.TestCase): self.assertIn('## Runtime Context Reduction', summary) self.assertIn('- compact_boundaries=1', summary) self.assertIn('- compacted_lineages=', summary) + self.assertIn('- max_source_mutation_serial=', summary) self.assertIn('## Runtime Lineage', summary) self.assertIn('- seen_lineages=', summary) self.assertIn('- compacted_lineages=', summary) @@ -1148,3 +1149,190 @@ class QueryEngineRuntimeTests(unittest.TestCase): self.assertEqual(turn.output, 'Parent completed after resumed child.') self.assertIn('## Runtime Orchestration', summary) self.assertIn('- resumed_children=1', summary) + + def test_query_engine_runtime_summary_tracks_topological_delegate_batches(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Delegating batched subtasks.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'delegate_agent', + 'arguments': json.dumps( + { + 'subtasks': [ + { + 'label': 'summarize', + 'prompt': 'Summarize the project.', + 'depends_on': ['scan'], + }, + { + 'label': 'scan', + 'prompt': 'Scan the project.', + }, + ], + 'strategy': 'topological', + 'max_turns': 2, + } + ), + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Child scan result.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Child summary result.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Parent completed after topological delegation.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_urlopen_side_effect(responses), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + engine = QueryEnginePort.from_runtime_agent(agent) + turn = engine.submit_message('Use topological delegated subtasks') + summary = engine.render_summary() + + self.assertEqual(turn.output, 'Parent completed after topological delegation.') + self.assertIn('## Runtime Events', summary) + self.assertIn('- delegate_batch_result=2', summary) + self.assertIn('## Agent Manager', summary) + self.assertIn('strategy=topological', summary) + self.assertIn('batches=2', summary) + + def test_query_engine_runtime_summary_tracks_dependency_skips(self) -> None: + responses = [ + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Delegating dependency-aware subtasks.', + 'tool_calls': [ + { + 'id': 'call_1', + 'type': 'function', + 'function': { + 'name': 'delegate_agent', + 'arguments': json.dumps( + { + 'subtasks': [ + { + 'label': 'summarize', + 'prompt': 'Summarize the project.', + 'depends_on': ['scan'], + }, + { + 'label': 'scan', + 'prompt': 'Scan the project.', + }, + ], + 'max_turns': 2, + } + ), + }, + } + ], + }, + 'finish_reason': 'tool_calls', + } + ], + 'usage': {'prompt_tokens': 8, 'completion_tokens': 3}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Child scan result.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 5, 'completion_tokens': 2}, + }, + { + 'choices': [ + { + 'message': { + 'role': 'assistant', + 'content': 'Parent completed after dependency-aware delegation.', + }, + 'finish_reason': 'stop', + } + ], + 'usage': {'prompt_tokens': 7, 'completion_tokens': 2}, + }, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir) + with patch( + 'src.openai_compat.request.urlopen', + side_effect=make_urlopen_side_effect(responses), + ): + agent = LocalCodingAgent( + model_config=ModelConfig( + model='Qwen/Qwen3-Coder-30B-A3B-Instruct', + base_url='http://127.0.0.1:8000/v1', + ), + runtime_config=AgentRuntimeConfig(cwd=workspace), + ) + engine = QueryEnginePort.from_runtime_agent(agent) + turn = engine.submit_message('Use dependency-aware delegated subtasks') + summary = engine.render_summary() + + self.assertEqual(turn.output, 'Parent completed after dependency-aware delegation.') + self.assertIn('## Runtime Orchestration', summary) + self.assertIn('- child_stop:pending_dependency=1', summary) + self.assertIn('- child_stop:stop=1', summary)