Cyrene_v2/internal/runtime/loop/resume_test.go

package loop

import (
	"context"
	"encoding/json"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/persist"
)

// fakeCheckpointStore is a seam for asserting the resume path without a real
// database. It records every LoadLatest call and returns a scripted checkpoint or
// error, so a test can prove the loop asks the store for the session's checkpoint
// and resumes (or stays fresh) based on what comes back.
type fakeCheckpointStore struct {
	latest    *persist.Checkpoint
	latestErr error
	loadCalls []string
}

func (f *fakeCheckpointStore) Save(context.Context, *persist.Checkpoint) error { return nil }

func (f *fakeCheckpointStore) LoadLatest(_ context.Context, agentID string) (*persist.Checkpoint, error) {
	f.loadCalls = append(f.loadCalls, agentID)
	return f.latest, f.latestErr
}

// TestRun_ResumeFromCheckpoint_ContinuesNoRepeat simulates a long task that was
// interrupted mid-tool-loop: the loop left a tool_calls checkpoint at step 2 with
// the message window as it stood after folding two tool batches in, then the
// process died before step 3 could run. A fresh Run on the same sessionID must
// resume at step 3 — rebuilding the context from the checkpoint rather than
// restarting — and must not re-execute the already-completed steps.
func TestRun_ResumeFromCheckpoint_ContinuesNoRepeat(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	ad := &mockAdapter{chat: []mgResp{{resp: model.Response{Content: "final"}}}}
	a := newTestAgent(t, ad, nil)
	cp := newTestCheckpointStore(t)
	l := New(Options{Agent: a, Checkpoint: cp, MaxRounds: 8})

	const sess = "sess-resume"
	interrupted := []model.Message{
		model.TextMessage(model.RoleUser, "do the task"),
		model.TextMessage(model.RoleTool, "tool aa result: va"),
		model.TextMessage(model.RoleTool, "tool bb result: vb"),
	}
	// Seed exactly what a crashed run leaves behind: a tool_calls checkpoint at
	// step 2 with the window computed after step 2's tool batch.
	l.saveCheckpoint(ctx, sess, 2, interrupted, agent.StopReasonToolCalls)

	res, err := l.Run(ctx, sess, []model.Message{model.TextMessage(model.RoleUser, "do the task")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}

	// Resume continues at step 3 (the interrupted step), not step 1.
	if res.Steps != 3 {
		t.Fatalf("Steps = %d, want 3 (resumed at step 3)", res.Steps)
	}
	// No completed step's tool call is re-executed.
	if res.ToolCallsExecuted != 0 {
		t.Fatalf("ToolCallsExecuted = %d, want 0 (completed steps not re-run)", res.ToolCallsExecuted)
	}
	if res.Text != "final" {
		t.Fatalf("Text = %q, want final", res.Text)
	}

	// The resumed run makes exactly one model call (step 3), fed the restored
	// window — not a restart with a fresh conversation.
	if len(ad.reqs) != 1 {
		t.Fatalf("resumed run made %d model calls, want 1", len(ad.reqs))
	}
	got := ad.reqs[0].Messages
	if len(got) != len(interrupted)+1 {
		t.Fatalf("resumed request messages = %d, want %d (system + restored window)", len(got), len(interrupted)+1)
	}
	if got[0].Role != model.RoleSystem {
		t.Fatalf("leading message role = %q, want system (instruction)", got[0].Role)
	}
	for i, m := range interrupted {
		if got[i+1].Role != m.Role || messageText(got[i+1]) != messageText(m) {
			t.Fatalf("restored message %d mismatch: got %+v, want %+v", i, got[i+1], m)
		}
	}
}

// TestRun_NewSession_NotResumed proves a different sessionID never resurrects an
// old checkpoint: the caller opens a brand-new session, LoadLatest misses, and the
// loop runs fresh from step 1 with the caller's own messages.
func TestRun_NewSession_NotResumed(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	call := `{"tool_calls":[{"name":"echo","args":{"q":"x"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "done"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)
	cp := newTestCheckpointStore(t)
	l := New(Options{Agent: a, Checkpoint: cp, MaxRounds: 8})

	// sess-a holds an interrupted run's checkpoint.
	l.saveCheckpoint(ctx, "sess-a", 2, []model.Message{
		model.TextMessage(model.RoleUser, "a task"),
		model.TextMessage(model.RoleTool, "tool echo result: hi"),
	}, agent.StopReasonToolCalls)

	res, err := l.Run(ctx, "sess-b", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Steps != 2 {
		t.Fatalf("Steps = %d, want 2 (fresh two-step run)", res.Steps)
	}
	if res.ToolCallsExecuted != 1 {
		t.Fatalf("ToolCallsExecuted = %d, want 1", res.ToolCallsExecuted)
	}
	// The first model call must start from the caller's fresh message, never from
	// sess-a's restored window.
	first := ad.reqs[0].Messages
	if len(first) != 2 {
		t.Fatalf("fresh run first request messages = %d, want 2 (system + user)", len(first))
	}
	if messageText(first[1]) != "go" {
		t.Fatalf("fresh run first user message = %q, want 'go'", messageText(first[1]))
	}
}

// TestRun_Resume_LoadLatestCalled asserts Run asks the store for the session's
// latest checkpoint and resumes from it (LoadLatest → tool_calls checkpoint →
// continue at the next step).
func TestRun_Resume_LoadLatestCalled(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	ad := &mockAdapter{chat: []mgResp{{resp: model.Response{Content: "final"}}}}
	a := newTestAgent(t, ad, nil)

	buf, err := json.Marshal(checkpointState{
		Step:       2,
		Messages:   []model.Message{model.TextMessage(model.RoleUser, "keep going")},
		StopReason: agent.StopReasonToolCalls,
	})
	if err != nil {
		t.Fatalf("marshal checkpoint: %v", err)
	}
	fake := &fakeCheckpointStore{latest: &persist.Checkpoint{ID: "sess:2", AgentID: "sess", Step: 2, State: buf}}
	l := New(Options{Agent: a, Checkpoint: fake, MaxRounds: 8})

	res, err := l.Run(ctx, "sess", []model.Message{model.TextMessage(model.RoleUser, "new")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if len(fake.loadCalls) != 1 || fake.loadCalls[0] != "sess" {
		t.Fatalf("LoadLatest calls = %v, want exactly ['sess']", fake.loadCalls)
	}
	if res.Steps != 3 {
		t.Fatalf("Steps = %d, want 3 (resumed at step 3)", res.Steps)
	}
	got := ad.reqs[0].Messages
	if len(got) != 2 || messageText(got[1]) != "keep going" {
		t.Fatalf("resumed request = %+v, want [system, 'keep going']", got)
	}
}

// TestRun_Resume_NotFound_Fresh asserts LoadLatest is still called on the fresh
// path and that ErrNotFound yields a full restart from step 1 with the caller's
// messages (no checkpoint to resurrect).
func TestRun_Resume_NotFound_Fresh(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	ad := &mockAdapter{chat: []mgResp{{resp: model.Response{Content: "hello"}}}}
	a := newTestAgent(t, ad, nil)
	fake := &fakeCheckpointStore{latestErr: persist.ErrNotFound}
	l := New(Options{Agent: a, Checkpoint: fake, MaxRounds: 8})

	res, err := l.Run(ctx, "sess", []model.Message{model.TextMessage(model.RoleUser, "hi")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if len(fake.loadCalls) != 1 || fake.loadCalls[0] != "sess" {
		t.Fatalf("LoadLatest calls = %v, want exactly ['sess']", fake.loadCalls)
	}
	if res.Steps != 1 {
		t.Fatalf("Steps = %d, want 1 (fresh run)", res.Steps)
	}
	got := ad.reqs[0].Messages
	if len(got) != 2 || messageText(got[1]) != "hi" {
		t.Fatalf("fresh request = %+v, want [system, 'hi']", got)
	}
}