From ae242c507d53560bd4ba780869fe3c5655a01661 Mon Sep 17 00:00:00 2001 From: Greyson LaLonde Date: Mon, 20 Apr 2026 22:47:37 +0800 Subject: [PATCH] feat: add checkpoint and fork support to standalone agents Add fork classmethod, _restore_runtime, and _restore_event_scope to BaseAgent. Fix from_checkpoint to set runtime state on the event bus and restore event scopes. Store kickoff event ID across checkpoints to skip re-emission on resume. Handle agent entity type in checkpoint CLI and TUI. --- lib/crewai/src/crewai/agent/core.py | 89 +++++++++++++------ .../crewai/agents/agent_builder/base_agent.py | 75 ++++++++++++++-- lib/crewai/src/crewai/cli/checkpoint_cli.py | 7 ++ lib/crewai/src/crewai/cli/checkpoint_tui.py | 23 ++++- lib/crewai/src/crewai/state/runtime.py | 3 + lib/crewai/tests/test_checkpoint.py | 72 +++++++++++++++ 6 files changed, 233 insertions(+), 36 deletions(-) diff --git a/lib/crewai/src/crewai/agent/core.py b/lib/crewai/src/crewai/agent/core.py index b83752344..74a3e85de 100644 --- a/lib/crewai/src/crewai/agent/core.py +++ b/lib/crewai/src/crewai/agent/core.py @@ -29,7 +29,7 @@ from pydantic import ( model_validator, ) from pydantic.functional_serializers import PlainSerializer -from typing_extensions import Self +from typing_extensions import Self, TypeIs from crewai.agent.planning_config import PlanningConfig from crewai.agent.utils import ( @@ -133,6 +133,13 @@ _EXECUTOR_CLASS_MAP: dict[str, type] = { } +def _is_resuming_agent_executor( + executor: CrewAgentExecutor | AgentExecutor | None, +) -> TypeIs[AgentExecutor]: + """Type guard: True when the executor is resuming from a checkpoint.""" + return isinstance(executor, AgentExecutor) and executor._resuming + + def _validate_executor_class(value: Any) -> Any: if isinstance(value, str): cls = _EXECUTOR_CLASS_MAP.get(value) @@ -1366,24 +1373,42 @@ class Agent(BaseAgent): prompt, stop_words, rpm_limit_fn = self._build_execution_prompt(raw_tools) - executor = AgentExecutor( - llm=cast(BaseLLM, self.llm), - agent=self, - prompt=prompt, - max_iter=self.max_iter, - tools=parsed_tools, - tools_names=get_tool_names(parsed_tools), - stop_words=stop_words, - tools_description=render_text_description_and_args(parsed_tools), - tools_handler=self.tools_handler, - original_tools=raw_tools, - step_callback=self.step_callback, - function_calling_llm=self.function_calling_llm, - respect_context_window=self.respect_context_window, - request_within_rpm_limit=rpm_limit_fn, - callbacks=[TokenCalcHandler(self._token_process)], - response_model=response_format, - ) + if _is_resuming_agent_executor(self.agent_executor): + executor = self.agent_executor + executor.tools = parsed_tools + executor.tools_names = get_tool_names(parsed_tools) + executor.tools_description = render_text_description_and_args(parsed_tools) + executor.original_tools = raw_tools + executor.prompt = prompt + executor.response_model = response_format + executor.stop_words = stop_words + executor.tools_handler = self.tools_handler + executor.step_callback = self.step_callback + executor.function_calling_llm = cast( + BaseLLM | None, self.function_calling_llm + ) + executor.respect_context_window = self.respect_context_window + executor.request_within_rpm_limit = rpm_limit_fn + executor.callbacks = [TokenCalcHandler(self._token_process)] + else: + executor = AgentExecutor( + llm=cast(BaseLLM, self.llm), + agent=self, + prompt=prompt, + max_iter=self.max_iter, + tools=parsed_tools, + tools_names=get_tool_names(parsed_tools), + stop_words=stop_words, + tools_description=render_text_description_and_args(parsed_tools), + tools_handler=self.tools_handler, + original_tools=raw_tools, + step_callback=self.step_callback, + function_calling_llm=self.function_calling_llm, + respect_context_window=self.respect_context_window, + request_within_rpm_limit=rpm_limit_fn, + callbacks=[TokenCalcHandler(self._token_process)], + response_model=response_format, + ) all_files: dict[str, Any] = {} if isinstance(messages, str): @@ -1504,14 +1529,17 @@ class Agent(BaseAgent): ) try: - crewai_event_bus.emit( - self, - event=LiteAgentExecutionStartedEvent( + if self.checkpoint_kickoff_event_id is not None: + self._kickoff_event_id = self.checkpoint_kickoff_event_id + self.checkpoint_kickoff_event_id = None + else: + started_event = LiteAgentExecutionStartedEvent( agent_info=agent_info, tools=parsed_tools, messages=messages, - ), - ) + ) + crewai_event_bus.emit(self, event=started_event) + self._kickoff_event_id = started_event.event_id output = self._execute_and_build_output(executor, inputs, response_format) return self._finalize_kickoff( @@ -1808,14 +1836,17 @@ class Agent(BaseAgent): ) try: - crewai_event_bus.emit( - self, - event=LiteAgentExecutionStartedEvent( + if self.checkpoint_kickoff_event_id is not None: + self._kickoff_event_id = self.checkpoint_kickoff_event_id + self.checkpoint_kickoff_event_id = None + else: + started_event = LiteAgentExecutionStartedEvent( agent_info=agent_info, tools=parsed_tools, messages=messages, - ), - ) + ) + crewai_event_bus.emit(self, event=started_event) + self._kickoff_event_id = started_event.event_id output = await self._execute_and_build_output_async( executor, inputs, response_format diff --git a/lib/crewai/src/crewai/agents/agent_builder/base_agent.py b/lib/crewai/src/crewai/agents/agent_builder/base_agent.py index a00f9b49f..74d30e0b2 100644 --- a/lib/crewai/src/crewai/agents/agent_builder/base_agent.py +++ b/lib/crewai/src/crewai/agents/agent_builder/base_agent.py @@ -28,6 +28,9 @@ from crewai.agents.agent_builder.base_agent_executor import BaseAgentExecutor from crewai.agents.agent_builder.utilities.base_token_process import TokenProcess from crewai.agents.cache.cache_handler import CacheHandler from crewai.agents.tools_handler import ToolsHandler +from crewai.events.base_events import set_emission_counter +from crewai.events.event_bus import crewai_event_bus +from crewai.events.event_context import restore_event_scope, set_last_event_id from crewai.knowledge.knowledge import Knowledge from crewai.knowledge.knowledge_config import KnowledgeConfig from crewai.knowledge.source.base_knowledge_source import BaseKnowledgeSource @@ -51,6 +54,7 @@ from crewai.utilities.string_utils import interpolate_only if TYPE_CHECKING: from crewai.context import ExecutionContext from crewai.crew import Crew + from crewai.state.runtime import RuntimeState def _validate_crew_ref(value: Any) -> Any: @@ -219,6 +223,7 @@ class BaseAgent(BaseModel, ABC, metaclass=AgentMeta): _original_goal: str | None = PrivateAttr(default=None) _original_backstory: str | None = PrivateAttr(default=None) _token_process: TokenProcess = PrivateAttr(default_factory=TokenProcess) + _kickoff_event_id: str | None = PrivateAttr(default=None) id: UUID4 = Field(default_factory=uuid.uuid4, frozen=True) role: str = Field(description="Role of the agent") goal: str = Field(description="Objective of the agent") @@ -335,30 +340,90 @@ class BaseAgent(BaseModel, ABC, metaclass=AgentMeta): min_length=1, ) execution_context: ExecutionContext | None = Field(default=None) + checkpoint_kickoff_event_id: str | None = Field(default=None) @classmethod def from_checkpoint(cls, config: CheckpointConfig) -> Self: - """Restore an Agent from a checkpoint. + """Restore an Agent from a checkpoint, ready to resume via kickoff(). Args: - config: Checkpoint configuration with ``restore_from`` set. + config: Checkpoint configuration with ``restore_from`` set to + the path of the checkpoint to load. + + Returns: + An Agent instance. Call kickoff() to resume execution. """ from crewai.context import apply_execution_context from crewai.state.runtime import RuntimeState state = RuntimeState.from_checkpoint(config, context={"from_checkpoint": True}) + crewai_event_bus.set_runtime_state(state) for entity in state.root: if isinstance(entity, cls): if entity.execution_context is not None: apply_execution_context(entity.execution_context) - if entity.agent_executor is not None: - entity.agent_executor.agent = entity - entity.agent_executor._resuming = True + entity._restore_runtime(state) return entity raise ValueError( f"No {cls.__name__} found in checkpoint: {config.restore_from}" ) + @classmethod + def fork(cls, config: CheckpointConfig, branch: str | None = None) -> Self: + """Fork an Agent from a checkpoint, creating a new execution branch. + + Args: + config: Checkpoint configuration with ``restore_from`` set. + branch: Branch label for the fork. Auto-generated if not provided. + + Returns: + An Agent instance on the new branch. Call kickoff() to run. + """ + agent = cls.from_checkpoint(config) + state = crewai_event_bus._runtime_state + if state is None: + raise RuntimeError("Cannot fork: no runtime state on the event bus.") + state.fork(branch) + return agent + + def _restore_runtime(self, state: RuntimeState) -> None: + """Re-create runtime objects after restoring from a checkpoint. + + Args: + state: The RuntimeState containing the event record. + """ + if self.agent_executor is not None: + self.agent_executor.agent = self + self.agent_executor._resuming = True + if self.checkpoint_kickoff_event_id is not None: + self._kickoff_event_id = self.checkpoint_kickoff_event_id + self._restore_event_scope(state) + + def _restore_event_scope(self, state: RuntimeState) -> None: + """Rebuild the event scope stack from the checkpoint's event record. + + Args: + state: The RuntimeState containing the event record. + """ + stack: list[tuple[str, str]] = [] + kickoff_id = self._kickoff_event_id + if kickoff_id: + stack.append((kickoff_id, "lite_agent_execution_started")) + + restore_event_scope(tuple(stack)) + + last_event_id: str | None = None + max_seq = 0 + for node in state.event_record.nodes.values(): + seq = node.event.emission_sequence or 0 + if seq > max_seq: + max_seq = seq + last_event_id = node.event.event_id + if last_event_id is not None: + set_last_event_id(last_event_id) + if max_seq > 0: + set_emission_counter(max_seq) + @model_validator(mode="before") @classmethod def process_model_config(cls, values: Any) -> dict[str, Any]: diff --git a/lib/crewai/src/crewai/cli/checkpoint_cli.py b/lib/crewai/src/crewai/cli/checkpoint_cli.py index 5e8572c62..0b3139d7d 100644 --- a/lib/crewai/src/crewai/cli/checkpoint_cli.py +++ b/lib/crewai/src/crewai/cli/checkpoint_cli.py @@ -472,6 +472,8 @@ def _entity_type_from_meta(meta: dict[str, Any]) -> str: for ent in meta.get("entities", []): if ent.get("type") == "flow": return "flow" + if ent.get("type") == "agent": + return "agent" return "crew" @@ -505,6 +507,11 @@ def resume_checkpoint(location: str, checkpoint_id: str | None) -> None: flow = Flow.from_checkpoint(config) result = asyncio.run(flow.kickoff_async(inputs=inputs)) + elif entity_type == "agent": + from crewai.agent import Agent + + agent = Agent.from_checkpoint(config) + result = asyncio.run(agent.akickoff(messages="Resume execution.")) else: from crewai.crew import Crew diff --git a/lib/crewai/src/crewai/cli/checkpoint_tui.py b/lib/crewai/src/crewai/cli/checkpoint_tui.py index 0f945dd31..7cc1d6867 100644 --- a/lib/crewai/src/crewai/cli/checkpoint_tui.py +++ b/lib/crewai/src/crewai/cli/checkpoint_tui.py @@ -123,7 +123,7 @@ _TuiResult = ( str, dict[str, Any] | None, dict[int, str] | None, - Literal["crew", "flow"], + Literal["crew", "flow", "agent"], ] | None ) @@ -682,10 +682,14 @@ class CheckpointTUI(App[_TuiResult]): overrides[task_idx] = editor.text return overrides or None - def _detect_entity_type(self, entry: dict[str, Any]) -> Literal["crew", "flow"]: + def _detect_entity_type( + self, entry: dict[str, Any] + ) -> Literal["crew", "flow", "agent"]: for ent in entry.get("entities", []): if ent.get("type") == "flow": return "flow" + if ent.get("type") == "agent": + return "agent" return "crew" def _resolve_location(self, entry: dict[str, Any]) -> str: @@ -829,6 +833,21 @@ async def _run_checkpoint_tui_async(location: str) -> None: click.echo(f"\nResult: {getattr(result, 'raw', result)}") return + if entity_type == "agent": + from crewai.agent import Agent + + if action == "fork": + click.echo(f"\nForking agent from: {selected}\n") + agent = Agent.fork(config) + else: + click.echo(f"\nResuming agent from: {selected}\n") + agent = Agent.from_checkpoint(config) + + click.echo() + result = await agent.akickoff(messages="Resume execution.") + click.echo(f"\nResult: {getattr(result, 'raw', result)}") + return + from crewai.crew import Crew if action == "fork": diff --git a/lib/crewai/src/crewai/state/runtime.py b/lib/crewai/src/crewai/state/runtime.py index daae0620e..3243d4c19 100644 --- a/lib/crewai/src/crewai/state/runtime.py +++ b/lib/crewai/src/crewai/state/runtime.py @@ -44,9 +44,12 @@ def _sync_checkpoint_fields(entity: object) -> None: entity: The entity whose private runtime attributes will be copied into its public checkpoint fields. """ + from crewai.agents.agent_builder.base_agent import BaseAgent from crewai.crew import Crew from crewai.flow.flow import Flow + if isinstance(entity, BaseAgent): + entity.checkpoint_kickoff_event_id = entity._kickoff_event_id if isinstance(entity, Flow): entity.checkpoint_completed_methods = ( set(entity._completed_methods) if entity._completed_methods else None diff --git a/lib/crewai/tests/test_checkpoint.py b/lib/crewai/tests/test_checkpoint.py index b1ad9e2df..d92a24803 100644 --- a/lib/crewai/tests/test_checkpoint.py +++ b/lib/crewai/tests/test_checkpoint.py @@ -562,3 +562,75 @@ class TestKickoffFromCheckpoint: ) assert mock_restored.checkpoint.restore_from is None assert result == "flow_result" + + +# ---------- Agent checkpoint/fork ---------- + + +class TestAgentCheckpoint: + def _make_agent_state(self) -> RuntimeState: + agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini") + return RuntimeState(root=[agent]) + + def test_agent_from_checkpoint_sets_runtime_state(self) -> None: + state = self._make_agent_state() + state._provider = JsonProvider() + with tempfile.TemporaryDirectory() as d: + loc = state.checkpoint(d) + cfg = CheckpointConfig(restore_from=loc) + + from crewai.events.event_bus import crewai_event_bus + + crewai_event_bus._runtime_state = None + Agent.from_checkpoint(cfg) + assert crewai_event_bus._runtime_state is not None + + def test_agent_fork_sets_branch(self) -> None: + state = self._make_agent_state() + state._provider = JsonProvider() + with tempfile.TemporaryDirectory() as d: + loc = state.checkpoint(d) + cfg = CheckpointConfig(restore_from=loc) + + from crewai.events.event_bus import crewai_event_bus + + Agent.fork(cfg, branch="agent-experiment") + rt = crewai_event_bus._runtime_state + assert rt is not None + assert rt._branch == "agent-experiment" + + def test_agent_fork_auto_branch(self) -> None: + state = self._make_agent_state() + state._provider = JsonProvider() + with tempfile.TemporaryDirectory() as d: + loc = state.checkpoint(d) + cfg = CheckpointConfig(restore_from=loc) + + from crewai.events.event_bus import crewai_event_bus + + Agent.fork(cfg) + rt = crewai_event_bus._runtime_state + assert rt is not None + assert rt._branch.startswith("fork/") + + def test_sync_checkpoint_fields_agent(self) -> None: + from crewai.state.runtime import _sync_checkpoint_fields + + agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini") + agent._kickoff_event_id = "evt-123" + _sync_checkpoint_fields(agent) + assert agent.checkpoint_kickoff_event_id == "evt-123" + + def test_agent_restore_kickoff_event_id(self) -> None: + agent = Agent(role="r", goal="g", backstory="b", llm="gpt-4o-mini") + agent._kickoff_event_id = "evt-456" + state = RuntimeState(root=[agent]) + state._provider = JsonProvider() + with tempfile.TemporaryDirectory() as d: + from crewai.state.runtime import _prepare_entities + + _prepare_entities(state.root) + loc = state.checkpoint(d) + cfg = CheckpointConfig(restore_from=loc) + restored = Agent.from_checkpoint(cfg) + assert restored._kickoff_event_id == "evt-456"