/agent/checkpoint_test.go

package agent

import (
	"context"
	"errors"
	"strings"
	"testing"

	"go-micro.dev/v6/ai"
	"go-micro.dev/v6/flow"
	"go-micro.dev/v6/store"
)

func TestResumeCompletedCheckpointDoesNotReplayModel(t *testing.T) {
	ctx := context.Background()
	cp := flow.StoreCheckpoint(store.NewMemoryStore(), "durable-agent")
	calls := 0
	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
		calls++
		return &ai.Response{Reply: "done"}, nil
	}
	defer func() { fakeGen = nil }()

	a := newTestAgent(Name("durable-agent"), WithCheckpoint(cp))
	resp, err := a.Ask(ctx, "finish the work")
	if err != nil {
		t.Fatalf("Ask: %v", err)
	}
	if calls != 1 {
		t.Fatalf("model calls after Ask = %d, want 1", calls)
	}

	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
		calls++
		t.Fatal("Resume of a completed run replayed the model")
		return nil, nil
	}
	resumed, err := Resume(ctx, a, resp.RunID)
	if err != nil {
		t.Fatalf("Resume: %v", err)
	}
	if resumed.Reply != "done" {
		t.Fatalf("resumed reply = %q, want done", resumed.Reply)
	}
	if resumed.RunID != resp.RunID {
		t.Fatalf("resumed run id = %q, want %q", resumed.RunID, resp.RunID)
	}
	if calls != 1 {
		t.Fatalf("model calls after Resume = %d, want 1", calls)
	}
}

func TestResumeFailedCheckpointDoesNotReplayCompletedTool(t *testing.T) {
	ctx := context.Background()
	cp := flow.StoreCheckpoint(store.NewMemoryStore(), "tool-resume-agent")
	toolRuns := 0
	first := true
	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
		if opts.ToolHandler != nil {
			res := opts.ToolHandler(ctx, ai.ToolCall{ID: "call-1", Name: "external.charge", Input: map[string]any{"order": "42"}})
			if res.Content != "charged" {
				t.Fatalf("tool result = %q, want charged", res.Content)
			}
		}
		if first {
			first = false
			return nil, errors.New("model connection dropped after tool")
		}
		return &ai.Response{Reply: "finished from checkpoint"}, nil
	}
	defer func() { fakeGen = nil }()

	a := newTestAgent(Name("tool-resume-agent"), WithCheckpoint(cp),
		WithTool("external.charge", "charge once", nil, func(context.Context, map[string]any) (string, error) {
			toolRuns++
			return "charged", nil
		}))
	_, err := a.Ask(ctx, "charge order 42")
	if err == nil {
		t.Fatal("Ask succeeded, want simulated failure")
	}
	if toolRuns != 1 {
		t.Fatalf("tool executions after failed Ask = %d, want 1", toolRuns)
	}

	runs, err := Pending(ctx, a)
	if err != nil {
		t.Fatalf("Pending: %v", err)
	}
	if len(runs) != 1 {
		t.Fatalf("Pending returned %d runs, want 1", len(runs))
	}
	resp, err := Resume(ctx, a, runs[0].ID)
	if err != nil {
		t.Fatalf("Resume: %v", err)
	}
	if resp.Reply != "finished from checkpoint" {
		t.Fatalf("Resume reply = %q", resp.Reply)
	}
	if toolRuns != 1 {
		t.Fatalf("tool executions after Resume = %d, want completed tool was not replayed", toolRuns)
	}
}

func TestResumeFailedCheckpointAfterFreshAgentRestart(t *testing.T) {
	ctx := context.Background()
	cp := flow.StoreCheckpoint(store.NewMemoryStore(), "restart-resume-agent")
	toolRuns := 0
	modelCalls := 0
	failFirst := true
	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
		modelCalls++
		if opts.ToolHandler != nil {
			res := opts.ToolHandler(ctx, ai.ToolCall{ID: "call-1", Name: "external.provision", Input: map[string]any{"service": "api"}})
			if res.Content != "provisioned" {
				t.Fatalf("tool result = %q, want provisioned", res.Content)
			}
		}
		if failFirst {
			failFirst = false
			return nil, errors.New("process stopped after tool checkpoint")
		}
		return &ai.Response{Reply: "resumed after restart"}, nil
	}
	defer func() { fakeGen = nil }()

	newAgent := func() *agentImpl {
		return newTestAgent(Name("restart-resume-agent"), WithCheckpoint(cp),
			WithTool("external.provision", "provision service once", nil, func(context.Context, map[string]any) (string, error) {
				toolRuns++
				return "provisioned", nil
			}))
	}

	first := newAgent()
	_, err := first.Ask(ctx, "provision api")
	if err == nil {
		t.Fatal("Ask succeeded, want simulated process stop")
	}
	if toolRuns != 1 {
		t.Fatalf("tool executions after failed Ask = %d, want 1", toolRuns)
	}
	runs, err := Pending(ctx, first)
	if err != nil {
		t.Fatalf("Pending before restart: %v", err)
	}
	if len(runs) != 1 {
		t.Fatalf("Pending before restart returned %d runs, want 1", len(runs))
	}

	restarted := newAgent()
	resp, err := Resume(ctx, restarted, runs[0].ID)
	if err != nil {
		t.Fatalf("Resume after restart: %v", err)
	}
	if resp.Reply != "resumed after restart" || resp.RunID != runs[0].ID {
		t.Fatalf("response = %#v, want resumed reply on original run id", resp)
	}
	if toolRuns != 1 {
		t.Fatalf("tool executions after restart resume = %d, want checkpointed tool not replayed", toolRuns)
	}
	if modelCalls != 2 {
		t.Fatalf("model calls = %d, want initial call plus resumed call", modelCalls)
	}
	loaded, ok, err := cp.Load(ctx, runs[0].ID)
	if err != nil || !ok {
		t.Fatalf("Load resumed run ok=%v err=%v", ok, err)
	}
	if loaded.Status != "done" || loaded.ParentID != runs[0].ParentID {
		t.Fatalf("loaded run status/parent = %s/%s, want done/%s", loaded.Status, loaded.ParentID, runs[0].ParentID)
	}
}

func TestResumeFailedCheckpointDoesNotDuplicateCompactedMemory(t *testing.T) {
	ctx := context.Background()
	st := store.NewMemoryStore()
	cp := flow.StoreCheckpoint(st, "memory-resume-agent")
	failRetry := true
	var sawRecall bool
	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
		for _, msg := range req.Messages {
			if text, ok := msg.Content.(string); ok && strings.Contains(text, "alpha code is 42") {
				sawRecall = true
			}
		}
		if strings.Contains(req.Prompt, "use alpha code") && failRetry {
			failRetry = false
			return nil, errors.New("model connection dropped")
		}
		return &ai.Response{Reply: "ok"}, nil
	}
	defer func() { fakeGen = nil }()

	a := newTestAgent(Name("memory-resume-agent"), WithStore(st), WithCheckpoint(cp), CompactMemory(4, 1), MemoryRecallLimit(2))
	for _, msg := range []string{"alpha code is 42", "beta note", "gamma note"} {
		if _, err := a.Ask(ctx, msg); err != nil {
			t.Fatalf("Ask(%q): %v", msg, err)
		}
	}

	_, err := a.Ask(ctx, "use alpha code now")
	if err == nil {
		t.Fatal("Ask succeeded, want simulated provider failure")
	}
	if got := countMemoryContent(a.mem.Messages(), "use alpha code now"); got != 1 {
		t.Fatalf("failed Ask stored prompt %d times, want 1", got)
	}

	runs, err := Pending(ctx, a)
	if err != nil {
		t.Fatalf("Pending: %v", err)
	}
	if len(runs) != 1 {
		t.Fatalf("Pending returned %d runs, want 1", len(runs))
	}
	if _, err := Resume(ctx, a, runs[0].ID); err != nil {
		t.Fatalf("Resume: %v", err)
	}
	if got := countMemoryContent(a.mem.Messages(), "use alpha code now"); got != 1 {
		t.Fatalf("resumed failed Ask stored prompt %d times, want no duplicate", got)
	}
	if !sawRecall {
		t.Fatal("resume did not retrieve archived compacted memory")
	}
	if got := len(a.mem.Messages()); got > 4 {
		t.Fatalf("compacted memory retained %d messages after resume, want <= 4", got)
	}
}

func countMemoryContent(messages []ai.Message, needle string) int {
	var count int
	for _, msg := range messages {
		if text, ok := msg.Content.(string); ok && strings.Contains(text, needle) {
			count++
		}
	}
	return count
}

func TestPendingReturnsUnfinishedAgentRuns(t *testing.T) {
	ctx := context.Background()
	cp := flow.StoreCheckpoint(store.NewMemoryStore(), "pending-agent")
	run := flow.Run{ID: "run-1", Flow: "pending-agent", Status: "failed", State: flow.State{Stage: agentAskStep, Data: []byte("retry me")}}
	if err := cp.Save(ctx, run); err != nil {
		t.Fatalf("Save: %v", err)
	}
	a := newTestAgent(Name("pending-agent"), WithCheckpoint(cp))
	runs, err := Pending(ctx, a)
	if err != nil {
		t.Fatalf("Pending: %v", err)
	}
	if len(runs) != 1 || runs[0].ID != "run-1" {
		t.Fatalf("Pending = %#v, want run-1", runs)
	}
}

func TestPendingSkipsTerminalCanceledAndExpiredAgentRuns(t *testing.T) {
	ctx := context.Background()
	cp := flow.StoreCheckpoint(store.NewMemoryStore(), "terminal-agent")
	for _, run := range []flow.Run{
		{ID: "active", Flow: "terminal-agent", Status: "failed", State: flow.State{Stage: agentAskStep, Data: []byte("retry me")}},
		{ID: "done", Flow: "terminal-agent", Status: "done", State: flow.State{Stage: agentAskStep, Data: []byte("done")}},
		{ID: "canceled", Flow: "terminal-agent", Status: "canceled", State: flow.State{Stage: agentAskStep, Data: []byte("canceled")}},
		{ID: "expired", Flow: "terminal-agent", Status: "expired", State: flow.State{Stage: agentAskStep, Data: []byte("expired")}},
	} {
		if err := cp.Save(ctx, run); err != nil {
			t.Fatalf("Save(%s): %v", run.ID, err)
		}
	}

	a := newTestAgent(Name("terminal-agent"), WithCheckpoint(cp))
	runs, err := Pending(ctx, a)
	if err != nil {
		t.Fatalf("Pending: %v", err)
	}
	if len(runs) != 1 || runs[0].ID != "active" {
		t.Fatalf("Pending = %#v, want only active failed run", runs)
	}
	for _, id := range []string{"canceled", "expired"} {
		if _, err := Resume(ctx, a, id); err == nil || !strings.Contains(err.Error(), "terminal") {
			t.Fatalf("Resume(%s) err = %v, want terminal status error", id, err)
		}
	}
}

func TestHumanInputPauseResumesSameRunWithInput(t *testing.T) {
	ctx := context.Background()
	cp := flow.StoreCheckpoint(store.NewMemoryStore(), "input-agent")
	calls := 0
	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
		calls++
		if calls == 1 {
			if opts.ToolHandler != nil {
				opts.ToolHandler(ctx, ai.ToolCall{ID: "input-1", Name: toolHumanInput, Input: map[string]any{"prompt": "Which region should I deploy to?"}})
			}
			return &ai.Response{Reply: "waiting"}, nil
		}
		if !strings.Contains(req.Prompt, "Human input: us-east-1") {
			t.Fatalf("resumed prompt = %q, want human input", req.Prompt)
		}
		return &ai.Response{Reply: "deploying to us-east-1"}, nil
	}
	defer func() { fakeGen = nil }()

	a := newTestAgent(Name("input-agent"), WithCheckpoint(cp))
	_, err := a.Ask(ctx, "deploy the service")
	if err == nil {
		t.Fatal("Ask succeeded, want input-required pause")
	}
	runs, err := Pending(ctx, a)
	if err != nil {
		t.Fatalf("Pending: %v", err)
	}
	if len(runs) != 1 || runs[0].Status != "paused" || runs[0].State.Stage != agentInputStep {
		t.Fatalf("paused runs = %#v, want one input-required run", runs)
	}
	var pause inputPause
	if err := runs[0].State.Scan(&pause); err != nil {
		t.Fatalf("Scan pause: %v", err)
	}
	if pause.OriginalMessage != "deploy the service" || pause.Prompt != "Which region should I deploy to?" {
		t.Fatalf("pause = %#v", pause)
	}

	if _, err := Resume(ctx, a, runs[0].ID); err == nil || !strings.Contains(err.Error(), "ResumeInput") {
		t.Fatalf("Resume input-required err = %v, want guidance", err)
	}
	resp, err := ResumeInput(ctx, a, runs[0].ID, "us-east-1")
	if err != nil {
		t.Fatalf("ResumeInput: %v", err)
	}
	if resp.RunID != runs[0].ID || resp.Reply != "deploying to us-east-1" {
		t.Fatalf("response = %#v", resp)
	}
	loaded, ok, err := cp.Load(ctx, runs[0].ID)
	if err != nil || !ok {
		t.Fatalf("Load resumed run ok=%v err=%v", ok, err)
	}
	if loaded.Status != "done" {
		t.Fatalf("resumed run status = %q, want done", loaded.Status)
	}
}

func TestHumanInputResumeHonorsCanceledContextAndLeavesRunPending(t *testing.T) {
	ctx := context.Background()
	cp := flow.StoreCheckpoint(store.NewMemoryStore(), "input-cancel-agent")
	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
		if opts.ToolHandler != nil {
			opts.ToolHandler(ctx, ai.ToolCall{ID: "input-1", Name: toolHumanInput, Input: map[string]any{"prompt": "Approve deploy?"}})
		}
		return &ai.Response{Reply: "waiting"}, nil
	}
	defer func() { fakeGen = nil }()

	a := newTestAgent(Name("input-cancel-agent"), WithCheckpoint(cp))
	if _, err := a.Ask(ctx, "deploy the service"); err == nil {
		t.Fatal("Ask succeeded, want input-required pause")
	}
	runs, err := Pending(ctx, a)
	if err != nil {
		t.Fatalf("Pending: %v", err)
	}
	if len(runs) != 1 {
		t.Fatalf("Pending returned %d runs, want 1: %#v", len(runs), runs)
	}

	canceled, cancel := context.WithCancel(ctx)
	cancel()
	if _, err := ResumeInput(canceled, a, runs[0].ID, "yes"); !errors.Is(err, context.Canceled) {
		t.Fatalf("ResumeInput canceled err = %v, want context.Canceled", err)
	}

	loaded, ok, err := cp.Load(ctx, runs[0].ID)
	if err != nil || !ok {
		t.Fatalf("Load paused run ok=%v err=%v", ok, err)
	}
	if loaded.Status != "paused" || loaded.State.Stage != agentInputStep {
		t.Fatalf("run status/stage after canceled resume = %s/%s, want paused/%s", loaded.Status, loaded.State.Stage, agentInputStep)
	}
	var pause inputPause
	if err := loaded.State.Scan(&pause); err != nil {
		t.Fatalf("Scan pause after canceled resume: %v", err)
	}
	if pause.OriginalMessage != "deploy the service" || pause.Prompt != "Approve deploy?" {
		t.Fatalf("pause after canceled resume = %#v", pause)
	}
}

func TestApprovalDenialPausesCheckpointedRunAndResumeContinues(t *testing.T) {
	ctx := context.Background()
	cp := flow.StoreCheckpoint(store.NewMemoryStore(), "approval-agent")
	calls := 0
	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
		calls++
		if opts.ToolHandler != nil {
			opts.ToolHandler(ctx, ai.ToolCall{ID: "call-1", Name: "external.approve", Input: map[string]any{"id": "42"}})
		}
		return &ai.Response{Reply: "model saw approval result"}, nil
	}
	defer func() { fakeGen = nil }()

	approved := false
	a := newTestAgent(Name("approval-agent"), WithCheckpoint(cp),
		WithTool("external.approve", "guarded external action", nil, func(context.Context, map[string]any) (string, error) { return "ok", nil }),
		ApproveTool(func(tool string, input map[string]any) (bool, string) {
			return approved, "waiting for operator"
		}))
	_, err := a.Ask(ctx, "send the guarded update")
	if err == nil {
		t.Fatal("Ask succeeded, want paused approval error")
	}

	runs, err := Pending(ctx, a)
	if err != nil {
		t.Fatalf("Pending: %v", err)
	}
	if len(runs) != 1 {
		t.Fatalf("Pending returned %d runs, want 1: %#v", len(runs), runs)
	}
	if runs[0].Status != "paused" || runs[0].State.Stage != agentApprovalStep {
		t.Fatalf("run status/stage = %s/%s, want paused/%s", runs[0].Status, runs[0].State.Stage, agentApprovalStep)
	}
	if got := string(runs[0].State.Data); got != "send the guarded update" {
		t.Fatalf("paused run data = %q", got)
	}

	approved = true
	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
		calls++
		if opts.ToolHandler != nil {
			res := opts.ToolHandler(ctx, ai.ToolCall{ID: "call-2", Name: "external.approve", Input: map[string]any{"id": "42"}})
			if res.Refused != "" {
				t.Fatalf("resumed call was refused: %#v", res)
			}
		}
		return &ai.Response{Reply: "done after approval"}, nil
	}
	resp, err := Resume(ctx, a, runs[0].ID)
	if err != nil {
		t.Fatalf("Resume: %v", err)
	}
	if resp.Reply != "done after approval" {
		t.Fatalf("Resume reply = %q", resp.Reply)
	}
	loaded, ok, err := cp.Load(ctx, runs[0].ID)
	if err != nil || !ok {
		t.Fatalf("Load resumed run ok=%v err=%v", ok, err)
	}
	if loaded.Status != "done" {
		t.Fatalf("resumed run status = %q, want done", loaded.Status)
	}
	if calls != 2 {
		t.Fatalf("model calls = %d, want 2", calls)
	}
}