1package tools
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "io"
8 "net"
9 "net/http"
10 "net/http/httptest"
11 "net/url"
12 "strings"
13 "sync/atomic"
14 "testing"
15 "time"
16
17 "github.com/bnema/zut/packages/provider"
18)
19
20type webSearchRoundTripper func(*http.Request) (*http.Response, error)
21
22func (f webSearchRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
23 return f(r)
24}
25
26func webSearchDirectTestTransport(t *testing.T) http.RoundTripper {
27 t.Helper()
28 transport := &http.Transport{Proxy: nil}
29 t.Cleanup(transport.CloseIdleConnections)
30 return transport
31}
32
33func webSearchText(t *testing.T, resultText any) string {
34 t.Helper()
35 result, ok := resultText.(provider.TextBlock)
36 if !ok {
37 t.Fatalf("content block = %T, want provider.TextBlock", resultText)
38 }
39 return result.Text
40}
41
42func executeWebSearch(t *testing.T, tool *WebSearchTool, ctx context.Context, args any) (string, bool, any) {
43 t.Helper()
44 raw, err := json.Marshal(args)
45 if err != nil {
46 t.Fatal(err)
47 }
48 return executeWebSearchRaw(t, tool, ctx, raw)
49}
50
51func executeWebSearchRaw(t *testing.T, tool *WebSearchTool, ctx context.Context, raw json.RawMessage) (string, bool, any) {
52 t.Helper()
53 result, err := tool.Execute(ctx, raw, nil)
54 if err != nil {
55 t.Fatal(err)
56 }
57 if len(result.Content) != 1 {
58 t.Fatalf("content = %#v", result.Content)
59 }
60 return webSearchText(t, result.Content[0]), result.IsError, result.Details
61}
62
63func webSearchHTMLResponse(request *http.Request, body string) *http.Response {
64 return &http.Response{
65 StatusCode: http.StatusOK,
66 Status: "200 OK",
67 Header: http.Header{"Content-Type": []string{"text/html; charset=utf-8"}},
68 Body: io.NopCloser(strings.NewReader(body)),
69 Request: request,
70 }
71}
72
73const webSearchFixture = `<!doctype html><html><body>
74<div class="result"><a class="result__a" href="/l/?uddg=https%3A%2F%2Fexample.com%2Fone">One <b>result</b></a><a class="result__snippet">First <em>snippet</em>.</a></div>
75<div class="result"><h2><a href="https://EXAMPLE.com/one">Duplicate</a></h2><span class="result__snippet">Ignored duplicate.</span></div>
76<div class="result"><a class="result__a" href="https://example.org/two#fragment">Two</a><div class="result__snippet">Second snippet.</div></div>
77</body></html>`
78
79func TestWebSearchRequestAndResultContract(t *testing.T) {
80 var calls atomic.Int64
81 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
82 calls.Add(1)
83 if r.Method != http.MethodGet {
84 t.Errorf("method = %s, want GET", r.Method)
85 }
86 if r.URL.Path != "/html/" {
87 t.Errorf("path = %q, want /html/", r.URL.Path)
88 }
89 if got, want := r.URL.RawQuery, "q=Go+JSON-RPC+%E2%9C%93"; got != want {
90 t.Errorf("query = %q, want %q", got, want)
91 }
92 if got := r.Header.Get("User-Agent"); got != webSearchUserAgent {
93 t.Errorf("User-Agent = %q", got)
94 }
95 if got := r.Header.Get("Cookie"); got != "" {
96 t.Errorf("Cookie = %q, want none", got)
97 }
98 w.Header().Set("Set-Cookie", "session=should-not-be-reused")
99 w.Header().Set("Content-Type", "text/html; charset=utf-8")
100 _, _ = w.Write([]byte(webSearchFixture))
101 }))
102 defer server.Close()
103
104 tool := &WebSearchTool{transport: webSearchDirectTestTransport(t), endpointOverride: server.URL + "/html/?ignored=value"}
105 text, isError, details := executeWebSearch(t, tool, context.Background(), map[string]any{
106 "query": " Go JSON-RPC ✓ ",
107 "max_results": 2,
108 })
109 if isError {
110 t.Fatalf("unexpected error: %s", text)
111 }
112 if got := calls.Load(); got != 1 {
113 t.Fatalf("calls = %d, want 1", got)
114 }
115 if !strings.HasPrefix(text, webSearchDisclosure+"\n\n[1] One result\n https://example.com/one\n First snippet.") {
116 t.Fatalf("unexpected result text:\n%s", text)
117 }
118 if strings.Contains(text, "#fragment") || !strings.Contains(text, "https://example.org/two") {
119 t.Fatalf("result URLs were not normalized:\n%s", text)
120 }
121 metadata, ok := details.(map[string]any)
122 if !ok || metadata["backend"] != "DuckDuckGo HTML" || metadata["query"] != "Go JSON-RPC ✓" {
123 t.Fatalf("details = %#v", details)
124 }
125 if count, ok := metadata["result_count"].(int); !ok || count != 2 {
126 t.Fatalf("result_count = %#v", metadata["result_count"])
127 }
128}
129
130func TestWebSearchEndpointIsFixed(t *testing.T) {
131 var gotURL string
132 tool := NewWebSearchTool()
133 tool.transport = webSearchRoundTripper(func(request *http.Request) (*http.Response, error) {
134 gotURL = request.URL.String()
135 return webSearchHTMLResponse(request, webSearchFixture), nil
136 })
137 text, isError, _ := executeWebSearch(t, tool, context.Background(), map[string]any{"query": "fixed"})
138 if isError {
139 t.Fatalf("unexpected error: %s", text)
140 }
141 if gotURL != webSearchEndpoint+"?q=fixed" {
142 t.Fatalf("request URL = %q, want %q", gotURL, webSearchEndpoint+"?q=fixed")
143 }
144}
145
146func TestWebSearchUsesConfiguredProxyTransport(t *testing.T) {
147 var proxyCalls atomic.Int64
148 proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
149 proxyCalls.Add(1)
150 if r.URL.Scheme != "http" || r.URL.Host != "web-search.invalid" {
151 t.Errorf("proxy request URL = %q, want absolute web-search.invalid URL", r.URL)
152 }
153 if got, want := r.URL.Path, "/html/"; got != want {
154 t.Errorf("proxy request path = %q, want %q", got, want)
155 }
156 if got, want := r.URL.RawQuery, "q=environment"; got != want {
157 t.Errorf("proxy request query = %q, want %q", got, want)
158 }
159 w.Header().Set("Content-Type", "text/html")
160 _, _ = w.Write([]byte(webSearchFixture))
161 }))
162 defer proxy.Close()
163
164 proxyURL, err := url.Parse(proxy.URL)
165 if err != nil {
166 t.Fatal(err)
167 }
168 transport := &http.Transport{Proxy: http.ProxyURL(proxyURL)}
169 t.Cleanup(transport.CloseIdleConnections)
170 tool := &WebSearchTool{
171 endpointOverride: "http://web-search.invalid/html/",
172 transport: transport,
173 }
174 text, isError, _ := executeWebSearch(t, tool, context.Background(), map[string]any{"query": "environment"})
175 if isError {
176 t.Fatalf("unexpected error through configured proxy: %s", text)
177 }
178 if got := proxyCalls.Load(); got != 1 {
179 t.Fatalf("proxy calls = %d, want 1", got)
180 }
181}
182
183func TestWebSearchExecuteClosesOwnedTransportConnections(t *testing.T) {
184 const calls = 3
185 closed := make(chan struct{}, calls)
186 server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
187 w.Header().Set("Content-Type", "text/html")
188 _, _ = w.Write([]byte(webSearchFixture))
189 }))
190 server.Config.ConnState = func(_ net.Conn, state http.ConnState) {
191 if state == http.StateClosed {
192 closed <- struct{}{}
193 }
194 }
195 server.Start()
196 defer server.Close()
197
198 // ProxyFromEnvironment always bypasses localhost. This keeps this regression
199 // test local even after the proxy contract test initializes its cached proxy.
200 endpoint := strings.Replace(server.URL, "127.0.0.1", "localhost", 1)
201 tool := &WebSearchTool{endpointOverride: endpoint}
202 for i := 0; i < calls; i++ {
203 text, isError, _ := executeWebSearch(t, tool, context.Background(), map[string]any{"query": "lifecycle"})
204 if isError {
205 t.Fatalf("call %d returned an error: %s", i+1, text)
206 }
207 }
208 for i := 0; i < calls; i++ {
209 select {
210 case <-closed:
211 case <-time.After(time.Second):
212 t.Fatalf("call %d left its transport connection open", i+1)
213 }
214 }
215}
216
217func TestWebSearchSanitizesResultsAndDetails(t *testing.T) {
218 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
219 w.Header().Set("Content-Type", "text/html")
220 _, _ = w.Write([]byte("<div class=\"result\"><a class=\"result__a\" href=\"https://example.com\">\x1b[31mRed\x1b[0m\x00 title</a><span class=\"result__snippet\">\x1b]0;terminal title\x07Clean\x01 snippet</span></div>"))
221 }))
222 defer server.Close()
223
224 tool := &WebSearchTool{transport: webSearchDirectTestTransport(t), endpointOverride: server.URL}
225 text, isError, details := executeWebSearch(t, tool, context.Background(), map[string]any{"query": "safe\x1b[31m query\x00"})
226 if isError {
227 t.Fatalf("unexpected error: %s", text)
228 }
229 if want := webSearchDisclosure + "\n\n[1] Red title\n https://example.com\n Clean snippet\n"; text != want {
230 t.Fatalf("sanitized text = %q, want %q", text, want)
231 }
232 metadata, ok := details.(map[string]any)
233 if !ok {
234 t.Fatalf("details = %#v", details)
235 }
236 if metadata["query"] != "safe query" {
237 t.Fatalf("sanitized query details = %#v", metadata["query"])
238 }
239 results, ok := metadata["results"].([]webSearchResult)
240 if !ok || len(results) != 1 {
241 t.Fatalf("sanitized result details = %#v", metadata["results"])
242 }
243 if results[0].Title != "Red title" || results[0].Snippet != "Clean snippet" {
244 t.Fatalf("sanitized result = %#v", results[0])
245 }
246 if strings.ContainsAny(text+results[0].Title+results[0].Snippet+metadata["query"].(string), "\x00\x01\x1b") {
247 t.Fatalf("control character remained in model-visible output")
248 }
249}
250
251func TestWebSearchRejectsInvalidArgumentsWithoutRequest(t *testing.T) {
252 tool := &WebSearchTool{transport: webSearchRoundTripper(func(*http.Request) (*http.Response, error) {
253 t.Fatal("unexpected request")
254 return nil, nil
255 })}
256 cases := []struct {
257 name string
258 args any
259 want string
260 }{
261 {name: "missing query", args: map[string]any{}, want: "web search: invalid query"},
262 {name: "empty query", args: map[string]any{"query": " \t "}, want: "web search: invalid query"},
263 {name: "long query", args: map[string]any{"query": strings.Repeat("界", webSearchMaxQueryRunes+1)}, want: "web search: invalid query"},
264 {name: "zero results", args: map[string]any{"query": "x", "max_results": 0}, want: "web search: invalid max_results"},
265 {name: "too many results", args: map[string]any{"query": "x", "max_results": 11}, want: "web search: invalid max_results"},
266 {name: "noninteger results", args: map[string]any{"query": "x", "max_results": "2"}, want: "web search: invalid max_results"},
267 {name: "unknown key", args: map[string]any{"query": "x", "unexpected": true}, want: "web search: invalid query"},
268 }
269
270 for _, tc := range cases {
271 t.Run(tc.name, func(t *testing.T) {
272 text, isError, _ := executeWebSearch(t, tool, context.Background(), tc.args)
273 if !isError || text != tc.want {
274 t.Fatalf("result = (%q, %v), want (%q, true)", text, isError, tc.want)
275 }
276 })
277 }
278
279 largeRaw := json.RawMessage(strings.Repeat(" ", webSearchMaxInputBytes) + `{"query":"x"}`)
280 text, isError, _ := executeWebSearchRaw(t, tool, context.Background(), largeRaw)
281 if !isError || text != "web search: invalid query" {
282 t.Fatalf("large raw input result = (%q, %v)", text, isError)
283 }
284}
285
286func TestWebSearchBackendFailuresAreSanitized(t *testing.T) {
287 cases := []struct {
288 name string
289 status int
290 content string
291 body string
292 want string
293 }{
294 {name: "rate limited", status: http.StatusTooManyRequests, content: "text/html", body: "private backend diagnostics", want: "web search: DuckDuckGo returned an unavailable response"},
295 {name: "wrong content type", status: http.StatusOK, content: "application/json", body: `{}`, want: "web search: DuckDuckGo returned an unavailable response"},
296 {name: "challenge", status: http.StatusOK, content: "text/html", body: "<html>unusual traffic captcha</html>", want: "web search: DuckDuckGo returned an unavailable response"},
297 {name: "unparseable", status: http.StatusOK, content: "text/html", body: "<html><body>ordinary page</body></html>", want: "web search: DuckDuckGo returned an unparseable response"},
298 {name: "no usable", status: http.StatusOK, content: "text/html", body: `<div class="result"><a class="result__a" href="javascript:alert(1)">bad</a></div>`, want: "web search: no usable HTTP(S) results"},
299 {name: "malformed redirect", status: http.StatusOK, content: "text/html", body: `<div class="result"><a class="result__a" href="/l/?uddg=%">bad</a></div>`, want: "web search: no usable HTTP(S) results"},
300 }
301 for _, tc := range cases {
302 t.Run(tc.name, func(t *testing.T) {
303 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
304 w.Header().Set("Content-Type", tc.content)
305 w.WriteHeader(tc.status)
306 _, _ = w.Write([]byte(tc.body))
307 }))
308 defer server.Close()
309 tool := &WebSearchTool{transport: webSearchDirectTestTransport(t), endpointOverride: server.URL}
310 text, isError, _ := executeWebSearch(t, tool, context.Background(), map[string]any{"query": "test"})
311 if !isError || text != tc.want {
312 t.Fatalf("result = (%q, %v), want (%q, true)", text, isError, tc.want)
313 }
314 })
315 }
316}
317
318func TestWebSearchDoesNotUseCookies(t *testing.T) {
319 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
320 if got := r.Header.Get("Cookie"); got != "" {
321 t.Errorf("Cookie = %q, want none", got)
322 }
323 w.Header().Set("Set-Cookie", "session=private")
324 w.Header().Set("Content-Type", "text/html")
325 _, _ = w.Write([]byte(webSearchFixture))
326 }))
327 defer server.Close()
328
329 tool := &WebSearchTool{transport: webSearchDirectTestTransport(t), endpointOverride: server.URL}
330 for i := 0; i < 2; i++ {
331 text, isError, _ := executeWebSearch(t, tool, context.Background(), map[string]any{"query": "cookie"})
332 if isError {
333 t.Fatalf("unexpected error: %s", text)
334 }
335 }
336 if client := tool.client(); client.Jar != nil {
337 t.Fatalf("client jar = %v, want nil", client.Jar)
338 }
339}
340
341func TestWebSearchRejectsRedirectAndOversizedResponse(t *testing.T) {
342 var followed atomic.Bool
343 redirect := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
344 if r.URL.Path == "/next" {
345 followed.Store(true)
346 }
347 http.Redirect(w, r, "/next", http.StatusFound)
348 }))
349 defer redirect.Close()
350 tool := &WebSearchTool{transport: webSearchDirectTestTransport(t), endpointOverride: redirect.URL}
351 text, isError, _ := executeWebSearch(t, tool, context.Background(), map[string]any{"query": "test"})
352 if !isError || text != "web search: DuckDuckGo request failed" {
353 t.Fatalf("redirect result = (%q, %v)", text, isError)
354 }
355 if followed.Load() {
356 t.Fatal("redirect was followed")
357 }
358
359 oversized := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
360 w.Header().Set("Content-Type", "text/html")
361 _, _ = w.Write([]byte(strings.Repeat("x", webSearchMaxBodyBytes+1)))
362 }))
363 defer oversized.Close()
364 tool = &WebSearchTool{transport: webSearchDirectTestTransport(t), endpointOverride: oversized.URL}
365 text, isError, _ = executeWebSearch(t, tool, context.Background(), map[string]any{"query": "test"})
366 if !isError || text != "web search: DuckDuckGo returned an unavailable response" {
367 t.Fatalf("oversized result = (%q, %v)", text, isError)
368 }
369}
370
371func TestWebSearchCancellationAndTransportFailure(t *testing.T) {
372 cancelled, cancel := context.WithCancel(context.Background())
373 cancel()
374 tool := &WebSearchTool{}
375 text, isError, _ := executeWebSearch(t, tool, cancelled, map[string]any{"query": "test"})
376 if !isError || text != "web search: request cancelled" {
377 t.Fatalf("cancelled result = (%q, %v)", text, isError)
378 }
379
380 deadline, deadlineCancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second))
381 defer deadlineCancel()
382 text, isError, _ = executeWebSearch(t, tool, deadline, map[string]any{"query": "test"})
383 if !isError || text != "web search: timed out" {
384 t.Fatalf("deadline result = (%q, %v)", text, isError)
385 }
386
387 tool = &WebSearchTool{transport: webSearchRoundTripper(func(*http.Request) (*http.Response, error) {
388 return nil, errors.New("dial private proxy failed")
389 })}
390 text, isError, _ = executeWebSearch(t, tool, context.Background(), map[string]any{"query": "test"})
391 if !isError || text != "web search: DuckDuckGo request failed" || strings.Contains(text, "proxy") {
392 t.Fatalf("transport result = (%q, %v)", text, isError)
393 }
394}
395
396func TestWebSearchActiveCancellationAndRequestDeadline(t *testing.T) {
397 type asyncResult struct {
398 text string
399 isError bool
400 }
401 call := func(tool *WebSearchTool, ctx context.Context) <-chan asyncResult {
402 done := make(chan asyncResult, 1)
403 go func() {
404 result, err := tool.Execute(ctx, json.RawMessage(`{"query":"active"}`), nil)
405 if err != nil || len(result.Content) != 1 {
406 done <- asyncResult{text: "unexpected tool error", isError: true}
407 return
408 }
409 block, ok := result.Content[0].(provider.TextBlock)
410 if !ok {
411 done <- asyncResult{text: "unexpected content", isError: true}
412 return
413 }
414 done <- asyncResult{text: block.Text, isError: result.IsError}
415 }()
416 return done
417 }
418 wait := func(t *testing.T, done <-chan asyncResult) asyncResult {
419 t.Helper()
420 select {
421 case result := <-done:
422 return result
423 case <-time.After(time.Second):
424 t.Fatal("web search did not finish after context termination")
425 return asyncResult{}
426 }
427 }
428
429 started := make(chan struct{})
430 tool := &WebSearchTool{transport: webSearchRoundTripper(func(request *http.Request) (*http.Response, error) {
431 close(started)
432 <-request.Context().Done()
433 return nil, request.Context().Err()
434 })}
435 ctx, cancel := context.WithCancel(context.Background())
436 done := call(tool, ctx)
437 select {
438 case <-started:
439 case <-time.After(time.Second):
440 t.Fatal("web search request did not start")
441 }
442 cancel()
443 result := wait(t, done)
444 if !result.isError || result.text != "web search: request cancelled" {
445 t.Fatalf("active cancellation result = (%q, %v)", result.text, result.isError)
446 }
447
448 started = make(chan struct{})
449 tool = &WebSearchTool{transport: webSearchRoundTripper(func(request *http.Request) (*http.Response, error) {
450 close(started)
451 <-request.Context().Done()
452 return nil, request.Context().Err()
453 })}
454 deadline, deadlineCancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
455 defer deadlineCancel()
456 done = call(tool, deadline)
457 select {
458 case <-started:
459 case <-time.After(time.Second):
460 t.Fatal("deadline request did not start")
461 }
462 result = wait(t, done)
463 if !result.isError || result.text != "web search: timed out" {
464 t.Fatalf("request deadline result = (%q, %v)", result.text, result.isError)
465 }
466}
467
468func TestWebSearchPreviewIsLocalAndOutputIsBounded(t *testing.T) {
469 called := false
470 tool := &WebSearchTool{transport: webSearchRoundTripper(func(*http.Request) (*http.Response, error) {
471 called = true
472 return nil, errors.New("unexpected request")
473 })}
474 raw := json.RawMessage(`{"query":"current Go release"}`)
475 preview, err := tool.Preview(context.Background(), raw)
476 if err != nil {
477 t.Fatal(err)
478 }
479 if preview.IsError || called {
480 t.Fatalf("preview error=%v called=%v", preview.IsError, called)
481 }
482 if got := webSearchText(t, preview.Content[0]); !strings.Contains(got, "DuckDuckGo HTML") || !strings.Contains(got, "current Go release") {
483 t.Fatalf("preview = %q", got)
484 }
485
486 results := make([]webSearchResult, webSearchMaxResults)
487 for i := range results {
488 results[i] = webSearchResult{
489 Title: strings.Repeat("界", webSearchMaxTitleRunes+20),
490 URL: "https://example.com/" + strings.Repeat("a", webSearchMaxURLBytes-32),
491 Snippet: strings.Repeat("界", webSearchMaxSnippetRunes+20),
492 }
493 }
494 output, visible, truncated := formatWebSearchResults(results)
495 if len(output) > webSearchMaxOutputBytes || visible == 0 || !truncated {
496 t.Fatalf("output bounds = len %d, visible %d, truncated %v", len(output), visible, truncated)
497 }
498 if got := boundWebSearchText(strings.Repeat("界", webSearchMaxTitleRunes+1), webSearchMaxTitleRunes); len([]rune(got)) != webSearchMaxTitleRunes || !strings.HasSuffix(got, "…") {
499 t.Fatalf("rune bound = %q", got)
500 }
501}