1package main
2
3import (
4 "bufio"
5 "encoding/json"
6 "io"
7 "log"
8 "os"
9 "reflect"
10 "testing"
11 "time"
12
13 "github.com/bnema/zut/packages/agent/ext"
14 "github.com/bnema/zut/packages/agent/extproto"
15)
16
17type bridgeExtensionHarness struct {
18 ext *ext.Extension
19 hostW *os.File
20 frames chan []byte
21}
22
23func newBridgeExtensionHarness(t *testing.T) *bridgeExtensionHarness {
24 t.Helper()
25 extIn, hostW, err := os.Pipe()
26 if err != nil {
27 t.Fatal(err)
28 }
29 hostR, extOut, err := os.Pipe()
30 if err != nil {
31 extIn.Close()
32 hostW.Close()
33 t.Fatal(err)
34 }
35
36 oldIn, oldOut := os.Stdin, os.Stdout
37 os.Stdin, os.Stdout = extIn, extOut
38 e := ext.New("mcp-test", "0.0.0-test")
39 os.Stdin, os.Stdout = oldIn, oldOut
40
41 h := &bridgeExtensionHarness{ext: e, hostW: hostW, frames: make(chan []byte, 32)}
42 go func() {
43 scanner := bufio.NewScanner(hostR)
44 for scanner.Scan() {
45 h.frames <- append([]byte(nil), scanner.Bytes()...)
46 }
47 close(h.frames)
48 }()
49 t.Cleanup(func() {
50 hostW.Close()
51 extIn.Close()
52 extOut.Close()
53 hostR.Close()
54 })
55 return h
56}
57
58func (h *bridgeExtensionHarness) next(t *testing.T) []byte {
59 t.Helper()
60 select {
61 case frame, ok := <-h.frames:
62 if !ok {
63 t.Fatal("extension output closed")
64 }
65 return frame
66 case <-time.After(2 * time.Second):
67 t.Fatal("timed out waiting for extension frame")
68 return nil
69 }
70}
71
72func (h *bridgeExtensionHarness) send(t *testing.T, frame any) {
73 t.Helper()
74 data, err := extproto.Encode(frame)
75 if err != nil {
76 t.Fatal(err)
77 }
78 if _, err := h.hostW.Write(data); err != nil {
79 t.Fatal(err)
80 }
81}
82
83func (h *bridgeExtensionHarness) startAndCollectTools(t *testing.T) []extproto.RegisterToolFromExt {
84 t.Helper()
85 go func() { _ = h.ext.Run() }()
86
87 var hello extproto.Frame
88 if err := json.Unmarshal(h.next(t), &hello); err != nil {
89 t.Fatal(err)
90 }
91 if hello.Type != "hello" {
92 t.Fatalf("first extension frame = %q, want hello", hello.Type)
93 }
94 h.send(t, extproto.HelloAckFromHost{Type: "hello_ack", ProtocolVersion: extproto.ProtocolVersion})
95
96 var tools []extproto.RegisterToolFromExt
97 for {
98 frame := h.next(t)
99 var header extproto.Frame
100 if err := json.Unmarshal(frame, &header); err != nil {
101 t.Fatal(err)
102 }
103 switch header.Type {
104 case "register_tool":
105 var tool extproto.RegisterToolFromExt
106 if err := json.Unmarshal(frame, &tool); err != nil {
107 t.Fatal(err)
108 }
109 tools = append(tools, tool)
110 case "ready":
111 return tools
112 }
113 }
114}
115
116func (h *bridgeExtensionHarness) callTool(t *testing.T, id, name, args string) extproto.ToolResultFromExt {
117 t.Helper()
118 h.send(t, extproto.ToolCallFromHost{
119 Type: "tool_call",
120 ID: id,
121 Name: name,
122 Args: json.RawMessage(args),
123 })
124 for {
125 frame := h.next(t)
126 var header extproto.Frame
127 if err := json.Unmarshal(frame, &header); err != nil {
128 t.Fatal(err)
129 }
130 if header.Type != "tool_result" || header.ID != id {
131 continue
132 }
133 var result extproto.ToolResultFromExt
134 if err := json.Unmarshal(frame, &result); err != nil {
135 t.Fatal(err)
136 }
137 return result
138 }
139}
140
141func cachedDiscoveryFixture(t *testing.T, e *ext.Extension) (*bridge, toolCache) {
142 t.Helper()
143 b := newBridge(e, "/project", log.New(io.Discard, "", 0))
144 b.loadServers(Config{MCPServers: map[string]ServerConfig{
145 "alpha": {Transport: "stdio", Command: "alpha-server"},
146 "beta": {Transport: "stdio", Command: "beta-server"},
147 }})
148 return b, toolCache{Version: toolCacheVersion, Servers: map[string]cachedServer{
149 "alpha": {
150 Fingerprint: serverFingerprint(b.servers["alpha"].config),
151 Tools: []cachedTool{
152 {Name: "repo_lookup", Description: "Find repository metadata", Schema: json.RawMessage(`{"type":"object"}`)},
153 {Name: "issue_scan", Description: "Search repository issues", Schema: json.RawMessage(`{"type":"object"}`)},
154 {Name: "deploy", Description: "Ship a release", Schema: json.RawMessage(`{"type":"object"}`)},
155 },
156 },
157 "beta": {
158 Fingerprint: serverFingerprint(b.servers["beta"].config),
159 Tools: []cachedTool{
160 {Name: "repository_stats", Description: "Count projects", Schema: json.RawMessage(`{"type":"object"}`)},
161 {Name: "code_search", Description: "Find source snippets", Schema: json.RawMessage(`{"type":"object"}`)},
162 },
163 },
164 }}
165}
166
167func TestCachedMCPToolsAreDeferredWithOneActiveLoader(t *testing.T) {
168 h := newBridgeExtensionHarness(t)
169 b, cache := cachedDiscoveryFixture(t, h.ext)
170 b.registerToolSearch()
171 if got := b.registerCachedTools(cache); got != 5 {
172 t.Fatalf("registerCachedTools() = %d, want 5", got)
173 }
174
175 registrations := h.startAndCollectTools(t)
176 var active []string
177 deferred := map[string]bool{}
178 for _, registration := range registrations {
179 if registration.Deferred {
180 deferred[registration.Name] = true
181 } else {
182 active = append(active, registration.Name)
183 }
184 }
185 wantDeferred := []string{
186 "mcp__alpha__deploy",
187 "mcp__alpha__issue_scan",
188 "mcp__alpha__repo_lookup",
189 "mcp__beta__code_search",
190 "mcp__beta__repository_stats",
191 }
192 for _, name := range wantDeferred {
193 if !deferred[name] {
194 t.Errorf("cached MCP tool %q was not registered as deferred", name)
195 }
196 }
197 if !reflect.DeepEqual(active, []string{mcpSearchToolName}) {
198 t.Fatalf("active tools = %v, want only %s", active, mcpSearchToolName)
199 }
200}
201
202func TestMCPToolSearchMatchesNameAndDescriptionDeterministically(t *testing.T) {
203 h := newBridgeExtensionHarness(t)
204 b, cache := cachedDiscoveryFixture(t, h.ext)
205 b.registerToolSearch()
206 b.registerCachedTools(cache)
207 registrations := h.startAndCollectTools(t)
208
209 loader := ""
210 for _, registration := range registrations {
211 if !registration.Deferred {
212 loader = registration.Name
213 }
214 }
215 if loader == "" {
216 t.Fatal("no active MCP tool loader registered")
217 }
218
219 byName := h.callTool(t, "name", loader, `{"query":"repository_stats"}`)
220 if byName.IsError || !reflect.DeepEqual(byName.ActivateTools, []string{"mcp__beta__repository_stats"}) {
221 t.Fatalf("name search result: error=%v activate_tools=%v", byName.IsError, byName.ActivateTools)
222 }
223
224 byDescription := h.callTool(t, "description", loader, `{"query":"source snippets"}`)
225 if byDescription.IsError || !reflect.DeepEqual(byDescription.ActivateTools, []string{"mcp__beta__code_search"}) {
226 t.Fatalf("description search result: error=%v activate_tools=%v", byDescription.IsError, byDescription.ActivateTools)
227 }
228
229 wantLimited := []string{"mcp__beta__repository_stats", "mcp__alpha__issue_scan"}
230 first := h.callTool(t, "limited-1", loader, `{"query":"repository","limit":2}`)
231 second := h.callTool(t, "limited-2", loader, `{"query":"repository","limit":2}`)
232 if first.IsError || !reflect.DeepEqual(first.ActivateTools, wantLimited) {
233 t.Fatalf("limited search activate_tools = %v, want %v", first.ActivateTools, wantLimited)
234 }
235 if !reflect.DeepEqual(second.ActivateTools, first.ActivateTools) {
236 t.Fatalf("repeated search was not deterministic: first=%v second=%v", first.ActivateTools, second.ActivateTools)
237 }
238}
239
240func TestMCPToolSearchEmptyAndNoMatchAreSafe(t *testing.T) {
241 h := newBridgeExtensionHarness(t)
242 b, cache := cachedDiscoveryFixture(t, h.ext)
243 b.registerToolSearch()
244 b.registerCachedTools(cache)
245 registrations := h.startAndCollectTools(t)
246
247 loader := ""
248 for _, registration := range registrations {
249 if !registration.Deferred {
250 loader = registration.Name
251 }
252 }
253 if loader == "" {
254 t.Fatal("no active MCP tool loader registered")
255 }
256
257 empty := h.callTool(t, "empty", loader, `{}`)
258 if len(empty.ActivateTools) != 0 {
259 t.Errorf("empty search activated tools: %v", empty.ActivateTools)
260 }
261
262 noMatch := h.callTool(t, "no-match", loader, `{"query":"definitely-not-an-mcp-tool"}`)
263 if noMatch.IsError {
264 t.Error("no-match search returned an error")
265 }
266 if len(noMatch.ActivateTools) != 0 {
267 t.Errorf("no-match search activated tools: %v", noMatch.ActivateTools)
268 }
269}