diff --git a/model_tools.py b/model_tools.py index 36cea8f304..b7dbb340b7 100644 --- a/model_tools.py +++ b/model_tools.py @@ -366,7 +366,7 @@ def get_tool_definitions( # because they need agent-level state (TodoStore, MemoryStore, etc.). # The registry still holds their schemas; dispatch just returns a stub error # so if something slips through, the LLM sees a sensible message. -_AGENT_LOOP_TOOLS = {"todo", "memory", "session_search", "delegate_task"} +_AGENT_LOOP_TOOLS = {"todo", "memory", "session_search", "delegate_task", "checkpoint"} _READ_SEARCH_TOOLS = {"read_file", "search_files"} diff --git a/run_agent.py b/run_agent.py index 6770f568c0..017812ac48 100644 --- a/run_agent.py +++ b/run_agent.py @@ -1561,6 +1561,11 @@ class AIAgent: # In-memory todo list for task planning (one per agent/session) from tools.todo_tool import TodoStore self._todo_store = TodoStore() + + # Layer 4: Checkpoint store (session resumption across compression) + from agent.checkpoint_store import CheckpointStore + self._checkpoint_store = CheckpointStore() + self._checkpoint_store.garbage_collect() # prune stale checkpoints on startup # Load config once for memory, skills, and compression sections try: @@ -8351,6 +8356,19 @@ class AIAgent: ) elif function_name == "delegate_task": return self._dispatch_delegate_task(function_args) + elif function_name == "checkpoint": + from tools.checkpoint_tool import checkpoint_tool as _checkpoint_tool + return _checkpoint_tool( + action=function_args.get("action"), + task=function_args.get("task"), + progress=function_args.get("progress"), + state=function_args.get("state"), + decisions=function_args.get("decisions"), + blocked=function_args.get("blocked"), + unresolved=function_args.get("unresolved"), + store=self._checkpoint_store, + agent=self, + ) else: return handle_function_call( function_name, function_args, effective_task_id, @@ -8892,6 +8910,22 @@ class AIAgent: spinner.stop(cute_msg) elif self._should_emit_quiet_tool_messages(): self._vprint(f" {cute_msg}") + elif function_name == "checkpoint": + from tools.checkpoint_tool import checkpoint_tool as _checkpoint_tool + function_result = _checkpoint_tool( + action=function_args.get("action"), + task=function_args.get("task"), + progress=function_args.get("progress"), + state=function_args.get("state"), + decisions=function_args.get("decisions"), + blocked=function_args.get("blocked"), + unresolved=function_args.get("unresolved"), + store=self._checkpoint_store, + agent=self, + ) + tool_duration = time.time() - tool_start_time + if self._should_emit_quiet_tool_messages(): + self._vprint(f" {_get_cute_tool_message_impl('checkpoint', function_args, tool_duration, result=function_result)}") elif self._context_engine_tool_names and function_name in self._context_engine_tool_names: # Context engine tools (lcm_grep, lcm_describe, lcm_expand, etc.) spinner = None diff --git a/tests/test_checkpoint_agent_dispatch.py b/tests/test_checkpoint_agent_dispatch.py new file mode 100644 index 0000000000..6e48a8201d --- /dev/null +++ b/tests/test_checkpoint_agent_dispatch.py @@ -0,0 +1,16 @@ +"""Verify that checkpoint is routed through the agent loop, not the registry.""" +import json +import pytest +from model_tools import handle_function_call, _AGENT_LOOP_TOOLS + + +def test_checkpoint_is_agent_loop_tool(): + """checkpoint must be in _AGENT_LOOP_TOOLS so it gets agent-level state.""" + assert "checkpoint" in _AGENT_LOOP_TOOLS + + +def test_checkpoint_registry_dispatch_returns_error(): + """Calling handle_function_call for checkpoint should return agent-loop error.""" + result = handle_function_call("checkpoint", {"action": "read"}) + data = json.loads(result) + assert "must be handled" in data["error"].lower() \ No newline at end of file