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/checkpoint.go

@@ -16,21 +16,10 @@ package agent
 
 import (
 	"context"
-	"errors"
 
 	"go.probo.inc/probo/pkg/llm"
 )
 
-// ErrSuspendForCheckpoint is the recommended cancel cause for callers
-// who want the agent loop to gracefully suspend (build a checkpoint
-// and return *SuspendedError) rather than treat the cancellation as a
-// silent close. The agent loop only inspects ctx.Err(); any cancel
-// cause produces a graceful suspend, but using this sentinel makes
-// the intent explicit and lets supervisors distinguish a user-driven
-// cancel from infrastructure-level causes (lease loss, heartbeat
-// failure) when they inspect context.Cause(ctx).
-var ErrSuspendForCheckpoint = errors.New("agent run: graceful suspend requested")
-
 type (
 	AgentStatus string
 

pkg/agent/errors.go

@@ -15,11 +15,16 @@
 package agent
 
 import (
+	"errors"
 	"fmt"
 
 	"go.probo.inc/probo/pkg/llm"
 )
 
+// ErrSuspendForCheckpoint is the cancel cause to use when the caller
+// wants the agent loop to gracefully suspend.
+var ErrSuspendForCheckpoint = errors.New("agent run graceful suspend requested")
+
 type (
 	MaxTurnsExceededError struct {
 		MaxTurns int

pkg/agent/restore.go

@@ -25,11 +25,8 @@ import (
 
 // Restore continues a previously suspended or approval-interrupted agent run
 // from its last persisted checkpoint. The registry must contain all agents
-// that may have been active (including handoff targets).
-//
-// ctx follows the same graceful-suspend contract as Run: cancelling it
-// asks the resumed loop to checkpoint at the next safe boundary and
-// return a *SuspendedError with the new checkpoint. See Run for details.
+// that may have been active (including handoff targets). ctx follows Run's
+// graceful-suspend contract.
 func Restore(
 	ctx context.Context,
 	store Checkpointer,
@@ -138,13 +135,8 @@ func restoreNestedSuspended(
 	runID string,
 	registry AgentRegistry,
 ) (*Result, error) {
-	// Graceful-suspend contract: the saveProgress closure and the
-	// snapshot hook must survive a cancellation on the supervisor's
-	// runCtx so a partial save can land before we surface
-	// SuspendedError. Use a non-cancellable shadow for those sites
-	// only; the recursive restoreCheckpoint call keeps the original
-	// ctx so the nested coreLoop sees the cancel and suspends at its
-	// own next turn boundary.
+	// saveCtx survives an outer cancel so partial restore progress
+	// is persisted before SuspendedError surfaces.
 	saveCtx := context.WithoutCancel(ctx)
 
 	type nestedRestoreEntry struct {

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

pkg/agent/stream.go

@@ -61,9 +61,7 @@ func (sr *StreamedRun) Wait() (*Result, error) {
 
 // RunStreamed launches the agent loop and returns immediately with a
 // StreamedRun whose Events channel emits incremental progress. ctx
-// follows the same graceful-suspend contract as Run: cancellation
-// triggers a checkpoint at the next safe boundary and the run finishes
-// with a *SuspendedError surfaced via Wait. See Run for details.
+// follows Run's graceful-suspend contract.
 func (a *Agent) RunStreamed(ctx context.Context, messages []llm.Message, opts ...RunOption) *StreamedRun {
 	events := make(chan StreamEvent, 64)
 	sr := &StreamedRun{