diff --git a/grpc/mcp/mcp.go b/grpc/mcp/mcp.go index 62ac1a08..eb0ae129 100644 --- a/grpc/mcp/mcp.go +++ b/grpc/mcp/mcp.go @@ -8,12 +8,37 @@ import ( "google.golang.org/grpc/status" goumcp "github.com/yaoapp/gou/mcp" + "github.com/yaoapp/yao/grpc/auth" "github.com/yaoapp/yao/grpc/pb" ) // Handler implements the MCP gRPC methods. 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. func (h *Handler) MCPListTools(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPListResponse, error) { 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 { return nil, status.Errorf(codes.Internal, "CallTool failed: %v", err) } diff --git a/tools/webfetch/fetch.go b/tools/webfetch/fetch.go index 98690aaf..9d5fc909 100644 --- a/tools/webfetch/fetch.go +++ b/tools/webfetch/fetch.go @@ -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" ) +var brightdataEndpoint = "https://api.brightdata.com/request" + type fetchResult struct { Body []byte StatusCode int @@ -69,7 +71,7 @@ func brightdataFetch(targetURL, apiKey, zone string) ([]byte, error) { }) 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 { return nil, fmt.Errorf("build brightdata request: %w", err) } @@ -108,8 +110,22 @@ func headCheck(url string) (int, error) { 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 { + 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) if err == nil && res.StatusCode == 200 && len(res.Body) >= minDirectBody { 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) if err == nil { htmlStr := string(body) @@ -146,7 +162,6 @@ func fetchMarkdown(cfg *fetchConfig, targetURL string) *FetchResponse { lower := strings.ToLower(targetURL) ext := strings.ToLower(path.Ext(strings.TrimSuffix(lower, "/"))) - // Already a .md file if ext == ".md" || ext == ".mdx" { res, err := directFetch(targetURL, true) 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) if mdURL != "" { code, err := headCheck(mdURL) if err == nil && code == 200 { 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{ URL: targetURL, Content: string(res.Body), @@ -174,7 +188,6 @@ func fetchMarkdown(cfg *fetchConfig, targetURL string) *FetchResponse { } } - // Fallback: fetch HTML and convert htmlRes := fetchRawHTML(cfg, targetURL) if htmlRes == nil { 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 { + if cfg.Provider == "brightdata" && cfg.BrightdataKey != "" { + body, err := brightdataFetch(targetURL, cfg.BrightdataKey, cfg.BrightdataZone) + if err == nil { + return body + } + } + res, err := directFetch(targetURL, false) if err == nil && res.StatusCode == 200 && len(res.Body) >= minDirectBody { return res.Body } - if cfg.BrightdataKey != "" { + if cfg.Provider != "brightdata" && cfg.BrightdataKey != "" { body, err := brightdataFetch(targetURL, cfg.BrightdataKey, cfg.BrightdataZone) if err == nil { return body @@ -224,6 +245,22 @@ func buildMdURL(u string) string { 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, "Direct

" + strings.Repeat("direct-content ", 40) + "

" + brightdataHTML = "Brightdata

" + strings.Repeat("brightdata-content ", 40) + "

" +) + +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)]) + } +}