Tighten ctx-suspend plumbing and trim docs

Address review feedback:

- Move ErrSuspendForCheckpoint from checkpoint.go to errors.go
  next to the rest of the agent error declarations; drop the
  colon in the error string so it matches the existing
  `agent run <event>` style used by the supervisor sentinels.
- Replace the inline `outerCtx := ctx; ctx = context.WithoutCancel(ctx)`
  pattern with a small `suspendShield` helper in context.go used
  by coreLoop, resumeWithOpts, and resumeNested. Reads more
  cleanly and stops surfacing the WithoutCancel mechanism at
  every call site.
- Trim the doc comments on Run, RunStreamed, Resume, Restore, the
  ErrSuspendForCheckpoint declaration, and the saveCtx comment in
  restoreNestedSuspended down to the contract bullet.

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>

pkg/agent/run.go

@@ -115,14 +115,9 @@ func blockingCallLLM(ctx context.Context, agent *Agent, req *llm.ChatCompletionR
 	return acc.Response(), nil
 }
 
-// Run executes the agent loop. ctx is a graceful-stop signal: when it is
-// cancelled the loop checkpoints at its next safe boundary and returns
-// a *SuspendedError. The in-flight LLM call, the in-flight tool, and
-// the checkpoint write all run on a non-cancellable shadow of ctx so
-// they complete normally. ctx deadlines therefore become a "max
-// wall-clock budget, then suspend" rather than a hard cut-off; there
-// is no in-process hard-abort path. Callers that need to truly kill a
-// run must terminate the process.
+// Run executes the agent loop. Cancelling ctx triggers a graceful
+// suspend: the loop checkpoints at the next safe boundary and returns
+// *SuspendedError. There is no in-process hard-abort path.
 func (a *Agent) Run(ctx context.Context, messages []llm.Message, opts ...RunOption) (*Result, error) {
 	ro := runOpts{
 		callLLM: blockingCallLLM,
@@ -283,15 +278,11 @@ func (s *loopState) applyHandoff(ctx context.Context, handoffTarget *Handoff) er
 }
 
 func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Message, opts runOpts) (*Result, error) {
-	// Cancellation contract: ctx.Done() means "graceful suspend" — the
-	// loop checkpoints at its next safe boundary and returns
-	// SuspendedError. Every downstream call (LLM, tools, hooks,
-	// guardrails, save) runs on a non-cancellable shadow so an in-flight
-	// turn completes naturally and the checkpoint write itself isn't
-	// killed by the cancel that triggered the suspend. outerCtx is
-	// retained only for the at-boundary cancellation check.
-	outerCtx := ctx
-	ctx = context.WithoutCancel(ctx)
+	// outerCtx keeps the cancel signal for turn-boundary checkpointing;
+	// ctx survives cancellation so every downstream call (LLM, tools,
+	// hooks, save) carries through to completion once a checkpoint is
+	// requested.
+	outerCtx, ctx := ctx, context.WithoutCancel(ctx)
 
 	s := &loopState{
 		agent:         startAgent,
@@ -386,11 +377,12 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
 	exploring := structuredFormat != nil && len(s.toolDefs) > 0
 
 	for {
-		if err := outerCtx.Err(); err != nil {
+		select {
+		case <-outerCtx.Done():
 			cp := s.buildCheckpoint(AgentStatusSuspended)
 			se := &SuspendedError{RunID: s.opts.runID}
 
-			if s.opts.checkpointer != nil && s.opts.runID != "" {
+			if s.opts.checkpointer != nil {
 				if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil {
 					s.logger.ErrorCtx(ctx, "cannot save suspension checkpoint", log.Error(saveErr))
 					se.Checkpoint = cp
@@ -402,6 +394,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
 			}
 
 			return s.finishRun(ctx, nil, se)
+		default:
 		}
 
 		if s.turns >= s.agent.maxTurns {
@@ -568,7 +561,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
 						outerCP.InnerCheckpoints = se.Checkpoint.InnerCheckpoints
 						outerCP.CompletedCalls = se.Checkpoint.CompletedCalls
 					}
-					if s.opts.checkpointer != nil && s.opts.runID != "" {
+					if s.opts.checkpointer != nil {
 						if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, outerCP); saveErr != nil {
 							s.logger.ErrorCtx(ctx, "cannot save checkpoint", log.Error(saveErr))
 						} else {
@@ -588,7 +581,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
 					msgsCopy := make([]llm.Message, len(s.messages))
 					copy(msgsCopy, s.messages)
 
-					if s.opts.checkpointer != nil && s.opts.runID != "" {
+					if s.opts.checkpointer != nil {
 						cp := s.buildCheckpoint(AgentStatusAwaitingApproval)
 						cp.PendingToolCalls = nae.allToolCalls
 						cp.PendingApprovals = nae.pendingApprovals
@@ -624,7 +617,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
 					msgsCopy := make([]llm.Message, len(s.messages))
 					copy(msgsCopy, s.messages)
 
-					if s.opts.checkpointer != nil && s.opts.runID != "" {
+					if s.opts.checkpointer != nil {
 						cp := s.buildCheckpoint(AgentStatusAwaitingApproval)
 						cp.PendingToolCalls = nie.inner.ToolCalls
 						cp.PendingApprovals = nie.inner.PendingApprovals
@@ -709,7 +702,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
 			}
 
 			// Save incremental checkpoint after completed tool-call turn.
-			if s.opts.checkpointer != nil && s.opts.runID != "" {
+			if s.opts.checkpointer != nil {
 				cp := s.buildCheckpoint(AgentStatusSuspended)
 				if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil {
 					s.logger.ErrorCtx(ctx, "cannot save checkpoint", log.Error(saveErr))
@@ -1259,11 +1252,7 @@ func runOutputGuardrails(ctx context.Context, agent *Agent, message llm.Message)
 // been collected. It executes or denies each pending tool call according to
 // the provided ResumeInput, then re-enters the agent loop. Input guardrails
 // are not re-evaluated because the messages were already validated in the
-// original Run call.
-//
-// ctx follows the same graceful-suspend contract as Run: cancellation
-// produces a SuspendedError with a checkpoint covering progress up to
-// the next safe boundary. See Run for details.
+// original Run call. ctx follows Run's graceful-suspend contract.
 func Resume(ctx context.Context, interrupted *InterruptedError, input ResumeInput, opts ...RunOption) (*Result, error) {
 	ro := runOpts{
 		callLLM: blockingCallLLM,
@@ -1277,13 +1266,7 @@ func Resume(ctx context.Context, interrupted *InterruptedError, input ResumeInpu
 }
 
 func resumeWithOpts(ctx context.Context, interrupted *InterruptedError, input ResumeInput, ro runOpts) (*Result, error) {
-	// Same cancellation contract as Run: ctx.Done() means graceful
-	// suspend. Pre-loop tool execution and handoff callbacks run on the
-	// non-cancellable shadow so an in-flight tool completes; coreLoop
-	// receives the original ctx so it can detect cancellation at its
-	// next turn boundary and trigger the suspend path.
-	outerCtx := ctx
-	ctx = context.WithoutCancel(ctx)
+	outerCtx, ctx := ctx, context.WithoutCancel(ctx)
 
 	if interrupted.outerState != nil {
 		return resumeNested(outerCtx, interrupted, input, ro)
@@ -1446,11 +1429,7 @@ func resumeWithOpts(ctx context.Context, interrupted *InterruptedError, input Re
 }
 
 func resumeNested(ctx context.Context, interrupted *InterruptedError, input ResumeInput, ro runOpts) (*Result, error) {
-	// Same shadow as resumeWithOpts: pre-loop bridging work runs on
-	// the non-cancellable shadow; coreLoop receives outerCtx so it can
-	// detect a graceful-suspend cancel at its turn boundary.
-	outerCtx := ctx
-	ctx = context.WithoutCancel(ctx)
+	outerCtx, ctx := ctx, context.WithoutCancel(ctx)
 
 	outer := interrupted.outerState
 	logger := outer.agent.logger