diff --git a/pkg/agent/checkpoint.go b/pkg/agent/checkpoint.go index 076c1d393..d81b2ae4e 100644 --- a/pkg/agent/checkpoint.go +++ b/pkg/agent/checkpoint.go @@ -23,9 +23,25 @@ import ( type ( AgentStatus string + // AgentConfig captures the subset of agent options that must remain + // stable across a suspend/restore cycle to keep the run coherent. + // Currently that is only MaxTurns, because Checkpoint.Turns is a + // counter compared against it — if the live agent's bound were + // lowered below the saved counter we would either short-circuit the + // restored run or fail the warning at restoreSuspended. Other loop + // bounds (maxEmptyOutputRetries, maxToolDepth) reset per turn and + // are safe to change mid-suspension. Live references (tools, + // handoffs, hooks, LLM client, approval callbacks, guardrails) are + // intentionally not snapshotted so deploys can update behavior + // while runs are paused. + AgentConfig struct { + MaxTurns int + } + Checkpoint struct { Status AgentStatus AgentName string + Config AgentConfig Messages []llm.Message Usage llm.Usage Turns int diff --git a/pkg/agent/restore.go b/pkg/agent/restore.go index 3c82a7ffe..bd7912318 100644 --- a/pkg/agent/restore.go +++ b/pkg/agent/restore.go @@ -20,7 +20,6 @@ import ( "fmt" "sync" - "go.gearno.de/kit/log" "go.probo.inc/probo/pkg/llm" ) @@ -44,10 +43,23 @@ func Restore( if err != nil { return nil, fmt.Errorf("cannot resolve agent %q: %w", cp.AgentName, err) } + agent = applyCheckpointConfig(agent, cp.Config) return restoreCheckpoint(ctx, agent, cp, store, runID, registry) } +// applyCheckpointConfig returns a clone of agent with the bounds from +// the checkpoint snapshot overriding the live values. Zero values in +// cfg fall through to the live agent so older checkpoints written +// before the Config field existed, or test-constructed Checkpoint +// literals that omit Config, still resume correctly. +func applyCheckpointConfig(agent *Agent, cfg AgentConfig) *Agent { + if cfg.MaxTurns <= 0 { + return agent + } + return agent.Clone(WithMaxTurns(cfg.MaxTurns)) +} + func restoreCheckpoint( ctx context.Context, agent *Agent, @@ -96,16 +108,6 @@ func continueFromMessages( messagesCopy := make([]llm.Message, len(messages)) copy(messagesCopy, messages) - if cp.Turns >= agent.maxTurns { - agent.logger.WarnCtx( - ctx, - "restored agent run has already reached max turns", - log.String("agent", agent.name), - log.Int("turns", cp.Turns), - log.Int("max_turns", agent.maxTurns), - ) - } - return coreLoop( ctx, agent, @@ -169,6 +171,7 @@ func restoreNestedSuspended( entries[i].err = fmt.Errorf("cannot resolve inner agent %q: %w", innerCP.AgentName, err) continue } + innerAgent = applyCheckpointConfig(innerAgent, innerCP.Config) wg.Add(1) go func(i int, tc llm.ToolCall, innerAgent *Agent, innerCP *Checkpoint) { @@ -297,6 +300,7 @@ func restoreAwaitingApproval( if err != nil { return nil, fmt.Errorf("cannot resolve inner agent %q: %w", innerCP.AgentName, err) } + innerAgent = applyCheckpointConfig(innerAgent, innerCP.Config) innerIE := &InterruptedError{ ToolCalls: innerCP.PendingToolCalls, diff --git a/pkg/agent/restore_test.go b/pkg/agent/restore_test.go index 2dd663d54..274064945 100644 --- a/pkg/agent/restore_test.go +++ b/pkg/agent/restore_test.go @@ -461,4 +461,100 @@ func TestRestore(t *testing.T) { assert.Contains(t, err.Error(), "unknown checkpoint status") }, ) + + t.Run( + "buildCheckpoint captures MaxTurns in config snapshot", + func(t *testing.T) { + t.Parallel() + + ag := agent.New( + "producer-agent", + newTestClient(&mockProvider{}), + agent.WithModel("test-model"), + agent.WithMaxTurns(7), + ) + + store := newMemoryCheckpointer() + + stopCh := make(chan struct{}) + close(stopCh) + ctx := agent.WithStopSignal(context.Background(), stopCh) + + _, err := ag.RunWithOpts( + ctx, + []llm.Message{ + { + Role: llm.RoleUser, + Parts: []llm.Part{llm.TextPart{Text: "begin"}}, + }, + }, + agent.WithCheckpointer(store, "run-save-side"), + ) + + var se *agent.SuspendedError + require.ErrorAs(t, err, &se) + + cp, loadErr := store.Load(context.Background(), "run-save-side") + require.NoError(t, loadErr) + require.NotNil(t, cp) + assert.Equal(t, 7, cp.Config.MaxTurns) + }, + ) + + t.Run( + "checkpoint config supersedes live agent config on restore", + func(t *testing.T) { + t.Parallel() + + provider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + stopResponse("Completed after resume."), + }, + } + + // Agent registered at restore time has a bound tighter + // than the count of turns already taken. Without the + // checkpoint config snapshot, coreLoop would immediately + // trip MaxTurnsExceededError on its first iteration. + restoreAgent := agent.New( + "test-agent", + newTestClient(provider), + agent.WithInstructions("Test."), + agent.WithModel("test-model"), + agent.WithMaxTurns(5), + ) + + store := newMemoryCheckpointer() + err := store.Save(context.Background(), "run-config", &agent.Checkpoint{ + Status: agent.AgentStatusSuspended, + AgentName: "test-agent", + Config: agent.AgentConfig{MaxTurns: 20}, + Messages: []llm.Message{ + { + Role: llm.RoleUser, + Parts: []llm.Part{llm.TextPart{Text: "hi"}}, + }, + }, + Turns: 15, + }) + require.NoError(t, err) + + registry := &simpleRegistry{ + agents: map[string]*agent.Agent{ + "test-agent": restoreAgent, + }, + } + + result, err := agent.Restore( + context.Background(), + store, + "run-config", + registry, + ) + + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, "Completed after resume.", result.FinalMessage().Text()) + }, + ) } diff --git a/pkg/agent/run.go b/pkg/agent/run.go index a97fcb0fb..07de1f4cb 100644 --- a/pkg/agent/run.go +++ b/pkg/agent/run.go @@ -227,8 +227,11 @@ func (s *loopState) buildCheckpoint(status AgentStatus) *Checkpoint { copy(msgsCopy, s.messages) return &Checkpoint{ - Status: status, - AgentName: s.agent.name, + Status: status, + AgentName: s.agent.name, + Config: AgentConfig{ + MaxTurns: s.agent.maxTurns, + }, Messages: msgsCopy, Usage: s.totalUsage, Turns: s.turns,