rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Drive graceful agent suspend from ctx cancellation

Browse Source 
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>
This commit is contained in:
Aurélien Sibiril
2026-04-27 09:51:45 +02:00
parent 3acc3191ea
commit c4228e8e7c
5 changed files with 76 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
pkg/agent/restore.go
View File

@@ -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
}

84
pkg/agent/run.go
View File

@@ -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{

31
pkg/agent/stop.go
View File

@@ -1,31 +0,0 @@
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
//
// 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
}

63
pkg/agent/stop_test.go
View File

@@ -1,63 +0,0 @@
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
//
// 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")
}
})
}

5
pkg/agent/stream.go
View File

@@ -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{