Add checkpoint persistence to agent core loop

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

@@ -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)                    {}

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(
+	return a.RunWithOpts(ctx, messages)
-		ctx,
+}
-		a,
+
-		messages,
+func (a *Agent) RunWithOpts(ctx context.Context, messages []llm.Message, opts ...RunOption) (*Result, error) {
-		runOpts{
+	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{
+						CompletedCall{
-							toolCallID: toolCalls[j].ID,
+							ToolCallID: toolCalls[j].ID,
-							result:     results[j].Result,
+							Result:     results[j].Result,
 						},
 					)
 				}
@@ -803,9 +923,12 @@ func executeParallel(
 	wg.Wait()
 
 	for i, entry := range entries {
-		var ie *InterruptedError
+		ie, ok := errors.AsType[*InterruptedError](entry.err)
-		if entry.err != nil && errors.As(entry.err, &ie) {
+		if !ok {
-			var completed []completedCall
+			continue
+		}
+
+		var completed []CompletedCall
 		for j, other := range entries {
 			if j == i {
 				continue
@@ -813,9 +936,9 @@ func executeParallel(
 			if other.err != nil {
 				completed = append(
 					completed,
-						completedCall{
+					CompletedCall{
-							toolCallID: toolCalls[j].ID,
+						ToolCallID: toolCalls[j].ID,
-							result: ToolResult{
+						Result: ToolResult{
 							Content: fmt.Sprintf("Error: %s", other.err.Error()),
 							IsError: true,
 						},
@@ -825,9 +948,9 @@ func executeParallel(
 			}
 			completed = append(
 				completed,
-					completedCall{
+				CompletedCall{
-						toolCallID: toolCalls[j].ID,
+					ToolCallID: toolCalls[j].ID,
-						result:     other.result,
+					Result:     other.result,
 				},
 			)
 		}
@@ -838,6 +961,61 @@ func executeParallel(
 			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 {
+					completed = append(
+						completed,
+						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,
+							},
+						},
+					)
+				}
+			}
+
+			innerCheckpoints[toolCalls[i].ID] = se.Checkpoint
+
+			outerSE := &SuspendedError{
+				Checkpoint: &Checkpoint{
+					Version:          CheckpointVersion,
+					Status:           CheckpointStatusSuspended,
+					AllToolCalls:     toolCalls,
+					InnerCheckpoints: innerCheckpoints,
+					CompletedCalls:   completed,
+				},
+			}
+			return nil, nil, outerSE
+		}
 	}
 
 	var (
@@ -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,
+			callLLM:             ro.callLLM,
-			onEvent:             noopEvent,
+			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
+		innerIE, ok := errors.AsType[*InterruptedError](err)
-		if errors.As(err, &innerIE) {
+		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,
+			callLLM:             ro.callLLM,
-			onEvent:             noopEvent,
+			onEvent:             ro.onEvent,
 			skipInputGuardrails: true,
 			skipSessionLoad:     true,
 			initialUsage:        outer.usage,
 			initialTurns:        outer.turns,
+			checkpointStore:     ro.checkpointStore,
+			runID:               ro.runID,
+			toolUsedInRun:       ro.toolUsedInRun,
 		},
 	)
 }

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(
+		ro := runOpts{
-			ctx,
-			a,
-			messages,
-			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