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.
This commit is contained in:
Greyson LaLonde
2026-04-20 22:47:37 +08:00
committed by GitHub
parent 0b120fac90
commit ae242c507d
6 changed files with 233 additions and 36 deletions

View File

@@ -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

View File

@@ -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]:

View File

@@ -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

View File

@@ -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":

View File

@@ -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

View File

@@ -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"