rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add checkpoint persistence to agent core loop

Browse Source 
coreLoop now saves incremental checkpoints after each tool-call turn
and checks a cooperative stop signal at turn boundaries. SuspendedError
is handled in finishRun, executeParallel, and executeSingleTool.
Approval-interrupted checkpoints are persisted for both flat and
nested interruptions.

Introduce RunOption, WithCheckpointStore, RunWithOpts, ResumeWithOpts,
and RunStreamedWithOpts so callers can provide checkpoint storage.
Add StreamEventSuspended and OnRunRestore hook.

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
This commit is contained in:
Aurélien Sibiril
2026-04-10 18:53:06 +02:00
parent 51df618c93
commit 7d5b933aa4
3 changed files with 272 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
pkg/agent/hooks.go
View File

@@ -24,6 +24,7 @@ import (
type RunHooks interface {
OnRunStart(ctx context.Context, agent *Agent, messages []llm.Message)
OnRunEnd(ctx context.Context, agent *Agent, result *Result, err error)
OnRunRestore(ctx context.Context, agent *Agent, checkpoint *Checkpoint)
OnLLMStart(ctx context.Context, agent *Agent, messages []llm.Message)
OnLLMEnd(ctx context.Context, agent *Agent, response *llm.ChatCompletionResponse, err error)
OnToolStart(ctx context.Context, agent *Agent, tool Tool, arguments string)
@@ -39,6 +40,7 @@ var _ RunHooks = NoOpHooks{}
func (NoOpHooks) OnRunStart(context.Context, *Agent, []llm.Message) {}
func (NoOpHooks) OnRunEnd(context.Context, *Agent, *Result, error) {}
func (NoOpHooks) OnRunRestore(context.Context, *Agent, *Checkpoint) {}
func (NoOpHooks) OnLLMStart(context.Context, *Agent, []llm.Message) {}
func (NoOpHooks) OnLLMEnd(context.Context, *Agent, *llm.ChatCompletionResponse, error) {}
func (NoOpHooks) OnToolStart(context.Context, *Agent, Tool, string) {}

299
pkg/agent/run.go
View File

@@ -40,6 +40,8 @@ const (
type (
CallLLMFunc func(ctx context.Context, agent *Agent, req *llm.ChatCompletionRequest) (*llm.ChatCompletionResponse, error)
RunOption func(*runOpts)
runOpts struct {
callLLM CallLLMFunc
onEvent func(ctx context.Context, ev StreamEvent)
@@ -47,6 +49,9 @@ type (
skipSessionLoad bool
initialUsage llm.Usage
initialTurns int
checkpointStore CheckpointStore
runID string
toolUsedInRun bool
}
loopState struct {
@@ -72,6 +77,13 @@ type (
}
)
func WithCheckpointStore(store CheckpointStore, runID string) RunOption {
return func(o *runOpts) {
o.checkpointStore = store
o.runID = runID
}
}
func noopEvent(_ context.Context, _ StreamEvent) {}
func blockingCallLLM(ctx context.Context, agent *Agent, req *llm.ChatCompletionRequest) (*llm.ChatCompletionResponse, error) {
@@ -104,15 +116,19 @@ func blockingCallLLM(ctx context.Context, agent *Agent, req *llm.ChatCompletionR
}
func (a *Agent) Run(ctx context.Context, messages []llm.Message) (*Result, error) {
return coreLoop(
ctx,
a,
messages,
runOpts{
callLLM: blockingCallLLM,
onEvent: noopEvent,
},
)
return a.RunWithOpts(ctx, messages)
}
func (a *Agent) RunWithOpts(ctx context.Context, messages []llm.Message, opts ...RunOption) (*Result, error) {
ro := runOpts{
callLLM: blockingCallLLM,
onEvent: noopEvent,
}
for _, opt := range opts {
opt(&ro)
}
return coreLoop(ctx, a, messages, ro)
}
func (s *loopState) resolveAgentTools(ctx context.Context) error {
@@ -138,6 +154,20 @@ func (s *loopState) finishRun(ctx context.Context, result *Result, err error) (*
}()
if err != nil {
if _, ok := errors.AsType[*SuspendedError](err); ok {
s.runSpan.SetAttributes(attribute.Bool("agent.suspended", true))
s.opts.onEvent(ctx, StreamEvent{Type: StreamEventSuspended, Agent: s.agent})
s.logger.InfoCtx(
ctx,
"agent run suspended",
log.String("agent", s.agent.name),
log.Int("turns", s.turns),
)
return result, err
}
s.runSpan.RecordError(err)
s.runSpan.SetStatus(codes.Error, err.Error())
s.opts.onEvent(ctx, StreamEvent{Type: StreamEventError, Agent: s.agent, Err: err})
@@ -192,6 +222,21 @@ func (s *loopState) finishRun(ctx context.Context, result *Result, err error) (*
return result, err
}
func (s *loopState) buildCheckpoint(status CheckpointStatus) *Checkpoint {
msgsCopy := make([]llm.Message, len(s.messages))
copy(msgsCopy, s.messages)
return &Checkpoint{
Version: CheckpointVersion,
Status: status,
AgentName: s.agent.name,
Messages: msgsCopy,
Usage: s.totalUsage,
Turns: s.turns,
ToolUsedInRun: s.toolUsedInRun,
}
}
func (s *loopState) applyHandoff(ctx context.Context, handoffTarget *Handoff) error {
emitHook(s.agent, func(h RunHooks) { h.OnHandoff(ctx, s.agent, handoffTarget.Agent) })
emitAgentHook(handoffTarget.Agent, func(h AgentHooks) { h.OnHandoff(ctx, handoffTarget.Agent, s.agent) })
@@ -238,6 +283,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
systemPrompt: startAgent.buildSystemPrompt(ctx),
totalUsage: opts.initialUsage,
turns: opts.initialTurns,
toolUsedInRun: opts.toolUsedInRun,
tracer: otel.GetTracerProvider().Tracer(tracerName),
opts: opts,
logger: startAgent.logger,
@@ -328,6 +374,25 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
return s.finishRun(ctx, nil, fmt.Errorf("cannot complete: %w", err))
}
if ch := stopSignalFrom(ctx); ch != nil {
select {
case <-ch:
cp := s.buildCheckpoint(CheckpointStatusSuspended)
se := &SuspendedError{RunID: s.opts.runID}
if s.opts.checkpointStore != nil && s.opts.runID != "" {
if saveErr := s.opts.checkpointStore.Save(ctx, s.opts.runID, cp); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save suspension checkpoint", log.Error(saveErr))
}
} else {
se.Checkpoint = cp
}
return s.finishRun(ctx, nil, se)
default:
}
}
if s.turns >= s.agent.maxTurns {
return s.finishRun(ctx, nil, &MaxTurnsExceededError{MaxTurns: s.agent.maxTurns})
}
@@ -485,6 +550,21 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
s.messages = append(s.messages, toolMsgs...)
if err != nil {
if se, ok := errors.AsType[*SuspendedError](err); ok {
outerCP := s.buildCheckpoint(CheckpointStatusSuspended)
if se.Checkpoint != nil {
outerCP.AllToolCalls = se.Checkpoint.AllToolCalls
outerCP.InnerCheckpoints = se.Checkpoint.InnerCheckpoints
outerCP.CompletedCalls = se.Checkpoint.CompletedCalls
}
if s.opts.checkpointStore != nil && s.opts.runID != "" {
if saveErr := s.opts.checkpointStore.Save(ctx, s.opts.runID, outerCP); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save checkpoint", log.Error(saveErr))
}
}
return s.finishRun(ctx, nil, &SuspendedError{RunID: s.opts.runID})
}
if nae, ok := errors.AsType[*needsApprovalError](err); ok {
s.logger.InfoCtx(
ctx,
@@ -495,6 +575,15 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
msgsCopy := make([]llm.Message, len(s.messages))
copy(msgsCopy, s.messages)
if s.opts.checkpointStore != nil && s.opts.runID != "" {
cp := s.buildCheckpoint(CheckpointStatusAwaitingApproval)
cp.PendingToolCalls = nae.allToolCalls
cp.PendingApprovals = nae.pendingApprovals
if saveErr := s.opts.checkpointStore.Save(ctx, s.opts.runID, cp); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save approval checkpoint", log.Error(saveErr))
}
}
return s.finishRun(
ctx,
nil,
@@ -520,6 +609,29 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
msgsCopy := make([]llm.Message, len(s.messages))
copy(msgsCopy, s.messages)
if s.opts.checkpointStore != nil && s.opts.runID != "" {
cp := s.buildCheckpoint(CheckpointStatusAwaitingApproval)
cp.PendingToolCalls = nie.inner.ToolCalls
cp.PendingApprovals = nie.inner.PendingApprovals
cp.AllToolCalls = nie.allToolCalls
cp.CompletedCalls = nie.completedCalls
cp.InnerCheckpoints = map[string]*Checkpoint{
nie.toolCallID: {
Version: CheckpointVersion,
Status: CheckpointStatusAwaitingApproval,
AgentName: nie.inner.Agent.name,
Messages: nie.inner.Messages,
Usage: nie.inner.Usage,
Turns: nie.inner.Turns,
PendingToolCalls: nie.inner.ToolCalls,
PendingApprovals: nie.inner.PendingApprovals,
},
}
if saveErr := s.opts.checkpointStore.Save(ctx, s.opts.runID, cp); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save nested approval checkpoint", log.Error(saveErr))
}
}
return s.finishRun(
ctx,
nil,
@@ -580,6 +692,14 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
}
}
// Save incremental checkpoint after completed tool-call turn.
if s.opts.checkpointStore != nil && s.opts.runID != "" {
cp := s.buildCheckpoint(CheckpointStatusSuspended)
if saveErr := s.opts.checkpointStore.Save(ctx, s.opts.runID, cp); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save checkpoint", log.Error(saveErr))
}
}
case llm.FinishReasonContentFilter:
return s.finishRun(ctx, nil, fmt.Errorf("cannot complete: content was filtered by the provider"))
@@ -703,13 +823,13 @@ func executeWithHandoff(
tr, err := executeSingleTool(ctx, tracer, agent, toolCalls[i], descriptors[i].(Tool), onEvent, logger)
if err != nil {
if ie, ok := errors.AsType[*InterruptedError](err); ok {
var completed []completedCall
var completed []CompletedCall
for j := range results {
completed = append(
completed,
completedCall{
toolCallID: toolCalls[j].ID,
result: results[j].Result,
CompletedCall{
ToolCallID: toolCalls[j].ID,
Result: results[j].Result,
},
)
}
@@ -803,40 +923,98 @@ func executeParallel(
wg.Wait()
for i, entry := range entries {
var ie *InterruptedError
if entry.err != nil && errors.As(entry.err, &ie) {
var completed []completedCall
ie, ok := errors.AsType[*InterruptedError](entry.err)
if !ok {
continue
}
var completed []CompletedCall
for j, other := range entries {
if j == i {
continue
}
if other.err != nil {
completed = append(
completed,
CompletedCall{
ToolCallID: toolCalls[j].ID,
Result: ToolResult{
Content: fmt.Sprintf("Error: %s", other.err.Error()),
IsError: true,
},
},
)
continue
}
completed = append(
completed,
CompletedCall{
ToolCallID: toolCalls[j].ID,
Result: other.result,
},
)
}
return nil, nil, &nestedInterruptionError{
inner: ie,
toolCallID: toolCalls[i].ID,
allToolCalls: toolCalls,
completedCalls: completed,
}
}
// Check for suspended inner agents (stop signal propagated).
for i, entry := range entries {
if entry.err == nil {
continue
}
se, ok := errors.AsType[*SuspendedError](entry.err)
if ok && se.Checkpoint != nil {
innerCheckpoints := make(map[string]*Checkpoint)
var completed []CompletedCall
for j, other := range entries {
if j == i {
continue
}
if other.err != nil {
if other.err == nil {
completed = append(
completed,
completedCall{
toolCallID: toolCalls[j].ID,
result: ToolResult{
CompletedCall{
ToolCallID: toolCalls[j].ID,
Result: other.result,
},
)
continue
}
otherSE, ok := errors.AsType[*SuspendedError](other.err)
if ok && otherSE.Checkpoint != nil {
innerCheckpoints[toolCalls[j].ID] = otherSE.Checkpoint
} else {
completed = append(
completed,
CompletedCall{
ToolCallID: toolCalls[j].ID,
Result: ToolResult{
Content: fmt.Sprintf("Error: %s", other.err.Error()),
IsError: true,
},
},
)
continue
}
completed = append(
completed,
completedCall{
toolCallID: toolCalls[j].ID,
result: other.result,
},
)
}
return nil, nil, &nestedInterruptionError{
inner: ie,
toolCallID: toolCalls[i].ID,
allToolCalls: toolCalls,
completedCalls: completed,
innerCheckpoints[toolCalls[i].ID] = se.Checkpoint
outerSE := &SuspendedError{
Checkpoint: &Checkpoint{
Version: CheckpointVersion,
Status: CheckpointStatusSuspended,
AllToolCalls: toolCalls,
InnerCheckpoints: innerCheckpoints,
CompletedCalls: completed,
},
}
return nil, nil, outerSE
}
}
@@ -939,6 +1117,17 @@ func executeSingleTool(
return ToolResult{}, err
}
if _, ok := errors.AsType[*SuspendedError](err); ok {
toolSpan.SetAttributes(attribute.Bool("tool.suspended", true))
toolSpan.End()
onEvent(ctx, StreamEvent{Type: StreamEventToolEnd, Agent: agent, Tool: tool, Err: err})
emitHook(agent, func(h RunHooks) { h.OnToolEnd(ctx, agent, tool, ToolResult{}, err) })
emitAgentHook(agent, func(h AgentHooks) { h.OnToolEnd(ctx, agent, tool, ToolResult{}) })
return ToolResult{}, err
}
toolSpan.RecordError(err)
toolSpan.SetStatus(codes.Error, err.Error())
toolSpan.End()
@@ -1055,8 +1244,24 @@ func runOutputGuardrails(ctx context.Context, agent *Agent, message llm.Message)
// are not re-evaluated because the messages were already validated in the
// original Run call.
func Resume(ctx context.Context, interrupted *InterruptedError, input ResumeInput) (*Result, error) {
return ResumeWithOpts(ctx, interrupted, input)
}
func ResumeWithOpts(ctx context.Context, interrupted *InterruptedError, input ResumeInput, opts ...RunOption) (*Result, error) {
ro := runOpts{
callLLM: blockingCallLLM,
onEvent: noopEvent,
}
for _, opt := range opts {
opt(&ro)
}
return resumeWithOpts(ctx, interrupted, input, ro)
}
func resumeWithOpts(ctx context.Context, interrupted *InterruptedError, input ResumeInput, ro runOpts) (*Result, error) {
if interrupted.outerState != nil {
return resumeNested(ctx, interrupted, input)
return resumeNested(ctx, interrupted, input, ro)
}
tracer := otel.GetTracerProvider().Tracer(tracerName)
@@ -1153,7 +1358,7 @@ func Resume(ctx context.Context, interrupted *InterruptedError, input ResumeInpu
break
}
tr, execErr := executeSingleTool(ctx, tracer, agent, tc, desc.(Tool), noopEvent, logger)
tr, execErr := executeSingleTool(ctx, tracer, agent, tc, desc.(Tool), ro.onEvent, logger)
if execErr != nil {
return nil, execErr
}
@@ -1202,17 +1407,20 @@ func Resume(ctx context.Context, interrupted *InterruptedError, input ResumeInpu
resumeAgent,
messages,
runOpts{
callLLM: blockingCallLLM,
onEvent: noopEvent,
callLLM: ro.callLLM,
onEvent: ro.onEvent,
skipInputGuardrails: true,
skipSessionLoad: true,
initialUsage: interrupted.Usage,
initialTurns: interrupted.Turns,
checkpointStore: ro.checkpointStore,
runID: ro.runID,
toolUsedInRun: ro.toolUsedInRun,
},
)
}
func resumeNested(ctx context.Context, interrupted *InterruptedError, input ResumeInput) (*Result, error) {
func resumeNested(ctx context.Context, interrupted *InterruptedError, input ResumeInput, ro runOpts) (*Result, error) {
outer := interrupted.outerState
logger := outer.agent.logger
@@ -1223,10 +1431,10 @@ func resumeNested(ctx context.Context, interrupted *InterruptedError, input Resu
log.String("inner_agent", interrupted.Agent.name),
)
innerResult, err := Resume(ctx, outer.innerInterrupt, input)
innerResult, err := resumeWithOpts(ctx, outer.innerInterrupt, input, ro)
if err != nil {
var innerIE *InterruptedError
if errors.As(err, &innerIE) {
innerIE, ok := errors.AsType[*InterruptedError](err)
if ok {
return nil, &InterruptedError{
ToolCalls: innerIE.ToolCalls,
PendingApprovals: innerIE.PendingApprovals,
@@ -1251,7 +1459,7 @@ func resumeNested(ctx context.Context, interrupted *InterruptedError, input Resu
completedMap := make(map[string]ToolResult, len(outer.completedCalls))
for _, cc := range outer.completedCalls {
completedMap[cc.toolCallID] = cc.result
completedMap[cc.ToolCallID] = cc.Result
}
messages := make([]llm.Message, len(outer.messages))
@@ -1282,12 +1490,15 @@ func resumeNested(ctx context.Context, interrupted *InterruptedError, input Resu
outer.agent,
messages,
runOpts{
callLLM: blockingCallLLM,
onEvent: noopEvent,
callLLM: ro.callLLM,
onEvent: ro.onEvent,
skipInputGuardrails: true,
skipSessionLoad: true,
initialUsage: outer.usage,
initialTurns: outer.turns,
checkpointStore: ro.checkpointStore,
runID: ro.runID,
toolUsedInRun: ro.toolUsedInRun,
},
)
}

25
pkg/agent/stream.go
View File

@@ -49,6 +49,7 @@ const (
StreamEventToolEnd StreamEventType = "tool_end"
StreamEventHandoff StreamEventType = "handoff"
StreamEventComplete StreamEventType = "complete"
StreamEventSuspended StreamEventType = "suspended"
StreamEventError StreamEventType = "error"
)
@@ -59,6 +60,10 @@ func (sr *StreamedRun) Wait() (*Result, error) {
}
func (a *Agent) RunStreamed(ctx context.Context, messages []llm.Message) *StreamedRun {
return a.RunStreamedWithOpts(ctx, messages)
}
func (a *Agent) RunStreamedWithOpts(ctx context.Context, messages []llm.Message, opts ...RunOption) *StreamedRun {
events := make(chan StreamEvent, 64)
sr := &StreamedRun{
Events: events,
@@ -69,17 +74,17 @@ func (a *Agent) RunStreamed(ctx context.Context, messages []llm.Message) *Stream
defer close(sr.done)
defer close(events)
result, err := coreLoop(
ctx,
a,
messages,
runOpts{
callLLM: streamingCallLLM(events),
onEvent: func(ctx context.Context, ev StreamEvent) {
trySendEvent(ctx, events, ev)
},
ro := runOpts{
callLLM: streamingCallLLM(events),
onEvent: func(ctx context.Context, ev StreamEvent) {
trySendEvent(ctx, events, ev)
},
)
}
for _, opt := range opts {
opt(&ro)
}
result, err := coreLoop(ctx, a, messages, ro)
sr.result = result
sr.err = err