| 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 | func marshalMessages(t *testing.T, msgs []provider.Message) []byte { |
| 56 | t.Helper() |
| 57 | b, err := json.Marshal(msgs) |
| 58 | if err != nil { |
| 59 | t.Fatalf("marshal messages: %v", err) |
| 60 | } |
| 61 | return b |
| 62 | } |
| 63 | |
| 64 | // copySessionFiles clones a saved transcript (checkpoint anchor, event log, |
| 65 | // meta sidecar) to an independent path, so a rebind can load the state saved |
| 66 | // at that moment while the original controller keeps running and autosaving. |
| 67 | func copySessionFiles(t *testing.T, from, to string) { |
| 68 | t.Helper() |
| 69 | copied := false |
| 70 | for _, suffix := range []string{"", ".events.jsonl", ".meta"} { |
| 71 | b, err := os.ReadFile(from + suffix) |
| 72 | if err != nil { |
| 73 | continue |
| 74 | } |
| 75 | if err := os.WriteFile(to+suffix, b, 0o644); err != nil { |
| 76 | t.Fatalf("copy session file %s: %v", suffix, err) |
| 77 | } |
| 78 | copied = true |
| 79 | } |
| 80 | if !copied { |
| 81 | t.Fatalf("no session files found at %s", from) |
| 82 | } |
| 83 | } |
| 84 | |
| 85 | // TestRebindReproducesRequestBytes is the desktop-level byte-stability guard |
| 86 | // for the provider prefix cache. It builds the strongest comparison available: |
| 87 | // from ONE saved transcript, run the same follow-up turn twice — once on the |
| 88 | // original controller (no rebind, the provider-cache-warm baseline) and once |
| 89 | // after the desktop rebind path (agent.LoadSession + sessionWithFreshSystemPrompt |
| 90 | // + Resume on a freshly built controller, the shape of tabs.go's restore). The |
| 91 | // two requests must be byte-identical END TO END — system prompt, prior user |
| 92 | // AND assistant turns, and the composed follow-up. Any divergence means a |
| 93 | // desktop rebuild cold-starts the conversation's provider cache at 10x miss |
| 94 | // pricing (#2945, #5614). |
| 95 | func TestRebindReproducesRequestBytes(t *testing.T) { |
| 96 | isolateDesktopUserDirs(t) |
| 97 | dir := t.TempDir() |
| 98 | path := filepath.Join(dir, "session.jsonl") |
| 99 | const systemPrompt = "SYSPROMPT stable bytes" |
| 100 | |
| 101 | prov := &capturingProvider{} |
| 102 | exec := agent.New(prov, tool.NewRegistry(), agent.NewSession(systemPrompt), agent.Options{}, event.Discard) |
| 103 | ctrl := control.New(control.Options{Runner: exec, Executor: exec, SystemPrompt: systemPrompt, SessionDir: dir, SessionPath: path, Label: "test", Sink: event.Discard}) |
| 104 | |
| 105 | if err := ctrl.RunTurn(context.Background(), "first question"); err != nil { |
| 106 | t.Fatalf("first turn: %v", err) |
| 107 | } |
| 108 | if err := ctrl.Snapshot(); err != nil { |
| 109 | t.Fatalf("Snapshot: %v", err) |
| 110 | } |
| 111 | // Freeze the after-turn-one transcript on an independent path: the |
| 112 | // baseline turn below autosaves onto the original path, and the rebind |
| 113 | // must load the state as saved at this moment. |
| 114 | rebindPath := filepath.Join(dir, "rebind.jsonl") |
| 115 | copySessionFiles(t, path, rebindPath) |
| 116 | |
| 117 | // Baseline: the follow-up turn on the ORIGINAL controller — the exact |
| 118 | // request an uninterrupted (cache-warm) session would send. |
| 119 | if err := ctrl.RunTurn(context.Background(), "second question"); err != nil { |
| 120 | t.Fatalf("baseline second turn: %v", err) |
| 121 | } |
| 122 | baseline := prov.lastRequestMessages(t) |
| 123 | baselineBytes := marshalMessages(t, baseline) |
| 124 | if len(baseline) < 4 { |
| 125 | t.Fatalf("baseline request has %d messages, want system + first exchange + follow-up", len(baseline)) |
| 126 | } |
| 127 | |
| 128 | // Rebind from the transcript saved after turn one: a NEW controller |
| 129 | // composes its (identical) system prompt, the persisted transcript is |
| 130 | // loaded, the fresh prompt is swapped in, and the controller resumes — |
| 131 | // then sends the same follow-up. |
| 132 | prov2 := &capturingProvider{} |
| 133 | exec2 := agent.New(prov2, tool.NewRegistry(), agent.NewSession(systemPrompt), agent.Options{}, event.Discard) |
| 134 | ctrl2 := control.New(control.Options{Runner: exec2, Executor: exec2, SystemPrompt: systemPrompt, SessionDir: dir, SessionPath: rebindPath, Label: "test", Sink: event.Discard}) |
| 135 | loaded, err := agent.LoadSession(rebindPath) |
| 136 | if err != nil { |
| 137 | t.Fatalf("LoadSession: %v", err) |
| 138 | } |
| 139 | ctrl2.Resume(sessionWithFreshSystemPrompt(loaded, systemPromptFrom(ctrl2.History())), rebindPath) |
| 140 | |
| 141 | if err := ctrl2.RunTurn(context.Background(), "second question"); err != nil { |
| 142 | t.Fatalf("post-rebind second turn: %v", err) |
| 143 | } |
| 144 | rebound := prov2.lastRequestMessages(t) |
| 145 | reboundBytes := marshalMessages(t, rebound) |
| 146 | if string(reboundBytes) != string(baselineBytes) { |
| 147 | t.Fatalf("rebind changed the request bytes — the provider prefix cache is invalidated:\nbaseline: %s\nrebound: %s", baselineBytes, reboundBytes) |
| 148 | } |
| 149 | } |
| 150 | |
| 151 | // TestRebindWithDriftedPromptBreaksRequestPrefix pins the failure mode the |
| 152 | // guard above protects against: when the freshly composed prompt differs from |
| 153 | // the one the transcript was recorded with, the swap rewrites the first |
| 154 | // message and the request diverges from the no-rebind baseline. If a future |
| 155 | // change moves the swap policy to keep the persisted prompt for resumed |
| 156 | // conversations, this test should be updated to assert the bytes survive |
| 157 | // instead. |
| 158 | func TestRebindWithDriftedPromptBreaksRequestPrefix(t *testing.T) { |
| 159 | isolateDesktopUserDirs(t) |
| 160 | dir := t.TempDir() |
| 161 | path := filepath.Join(dir, "session.jsonl") |
| 162 | |
| 163 | prov := &capturingProvider{} |
| 164 | exec := agent.New(prov, tool.NewRegistry(), agent.NewSession("SYSPROMPT v1"), agent.Options{}, event.Discard) |
| 165 | ctrl := control.New(control.Options{Runner: exec, Executor: exec, SystemPrompt: "SYSPROMPT v1", SessionDir: dir, SessionPath: path, Label: "test", Sink: event.Discard}) |
| 166 | if err := ctrl.RunTurn(context.Background(), "first question"); err != nil { |
| 167 | t.Fatalf("first turn: %v", err) |
| 168 | } |
| 169 | if err := ctrl.Snapshot(); err != nil { |
| 170 | t.Fatalf("Snapshot: %v", err) |
| 171 | } |
| 172 | rebindPath := filepath.Join(dir, "rebind.jsonl") |
| 173 | copySessionFiles(t, path, rebindPath) |
| 174 | if err := ctrl.RunTurn(context.Background(), "second question"); err != nil { |
| 175 | t.Fatalf("baseline second turn: %v", err) |
| 176 | } |
| 177 | baseline := prov.lastRequestMessages(t) |
| 178 | baselineBytes := marshalMessages(t, baseline) |
| 179 | |
| 180 | prov2 := &capturingProvider{} |
| 181 | exec2 := agent.New(prov2, tool.NewRegistry(), agent.NewSession("SYSPROMPT v2 drifted"), agent.Options{}, event.Discard) |
| 182 | ctrl2 := control.New(control.Options{Runner: exec2, Executor: exec2, SystemPrompt: "SYSPROMPT v2 drifted", SessionDir: dir, SessionPath: rebindPath, Label: "test", Sink: event.Discard}) |
| 183 | loaded, err := agent.LoadSession(rebindPath) |
| 184 | if err != nil { |
| 185 | t.Fatalf("LoadSession: %v", err) |
| 186 | } |
| 187 | ctrl2.Resume(sessionWithFreshSystemPrompt(loaded, systemPromptFrom(ctrl2.History())), rebindPath) |
| 188 | if err := ctrl2.RunTurn(context.Background(), "second question"); err != nil { |
| 189 | t.Fatalf("post-rebind turn: %v", err) |
| 190 | } |
| 191 | rebound := prov2.lastRequestMessages(t) |
| 192 | if string(marshalMessages(t, rebound)) == string(baselineBytes) { |
| 193 | t.Fatal("drifted prompt unexpectedly reproduced the baseline request — the swap policy changed; update these guards") |
| 194 | } |
| 195 | if len(rebound) == 0 || rebound[0].Content == baseline[0].Content { |
| 196 | t.Fatalf("drift should surface in the leading system message; got %q", rebound[0].Content) |
| 197 | } |
| 198 | } |
| 199 |