| 1 | package websearch |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "strings" |
| 8 | "sync" |
| 9 | "testing" |
| 10 | "unicode/utf8" |
| 11 | |
| 12 | "reasonix/internal/provider" |
| 13 | ) |
| 14 | |
| 15 | type fakeProvider struct { |
| 16 | stream func(context.Context, provider.Request) (<-chan provider.Chunk, error) |
| 17 | } |
| 18 | |
| 19 | func (fakeProvider) Name() string { return "search" } |
| 20 | func (p fakeProvider) Stream(ctx context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 21 | return p.stream(ctx, req) |
| 22 | } |
| 23 | |
| 24 | func chunks(values ...provider.Chunk) <-chan provider.Chunk { |
| 25 | ch := make(chan provider.Chunk, len(values)) |
| 26 | for _, value := range values { |
| 27 | ch <- value |
| 28 | } |
| 29 | close(ch) |
| 30 | return ch |
| 31 | } |
| 32 | |
| 33 | func TestSearchIsolatedRequestsAndBoundedResults(t *testing.T) { |
| 34 | var mu sync.Mutex |
| 35 | var requests []provider.Request |
| 36 | var usages []*provider.Usage |
| 37 | tool := &Tool{Factory: func() (provider.Provider, error) { |
| 38 | return fakeProvider{stream: func(ctx context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 39 | mu.Lock() |
| 40 | requests = append(requests, req) |
| 41 | mu.Unlock() |
| 42 | return chunks( |
| 43 | provider.Chunk{Type: provider.ChunkReasoning, Text: "PRIVATE REASONING"}, |
| 44 | provider.Chunk{Type: provider.ChunkResponsesItem, ResponsesItem: json.RawMessage(`{"secret":"OPAQUE REPLAY"}`)}, |
| 45 | provider.Chunk{Type: provider.ChunkServerSearch, ServerSearch: &provider.ServerSearchCall{ID: "one", Raw: json.RawMessage(`[]`), Results: []provider.ServerSearchHit{{Title: "Source", URL: "https://example.com"}, {Title: "Duplicate", URL: "https://example.com"}, {URL: "javascript:alert(1)"}}}}, |
| 46 | provider.Chunk{Type: provider.ChunkText, Text: strings.Repeat("搜索", 4000)}, |
| 47 | provider.Chunk{Type: provider.ChunkUsage, Usage: &provider.Usage{PromptTokens: 10, CompletionTokens: 20}}, |
| 48 | provider.Chunk{Type: provider.ChunkDone}, |
| 49 | ), nil |
| 50 | }}, nil |
| 51 | }, ReportUsage: func(u *provider.Usage) { mu.Lock(); usages = append(usages, u); mu.Unlock() }} |
| 52 | var wg sync.WaitGroup |
| 53 | for _, q := range []string{"first", "second"} { |
| 54 | wg.Go(func() { |
| 55 | output, err := tool.Execute(context.Background(), json.RawMessage(`{"query":"`+q+`"}`)) |
| 56 | if err != nil { |
| 57 | t.Error(err) |
| 58 | return |
| 59 | } |
| 60 | var result Result |
| 61 | if json.Unmarshal([]byte(output), &result) != nil || len(result.Sources) != 1 || !utf8.ValidString(result.Summary) || len(result.Summary) > maxSummaryBytes { |
| 62 | t.Errorf("bad result: %s", output) |
| 63 | } |
| 64 | if strings.Contains(output, "PRIVATE") || strings.Contains(output, "OPAQUE") { |
| 65 | t.Error("reasoning or replay escaped search") |
| 66 | } |
| 67 | }) |
| 68 | } |
| 69 | wg.Wait() |
| 70 | if len(requests) != 2 || len(usages) != 2 { |
| 71 | t.Fatalf("requests=%d usage=%d", len(requests), len(usages)) |
| 72 | } |
| 73 | for _, req := range requests { |
| 74 | if len(req.Messages) != 1 || req.Messages[0].Role != provider.RoleUser || len(req.Tools) != 0 || req.MaxTokens != maxOutputTokens { |
| 75 | t.Fatalf("unexpected request: %+v", req) |
| 76 | } |
| 77 | } |
| 78 | } |
| 79 | |
| 80 | func TestSearchRejectsIncompleteOrInventedResults(t *testing.T) { |
| 81 | for _, tc := range []struct { |
| 82 | name string |
| 83 | values []provider.Chunk |
| 84 | }{ |
| 85 | {"native error", []provider.Chunk{{Type: provider.ChunkServerSearch, ServerSearch: &provider.ServerSearchCall{Raw: json.RawMessage(`{"type":"web_search_tool_result_error","error_code":"unavailable"}`)}}, {Type: provider.ChunkDone}}}, |
| 86 | {"plain prose", []provider.Chunk{{Type: provider.ChunkText, Text: "I searched"}, {Type: provider.ChunkDone}}}, |
| 87 | {"interrupted", []provider.Chunk{{Type: provider.ChunkServerSearch, ServerSearch: &provider.ServerSearchCall{Raw: json.RawMessage(`[]`)}}}}, |
| 88 | {"start only", []provider.Chunk{{Type: provider.ChunkServerSearch, ServerSearch: &provider.ServerSearchCall{ID: "start"}}, {Type: provider.ChunkDone}}}, |
| 89 | {"client tool", []provider.Chunk{{Type: provider.ChunkToolCall}}}, |
| 90 | {"provider failure", []provider.Chunk{{Type: provider.ChunkError, Err: errors.New("upstream failed")}}}, |
| 91 | } { |
| 92 | t.Run(tc.name, func(t *testing.T) { |
| 93 | tool := &Tool{Factory: func() (provider.Provider, error) { |
| 94 | return fakeProvider{stream: func(context.Context, provider.Request) (<-chan provider.Chunk, error) { |
| 95 | return chunks(tc.values...), nil |
| 96 | }}, nil |
| 97 | }} |
| 98 | if _, err := tool.Execute(context.Background(), json.RawMessage(`{"query":"test"}`)); err == nil { |
| 99 | t.Fatal("expected search failure") |
| 100 | } |
| 101 | }) |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | func TestSearchCancellationAndValidation(t *testing.T) { |
| 106 | started := make(chan struct{}) |
| 107 | cancelled := make(chan struct{}) |
| 108 | tool := &Tool{Factory: func() (provider.Provider, error) { |
| 109 | return fakeProvider{stream: func(ctx context.Context, _ provider.Request) (<-chan provider.Chunk, error) { |
| 110 | close(started) |
| 111 | ch := make(chan provider.Chunk) |
| 112 | go func() { <-ctx.Done(); close(ch); close(cancelled) }() |
| 113 | return ch, nil |
| 114 | }}, nil |
| 115 | }} |
| 116 | for _, args := range []string{`{`, `{}`, `{"query":" "}`} { |
| 117 | if _, err := tool.Execute(context.Background(), json.RawMessage(args)); err == nil { |
| 118 | t.Fatal("invalid query accepted") |
| 119 | } |
| 120 | } |
| 121 | ctx, cancel := context.WithCancel(context.Background()) |
| 122 | finished := make(chan error, 1) |
| 123 | go func() { _, err := tool.Execute(ctx, json.RawMessage(`{"query":"test"}`)); finished <- err }() |
| 124 | <-started |
| 125 | cancel() |
| 126 | if err := <-finished; !errors.Is(err, context.Canceled) { |
| 127 | t.Fatalf("cancel error = %v", err) |
| 128 | } |
| 129 | <-cancelled |
| 130 | } |
| 131 | |
| 132 | func TestSearchEncodedResultBound(t *testing.T) { |
| 133 | result := Result{Summary: strings.Repeat("\x00", maxSummaryBytes)} |
| 134 | for range maxSources { |
| 135 | result.Sources = append(result.Sources, provider.ServerSearchHit{Title: strings.Repeat("\x00", maxSourceBytes), URL: "https://example.com/" + strings.Repeat("a", 1900)}) |
| 136 | } |
| 137 | output, err := encodeResult(result) |
| 138 | if err != nil || len(output) > 24000 || !json.Valid([]byte(output)) { |
| 139 | t.Fatalf("invalid bounded output: size=%d err=%v", len(output), err) |
| 140 | } |
| 141 | var got Result |
| 142 | if err := json.Unmarshal([]byte(output), &got); err != nil || !got.Truncated || len(got.Sources) == 0 { |
| 143 | t.Fatalf("lost sources or truncation marker: %+v", got) |
| 144 | } |
| 145 | } |
| 146 | |
| 147 | func TestSearchCompletedWithoutStructuredSources(t *testing.T) { |
| 148 | calls := 0 |
| 149 | status := "" |
| 150 | tool := &Tool{ReportSourcesStatus: func(value string) { status = value }, Factory: func() (provider.Provider, error) { |
| 151 | return fakeProvider{stream: func(context.Context, provider.Request) (<-chan provider.Chunk, error) { |
| 152 | calls++ |
| 153 | return chunks(provider.Chunk{Type: provider.ChunkServerSearch, ServerSearch: &provider.ServerSearchCall{ID: "s", Raw: json.RawMessage(`{"type":"web_search_call","status":"completed"}`)}}, provider.Chunk{Type: provider.ChunkText, Text: "Summary mentioning https://unverified.invalid is still only prose."}, provider.Chunk{Type: provider.ChunkDone}), nil |
| 154 | }}, nil |
| 155 | }} |
| 156 | output, err := tool.Execute(context.Background(), json.RawMessage(`{"query":"test"}`)) |
| 157 | if err != nil { |
| 158 | t.Fatal(err) |
| 159 | } |
| 160 | var result Result |
| 161 | if err := json.Unmarshal([]byte(output), &result); err != nil { |
| 162 | t.Fatal(err) |
| 163 | } |
| 164 | if calls != 1 || status != provider.SourcesNotProvided || result.SourcesStatus != status || len(result.Sources) != 0 || !strings.Contains(result.Summary, "Summary") { |
| 165 | t.Fatalf("calls=%d result=%+v status=%s", calls, result, status) |
| 166 | } |
| 167 | } |
| 168 |