1package core
2
3import (
4 "context"
5 "testing"
6
7 "github.com/bnema/zut/packages/provider"
8)
9
10func TestSessionMessageOriginsAcrossSwitch(t *testing.T) {
11 root, cwd := t.TempDir(), t.TempDir()
12 session, err := NewSession(root, cwd, "opencode", "model-a", "test")
13 if err != nil {
14 t.Fatal(err)
15 }
16 defer session.Close()
17 msg := provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{provider.ReasoningBlock{ID: "rs_a:rs_b", Encrypted: "opaque"}, provider.TextBlock{Text: "answer"}}}
18 if err := session.AppendMessage(msg); err != nil {
19 t.Fatal(err)
20 }
21 if err := session.UpdateModel("openai-codex", "model-b"); err != nil {
22 t.Fatal(err)
23 }
24 if err := session.AppendMessage(msg); err != nil {
25 t.Fatal(err)
26 }
27 snapshot, err := ReadSessionSnapshot(session.Path)
28 if err != nil {
29 t.Fatal(err)
30 }
31 for i, want := range []string{"opencode", "openai-codex"} {
32 if got := snapshot.Messages[i].Meta["provider"]; got != want {
33 t.Fatalf("message %d origin = %q, want %q", i, got, want)
34 }
35 }
36 if err := session.AppendCompaction(snapshot.Messages); err != nil {
37 t.Fatal(err)
38 }
39 if err := session.UpdateModel("google", "model-c"); err != nil {
40 t.Fatal(err)
41 }
42 snapshot, err = ReadSessionSnapshot(session.Path)
43 if err != nil {
44 t.Fatal(err)
45 }
46 if snapshot.Messages[0].Meta["provider"] != "opencode" || snapshot.Messages[1].Meta["model"] != "model-b" {
47 t.Fatalf("compaction lost origins: %+v", snapshot.Messages)
48 }
49 history, err := ReadSessionHistory(session.Path)
50 if err != nil {
51 t.Fatal(err)
52 }
53 branchPath, err := BranchSessionHiddenFromHistory(session.Path, root, cwd, "test", history.Segments[0], len(history.Segments[0].Messages))
54 if err != nil {
55 t.Fatal(err)
56 }
57 branch, err := ReadSessionSnapshot(branchPath)
58 if err != nil {
59 t.Fatal(err)
60 }
61 if len(branch.Messages) != 2 || branch.Messages[0].Meta["provider"] != "opencode" || branch.Messages[0].Meta["model"] != "model-a" || branch.Messages[1].Meta["provider"] != "openai-codex" || branch.Messages[1].Meta["model"] != "model-b" {
62 t.Fatalf("historical branch lost origins: %+v", branch.Messages)
63 }
64}
65
66func TestAgentRecordsMessageOriginAcrossModelSwitch(t *testing.T) {
67 client := &identityRecordingClient{}
68 agent := NewAgent(client, "model-a", "system", Registry{})
69 if err := agent.Prompt(context.Background(), "first", nil, nil); err != nil {
70 t.Fatal(err)
71 }
72 agent.Model = "model-b"
73 if err := agent.Prompt(context.Background(), "second", nil, nil); err != nil {
74 t.Fatal(err)
75 }
76 var models []string
77 for _, msg := range agent.Messages() {
78 if msg.Role == provider.RoleAssistant {
79 if msg.Meta["provider"] != client.Name() {
80 t.Fatal("missing provider origin")
81 }
82 models = append(models, msg.Meta["model"])
83 }
84 }
85 if len(models) != 2 || models[0] != "model-a" || models[1] != "model-b" {
86 t.Fatalf("origins = %v", models)
87 }
88}