feat(project): 项目完结

2026-05-30 23:16:49 +08:00
commit 60e7777cbd
97 changed files with 15027 additions and 0 deletions
--- a/internal/infrastructure/ai/mcp_sse_client.go
+++ b/internal/infrastructure/ai/mcp_sse_client.go
@@ -0,0 +1,509 @@
+package ai
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"ai-agent-scaffold-go/internal/domain/agent/ports"
+)
+
+const (
+	mcpProtocolVersion = "2024-11-05"
+	mcpClientName      = "ai-agent-scaffold-go"
+	mcpClientVersion   = "0.1.0"
+)
+
+type MCPSSEClient struct {
+	baseURI    string
+	endpoint   string
+	httpClient *http.Client
+	timeout    time.Duration
+
+	mu          sync.Mutex
+	started     bool
+	postURL     string
+	cancelFn    context.CancelFunc
+	pending     map[uint64]chan json.RawMessage
+	nextID      atomic.Uint64
+	streamErr   chan error
+	endpointSig chan struct{}
+}
+
+func NewMCPSSEClient(baseURI, endpoint string, requestTimeoutMillis int) *MCPSSEClient {
+	timeout := time.Duration(requestTimeoutMillis) * time.Millisecond
+	if timeout <= 0 {
+		timeout = 120 * time.Second
+	}
+	return &MCPSSEClient{
+		baseURI:    strings.TrimRight(baseURI, "/"),
+		endpoint:   endpoint,
+		httpClient: &http.Client{Timeout: 0},
+		timeout:    timeout,
+	}
+}
+
+func (c *MCPSSEClient) ensureStarted(ctx context.Context) error {
+	c.mu.Lock()
+	if c.started {
+		c.mu.Unlock()
+		return nil
+	}
+	c.pending = make(map[uint64]chan json.RawMessage)
+	c.endpointSig = make(chan struct{})
+	c.streamErr = make(chan error, 1)
+	streamCtx, cancel := context.WithCancel(context.Background())
+	c.cancelFn = cancel
+	c.started = true
+	c.mu.Unlock()
+
+	if err := c.openSSE(streamCtx); err != nil {
+		c.shutdown()
+		return err
+	}
+
+	waitCtx, waitCancel := context.WithTimeout(ctx, c.timeout)
+	defer waitCancel()
+	select {
+	case <-c.endpointSig:
+	case err := <-c.streamErr:
+		c.shutdown()
+		return err
+	case <-waitCtx.Done():
+		c.shutdown()
+		return fmt.Errorf("mcp sse endpoint event timeout: %w", waitCtx.Err())
+	}
+
+	if err := c.initialize(ctx); err != nil {
+		c.shutdown()
+		return err
+	}
+	return nil
+}
+
+func (c *MCPSSEClient) openSSE(ctx context.Context) error {
+	target := c.baseURI + c.endpoint
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
+	if err != nil {
+		return fmt.Errorf("mcp sse build request: %w", err)
+	}
+	req.Header.Set("Accept", "text/event-stream")
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return fmt.Errorf("mcp sse open: %w", err)
+	}
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		raw, _ := io.ReadAll(resp.Body)
+		_ = resp.Body.Close()
+		return fmt.Errorf("mcp sse status %d: %s", resp.StatusCode, truncate(string(raw), 200))
+	}
+	go c.readLoop(resp)
+	return nil
+}
+
+func (c *MCPSSEClient) readLoop(resp *http.Response) {
+	defer resp.Body.Close()
+	reader := bufio.NewReader(resp.Body)
+	var event string
+	var dataBuf strings.Builder
+	dispatch := func() {
+		defer func() {
+			event = ""
+			dataBuf.Reset()
+		}()
+		data := dataBuf.String()
+		if data == "" {
+			return
+		}
+		switch event {
+		case "endpoint", "":
+			c.handleEndpoint(data, event)
+		case "message":
+			c.handleMessage(data)
+		}
+	}
+	for {
+		line, err := reader.ReadString('\n')
+		if err != nil {
+			if err != io.EOF {
+				c.signalStreamErr(fmt.Errorf("mcp sse read: %w", err))
+			} else {
+				c.signalStreamErr(fmt.Errorf("mcp sse stream closed"))
+			}
+			return
+		}
+		line = strings.TrimRight(line, "\r\n")
+		if line == "" {
+			dispatch()
+			continue
+		}
+		switch {
+		case strings.HasPrefix(line, "event:"):
+			event = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
+		case strings.HasPrefix(line, "data:"):
+			if dataBuf.Len() > 0 {
+				dataBuf.WriteByte('\n')
+			}
+			dataBuf.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
+		}
+	}
+}
+
+func (c *MCPSSEClient) handleEndpoint(data, eventName string) {
+	if c.postURLAlreadySet() {
+		if eventName == "" {
+			c.handleMessage(data)
+		}
+		return
+	}
+	target := c.resolvePostURL(data)
+	c.mu.Lock()
+	if c.postURL == "" {
+		c.postURL = target
+		close(c.endpointSig)
+	}
+	c.mu.Unlock()
+}
+
+func (c *MCPSSEClient) postURLAlreadySet() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.postURL != ""
+}
+
+func (c *MCPSSEClient) resolvePostURL(raw string) string {
+	parsed, err := url.Parse(raw)
+	if err != nil || !parsed.IsAbs() {
+		base, baseErr := url.Parse(c.baseURI)
+		if baseErr == nil {
+			ref, refErr := url.Parse(raw)
+			if refErr == nil {
+				return base.ResolveReference(ref).String()
+			}
+		}
+	}
+	return raw
+}
+
+func (c *MCPSSEClient) handleMessage(data string) {
+	var resp struct {
+		ID     json.Number     `json:"id"`
+		Result json.RawMessage `json:"result"`
+		Error  *struct {
+			Code    int    `json:"code"`
+			Message string `json:"message"`
+		} `json:"error"`
+	}
+	if err := json.Unmarshal([]byte(data), &resp); err != nil {
+		return
+	}
+	if resp.ID == "" {
+		return
+	}
+	id, err := resp.ID.Int64()
+	if err != nil || id <= 0 {
+		return
+	}
+	c.mu.Lock()
+	ch, ok := c.pending[uint64(id)]
+	if ok {
+		delete(c.pending, uint64(id))
+	}
+	c.mu.Unlock()
+	if !ok {
+		return
+	}
+	if resp.Error != nil {
+		ch <- mustJSON(map[string]any{"__error__": resp.Error.Message, "code": resp.Error.Code})
+		close(ch)
+		return
+	}
+	ch <- resp.Result
+	close(ch)
+}
+
+func mustJSON(v any) json.RawMessage {
+	raw, _ := json.Marshal(v)
+	return raw
+}
+
+func (c *MCPSSEClient) signalStreamErr(err error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	select {
+	case c.streamErr <- err:
+	default:
+	}
+	for id, ch := range c.pending {
+		ch <- mustJSON(map[string]any{"__error__": err.Error()})
+		close(ch)
+		delete(c.pending, id)
+	}
+}
+
+func (c *MCPSSEClient) shutdown() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.cancelFn != nil {
+		c.cancelFn()
+	}
+	c.started = false
+	c.postURL = ""
+	c.pending = nil
+}
+
+func (c *MCPSSEClient) initialize(ctx context.Context) error {
+	_, err := c.callRPC(ctx, "initialize", map[string]any{
+		"protocolVersion": mcpProtocolVersion,
+		"capabilities":    map[string]any{},
+		"clientInfo": map[string]any{
+			"name":    mcpClientName,
+			"version": mcpClientVersion,
+		},
+	})
+	if err != nil {
+		return fmt.Errorf("mcp initialize: %w", err)
+	}
+	if err := c.notify(ctx, "notifications/initialized", map[string]any{}); err != nil {
+		return fmt.Errorf("mcp notifications/initialized: %w", err)
+	}
+	return nil
+}
+
+func (c *MCPSSEClient) CallTool(ctx context.Context, name, arguments string) (string, error) {
+	if err := c.ensureStarted(ctx); err != nil {
+		return "", err
+	}
+	args := map[string]any{}
+	trimmed := strings.TrimSpace(arguments)
+	if trimmed != "" {
+		if err := json.Unmarshal([]byte(trimmed), &args); err != nil {
+			return "", fmt.Errorf("mcp tool %q arguments not valid json: %w", name, err)
+		}
+	}
+	result, err := c.callRPC(ctx, "tools/call", map[string]any{
+		"name":      name,
+		"arguments": args,
+	})
+	if err != nil {
+		return "", fmt.Errorf("mcp tools/call %q: %w", name, err)
+	}
+	return extractToolText(result), nil
+}
+
+func (c *MCPSSEClient) ListTools(ctx context.Context) ([]string, error) {
+	if err := c.ensureStarted(ctx); err != nil {
+		return nil, err
+	}
+	result, err := c.callRPC(ctx, "tools/list", map[string]any{})
+	if err != nil {
+		return nil, err
+	}
+	var parsed struct {
+		Tools []struct {
+			Name string `json:"name"`
+		} `json:"tools"`
+	}
+	if err := json.Unmarshal(result, &parsed); err != nil {
+		return nil, err
+	}
+	names := make([]string, 0, len(parsed.Tools))
+	for _, t := range parsed.Tools {
+		names = append(names, t.Name)
+	}
+	return names, nil
+}
+
+type toolCallResult struct {
+	Content []struct {
+		Type string `json:"type"`
+		Text string `json:"text"`
+	} `json:"content"`
+	IsError bool `json:"isError"`
+}
+
+func extractToolText(raw json.RawMessage) string {
+	var parsed toolCallResult
+	if err := json.Unmarshal(raw, &parsed); err != nil {
+		return string(raw)
+	}
+	parts := make([]string, 0, len(parsed.Content))
+	for _, item := range parsed.Content {
+		if item.Type == "text" && item.Text != "" {
+			parts = append(parts, item.Text)
+		}
+	}
+	if len(parts) == 0 {
+		return string(raw)
+	}
+	return strings.Join(parts, "\n")
+}
+
+func (c *MCPSSEClient) callRPC(ctx context.Context, method string, params any) (json.RawMessage, error) {
+	id := c.nextID.Add(1)
+	ch := make(chan json.RawMessage, 1)
+	c.mu.Lock()
+	postURL := c.postURL
+	c.pending[id] = ch
+	c.mu.Unlock()
+	if postURL == "" {
+		return nil, fmt.Errorf("mcp post url is not set")
+	}
+	body, err := json.Marshal(map[string]any{
+		"jsonrpc": "2.0",
+		"id":      id,
+		"method":  method,
+		"params":  params,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("mcp marshal request: %w", err)
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, postURL, bytes.NewReader(body))
+	if err != nil {
+		return nil, fmt.Errorf("mcp build post: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("mcp post: %w", err)
+	}
+	io.Copy(io.Discard, resp.Body)
+	resp.Body.Close()
+	if resp.StatusCode >= 400 {
+		return nil, fmt.Errorf("mcp post status %d", resp.StatusCode)
+	}
+
+	waitCtx, cancel := context.WithTimeout(ctx, c.timeout)
+	defer cancel()
+	select {
+	case msg, ok := <-ch:
+		if !ok {
+			return nil, fmt.Errorf("mcp call %s closed unexpectedly", method)
+		}
+		var probe struct {
+			Err string `json:"__error__"`
+		}
+		if err := json.Unmarshal(msg, &probe); err == nil && probe.Err != "" {
+			return nil, fmt.Errorf("%s", probe.Err)
+		}
+		return msg, nil
+	case <-waitCtx.Done():
+		return nil, fmt.Errorf("mcp call %s timeout: %w", method, waitCtx.Err())
+	}
+}
+
+func (c *MCPSSEClient) notify(ctx context.Context, method string, params any) error {
+	c.mu.Lock()
+	postURL := c.postURL
+	c.mu.Unlock()
+	if postURL == "" {
+		return fmt.Errorf("mcp post url is not set")
+	}
+	body, err := json.Marshal(map[string]any{
+		"jsonrpc": "2.0",
+		"method":  method,
+		"params":  params,
+	})
+	if err != nil {
+		return err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, postURL, bytes.NewReader(body))
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return err
+	}
+	io.Copy(io.Discard, resp.Body)
+	resp.Body.Close()
+	return nil
+}
+
+type MCPToolRouter struct {
+	clients      map[string]*MCPSSEClient
+	tools        []ports.Tool
+	registered   map[string]struct{}
+	listTimeout  time.Duration
+	listFailures map[string]error
+}
+
+func NewMCPToolRouter() *MCPToolRouter {
+	return &MCPToolRouter{
+		clients:      make(map[string]*MCPSSEClient),
+		registered:   make(map[string]struct{}),
+		listTimeout:  10 * time.Second,
+		listFailures: make(map[string]error),
+	}
+}
+
+func (r *MCPToolRouter) Register(tool ports.Tool) {
+	r.RegisterAndExpand(tool)
+}
+
+func (r *MCPToolRouter) RegisterAndExpand(tool ports.Tool) []ports.Tool {
+	mcp, ok := tool.(MCPTool)
+	if !ok {
+		r.tools = append(r.tools, tool)
+		return []ports.Tool{tool}
+	}
+	if mcp.TransportType != "sse" {
+		r.tools = append(r.tools, mcp)
+		return []ports.Tool{mcp}
+	}
+	if _, exists := r.clients[mcp.ToolName]; exists {
+		return nil
+	}
+	client := NewMCPSSEClient(mcp.BaseURI, mcp.SSEEndpoint, mcp.RequestTimeout)
+	r.clients[mcp.ToolName] = client
+
+	ctx, cancel := context.WithTimeout(context.Background(), r.listTimeout)
+	defer cancel()
+	names, err := client.ListTools(ctx)
+	if err != nil {
+		r.listFailures[mcp.ToolName] = err
+		r.tools = append(r.tools, mcp)
+		return []ports.Tool{mcp}
+	}
+	expanded := make([]ports.Tool, 0, len(names))
+	for _, n := range names {
+		r.clients[n] = client
+		r.registered[n] = struct{}{}
+		t := EinoTool{ToolName: n}
+		r.tools = append(r.tools, t)
+		expanded = append(expanded, t)
+	}
+	return expanded
+}
+
+func (r *MCPToolRouter) Tools() []ports.Tool {
+	return r.tools
+}
+
+func (r *MCPToolRouter) CallTool(ctx context.Context, name, arguments string) (string, error) {
+	if client, ok := r.clients[name]; ok {
+		return client.CallTool(ctx, name, arguments)
+	}
+	for _, tool := range r.tools {
+		if tool.Name() != name {
+			continue
+		}
+		switch tool.(type) {
+		case MCPTool:
+			return "", fmt.Errorf("mcp tool %q transport not supported in runtime", name)
+		default:
+			return "", fmt.Errorf("tool %q is not callable in current runtime", name)
+		}
+	}
+	return "", fmt.Errorf("tool %q is not registered", name)
+}