| 1 | package main |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "os" |
| 7 | "path/filepath" |
| 8 | "sync" |
| 9 | "testing" |
| 10 | |
| 11 | "reasonix/internal/agent" |
| 12 | "reasonix/internal/control" |
| 13 | "reasonix/internal/event" |
| 14 | "reasonix/internal/provider" |
| 15 | "reasonix/internal/tool" |
| 16 | ) |
| 17 | |
| 18 | // capturingProvider records the exact message list of every request it |
| 19 | // receives, marshaled at capture time, so tests can compare request bytes. |
| 20 | type capturingProvider struct { |
| 21 | mu sync.Mutex |
| 22 | requests [][]byte |
| 23 | } |
| 24 | |
| 25 | func (p *capturingProvider) Name() string { return "capturing" } |
| 26 | |
| 27 | func (p *capturingProvider) Stream(_ context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 28 | b, err := json.Marshal(req.Messages) |
| 29 | if err != nil { |
| 30 | return nil, err |
| 31 | } |
| 32 | p.mu.Lock() |
| 33 | p.requests = append(p.requests, b) |
| 34 | p.mu.Unlock() |
| 35 | ch := make(chan provider.Chunk, 1) |
| 36 | ch <- provider.Chunk{Type: provider.ChunkText, Text: "ok"} |
| 37 | close(ch) |
| 38 | return ch, nil |
| 39 | } |
| 40 | |
| 41 | func (p *capturingProvider) lastRequestMessages(t *testing.T) []provider.Message { |
| 42 | t.Helper() |
| 43 | p.mu.Lock() |
| 44 | defer p.mu.Unlock() |
| 45 | if len(p.requests) == 0 { |
| 46 | t.Fatal("provider captured no requests") |
| 47 | } |
| 48 | var msgs []provider.Message |
| 49 | if err := json.Unmarshal(p.requests[len(p.requests)-1], &msgs); err != nil { |
| 50 | t.Fatalf("unmarshal captured request: %v", err) |
| 51 | } |
| 52 | return msgs |
| 53 | } |
| 54 | |
| 55 | // marshalMessages drops message ids first: they are local transcript identity |
| 56 | // that provider adapters never copy to the wire, and a freshly composed |
| 57 | // follow-up legitimately mints a new one on every run. |
| 58 | func marshalMessages(t *testing.T, msgs []provider.Message) []byte { |
| 59 | t.Helper() |
| 60 | msgs = append([]provider.Message(nil), msgs...) |
| 61 | for i := range msgs { |
| 62 | msgs[i].ID = "" |
| 63 | } |
| 64 | b, err := json.Marshal(msgs) |
| 65 | if err != nil { |
| 66 | t.Fatalf("marshal messages: %v", err) |
| 67 | } |
| 68 | return b |
| 69 | } |
| 70 | |
| 71 | // copySessionFiles clones a saved transcript (checkpoint anchor, event log, |
| 72 | // meta sidecar) to an independent path, so a rebind can load the state saved |
| 73 | // at that moment while the original controller keeps running and autosaving. |
| 74 | func copySessionFiles(t *testing.T, from, to string) { |
| 75 | t.Helper() |
| 76 | copied := false |
| 77 | for _, suffix := range []string{"", ".events.jsonl", ".meta"} { |
| 78 | b, err := os.ReadFile(from + suffix) |
| 79 | if err != nil { |
| 80 | continue |
| 81 | } |
| 82 | if err := os.WriteFile(to+suffix, b, 0o644); err != nil { |
| 83 | t.Fatalf("copy session file %s: %v", suffix, err) |
| 84 | } |
| 85 | copied = true |
| 86 | } |
| 87 | if !copied { |
| 88 | t.Fatalf("no session files found at %s", from) |
| 89 | } |
| 90 | } |
| 91 | |
| 92 | // TestRebindReproducesRequestBytes is the desktop-level byte-stability guard |
| 93 | // for the provider prefix cache. It builds the strongest comparison available: |
| 94 | // from ONE saved transcript, run the same follow-up turn twice — once on the |
| 95 | // original controller (no rebind, the provider-cache-warm baseline) and once |
| 96 | // after the desktop rebind path (agent.LoadSession + sessionWithFreshSystemPrompt |
| 97 | // + Resume on a freshly built controller, the shape of tabs.go's restore). The |
| 98 | // two requests must be byte-identical END TO END — system prompt, prior user |
| 99 | // AND assistant turns, and the composed follow-up. Any divergence means a |
| 100 | // desktop rebuild cold-starts the conversation's provider cache at 10x miss |
| 101 | // pricing (#2945, #5614). |
| 102 | func TestRebindReproducesRequestBytes(t *testing.T) { |
| 103 | isolateDesktopUserDirs(t) |
| 104 | dir := t.TempDir() |
| 105 | path := filepath.Join(dir, "session.jsonl") |
| 106 | const systemPrompt = "SYSPROMPT stable bytes" |
| 107 | |
| 108 | prov := &capturingProvider{} |
| 109 | exec := agent.New(prov, tool.NewRegistry(), agent.NewSession(systemPrompt), agent.Options{}, event.Discard) |
| 110 | ctrl := newFixtureController(t, control.Options{Runner: exec, Executor: exec, SystemPrompt: systemPrompt, SessionDir: dir, SessionPath: path, Label: "test", Sink: event.Discard}) |
| 111 | |
| 112 | if err := ctrl.RunTurn(context.Background(), "first question"); err != nil { |
| 113 | t.Fatalf("first turn: %v", err) |
| 114 | } |
| 115 | if err := ctrl.Snapshot(); err != nil { |
| 116 | t.Fatalf("Snapshot: %v", err) |
| 117 | } |
| 118 | // Freeze the after-turn-one transcript on an independent path: the |
| 119 | // baseline turn below autosaves onto the original path, and the rebind |
| 120 | // must load the state as saved at this moment. |
| 121 | rebindPath := filepath.Join(dir, "rebind.jsonl") |
| 122 | copySessionFiles(t, path, rebindPath) |
| 123 | |
| 124 | // Baseline: the follow-up turn on the ORIGINAL controller — the exact |
| 125 | // request an uninterrupted (cache-warm) session would send. |
| 126 | if err := ctrl.RunTurn(context.Background(), "second question"); err != nil { |
| 127 | t.Fatalf("baseline second turn: %v", err) |
| 128 | } |
| 129 | baseline := prov.lastRequestMessages(t) |
| 130 | baselineBytes := marshalMessages(t, baseline) |
| 131 | if len(baseline) < 4 { |
| 132 | t.Fatalf("baseline request has %d messages, want system + first exchange + follow-up", len(baseline)) |
| 133 | } |
| 134 | |
| 135 | // Rebind from the transcript saved after turn one: a NEW controller |
| 136 | // composes its (identical) system prompt, the persisted transcript is |
| 137 | // loaded, the fresh prompt is swapped in, and the controller resumes — |
| 138 | // then sends the same follow-up. |
| 139 | prov2 := &capturingProvider{} |
| 140 | exec2 := agent.New(prov2, tool.NewRegistry(), agent.NewSession(systemPrompt), agent.Options{}, event.Discard) |
| 141 | ctrl2 := newFixtureController(t, control.Options{Runner: exec2, Executor: exec2, SystemPrompt: systemPrompt, SessionDir: dir, SessionPath: rebindPath, Label: "test", Sink: event.Discard}) |
| 142 | loaded, err := agent.LoadSession(rebindPath) |
| 143 | if err != nil { |
| 144 | t.Fatalf("LoadSession: %v", err) |
| 145 | } |
| 146 | ctrl2.Resume(sessionWithFreshSystemPrompt(loaded, systemPromptFrom(ctrl2.History())), rebindPath) |
| 147 | |
| 148 | if err := ctrl2.RunTurn(context.Background(), "second question"); err != nil { |
| 149 | t.Fatalf("post-rebind second turn: %v", err) |
| 150 | } |
| 151 | rebound := prov2.lastRequestMessages(t) |
| 152 | reboundBytes := marshalMessages(t, rebound) |
| 153 | if string(reboundBytes) != string(baselineBytes) { |
| 154 | t.Fatalf("rebind changed the request bytes — the provider prefix cache is invalidated:\nbaseline: %s\nrebound: %s", baselineBytes, reboundBytes) |
| 155 | } |
| 156 | } |
| 157 | |
| 158 | // TestRebindWithDriftedPromptBreaksRequestPrefix pins the failure mode the |
| 159 | // guard above protects against: when the freshly composed prompt differs from |
| 160 | // the one the transcript was recorded with, the swap rewrites the first |
| 161 | // message and the request diverges from the no-rebind baseline. If a future |
| 162 | // change moves the swap policy to keep the persisted prompt for resumed |
| 163 | // conversations, this test should be updated to assert the bytes survive |
| 164 | // instead. |
| 165 | func TestRebindWithDriftedPromptBreaksRequestPrefix(t *testing.T) { |
| 166 | isolateDesktopUserDirs(t) |
| 167 | dir := t.TempDir() |
| 168 | path := filepath.Join(dir, "session.jsonl") |
| 169 | |
| 170 | prov := &capturingProvider{} |
| 171 | exec := agent.New(prov, tool.NewRegistry(), agent.NewSession("SYSPROMPT v1"), agent.Options{}, event.Discard) |
| 172 | ctrl := newFixtureController(t, control.Options{Runner: exec, Executor: exec, SystemPrompt: "SYSPROMPT v1", SessionDir: dir, SessionPath: path, Label: "test", Sink: event.Discard}) |
| 173 | if err := ctrl.RunTurn(context.Background(), "first question"); err != nil { |
| 174 | t.Fatalf("first turn: %v", err) |
| 175 | } |
| 176 | if err := ctrl.Snapshot(); err != nil { |
| 177 | t.Fatalf("Snapshot: %v", err) |
| 178 | } |
| 179 | rebindPath := filepath.Join(dir, "rebind.jsonl") |
| 180 | copySessionFiles(t, path, rebindPath) |
| 181 | if err := ctrl.RunTurn(context.Background(), "second question"); err != nil { |
| 182 | t.Fatalf("baseline second turn: %v", err) |
| 183 | } |
| 184 | baseline := prov.lastRequestMessages(t) |
| 185 | baselineBytes := marshalMessages(t, baseline) |
| 186 | |
| 187 | prov2 := &capturingProvider{} |
| 188 | exec2 := agent.New(prov2, tool.NewRegistry(), agent.NewSession("SYSPROMPT v2 drifted"), agent.Options{}, event.Discard) |
| 189 | ctrl2 := newFixtureController(t, control.Options{Runner: exec2, Executor: exec2, SystemPrompt: "SYSPROMPT v2 drifted", SessionDir: dir, SessionPath: rebindPath, Label: "test", Sink: event.Discard}) |
| 190 | loaded, err := agent.LoadSession(rebindPath) |
| 191 | if err != nil { |
| 192 | t.Fatalf("LoadSession: %v", err) |
| 193 | } |
| 194 | ctrl2.Resume(sessionWithFreshSystemPrompt(loaded, systemPromptFrom(ctrl2.History())), rebindPath) |
| 195 | if err := ctrl2.RunTurn(context.Background(), "second question"); err != nil { |
| 196 | t.Fatalf("post-rebind turn: %v", err) |
| 197 | } |
| 198 | rebound := prov2.lastRequestMessages(t) |
| 199 | if string(marshalMessages(t, rebound)) == string(baselineBytes) { |
| 200 | t.Fatal("drifted prompt unexpectedly reproduced the baseline request — the swap policy changed; update these guards") |
| 201 | } |
| 202 | if len(rebound) == 0 || rebound[0].Content == baseline[0].Content { |
| 203 | t.Fatalf("drift should surface in the leading system message; got %q", rebound[0].Content) |
| 204 | } |
| 205 | } |
| 206 |