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{