| 1 | package sidecar |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "sync" |
| 6 | "testing" |
| 7 | "time" |
| 8 | |
| 9 | "reasonix/internal/extension/protocol" |
| 10 | "reasonix/internal/pluginpkg" |
| 11 | ) |
| 12 | |
| 13 | // recordingRouter captures routed provider stream notifications. |
| 14 | type recordingRouter struct { |
| 15 | mu sync.Mutex |
| 16 | chunks []protocol.StreamChunkParams |
| 17 | ends []protocol.StreamEndParams |
| 18 | } |
| 19 | |
| 20 | func (r *recordingRouter) RouteStreamChunk(p protocol.StreamChunkParams) { |
| 21 | r.mu.Lock() |
| 22 | defer r.mu.Unlock() |
| 23 | r.chunks = append(r.chunks, p) |
| 24 | } |
| 25 | |
| 26 | func (r *recordingRouter) RouteStreamEnd(p protocol.StreamEndParams) { |
| 27 | r.mu.Lock() |
| 28 | defer r.mu.Unlock() |
| 29 | r.ends = append(r.ends, p) |
| 30 | } |
| 31 | |
| 32 | func (r *recordingRouter) counts() (int, int) { |
| 33 | r.mu.Lock() |
| 34 | defer r.mu.Unlock() |
| 35 | return len(r.chunks), len(r.ends) |
| 36 | } |
| 37 | |
| 38 | func openProviderStream(t *testing.T, client *Client, streamID string) { |
| 39 | t.Helper() |
| 40 | opened, err := client.ProviderStreamOpen(context.Background(), protocol.StreamOpenParams{ |
| 41 | StreamID: streamID, |
| 42 | ProviderRef: "plugin/fakeplugin/fake/x", |
| 43 | Request: protocol.ProviderRequest{Messages: []protocol.ProviderMessage{}, Tools: []protocol.ProviderToolSchema{}}, |
| 44 | SeqBase: 1, |
| 45 | }) |
| 46 | if err != nil { |
| 47 | t.Fatalf("ProviderStreamOpen: %v", err) |
| 48 | } |
| 49 | if !opened.Accepted { |
| 50 | t.Fatal("ProviderStreamOpen was declined") |
| 51 | } |
| 52 | } |
| 53 | |
| 54 | // TestStreamRouterReceivesWireNotifications pins the stage 7 seam: inbound |
| 55 | // stream/chunk and stream/end notifications reach the installed router, |
| 56 | // decoded and addressed by stream ID. |
| 57 | func TestStreamRouterReceivesWireNotifications(t *testing.T) { |
| 58 | recorder := &recordingRouter{} |
| 59 | client := startFakeClient(t, func(rt *pluginpkg.RuntimeSpec) { |
| 60 | rt.Env[fakeEnvMode] = "provider_stream" |
| 61 | }, func(opts *ClientOptions) { |
| 62 | opts.Streams = recorder |
| 63 | }) |
| 64 | |
| 65 | openProviderStream(t, client, "es_route") |
| 66 | waitFor(t, "chunk and end routed", 5*time.Second, func() bool { |
| 67 | chunks, ends := recorder.counts() |
| 68 | return chunks == 1 && ends == 1 |
| 69 | }) |
| 70 | recorder.mu.Lock() |
| 71 | defer recorder.mu.Unlock() |
| 72 | if recorder.chunks[0].StreamID != "es_route" || recorder.chunks[0].Seq != 1 || |
| 73 | recorder.chunks[0].Chunk.Type != protocol.ChunkText || recorder.chunks[0].Chunk.Text != "wired" { |
| 74 | t.Fatalf("routed chunk = %+v", recorder.chunks[0]) |
| 75 | } |
| 76 | if recorder.ends[0].StreamID != "es_route" || recorder.ends[0].LastSeq != 1 { |
| 77 | t.Fatalf("routed end = %+v", recorder.ends[0]) |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | // TestRollbackAfterStageKeepsOldRouterConsumingRealStream pins the fail-atomic |
| 82 | // Unchanged-sidecar contract: after a narrow-reload stage adopts a live client |
| 83 | // and then fails before commit (no SetStreamRouter / installSidecarStreamRouters), |
| 84 | // RollbackPlanStart reattaches the client and the pre-stage StreamRouter still |
| 85 | // receives real wire stream/chunk and stream/end notifications. |
| 86 | func TestRollbackAfterStageKeepsOldRouterConsumingRealStream(t *testing.T) { |
| 87 | old := &recordingRouter{} |
| 88 | client := startFakeClient(t, func(rt *pluginpkg.RuntimeSpec) { |
| 89 | rt.Env[fakeEnvMode] = "provider_stream" |
| 90 | }, func(opts *ClientOptions) { |
| 91 | opts.Streams = old |
| 92 | }) |
| 93 | |
| 94 | prev := &Manager{} |
| 95 | if err := prev.Adopt("fakeplugin", client); err != nil { |
| 96 | t.Fatal(err) |
| 97 | } |
| 98 | |
| 99 | // Stage: Unchanged adopt into next (as StartPackagesWithPlan). Deliberately |
| 100 | // do not install a next-gen router — that only happens on commit. |
| 101 | next := &Manager{planAdopted: map[string]*Client{}} |
| 102 | if c := prev.Detach("fakeplugin"); c == nil { |
| 103 | t.Fatal("expected live client on previous manager") |
| 104 | } else { |
| 105 | if err := next.Adopt("fakeplugin", c); err != nil { |
| 106 | t.Fatal(err) |
| 107 | } |
| 108 | next.planAdopted["fakeplugin"] = c |
| 109 | } |
| 110 | discarded := &recordingRouter{} // would be next-gen resolver if wrongly installed |
| 111 | if client.streamRouter() != old { |
| 112 | t.Fatal("stage adopt must not rewrite StreamRouter") |
| 113 | } |
| 114 | |
| 115 | next.RollbackPlanStart(prev) |
| 116 | if prev.Client("fakeplugin") != client { |
| 117 | t.Fatal("client must be reattached to previous manager") |
| 118 | } |
| 119 | if client.streamRouter() != old { |
| 120 | t.Fatal("after rollback StreamRouter must still be the old generation") |
| 121 | } |
| 122 | |
| 123 | openProviderStream(t, client, "es_after_rollback") |
| 124 | waitFor(t, "old router receives chunk and end after rollback", 5*time.Second, func() bool { |
| 125 | chunks, ends := old.counts() |
| 126 | return chunks == 1 && ends == 1 |
| 127 | }) |
| 128 | if chunks, ends := discarded.counts(); chunks != 0 || ends != 0 { |
| 129 | t.Fatalf("discarded next-gen router saw traffic: chunks=%d ends=%d", chunks, ends) |
| 130 | } |
| 131 | old.mu.Lock() |
| 132 | defer old.mu.Unlock() |
| 133 | if old.chunks[0].StreamID != "es_after_rollback" || old.chunks[0].Seq != 1 || |
| 134 | old.chunks[0].Chunk.Type != protocol.ChunkText || old.chunks[0].Chunk.Text != "wired" { |
| 135 | t.Fatalf("routed chunk after rollback = %+v", old.chunks[0]) |
| 136 | } |
| 137 | if old.ends[0].StreamID != "es_after_rollback" || old.ends[0].LastSeq != 1 { |
| 138 | t.Fatalf("routed end after rollback = %+v", old.ends[0]) |
| 139 | } |
| 140 | } |
| 141 | |
| 142 | // TestSetStreamRouterSwapsMidFlight: a router installed after start receives |
| 143 | // later notifications; the replaced one stops seeing them. |
| 144 | func TestSetStreamRouterSwapsMidFlight(t *testing.T) { |
| 145 | first := &recordingRouter{} |
| 146 | client := startFakeClient(t, func(rt *pluginpkg.RuntimeSpec) { |
| 147 | rt.Env[fakeEnvMode] = "provider_stream" |
| 148 | }, func(opts *ClientOptions) { |
| 149 | opts.Streams = first |
| 150 | }) |
| 151 | |
| 152 | second := &recordingRouter{} |
| 153 | client.SetStreamRouter(second) |
| 154 | openProviderStream(t, client, "es_swapped") |
| 155 | waitFor(t, "chunk routed to the swapped router", 5*time.Second, func() bool { |
| 156 | chunks, _ := second.counts() |
| 157 | return chunks == 1 |
| 158 | }) |
| 159 | if chunks, _ := first.counts(); chunks != 0 { |
| 160 | t.Fatalf("replaced router saw %d chunks after the swap", chunks) |
| 161 | } |
| 162 | } |
| 163 | |
| 164 | func TestSetStreamRouterNilRestoresDropDefault(t *testing.T) { |
| 165 | c := &Client{pluginID: "p", streams: dropStreamRouter{pluginID: "p"}} |
| 166 | first := &recordingRouter{} |
| 167 | c.SetStreamRouter(first) |
| 168 | if c.streamRouter() != first { |
| 169 | t.Fatal("SetStreamRouter did not install the router") |
| 170 | } |
| 171 | c.SetStreamRouter(nil) |
| 172 | if _, ok := c.streamRouter().(dropStreamRouter); !ok { |
| 173 | t.Fatalf("SetStreamRouter(nil) restored %T, want the drop default", c.streamRouter()) |
| 174 | } |
| 175 | } |
| 176 | |
| 177 | // TestDisconnectedClosesWithServeLoop: the provider stream watchers' signal |
| 178 | // fires on an orderly shutdown too, so in-flight streams never hang. |
| 179 | func TestDisconnectedClosesWithServeLoop(t *testing.T) { |
| 180 | client := startFakeClient(t, nil, nil) |
| 181 | select { |
| 182 | case <-client.Disconnected(): |
| 183 | t.Fatal("Disconnected closed on a live client") |
| 184 | default: |
| 185 | } |
| 186 | if err := client.Close(); err != nil { |
| 187 | t.Fatalf("Close: %v", err) |
| 188 | } |
| 189 | select { |
| 190 | case <-client.Disconnected(): |
| 191 | case <-time.After(5 * time.Second): |
| 192 | t.Fatal("Disconnected did not close after shutdown") |
| 193 | } |
| 194 | } |
| 195 | |
| 196 | func TestProviderCatalogRoundTrip(t *testing.T) { |
| 197 | client := startFakeClient(t, nil, nil) |
| 198 | providers, err := client.ProviderCatalog(context.Background()) |
| 199 | if err != nil { |
| 200 | t.Fatalf("ProviderCatalog: %v", err) |
| 201 | } |
| 202 | if len(providers) != 0 { |
| 203 | t.Fatalf("providers = %v, want the fake's empty catalog", providers) |
| 204 | } |
| 205 | } |
| 206 | |
| 207 | func TestProviderStreamCancelBestEffort(t *testing.T) { |
| 208 | client := startFakeClient(t, nil, nil) |
| 209 | // The fake answers {"cancelled":true}; the call must simply not wedge. |
| 210 | client.ProviderStreamCancel("es_test") |
| 211 | } |
| 212 |