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

Rename CheckpointStatus to AgentStatus

Browse Source 
The status values describe the agent state, not the
checkpoint data state.

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
This commit is contained in:
Aurélien Sibiril
2026-04-13 23:01:37 +02:00
parent ff18a5fc23
commit 71e2d00b3a
7 changed files with 1060 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
pkg/agent/restore.go
View File

@@ -29,7 +29,7 @@ import (
// that may have been active (including handoff targets).
func Restore(
ctx context.Context,
store CheckpointStore,
store Checkpointer,
runID string,
registry AgentRegistry,
) (*Result, error) {
@@ -40,10 +40,6 @@ func Restore(
if cp == nil {
return nil, fmt.Errorf("cannot restore: no checkpoint for run %s", runID)
}
if cp.Version != CheckpointVersion {
return nil, fmt.Errorf("cannot restore: unsupported checkpoint version %d", cp.Version)
}
agent, err := registry.Agent(cp.AgentName)
if err != nil {
return nil, fmt.Errorf("cannot resolve agent %q: %w", cp.AgentName, err)
@@ -56,17 +52,17 @@ func restoreCheckpoint(
ctx context.Context,
agent *Agent,
cp *Checkpoint,
store CheckpointStore,
store Checkpointer,
runID string,
registry AgentRegistry,
) (*Result, error) {
emitHook(agent, func(h RunHooks) { h.OnRunRestore(ctx, agent, cp) })
switch cp.Status {
case CheckpointStatusSuspended:
case AgentStatusSuspended:
return restoreSuspended(ctx, agent, cp, store, runID, registry)
case CheckpointStatusAwaitingApproval:
case AgentStatusAwaitingApproval:
return restoreAwaitingApproval(ctx, agent, cp, store, runID, registry)
default:
@@ -78,7 +74,7 @@ func restoreSuspended(
ctx context.Context,
agent *Agent,
cp *Checkpoint,
store CheckpointStore,
store Checkpointer,
runID string,
registry AgentRegistry,
) (*Result, error) {
@@ -94,7 +90,7 @@ func continueFromMessages(
agent *Agent,
messages []llm.Message,
cp *Checkpoint,
store CheckpointStore,
store Checkpointer,
runID string,
) (*Result, error) {
messagesCopy := make([]llm.Message, len(messages))
@@ -121,7 +117,7 @@ func continueFromMessages(
skipSessionLoad: true,
initialUsage: cp.Usage,
initialTurns: cp.Turns,
checkpointStore: store,
checkpointer: store,
runID: runID,
toolUsedInRun: cp.ToolUsedInRun,
},
@@ -132,7 +128,7 @@ func restoreNestedSuspended(
ctx context.Context,
agent *Agent,
cp *Checkpoint,
store CheckpointStore,
store Checkpointer,
runID string,
registry AgentRegistry,
) (*Result, error) {
@@ -277,7 +273,7 @@ func restoreAwaitingApproval(
ctx context.Context,
agent *Agent,
cp *Checkpoint,
store CheckpointStore,
store Checkpointer,
runID string,
registry AgentRegistry,
) (*Result, error) {
@@ -338,11 +334,11 @@ func restoreAwaitingApproval(
ie,
ResumeInput{Approvals: cp.ApprovalInput},
runOpts{
callLLM: blockingCallLLM,
onEvent: noopEvent,
checkpointStore: store,
runID: runID,
toolUsedInRun: cp.ToolUsedInRun,
callLLM: blockingCallLLM,
onEvent: noopEvent,
checkpointer: store,
runID: runID,
toolUsedInRun: cp.ToolUsedInRun,
},
)
}

91
pkg/agent/restore_test.go
View File

@@ -26,18 +26,18 @@ import (
"go.probo.inc/probo/pkg/llm"
)
type memoryCheckpointStore struct {
type memoryCheckpointer struct {
mu sync.Mutex
checkpoints map[string]*agent.Checkpoint
}
func newMemoryCheckpointStore() *memoryCheckpointStore {
return &memoryCheckpointStore{
func newMemoryCheckpointer() *memoryCheckpointer {
return &memoryCheckpointer{
checkpoints: make(map[string]*agent.Checkpoint),
}
}
func (s *memoryCheckpointStore) Save(_ context.Context, runID string, cp *agent.Checkpoint) error {
func (s *memoryCheckpointer) Save(_ context.Context, runID string, cp *agent.Checkpoint) error {
s.mu.Lock()
defer s.mu.Unlock()
@@ -46,7 +46,7 @@ func (s *memoryCheckpointStore) Save(_ context.Context, runID string, cp *agent.
return nil
}
func (s *memoryCheckpointStore) Load(_ context.Context, runID string) (*agent.Checkpoint, error) {
func (s *memoryCheckpointer) Load(_ context.Context, runID string) (*agent.Checkpoint, error) {
s.mu.Lock()
defer s.mu.Unlock()
@@ -59,14 +59,6 @@ func (s *memoryCheckpointStore) Load(_ context.Context, runID string) (*agent.Ch
return &clone, nil
}
func (s *memoryCheckpointStore) Delete(_ context.Context, runID string) error {
s.mu.Lock()
defer s.mu.Unlock()
delete(s.checkpoints, runID)
return nil
}
type simpleRegistry struct {
agents map[string]*agent.Agent
}
@@ -87,7 +79,7 @@ func TestRestore(t *testing.T) {
func(t *testing.T) {
t.Parallel()
store := newMemoryCheckpointStore()
store := newMemoryCheckpointer()
registry := &simpleRegistry{agents: map[string]*agent.Agent{}}
_, err := agent.Restore(
@@ -102,41 +94,6 @@ func TestRestore(t *testing.T) {
},
)
t.Run(
"unsupported checkpoint version returns error",
func(t *testing.T) {
t.Parallel()
store := newMemoryCheckpointStore()
err := store.Save(context.Background(), "run-1", &agent.Checkpoint{
Version: 999,
Status: agent.CheckpointStatusSuspended,
AgentName: "test-agent",
})
require.NoError(t, err)
registry := &simpleRegistry{
agents: map[string]*agent.Agent{
"test-agent": agent.New(
"test-agent",
newTestClient(&mockProvider{}),
agent.WithModel("test-model"),
),
},
}
_, err = agent.Restore(
context.Background(),
store,
"run-1",
registry,
)
require.Error(t, err)
assert.Contains(t, err.Error(), "unsupported checkpoint version")
},
)
t.Run(
"suspended checkpoint restores and completes",
func(t *testing.T) {
@@ -155,10 +112,9 @@ func TestRestore(t *testing.T) {
agent.WithModel("test-model"),
)
store := newMemoryCheckpointStore()
store := newMemoryCheckpointer()
err := store.Save(context.Background(), "run-suspended", &agent.Checkpoint{
Version: 1,
Status: agent.CheckpointStatusSuspended,
Status: agent.AgentStatusSuspended,
AgentName: "test-agent",
Messages: []llm.Message{
{
@@ -217,10 +173,9 @@ func TestRestore(t *testing.T) {
}),
)
store := newMemoryCheckpointStore()
store := newMemoryCheckpointer()
err := store.Save(context.Background(), "run-approval", &agent.Checkpoint{
Version: 1,
Status: agent.CheckpointStatusAwaitingApproval,
Status: agent.AgentStatusAwaitingApproval,
AgentName: "test-agent",
Messages: []llm.Message{
{
@@ -303,10 +258,9 @@ func TestRestore(t *testing.T) {
}),
)
store := newMemoryCheckpointStore()
store := newMemoryCheckpointer()
err := store.Save(context.Background(), "run-approved", &agent.Checkpoint{
Version: 1,
Status: agent.CheckpointStatusAwaitingApproval,
Status: agent.AgentStatusAwaitingApproval,
AgentName: "test-agent",
Messages: []llm.Message{
{
@@ -370,10 +324,9 @@ func TestRestore(t *testing.T) {
agent.WithModel("test-model"),
)
store := newMemoryCheckpointStore()
store := newMemoryCheckpointer()
err := store.Save(context.Background(), "run-nested", &agent.Checkpoint{
Version: 1,
Status: agent.CheckpointStatusAwaitingApproval,
Status: agent.AgentStatusAwaitingApproval,
AgentName: "test-agent",
Messages: []llm.Message{
{
@@ -401,13 +354,11 @@ func TestRestore(t *testing.T) {
},
InnerCheckpoints: map[string]*agent.Checkpoint{
"tc_inner_1": {
Version: 1,
Status: agent.CheckpointStatusAwaitingApproval,
Status: agent.AgentStatusAwaitingApproval,
AgentName: "inner-agent-1",
},
"tc_inner_2": {
Version: 1,
Status: agent.CheckpointStatusAwaitingApproval,
Status: agent.AgentStatusAwaitingApproval,
AgentName: "inner-agent-2",
},
},
@@ -452,10 +403,9 @@ func TestRestore(t *testing.T) {
func(t *testing.T) {
t.Parallel()
store := newMemoryCheckpointStore()
store := newMemoryCheckpointer()
err := store.Save(context.Background(), "run-unknown", &agent.Checkpoint{
Version: 1,
Status: agent.CheckpointStatusSuspended,
Status: agent.AgentStatusSuspended,
AgentName: "missing-agent",
})
require.NoError(t, err)
@@ -487,10 +437,9 @@ func TestRestore(t *testing.T) {
agent.WithModel("test-model"),
)
store := newMemoryCheckpointStore()
store := newMemoryCheckpointer()
err := store.Save(context.Background(), "run-bad-status", &agent.Checkpoint{
Version: 1,
Status: agent.CheckpointStatus("bogus"),
Status: agent.AgentStatus("bogus"),
AgentName: "test-agent",
})
require.NoError(t, err)

49
pkg/agent/run.go
View File

@@ -49,7 +49,7 @@ type (
skipSessionLoad bool
initialUsage llm.Usage
initialTurns int
checkpointStore CheckpointStore
checkpointer Checkpointer
runID string
toolUsedInRun bool
}
@@ -77,9 +77,9 @@ type (
}
)
func WithCheckpointStore(store CheckpointStore, runID string) RunOption {
func WithCheckpointer(cp Checkpointer, runID string) RunOption {
return func(o *runOpts) {
o.checkpointStore = store
o.checkpointer = cp
o.runID = runID
}
}
@@ -222,12 +222,11 @@ func (s *loopState) finishRun(ctx context.Context, result *Result, err error) (*
return result, err
}
func (s *loopState) buildCheckpoint(status CheckpointStatus) *Checkpoint {
func (s *loopState) buildCheckpoint(status AgentStatus) *Checkpoint {
msgsCopy := make([]llm.Message, len(s.messages))
copy(msgsCopy, s.messages)
return &Checkpoint{
Version: CheckpointVersion,
Status: status,
AgentName: s.agent.name,
Messages: msgsCopy,
@@ -381,11 +380,11 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
if ch := stopSignalFrom(ctx); ch != nil {
select {
case <-ch:
cp := s.buildCheckpoint(CheckpointStatusSuspended)
cp := s.buildCheckpoint(AgentStatusSuspended)
se := &SuspendedError{RunID: s.opts.runID}
if s.opts.checkpointStore != nil && s.opts.runID != "" {
if saveErr := s.opts.checkpointStore.Save(ctx, s.opts.runID, cp); saveErr != nil {
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
}
@@ -552,14 +551,14 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
if err != nil {
if se, ok := errors.AsType[*SuspendedError](err); ok {
outerCP := s.buildCheckpoint(CheckpointStatusSuspended)
outerCP := s.buildCheckpoint(AgentStatusSuspended)
if se.Checkpoint != nil {
outerCP.AllToolCalls = se.Checkpoint.AllToolCalls
outerCP.InnerCheckpoints = se.Checkpoint.InnerCheckpoints
outerCP.CompletedCalls = se.Checkpoint.CompletedCalls
}
if s.opts.checkpointStore != nil && s.opts.runID != "" {
if saveErr := s.opts.checkpointStore.Save(ctx, s.opts.runID, outerCP); saveErr != nil {
if s.opts.checkpointer != nil && s.opts.runID != "" {
if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, outerCP); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save checkpoint", log.Error(saveErr))
}
}
@@ -576,11 +575,11 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
msgsCopy := make([]llm.Message, len(s.messages))
copy(msgsCopy, s.messages)
if s.opts.checkpointStore != nil && s.opts.runID != "" {
cp := s.buildCheckpoint(CheckpointStatusAwaitingApproval)
if s.opts.checkpointer != nil && s.opts.runID != "" {
cp := s.buildCheckpoint(AgentStatusAwaitingApproval)
cp.PendingToolCalls = nae.allToolCalls
cp.PendingApprovals = nae.pendingApprovals
if saveErr := s.opts.checkpointStore.Save(ctx, s.opts.runID, cp); saveErr != nil {
if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save approval checkpoint", log.Error(saveErr))
}
}
@@ -610,16 +609,15 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
msgsCopy := make([]llm.Message, len(s.messages))
copy(msgsCopy, s.messages)
if s.opts.checkpointStore != nil && s.opts.runID != "" {
cp := s.buildCheckpoint(CheckpointStatusAwaitingApproval)
if s.opts.checkpointer != nil && s.opts.runID != "" {
cp := s.buildCheckpoint(AgentStatusAwaitingApproval)
cp.PendingToolCalls = nie.inner.ToolCalls
cp.PendingApprovals = nie.inner.PendingApprovals
cp.AllToolCalls = nie.allToolCalls
cp.CompletedCalls = nie.completedCalls
cp.InnerCheckpoints = map[string]*Checkpoint{
nie.toolCallID: {
Version: CheckpointVersion,
Status: CheckpointStatusAwaitingApproval,
Status: AgentStatusAwaitingApproval,
AgentName: nie.inner.Agent.name,
Messages: nie.inner.Messages,
Usage: nie.inner.Usage,
@@ -628,7 +626,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
PendingApprovals: nie.inner.PendingApprovals,
},
}
if saveErr := s.opts.checkpointStore.Save(ctx, s.opts.runID, cp); saveErr != nil {
if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save nested approval checkpoint", log.Error(saveErr))
}
}
@@ -694,9 +692,9 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
}
// Save incremental checkpoint after completed tool-call turn.
if s.opts.checkpointStore != nil && s.opts.runID != "" {
cp := s.buildCheckpoint(CheckpointStatusSuspended)
if saveErr := s.opts.checkpointStore.Save(ctx, s.opts.runID, cp); saveErr != nil {
if s.opts.checkpointer != nil && s.opts.runID != "" {
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))
}
}
@@ -1008,8 +1006,7 @@ func executeParallel(
outerSE := &SuspendedError{
Checkpoint: &Checkpoint{
Version: CheckpointVersion,
Status: CheckpointStatusSuspended,
Status: AgentStatusSuspended,
AllToolCalls: toolCalls,
InnerCheckpoints: innerCheckpoints,
CompletedCalls: completed,
@@ -1414,7 +1411,7 @@ func resumeWithOpts(ctx context.Context, interrupted *InterruptedError, input Re
skipSessionLoad: true,
initialUsage: interrupted.Usage,
initialTurns: interrupted.Turns,
checkpointStore: ro.checkpointStore,
checkpointer: ro.checkpointer,
runID: ro.runID,
toolUsedInRun: ro.toolUsedInRun,
},
@@ -1497,7 +1494,7 @@ func resumeNested(ctx context.Context, interrupted *InterruptedError, input Resu
skipSessionLoad: true,
initialUsage: outer.usage,
initialTurns: outer.turns,
checkpointStore: ro.checkpointStore,
checkpointer: ro.checkpointer,
runID: ro.runID,
toolUsedInRun: ro.toolUsedInRun,
},