feat(grpc): enhance MCP handler with gRPC authorization support
- Introduced grpcAuthProvider to adapt gRPC AuthorizedInfo for use in MCP process calls, enabling better authorization handling. - Updated MCPListTools and MCPCallTool methods to include authorization information in tool calls, improving security and context awareness. - Refactored fetch functions in webfetch to support Brightdata API endpoint configuration, enhancing flexibility in fetching HTML content. - Improved fetchHTML and fetchRawHTML methods to prioritize Brightdata when configured, streamlining content retrieval processes. - Added isHTMLContent utility function to determine response content type, enhancing the robustness of content handling in fetch operations.
This commit is contained in:
parent
194faac9b7
commit
960f47c238
3 changed files with 356 additions and 10 deletions
|
|
@ -8,12 +8,37 @@ import (
|
||||||
"google.golang.org/grpc/status"
|
"google.golang.org/grpc/status"
|
||||||
|
|
||||||
goumcp "github.com/yaoapp/gou/mcp"
|
goumcp "github.com/yaoapp/gou/mcp"
|
||||||
|
"github.com/yaoapp/yao/grpc/auth"
|
||||||
"github.com/yaoapp/yao/grpc/pb"
|
"github.com/yaoapp/yao/grpc/pb"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Handler implements the MCP gRPC methods.
|
// Handler implements the MCP gRPC methods.
|
||||||
type Handler struct{}
|
type Handler struct{}
|
||||||
|
|
||||||
|
// grpcAuthProvider adapts gRPC AuthorizedInfo to the AuthorizedProvider
|
||||||
|
// interface expected by gou/mcp/process for propagating auth to process calls.
|
||||||
|
type grpcAuthProvider struct {
|
||||||
|
m map[string]interface{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *grpcAuthProvider) GetAuthorizedMap() map[string]interface{} { return p.m }
|
||||||
|
|
||||||
|
func authProviderFromCtx(ctx context.Context) *grpcAuthProvider {
|
||||||
|
info := auth.GetAuthorizedInfo(ctx)
|
||||||
|
if info == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return &grpcAuthProvider{m: map[string]interface{}{
|
||||||
|
"sub": info.Subject,
|
||||||
|
"client_id": info.ClientID,
|
||||||
|
"scope": info.Scope,
|
||||||
|
"session_id": info.SessionID,
|
||||||
|
"user_id": info.UserID,
|
||||||
|
"team_id": info.TeamID,
|
||||||
|
"tenant_id": info.TenantID,
|
||||||
|
}}
|
||||||
|
}
|
||||||
|
|
||||||
// MCPListTools lists all available MCP tools for a given session.
|
// MCPListTools lists all available MCP tools for a given session.
|
||||||
func (h *Handler) MCPListTools(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPListResponse, error) {
|
func (h *Handler) MCPListTools(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPListResponse, error) {
|
||||||
client, err := goumcp.Select(req.SessionId)
|
client, err := goumcp.Select(req.SessionId)
|
||||||
|
|
@ -48,7 +73,12 @@ func (h *Handler) MCPCallTool(ctx context.Context, req *pb.MCPCallRequest) (*pb.
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
resp, err := client.CallTool(ctx, req.Tool, args)
|
var extraArgs []interface{}
|
||||||
|
if ap := authProviderFromCtx(ctx); ap != nil {
|
||||||
|
extraArgs = append(extraArgs, ap)
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := client.CallTool(ctx, req.Tool, args, extraArgs...)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, status.Errorf(codes.Internal, "CallTool failed: %v", err)
|
return nil, status.Errorf(codes.Internal, "CallTool failed: %v", err)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -22,6 +22,8 @@ const (
|
||||||
browserUserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
|
browserUserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
var brightdataEndpoint = "https://api.brightdata.com/request"
|
||||||
|
|
||||||
type fetchResult struct {
|
type fetchResult struct {
|
||||||
Body []byte
|
Body []byte
|
||||||
StatusCode int
|
StatusCode int
|
||||||
|
|
@ -69,7 +71,7 @@ func brightdataFetch(targetURL, apiKey, zone string) ([]byte, error) {
|
||||||
})
|
})
|
||||||
|
|
||||||
client := &http.Client{Timeout: brightdataTimeout}
|
client := &http.Client{Timeout: brightdataTimeout}
|
||||||
req, err := http.NewRequest(http.MethodPost, "https://api.brightdata.com/request", bytes.NewReader(payload))
|
req, err := http.NewRequest(http.MethodPost, brightdataEndpoint, bytes.NewReader(payload))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("build brightdata request: %w", err)
|
return nil, fmt.Errorf("build brightdata request: %w", err)
|
||||||
}
|
}
|
||||||
|
|
@ -108,8 +110,22 @@ func headCheck(url string) (int, error) {
|
||||||
return resp.StatusCode, nil
|
return resp.StatusCode, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// fetchHTML tries direct GET first, falls back to Brightdata.
|
// fetchHTML fetches HTML content. When provider is "brightdata", Brightdata is
|
||||||
|
// tried first with direct fetch as fallback; otherwise direct fetch comes first.
|
||||||
func fetchHTML(cfg *fetchConfig, targetURL string) *FetchResponse {
|
func fetchHTML(cfg *fetchConfig, targetURL string) *FetchResponse {
|
||||||
|
if cfg.Provider == "brightdata" && cfg.BrightdataKey != "" {
|
||||||
|
body, err := brightdataFetch(targetURL, cfg.BrightdataKey, cfg.BrightdataZone)
|
||||||
|
if err == nil {
|
||||||
|
htmlStr := string(body)
|
||||||
|
return &FetchResponse{
|
||||||
|
URL: targetURL,
|
||||||
|
Title: ExtractTitle(htmlStr),
|
||||||
|
Content: ExtractContent(htmlStr),
|
||||||
|
Format: "html",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
res, err := directFetch(targetURL, false)
|
res, err := directFetch(targetURL, false)
|
||||||
if err == nil && res.StatusCode == 200 && len(res.Body) >= minDirectBody {
|
if err == nil && res.StatusCode == 200 && len(res.Body) >= minDirectBody {
|
||||||
htmlStr := string(res.Body)
|
htmlStr := string(res.Body)
|
||||||
|
|
@ -121,7 +137,7 @@ func fetchHTML(cfg *fetchConfig, targetURL string) *FetchResponse {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if cfg.BrightdataKey != "" {
|
if cfg.Provider != "brightdata" && cfg.BrightdataKey != "" {
|
||||||
body, err := brightdataFetch(targetURL, cfg.BrightdataKey, cfg.BrightdataZone)
|
body, err := brightdataFetch(targetURL, cfg.BrightdataKey, cfg.BrightdataZone)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
htmlStr := string(body)
|
htmlStr := string(body)
|
||||||
|
|
@ -146,7 +162,6 @@ func fetchMarkdown(cfg *fetchConfig, targetURL string) *FetchResponse {
|
||||||
lower := strings.ToLower(targetURL)
|
lower := strings.ToLower(targetURL)
|
||||||
ext := strings.ToLower(path.Ext(strings.TrimSuffix(lower, "/")))
|
ext := strings.ToLower(path.Ext(strings.TrimSuffix(lower, "/")))
|
||||||
|
|
||||||
// Already a .md file
|
|
||||||
if ext == ".md" || ext == ".mdx" {
|
if ext == ".md" || ext == ".mdx" {
|
||||||
res, err := directFetch(targetURL, true)
|
res, err := directFetch(targetURL, true)
|
||||||
if err == nil && res.StatusCode == 200 && len(res.Body) >= minMarkdownBody {
|
if err == nil && res.StatusCode == 200 && len(res.Body) >= minMarkdownBody {
|
||||||
|
|
@ -158,13 +173,12 @@ func fetchMarkdown(cfg *fetchConfig, targetURL string) *FetchResponse {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Probe for .md version
|
|
||||||
mdURL := buildMdURL(targetURL)
|
mdURL := buildMdURL(targetURL)
|
||||||
if mdURL != "" {
|
if mdURL != "" {
|
||||||
code, err := headCheck(mdURL)
|
code, err := headCheck(mdURL)
|
||||||
if err == nil && code == 200 {
|
if err == nil && code == 200 {
|
||||||
res, err := directFetch(mdURL, true)
|
res, err := directFetch(mdURL, true)
|
||||||
if err == nil && res.StatusCode == 200 && len(res.Body) >= minMarkdownBody {
|
if err == nil && res.StatusCode == 200 && len(res.Body) >= minMarkdownBody && !isHTMLContent(res.ContentType, res.Body) {
|
||||||
return &FetchResponse{
|
return &FetchResponse{
|
||||||
URL: targetURL,
|
URL: targetURL,
|
||||||
Content: string(res.Body),
|
Content: string(res.Body),
|
||||||
|
|
@ -174,7 +188,6 @@ func fetchMarkdown(cfg *fetchConfig, targetURL string) *FetchResponse {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Fallback: fetch HTML and convert
|
|
||||||
htmlRes := fetchRawHTML(cfg, targetURL)
|
htmlRes := fetchRawHTML(cfg, targetURL)
|
||||||
if htmlRes == nil {
|
if htmlRes == nil {
|
||||||
return &FetchResponse{
|
return &FetchResponse{
|
||||||
|
|
@ -199,14 +212,22 @@ func fetchMarkdown(cfg *fetchConfig, targetURL string) *FetchResponse {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// fetchRawHTML fetches HTML with direct -> brightdata fallback.
|
// fetchRawHTML fetches raw HTML. When provider is "brightdata", Brightdata is
|
||||||
|
// tried first with direct fetch as fallback; otherwise direct fetch comes first.
|
||||||
func fetchRawHTML(cfg *fetchConfig, targetURL string) []byte {
|
func fetchRawHTML(cfg *fetchConfig, targetURL string) []byte {
|
||||||
|
if cfg.Provider == "brightdata" && cfg.BrightdataKey != "" {
|
||||||
|
body, err := brightdataFetch(targetURL, cfg.BrightdataKey, cfg.BrightdataZone)
|
||||||
|
if err == nil {
|
||||||
|
return body
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
res, err := directFetch(targetURL, false)
|
res, err := directFetch(targetURL, false)
|
||||||
if err == nil && res.StatusCode == 200 && len(res.Body) >= minDirectBody {
|
if err == nil && res.StatusCode == 200 && len(res.Body) >= minDirectBody {
|
||||||
return res.Body
|
return res.Body
|
||||||
}
|
}
|
||||||
|
|
||||||
if cfg.BrightdataKey != "" {
|
if cfg.Provider != "brightdata" && cfg.BrightdataKey != "" {
|
||||||
body, err := brightdataFetch(targetURL, cfg.BrightdataKey, cfg.BrightdataZone)
|
body, err := brightdataFetch(targetURL, cfg.BrightdataKey, cfg.BrightdataZone)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
return body
|
return body
|
||||||
|
|
@ -224,6 +245,22 @@ func buildMdURL(u string) string {
|
||||||
return trimmed + ".md"
|
return trimmed + ".md"
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// isHTMLContent returns true if the response looks like HTML rather than
|
||||||
|
// plain text or markdown, based on Content-Type header and body sniffing.
|
||||||
|
func isHTMLContent(contentType string, body []byte) bool {
|
||||||
|
ct := strings.ToLower(contentType)
|
||||||
|
if strings.Contains(ct, "text/html") || strings.Contains(ct, "application/xhtml") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if len(body) > 0 {
|
||||||
|
prefix := strings.TrimSpace(strings.ToLower(string(body[:min(len(body), 256)])))
|
||||||
|
if strings.HasPrefix(prefix, "<!doctype") || strings.HasPrefix(prefix, "<html") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
func truncate(s string, n int) string {
|
func truncate(s string, n int) string {
|
||||||
if len(s) <= n {
|
if len(s) <= n {
|
||||||
return s
|
return s
|
||||||
|
|
|
||||||
279
tools/webfetch/fetch_order_test.go
Normal file
279
tools/webfetch/fetch_order_test.go
Normal file
|
|
@ -0,0 +1,279 @@
|
||||||
|
package webfetch
|
||||||
|
|
||||||
|
import (
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func newDirectServer(body string) *httptest.Server {
|
||||||
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/html")
|
||||||
|
w.Write([]byte(body))
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
func newBrightdataServer(body string) *httptest.Server {
|
||||||
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/html")
|
||||||
|
w.Write([]byte(body))
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
var (
|
||||||
|
directContent = "<html><head><title>Direct</title></head><body><p>" + strings.Repeat("direct-content ", 40) + "</p></body></html>"
|
||||||
|
brightdataHTML = "<html><head><title>Brightdata</title></head><body><p>" + strings.Repeat("brightdata-content ", 40) + "</p></body></html>"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestFetchHTML_BrightdataProvider_PrefersProxy(t *testing.T) {
|
||||||
|
directSrv := newDirectServer(directContent)
|
||||||
|
defer directSrv.Close()
|
||||||
|
|
||||||
|
bdSrv := newBrightdataServer(brightdataHTML)
|
||||||
|
defer bdSrv.Close()
|
||||||
|
|
||||||
|
origEndpoint := brightdataEndpoint
|
||||||
|
brightdataEndpoint = bdSrv.URL
|
||||||
|
defer func() { brightdataEndpoint = origEndpoint }()
|
||||||
|
|
||||||
|
cfg := &fetchConfig{
|
||||||
|
Provider: "brightdata",
|
||||||
|
BrightdataKey: "test-key",
|
||||||
|
BrightdataZone: "test-zone",
|
||||||
|
}
|
||||||
|
|
||||||
|
resp := fetchHTML(cfg, directSrv.URL)
|
||||||
|
if resp == nil {
|
||||||
|
t.Fatal("expected non-nil response")
|
||||||
|
}
|
||||||
|
if resp.Title != "Brightdata" {
|
||||||
|
t.Errorf("expected Brightdata content first, got title=%q", resp.Title)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFetchHTML_DefaultProvider_PrefersDirect(t *testing.T) {
|
||||||
|
directSrv := newDirectServer(directContent)
|
||||||
|
defer directSrv.Close()
|
||||||
|
|
||||||
|
bdSrv := newBrightdataServer(brightdataHTML)
|
||||||
|
defer bdSrv.Close()
|
||||||
|
|
||||||
|
origEndpoint := brightdataEndpoint
|
||||||
|
brightdataEndpoint = bdSrv.URL
|
||||||
|
defer func() { brightdataEndpoint = origEndpoint }()
|
||||||
|
|
||||||
|
cfg := &fetchConfig{
|
||||||
|
Provider: "",
|
||||||
|
BrightdataKey: "test-key",
|
||||||
|
BrightdataZone: "test-zone",
|
||||||
|
}
|
||||||
|
|
||||||
|
resp := fetchHTML(cfg, directSrv.URL)
|
||||||
|
if resp == nil {
|
||||||
|
t.Fatal("expected non-nil response")
|
||||||
|
}
|
||||||
|
if resp.Title != "Direct" {
|
||||||
|
t.Errorf("expected direct content first, got title=%q", resp.Title)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFetchHTML_BrightdataProvider_FallsBackToDirect(t *testing.T) {
|
||||||
|
directSrv := newDirectServer(directContent)
|
||||||
|
defer directSrv.Close()
|
||||||
|
|
||||||
|
// Brightdata returns 500 → should fall back to direct
|
||||||
|
bdSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusInternalServerError)
|
||||||
|
w.Write([]byte("error"))
|
||||||
|
}))
|
||||||
|
defer bdSrv.Close()
|
||||||
|
|
||||||
|
origEndpoint := brightdataEndpoint
|
||||||
|
brightdataEndpoint = bdSrv.URL
|
||||||
|
defer func() { brightdataEndpoint = origEndpoint }()
|
||||||
|
|
||||||
|
cfg := &fetchConfig{
|
||||||
|
Provider: "brightdata",
|
||||||
|
BrightdataKey: "test-key",
|
||||||
|
BrightdataZone: "test-zone",
|
||||||
|
}
|
||||||
|
|
||||||
|
resp := fetchHTML(cfg, directSrv.URL)
|
||||||
|
if resp == nil {
|
||||||
|
t.Fatal("expected non-nil response")
|
||||||
|
}
|
||||||
|
if resp.Title != "Direct" {
|
||||||
|
t.Errorf("expected fallback to direct, got title=%q", resp.Title)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFetchHTML_DefaultProvider_FallsBackToBrightdata(t *testing.T) {
|
||||||
|
// Direct returns 403 → should fall back to Brightdata
|
||||||
|
directSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusForbidden)
|
||||||
|
w.Write([]byte("forbidden"))
|
||||||
|
}))
|
||||||
|
defer directSrv.Close()
|
||||||
|
|
||||||
|
bdSrv := newBrightdataServer(brightdataHTML)
|
||||||
|
defer bdSrv.Close()
|
||||||
|
|
||||||
|
origEndpoint := brightdataEndpoint
|
||||||
|
brightdataEndpoint = bdSrv.URL
|
||||||
|
defer func() { brightdataEndpoint = origEndpoint }()
|
||||||
|
|
||||||
|
cfg := &fetchConfig{
|
||||||
|
Provider: "",
|
||||||
|
BrightdataKey: "test-key",
|
||||||
|
BrightdataZone: "test-zone",
|
||||||
|
}
|
||||||
|
|
||||||
|
resp := fetchHTML(cfg, directSrv.URL)
|
||||||
|
if resp == nil {
|
||||||
|
t.Fatal("expected non-nil response")
|
||||||
|
}
|
||||||
|
if resp.Title != "Brightdata" {
|
||||||
|
t.Errorf("expected fallback to brightdata, got title=%q", resp.Title)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFetchRawHTML_BrightdataProvider_PrefersProxy(t *testing.T) {
|
||||||
|
directSrv := newDirectServer(directContent)
|
||||||
|
defer directSrv.Close()
|
||||||
|
|
||||||
|
bdSrv := newBrightdataServer(brightdataHTML)
|
||||||
|
defer bdSrv.Close()
|
||||||
|
|
||||||
|
origEndpoint := brightdataEndpoint
|
||||||
|
brightdataEndpoint = bdSrv.URL
|
||||||
|
defer func() { brightdataEndpoint = origEndpoint }()
|
||||||
|
|
||||||
|
cfg := &fetchConfig{
|
||||||
|
Provider: "brightdata",
|
||||||
|
BrightdataKey: "test-key",
|
||||||
|
BrightdataZone: "test-zone",
|
||||||
|
}
|
||||||
|
|
||||||
|
body := fetchRawHTML(cfg, directSrv.URL)
|
||||||
|
if body == nil {
|
||||||
|
t.Fatal("expected non-nil body")
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(body), "Brightdata") {
|
||||||
|
t.Error("expected Brightdata content when provider is brightdata")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFetchRawHTML_DefaultProvider_PrefersDirect(t *testing.T) {
|
||||||
|
directSrv := newDirectServer(directContent)
|
||||||
|
defer directSrv.Close()
|
||||||
|
|
||||||
|
bdSrv := newBrightdataServer(brightdataHTML)
|
||||||
|
defer bdSrv.Close()
|
||||||
|
|
||||||
|
origEndpoint := brightdataEndpoint
|
||||||
|
brightdataEndpoint = bdSrv.URL
|
||||||
|
defer func() { brightdataEndpoint = origEndpoint }()
|
||||||
|
|
||||||
|
cfg := &fetchConfig{
|
||||||
|
Provider: "",
|
||||||
|
BrightdataKey: "test-key",
|
||||||
|
BrightdataZone: "test-zone",
|
||||||
|
}
|
||||||
|
|
||||||
|
body := fetchRawHTML(cfg, directSrv.URL)
|
||||||
|
if body == nil {
|
||||||
|
t.Fatal("expected non-nil body")
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(body), "Direct") {
|
||||||
|
t.Error("expected direct content when provider is empty")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFetchRawHTML_BrightdataProvider_FallsBackToDirect(t *testing.T) {
|
||||||
|
directSrv := newDirectServer(directContent)
|
||||||
|
defer directSrv.Close()
|
||||||
|
|
||||||
|
bdSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusInternalServerError)
|
||||||
|
}))
|
||||||
|
defer bdSrv.Close()
|
||||||
|
|
||||||
|
origEndpoint := brightdataEndpoint
|
||||||
|
brightdataEndpoint = bdSrv.URL
|
||||||
|
defer func() { brightdataEndpoint = origEndpoint }()
|
||||||
|
|
||||||
|
cfg := &fetchConfig{
|
||||||
|
Provider: "brightdata",
|
||||||
|
BrightdataKey: "test-key",
|
||||||
|
BrightdataZone: "test-zone",
|
||||||
|
}
|
||||||
|
|
||||||
|
body := fetchRawHTML(cfg, directSrv.URL)
|
||||||
|
if body == nil {
|
||||||
|
t.Fatal("expected non-nil body from direct fallback")
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(body), "Direct") {
|
||||||
|
t.Error("expected direct content as fallback")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFetchHTML_BrightdataProvider_NeverCallsDirect_WhenProxySucceeds(t *testing.T) {
|
||||||
|
var directCalls atomic.Int32
|
||||||
|
directSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
directCalls.Add(1)
|
||||||
|
w.Write([]byte(directContent))
|
||||||
|
}))
|
||||||
|
defer directSrv.Close()
|
||||||
|
|
||||||
|
bdSrv := newBrightdataServer(brightdataHTML)
|
||||||
|
defer bdSrv.Close()
|
||||||
|
|
||||||
|
origEndpoint := brightdataEndpoint
|
||||||
|
brightdataEndpoint = bdSrv.URL
|
||||||
|
defer func() { brightdataEndpoint = origEndpoint }()
|
||||||
|
|
||||||
|
cfg := &fetchConfig{
|
||||||
|
Provider: "brightdata",
|
||||||
|
BrightdataKey: "test-key",
|
||||||
|
BrightdataZone: "test-zone",
|
||||||
|
}
|
||||||
|
|
||||||
|
resp := fetchHTML(cfg, directSrv.URL)
|
||||||
|
if resp == nil {
|
||||||
|
t.Fatal("expected non-nil response")
|
||||||
|
}
|
||||||
|
if directCalls.Load() != 0 {
|
||||||
|
t.Errorf("direct server should not be called when brightdata succeeds, got %d calls", directCalls.Load())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFetchMarkdown_BrightdataProvider_UsesProxy(t *testing.T) {
|
||||||
|
directSrv := newDirectServer(directContent)
|
||||||
|
defer directSrv.Close()
|
||||||
|
|
||||||
|
bdSrv := newBrightdataServer(brightdataHTML)
|
||||||
|
defer bdSrv.Close()
|
||||||
|
|
||||||
|
origEndpoint := brightdataEndpoint
|
||||||
|
brightdataEndpoint = bdSrv.URL
|
||||||
|
defer func() { brightdataEndpoint = origEndpoint }()
|
||||||
|
|
||||||
|
cfg := &fetchConfig{
|
||||||
|
Provider: "brightdata",
|
||||||
|
BrightdataKey: "test-key",
|
||||||
|
BrightdataZone: "test-zone",
|
||||||
|
}
|
||||||
|
|
||||||
|
resp := fetchMarkdown(cfg, directSrv.URL)
|
||||||
|
if resp == nil {
|
||||||
|
t.Fatal("expected non-nil response")
|
||||||
|
}
|
||||||
|
if resp.Format != "markdown" {
|
||||||
|
t.Errorf("expected format 'markdown', got '%s'", resp.Format)
|
||||||
|
}
|
||||||
|
if !strings.Contains(resp.Content, "Brightdata") {
|
||||||
|
t.Errorf("expected Brightdata content in markdown, got: %s", resp.Content[:min(len(resp.Content), 200)])
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Add table
Reference in a new issue