| 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 | // TestSetStreamRouterSwapsMidFlight: a router installed after start receives |
| 82 | // later notifications; the replaced one stops seeing them. |
| 83 | func TestSetStreamRouterSwapsMidFlight(t *testing.T) { |
| 84 | first := &recordingRouter{} |
| 85 | client := startFakeClient(t, func(rt *pluginpkg.RuntimeSpec) { |
| 86 | rt.Env[fakeEnvMode] = "provider_stream" |
| 87 | }, func(opts *ClientOptions) { |
| 88 | opts.Streams = first |
| 89 | }) |
| 90 | |
| 91 | second := &recordingRouter{} |
| 92 | client.SetStreamRouter(second) |
| 93 | openProviderStream(t, client, "es_swapped") |
| 94 | waitFor(t, "chunk routed to the swapped router", 5*time.Second, func() bool { |
| 95 | chunks, _ := second.counts() |
| 96 | return chunks == 1 |
| 97 | }) |
| 98 | if chunks, _ := first.counts(); chunks != 0 { |
| 99 | t.Fatalf("replaced router saw %d chunks after the swap", chunks) |
| 100 | } |
| 101 | } |
| 102 | |
| 103 | func TestSetStreamRouterNilRestoresDropDefault(t *testing.T) { |
| 104 | c := &Client{pluginID: "p", streams: dropStreamRouter{pluginID: "p"}} |
| 105 | first := &recordingRouter{} |
| 106 | c.SetStreamRouter(first) |
| 107 | if c.streamRouter() != first { |
| 108 | t.Fatal("SetStreamRouter did not install the router") |
| 109 | } |
| 110 | c.SetStreamRouter(nil) |
| 111 | if _, ok := c.streamRouter().(dropStreamRouter); !ok { |
| 112 | t.Fatalf("SetStreamRouter(nil) restored %T, want the drop default", c.streamRouter()) |
| 113 | } |
| 114 | } |
| 115 | |
| 116 | // TestDisconnectedClosesWithServeLoop: the provider stream watchers' signal |
| 117 | // fires on an orderly shutdown too, so in-flight streams never hang. |
| 118 | func TestDisconnectedClosesWithServeLoop(t *testing.T) { |
| 119 | client := startFakeClient(t, nil, nil) |
| 120 | select { |
| 121 | case <-client.Disconnected(): |
| 122 | t.Fatal("Disconnected closed on a live client") |
| 123 | default: |
| 124 | } |
| 125 | if err := client.Close(); err != nil { |
| 126 | t.Fatalf("Close: %v", err) |
| 127 | } |
| 128 | select { |
| 129 | case <-client.Disconnected(): |
| 130 | case <-time.After(5 * time.Second): |
| 131 | t.Fatal("Disconnected did not close after shutdown") |
| 132 | } |
| 133 | } |
| 134 | |
| 135 | func TestProviderCatalogRoundTrip(t *testing.T) { |
| 136 | client := startFakeClient(t, nil, nil) |
| 137 | providers, err := client.ProviderCatalog(context.Background()) |
| 138 | if err != nil { |
| 139 | t.Fatalf("ProviderCatalog: %v", err) |
| 140 | } |
| 141 | if len(providers) != 0 { |
| 142 | t.Fatalf("providers = %v, want the fake's empty catalog", providers) |
| 143 | } |
| 144 | } |
| 145 | |
| 146 | func TestProviderStreamCancelBestEffort(t *testing.T) { |
| 147 | client := startFakeClient(t, nil, nil) |
| 148 | // The fake answers {"cancelled":true}; the call must simply not wedge. |
| 149 | client.ProviderStreamCancel("es_test") |
| 150 | } |
| 151 |