From 7d5b933aa4c66a2ecafced14fa8da7e212172d64 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aur=C3=A9lien=20Sibiril?= <81782+aureliensibiril@users.noreply.github.com> Date: Fri, 10 Apr 2026 18:53:06 +0200 Subject: [PATCH] Add checkpoint persistence to agent core loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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> --- pkg/agent/hooks.go | 2 + pkg/agent/run.go | 299 +++++++++++++++++++++++++++++++++++++------- pkg/agent/stream.go | 25 ++-- 3 files changed, 272 insertions(+), 54 deletions(-) diff --git a/pkg/agent/hooks.go b/pkg/agent/hooks.go index e1fd952bd..67b02fd83 100644 --- a/pkg/agent/hooks.go +++ b/pkg/agent/hooks.go @@ -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) {} diff --git a/pkg/agent/run.go b/pkg/agent/run.go index b6267be56..84b8eb363 100644 --- a/pkg/agent/run.go +++ b/pkg/agent/run.go @@ -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, }, ) } diff --git a/pkg/agent/stream.go b/pkg/agent/stream.go index b77171bf6..5743e4f0d 100644 --- a/pkg/agent/stream.go +++ b/pkg/agent/stream.go @@ -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