From c4228e8e7ce8f4899102a996bcba214423b5a82d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aur=C3=A9lien=20Sibiril?= <81782+aureliensibiril@users.noreply.github.com> Date: Mon, 27 Apr 2026 09:51:45 +0200 Subject: [PATCH] Drive graceful agent suspend from ctx cancellation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Collapse the dual-mechanism (ctx.Done() = abort + WithStopSignal = graceful suspend) into a single signal: ctx.Done() now means graceful suspend. coreLoop shadows the incoming ctx with context.WithoutCancel(ctx) on entry and uses the shadow for every downstream call (LLM, tools, hooks, guardrails, save), keeping the original ctx only for the at-boundary cancellation check. restoreNestedSuspended applies the same shadow to its saveProgress closure so partial nested-restore writes survive a graceful cancel. Resume and resumeNested mirror the pattern so their pre-loop tool dispatch is non-cancellable while coreLoop still detects the cancel at its first turn boundary. The dedicated stop signal API (WithStopSignal / stopSignalFrom) is removed. There is no longer an in-process hard-abort path; tool authors who need a deadline must derive it themselves. Document the new contract on Run, RunStreamed, Resume, and Restore. Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com> --- pkg/agent/restore.go | 17 ++++++++- pkg/agent/run.go | 84 ++++++++++++++++++++++++++++-------------- pkg/agent/stop.go | 31 ---------------- pkg/agent/stop_test.go | 63 ------------------------------- pkg/agent/stream.go | 5 +++ 5 files changed, 76 insertions(+), 124 deletions(-) delete mode 100644 pkg/agent/stop.go delete mode 100644 pkg/agent/stop_test.go diff --git a/pkg/agent/restore.go b/pkg/agent/restore.go index 50b6f6a59..8129c5202 100644 --- a/pkg/agent/restore.go +++ b/pkg/agent/restore.go @@ -26,6 +26,10 @@ 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. func Restore( ctx context.Context, store Checkpointer, @@ -134,6 +138,15 @@ 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 := context.WithoutCancel(ctx) + type nestedRestoreEntry struct { toolCall llm.ToolCall originalCheckpoint *Checkpoint @@ -247,10 +260,10 @@ func restoreNestedSuspended( next.InnerCheckpoints = remainingInner next.CompletedCalls = completedCalls if store != nil && runID != "" { - if err := store.Save(ctx, runID, &next); err != nil { + if err := store.Save(saveCtx, runID, &next); err != nil { return nil, fmt.Errorf("cannot save nested restore progress: %w", err) } - emitHook(agent, func(h RunHooks) { h.OnRunSnapshot(ctx, agent, &next) }) + emitHook(agent, func(h RunHooks) { h.OnRunSnapshot(saveCtx, agent, &next) }) } return &next, nil } diff --git a/pkg/agent/run.go b/pkg/agent/run.go index 4e673b915..bf7f55227 100644 --- a/pkg/agent/run.go +++ b/pkg/agent/run.go @@ -115,6 +115,14 @@ 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. func (a *Agent) Run(ctx context.Context, messages []llm.Message, opts ...RunOption) (*Result, error) { ro := runOpts{ callLLM: blockingCallLLM, @@ -275,6 +283,16 @@ 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) + s := &loopState{ agent: startAgent, inputMessages: inputMessages, @@ -368,36 +386,28 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag exploring := structuredFormat != nil && len(s.toolDefs) > 0 for { - if err := ctx.Err(); err != nil { - return s.finishRun(ctx, nil, fmt.Errorf("cannot complete: %w", err)) + if err := outerCtx.Err(); err != nil { + cp := s.buildCheckpoint(AgentStatusSuspended) + se := &SuspendedError{RunID: s.opts.runID} + + if s.opts.checkpointer != nil && s.opts.runID != "" { + 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 + } else { + emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, cp) }) + } + } else { + se.Checkpoint = cp + } + + return s.finishRun(ctx, nil, se) } if s.turns >= s.agent.maxTurns { return s.finishRun(ctx, nil, &MaxTurnsExceededError{MaxTurns: s.agent.maxTurns}) } - if ch := stopSignalFrom(ctx); ch != nil { - select { - case <-ch: - cp := s.buildCheckpoint(AgentStatusSuspended) - se := &SuspendedError{RunID: s.opts.runID} - - if s.opts.checkpointer != nil && s.opts.runID != "" { - 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 - } else { - emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, cp) }) - } - } else { - se.Checkpoint = cp - } - - return s.finishRun(ctx, nil, se) - default: - } - } - fullMessages := buildFullMessages(s.systemPrompt, s.messages) var responseFormat *llm.ResponseFormat @@ -1250,6 +1260,10 @@ func runOutputGuardrails(ctx context.Context, agent *Agent, message llm.Message) // 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. func Resume(ctx context.Context, interrupted *InterruptedError, input ResumeInput, opts ...RunOption) (*Result, error) { ro := runOpts{ callLLM: blockingCallLLM, @@ -1263,8 +1277,16 @@ 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) + if interrupted.outerState != nil { - return resumeNested(ctx, interrupted, input, ro) + return resumeNested(outerCtx, interrupted, input, ro) } tracer := otel.GetTracerProvider().Tracer(tracerName) @@ -1406,7 +1428,7 @@ func resumeWithOpts(ctx context.Context, interrupted *InterruptedError, input Re } return coreLoop( - ctx, + outerCtx, resumeAgent, messages, runOpts{ @@ -1424,6 +1446,12 @@ 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) + outer := interrupted.outerState logger := outer.agent.logger @@ -1434,7 +1462,7 @@ func resumeNested(ctx context.Context, interrupted *InterruptedError, input Resu log.String("inner_agent", interrupted.Agent.name), ) - innerResult, err := resumeWithOpts(ctx, outer.innerInterrupt, input, ro) + innerResult, err := resumeWithOpts(outerCtx, outer.innerInterrupt, input, ro) if err != nil { innerIE, ok := errors.AsType[*InterruptedError](err) if ok { @@ -1489,7 +1517,7 @@ func resumeNested(ctx context.Context, interrupted *InterruptedError, input Resu } return coreLoop( - ctx, + outerCtx, outer.agent, messages, runOpts{ diff --git a/pkg/agent/stop.go b/pkg/agent/stop.go deleted file mode 100644 index 558a8bf64..000000000 --- a/pkg/agent/stop.go +++ /dev/null @@ -1,31 +0,0 @@ -// Copyright (c) 2026 Probo Inc . -// -// Permission to use, copy, modify, and/or distribute this software for any -// purpose with or without fee is hereby granted, provided that the above -// copyright notice and this permission notice appear in all copies. -// -// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH -// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY -// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, -// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM -// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR -// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR -// PERFORMANCE OF THIS SOFTWARE. - -package agent - -import "context" - -type stopSignalKey struct{} - -// WithStopSignal returns a derived context carrying a stop channel. -// The agent loop checks this channel at each turn boundary. -func WithStopSignal(ctx context.Context, ch <-chan struct{}) context.Context { - return context.WithValue(ctx, stopSignalKey{}, ch) -} - -// stopSignalFrom retrieves the stop channel from the context, or nil. -func stopSignalFrom(ctx context.Context) <-chan struct{} { - ch, _ := ctx.Value(stopSignalKey{}).(<-chan struct{}) - return ch -} diff --git a/pkg/agent/stop_test.go b/pkg/agent/stop_test.go deleted file mode 100644 index 49f8d860f..000000000 --- a/pkg/agent/stop_test.go +++ /dev/null @@ -1,63 +0,0 @@ -// Copyright (c) 2026 Probo Inc . -// -// Permission to use, copy, modify, and/or distribute this software for any -// purpose with or without fee is hereby granted, provided that the above -// copyright notice and this permission notice appear in all copies. -// -// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH -// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY -// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, -// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM -// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR -// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR -// PERFORMANCE OF THIS SOFTWARE. - -package agent - -import ( - "context" - "testing" - - "github.com/stretchr/testify/assert" -) - -func TestStopSignal(t *testing.T) { - t.Parallel() - - t.Run("absent returns nil", func(t *testing.T) { - t.Parallel() - assert.Nil(t, stopSignalFrom(context.Background())) - }) - - t.Run("round trip", func(t *testing.T) { - t.Parallel() - ch := make(chan struct{}) - ctx := WithStopSignal(context.Background(), ch) - assert.Equal(t, (<-chan struct{})(ch), stopSignalFrom(ctx)) - }) - - t.Run("non-blocking when open", func(t *testing.T) { - t.Parallel() - ch := make(chan struct{}) - ctx := WithStopSignal(context.Background(), ch) - sig := stopSignalFrom(ctx) - select { - case <-sig: - t.Fatal("should not fire") - default: - } - }) - - t.Run("fires when closed", func(t *testing.T) { - t.Parallel() - ch := make(chan struct{}) - ctx := WithStopSignal(context.Background(), ch) - close(ch) - sig := stopSignalFrom(ctx) - select { - case <-sig: - default: - t.Fatal("should have fired") - } - }) -} diff --git a/pkg/agent/stream.go b/pkg/agent/stream.go index eaee2eb33..50077a2d6 100644 --- a/pkg/agent/stream.go +++ b/pkg/agent/stream.go @@ -59,6 +59,11 @@ func (sr *StreamedRun) Wait() (*Result, error) { return sr.result, sr.err } +// 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. func (a *Agent) RunStreamed(ctx context.Context, messages []llm.Message, opts ...RunOption) *StreamedRun { events := make(chan StreamEvent, 64) sr := &StreamedRun{