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 是在没有真实数据库的情况下断言恢复路径的
// 接缝。它记录每次 LoadLatest 调用并返回脚本化的 checkpoint 或
// 错误，使测试能证明循环向 store 索要该会话的 checkpoint，
// 并据此恢复（或保持全新）。
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 模拟一个在工具循环中途
// 被中断的长任务：循环在步 2 留下了 tool_calls checkpoint，其消息
// 窗口是并入两批工具结果后的样子，随后进程在步 3 能运行前
// 死亡。同一 sessionID 上全新的 Run 必须从步 3 恢复 ——
// 由该 checkpoint 重建上下文而非重启 ——
// 且不得重跑已完成的步骤。
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"),
	}
	// 精确播种崩溃的 run 留下的东西：步 2 的 tool_calls checkpoint，
	// 其窗口在步 2 的工具批次之后算出。
	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)
	}

	// 恢复从步 3（被中断的那一步）继续，而不是步 1。
	if res.Steps != 3 {
		t.Fatalf("Steps = %d, want 3 (resumed at step 3)", res.Steps)
	}
	// 已完成步骤的工具调用不会重跑。
	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)
	}

	// 恢复后的 run 恰好做一次模型调用（步 3），喂入恢复的
	// 窗口 —— 而不是以全新对话重启。
	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 证明不同的 sessionID 绝不复活
// 旧的 checkpoint：调用方打开全新会话，LoadLatest 未命中，
// 循环以自己的消息从步 1 全新运行。
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 持有一个被中断 run 的 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)
	}
	// 第一次模型调用必须从调用方的新消息开始，绝不从
	// sess-a 恢复的窗口开始。
	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 断言 Run 向 store 索要该会话
// 最新的 checkpoint 并据此恢复（LoadLatest → tool_calls checkpoint →
// 从下一步继续）。
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 断言全新路径上仍会调用 LoadLatest，
// 且 ErrNotFound 会以调用方的消息从步 1 完全重启
// （没有可复活的 checkpoint）。
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)
	}
}