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"
|
||||
|
||||
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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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, "<!doctype") || strings.HasPrefix(prefix, "<html") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
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