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
}