probo/pkg/agent/cancel_test.go

// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.

package agent_test

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.probo.inc/probo/pkg/agent"
	"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()

	t.Run(
		"cancel before first turn suspends with empty checkpoint",
		func(t *testing.T) {
			t.Parallel()

			provider := &mockProvider{
				responses: []*llm.ChatCompletionResponse{
					stopResponse("never called"),
				},
			}

			ag := agent.New(
				"assistant",
				newTestClient(provider),
				agent.WithModel("test-model"),
			)

			store := newMemoryCheckpointer()

			ctx, cancel := context.WithCancel(context.Background())
			cancel()

			_, err := ag.Run(
				ctx,
				[]llm.Message{userMessage("hi")},
				agent.WithCheckpointer(store, "run-cancel"),
			)

			var se *agent.SuspendedError
			require.ErrorAs(t, err, &se)
			assert.Equal(t, 0, provider.calls, "LLM must not be invoked when ctx was already cancelled at entry")

			// When a checkpointer is configured, the persistent store
			// is the source of truth — the error itself doesn't carry a
			// Checkpoint. Load from the store to verify.
			cp, loadErr := store.Load(context.Background(), "run-cancel")
			require.NoError(t, loadErr)
			require.NotNil(t, cp, "checkpoint should be persisted before SuspendedError surfaces")
			assert.Equal(t, agent.AgentStatusSuspended, cp.Status)
		},
	)

	t.Run(
		"cancel mid-run preserves the just-completed turn",
		func(t *testing.T) {
			t.Parallel()

			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()

			provider := &mockProvider{
				responses: []*llm.ChatCompletionResponse{
					// First turn completes a tool call; the tool body
					// then cancels ctx so the next turn-boundary check
					// in coreLoop observes the cancellation.
					{
						Message: llm.Message{
							Role: llm.RoleAssistant,
							ToolCalls: []llm.ToolCall{{
								ID:       "tc_1",
								Function: llm.FunctionCall{Name: "noop", Arguments: `{}`},
							}},
						},
						FinishReason: llm.FinishReasonToolCalls,
					},
					stopResponse("never reached"),
				},
			}

			noopTool := agent.FunctionTool[struct{}](
				"noop",
				"no-op",
				func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
					cancel()
					return agent.ToolResult{Content: "ok"}, nil
				},
			)

			ag := agent.New(
				"assistant",
				newTestClient(provider),
				agent.WithModel("test-model"),
				agent.WithTools(noopTool),
			)

			store := newMemoryCheckpointer()

			_, err := ag.Run(
				ctx,
				[]llm.Message{userMessage("hi")},
				agent.WithCheckpointer(store, "run-mid"),
			)

			var se *agent.SuspendedError
			require.ErrorAs(t, err, &se)
			assert.Equal(t, 1, provider.calls, "second LLM call must not fire after cancel")

			cp, loadErr := store.Load(context.Background(), "run-mid")
			require.NoError(t, loadErr)
			require.NotNil(t, cp)
			assert.Equal(t, agent.AgentStatusSuspended, cp.Status)
			// The first LLM call completed; its output and the tool
			// reply must be in the checkpointed messages so a Restore
			// can resume from the next turn.
			assert.Equal(t, 1, cp.Turns)
			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)
		},
	)
}