From f715a806ee0e822652809ab4e912cfe0a35b3902 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Sibiril?=
 <81782+aureliensibiril@users.noreply.github.com>
Date: Mon, 27 Apr 2026 10:03:44 +0200
Subject: [PATCH] Strengthen ctx-cancel test coverage
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two additions:
- agent_test.go's "context cancellation triggers graceful suspend"
  now also asserts the input messages land in the suspension
  checkpoint — verifies the embedded-Checkpoint path that fires
  when no Checkpointer is configured.
- cancel_test.go gets a third subtest that parks the LLM provider
  inside ChatCompletion via a release channel, cancels ctx while
  the call is in flight, then confirms the LLM call still saw a
  non-cancelled ctx and the just-completed turn lands in the
  persisted checkpoint. Proves the framework's WithoutCancel
  shielding works end-to-end at the unit level.

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
---
 pkg/agent/agent_test.go  |   1 +
 pkg/agent/cancel_test.go | 119 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 120 insertions(+)

diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go
index 4c596dddd..a71270c49 100644
--- a/pkg/agent/agent_test.go
+++ b/pkg/agent/agent_test.go
@@ -504,6 +504,7 @@ func TestRun(t *testing.T) {
 			require.ErrorAs(t, err, &se)
 			require.NotNil(t, se.Checkpoint)
 			assert.Equal(t, agent.AgentStatusSuspended, se.Checkpoint.Status)
+			assert.NotEmpty(t, se.Checkpoint.Messages, "the input messages must land in the suspension checkpoint")
 			assert.Equal(t, 0, provider.calls)
 		},
 	)
diff --git a/pkg/agent/cancel_test.go b/pkg/agent/cancel_test.go
index 9ea4cd83b..0c25409a9 100644
--- a/pkg/agent/cancel_test.go
+++ b/pkg/agent/cancel_test.go
@@ -16,7 +16,9 @@ package agent_test
 
 import (
 	"context"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -24,6 +26,39 @@ import (
 	"go.probo.inc/probo/pkg/llm"
 )
 
+// blockingProvider holds the ChatCompletion call until a release
+// channel fires, so a test can race ctx cancellation against an
+// in-flight LLM call.
+type blockingProvider struct {
+	ready    chan struct{}
+	release  chan struct{}
+	response *llm.ChatCompletionResponse
+
+	mu       sync.Mutex
+	calls    int
+	ctxAtEnd error
+}
+
+func (p *blockingProvider) ChatCompletion(ctx context.Context, _ *llm.ChatCompletionRequest) (*llm.ChatCompletionResponse, error) {
+	p.mu.Lock()
+	p.calls++
+	first := p.calls == 1
+	p.mu.Unlock()
+
+	if first {
+		close(p.ready)
+		<-p.release
+		p.mu.Lock()
+		p.ctxAtEnd = ctx.Err()
+		p.mu.Unlock()
+	}
+	return p.response, nil
+}
+
+func (p *blockingProvider) ChatCompletionStream(_ context.Context, _ *llm.ChatCompletionRequest) (llm.ChatCompletionStream, error) {
+	return nil, assert.AnError
+}
+
 func TestRun_CtxCancelGracefulSuspend(t *testing.T) {
 	t.Parallel()
 
@@ -135,4 +170,88 @@ func TestRun_CtxCancelGracefulSuspend(t *testing.T) {
 			assert.GreaterOrEqual(t, len(cp.Messages), 3, "user + assistant tool-call + tool result")
 		},
 	)
+
+	t.Run(
+		"cancel during in-flight LLM call shields the call",
+		func(t *testing.T) {
+			t.Parallel()
+
+			// First response is a tool call so the loop iterates back
+			// to its turn-boundary cancel check after the LLM returns.
+			provider := &blockingProvider{
+				ready:   make(chan struct{}),
+				release: make(chan struct{}),
+				response: &llm.ChatCompletionResponse{
+					Message: llm.Message{
+						Role: llm.RoleAssistant,
+						ToolCalls: []llm.ToolCall{{
+							ID:       "tc_inflight",
+							Function: llm.FunctionCall{Name: "noop", Arguments: `{}`},
+						}},
+					},
+					FinishReason: llm.FinishReasonToolCalls,
+				},
+			}
+
+			noopTool := agent.FunctionTool[struct{}](
+				"noop",
+				"no-op",
+				func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
+					return agent.ToolResult{Content: "ok"}, nil
+				},
+			)
+
+			ag := agent.New(
+				"assistant",
+				newTestClient(provider),
+				agent.WithModel("test-model"),
+				agent.WithTools(noopTool),
+			)
+
+			store := newMemoryCheckpointer()
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			done := make(chan error, 1)
+			go func() {
+				_, err := ag.Run(
+					ctx,
+					[]llm.Message{userMessage("hi")},
+					agent.WithCheckpointer(store, "run-inflight"),
+				)
+				done <- err
+			}()
+
+			// Wait until the provider is parked inside ChatCompletion,
+			// then cancel ctx while the call is still in flight.
+			select {
+			case <-provider.ready:
+			case <-time.After(2 * time.Second):
+				t.Fatal("LLM call never started")
+			}
+			cancel()
+			close(provider.release)
+
+			var err error
+			select {
+			case err = <-done:
+			case <-time.After(2 * time.Second):
+				t.Fatal("agent.Run did not return after release")
+			}
+
+			var se *agent.SuspendedError
+			require.ErrorAs(t, err, &se)
+
+			provider.mu.Lock()
+			assert.NoError(t, provider.ctxAtEnd, "ctx passed to LLM must remain non-cancellable so the call completes")
+			assert.Equal(t, 1, provider.calls, "second LLM call must not fire after cancel")
+			provider.mu.Unlock()
+
+			cp, loadErr := store.Load(context.Background(), "run-inflight")
+			require.NoError(t, loadErr)
+			require.NotNil(t, cp)
+			assert.Equal(t, agent.AgentStatusSuspended, cp.Status)
+		},
+	)
 }