diff --git a/pkg/agent/checkpoint.go b/pkg/agent/checkpoint.go index d78876d3e..840d3d1b3 100644 --- a/pkg/agent/checkpoint.go +++ b/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 diff --git a/pkg/agent/errors.go b/pkg/agent/errors.go index 06ed0148a..41fdccd70 100644 --- a/pkg/agent/errors.go +++ b/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 diff --git a/pkg/agent/restore.go b/pkg/agent/restore.go index 8129c5202..2484890a0 100644 --- a/pkg/agent/restore.go +++ b/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 { diff --git a/pkg/agent/run.go b/pkg/agent/run.go index bf7f55227..5206004b0 100644 --- a/pkg/agent/run.go +++ b/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 diff --git a/pkg/agent/stream.go b/pkg/agent/stream.go index 50077a2d6..8f3a93f57 100644 --- a/pkg/agent/stream.go +++ b/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{