From fc21d65c63f6bbc8c512cbea3708109b0226db15 Mon Sep 17 00:00:00 2001
From: Bryan Frimin <bryan@probo.com>
Date: Sun, 7 Jun 2026 12:15:07 +0200
Subject: [PATCH] Cover partial restore progress in nested suspend test

The nested-restore test only checked that an unresolvable inner agent
left its tool call in the checkpoint for a later retry. It did not
verify what happens to a sibling inner agent that does resolve and
complete during the same restore.

Add a resolvable done-agent alongside the missing inner-agent and
assert that its progress is persisted: its inner checkpoint is dropped
and its result recorded as a completed call, so a later retry replays
only the still-unresolved branch instead of re-running finished work.

Signed-off-by: Bryan Frimin <bryan@probo.com>
---
 pkg/agent/restore_test.go | 42 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/pkg/agent/restore_test.go b/pkg/agent/restore_test.go
index 6cf23be1a..d6b8ee40c 100644
--- a/pkg/agent/restore_test.go
+++ b/pkg/agent/restore_test.go
@@ -577,7 +577,7 @@ func TestRestore(t *testing.T) {
 	)
 
 	t.Run(
-		"nested suspended restore keeps progress when inner agent missing",
+		"nested suspended restore persists progress when one inner agent missing",
 		func(t *testing.T) {
 			t.Parallel()
 
@@ -586,6 +586,18 @@ func TestRestore(t *testing.T) {
 				newTestClient(&mockProvider{}),
 				agent.WithModel("test-model"),
 			)
+			// done-agent resolves and completes on restore; its
+			// progress must be persisted even though inner-agent
+			// (below) cannot be resolved.
+			doneAgent := agent.New(
+				"done-agent",
+				newTestClient(&mockProvider{
+					responses: []*llm.ChatCompletionResponse{
+						stopResponse("inner done"),
+					},
+				}),
+				agent.WithModel("test-model"),
+			)
 
 			store := newMemoryCheckpointer()
 			err := store.Save(context.Background(), "run-nested-missing-inner", &agent.Checkpoint{
@@ -598,6 +610,13 @@ func TestRestore(t *testing.T) {
 					},
 				},
 				AllToolCalls: []llm.ToolCall{
+					{
+						ID: "tc_done",
+						Function: llm.FunctionCall{
+							Name:      "call_done",
+							Arguments: `{"input":"go"}`,
+						},
+					},
 					{
 						ID: "tc_missing",
 						Function: llm.FunctionCall{
@@ -607,6 +626,16 @@ func TestRestore(t *testing.T) {
 					},
 				},
 				InnerCheckpoints: map[string]*agent.Checkpoint{
+					"tc_done": {
+						Status:    agent.AgentStatusSuspended,
+						AgentName: "done-agent",
+						Messages: []llm.Message{
+							{
+								Role:  llm.RoleUser,
+								Parts: []llm.Part{llm.TextPart{Text: "go"}},
+							},
+						},
+					},
 					"tc_missing": {
 						Status:    agent.AgentStatusSuspended,
 						AgentName: "inner-agent",
@@ -618,6 +647,7 @@ func TestRestore(t *testing.T) {
 			registry := &simpleRegistry{
 				agents: map[string]*agent.Agent{
 					"outer-agent": outerAgent,
+					"done-agent":  doneAgent,
 				},
 			}
 
@@ -633,7 +663,17 @@ func TestRestore(t *testing.T) {
 			cp, loadErr := store.Load(context.Background(), "run-nested-missing-inner")
 			require.NoError(t, loadErr)
 			require.NotNil(t, cp)
+
+			// The unresolved tool call is retained for a future retry.
 			require.Contains(t, cp.InnerCheckpoints, "tc_missing")
+
+			// The resolved tool call's progress was persisted: its
+			// inner checkpoint is dropped and its result recorded as a
+			// completed call, so a later retry does not re-run it.
+			require.NotContains(t, cp.InnerCheckpoints, "tc_done")
+			require.Len(t, cp.CompletedCalls, 1)
+			require.Equal(t, "tc_done", cp.CompletedCalls[0].ToolCallID)
+			require.Equal(t, "inner done", cp.CompletedCalls[0].Result.Content)
 		},
 	)