robinhood-agentic-mcp/client/session.go

package client

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/modelcontextprotocol/go-sdk/auth"
	mcp "github.com/modelcontextprotocol/go-sdk/mcp"
	"github.com/modelcontextprotocol/go-sdk/oauthex"
	"golang.org/x/oauth2"
)

// Token holds the OAuth/bearer fields ConnectSession needs.
type Token struct {
	AccessToken, RefreshToken, TokenType string
	Expiry                               time.Time
	ClientID, ClientSecret               string
	AuthURL, TokenURL, RedirectURL       string
}

// ConnectSession opens a streamable MCP session using tok.
func ConnectSession(ctx context.Context, url string, tok Token, name, version string) (*Client, error) {
	sess, err := connectSession(ctx, url, tok, name, version)
	if err != nil {
		return nil, err
	}
	return &Client{URL: url, Token: tok.AccessToken, Name: name, Version: version, session: sess}, nil
}

// AttachSession sets the SDK session used by Call when Hook is nil.
func (c *Client) AttachSession(sess *mcp.ClientSession) {
	c.session = sess
}

func connectSession(ctx context.Context, mcpURL string, tok Token, name, version string) (*mcp.ClientSession, error) {
	var handler auth.OAuthHandler
	if tok.ClientID != "" && tok.TokenURL != "" {
		oc := &oauth2.Config{
			ClientID:     tok.ClientID,
			ClientSecret: tok.ClientSecret,
			RedirectURL:  tok.RedirectURL,
			Endpoint:     oauth2.Endpoint{AuthURL: tok.AuthURL, TokenURL: tok.TokenURL},
		}
		ot := &oauth2.Token{
			AccessToken:  tok.AccessToken,
			RefreshToken: tok.RefreshToken,
			TokenType:    tok.TokenType,
			Expiry:       tok.Expiry,
		}
		h, err := auth.NewAuthorizationCodeHandler(&auth.AuthorizationCodeHandlerConfig{
			RedirectURL: tok.RedirectURL,
			PreregisteredClient: &oauthex.ClientCredentials{
				ClientID: tok.ClientID,
			},
			InitialTokenSource: oc.TokenSource(ctx, ot),
			AuthorizationCodeFetcher: func(context.Context, *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) {
				return nil, fmt.Errorf("robinhood session expired; run Login again")
			},
		})
		if err != nil {
			return nil, err
		}
		handler = h
	}
	httpClient := &http.Client{Timeout: 60 * time.Second}
	if handler == nil && tok.AccessToken != "" {
		httpClient.Transport = bearerRT{token: tok.AccessToken, base: http.DefaultTransport}
	}
	t := &mcp.StreamableClientTransport{
		Endpoint:             mcpURL,
		HTTPClient:           httpClient,
		OAuthHandler:         handler,
		DisableStandaloneSSE: true,
	}
	cli := mcp.NewClient(&mcp.Implementation{Name: name, Version: version}, nil)
	return cli.Connect(ctx, t, nil)
}

func (c *Client) sessionCall(ctx context.Context, name string, args map[string]any) (json.RawMessage, error) {
	if c.session == nil {
		return nil, ToolErrorf(name, "session is nil")
	}
	res, err := c.session.CallTool(ctx, &mcp.CallToolParams{Name: name, Arguments: args})
	if err != nil {
		return nil, ToolErrorf(name, "%w", err)
	}
	raw, err := toolJSON(res)
	if err != nil {
		return nil, ToolErrorf(name, "%w", err)
	}
	return raw, nil
}

func toolJSON(res *mcp.CallToolResult) (json.RawMessage, error) {
	if res == nil {
		return nil, fmt.Errorf("empty tool result")
	}
	// GetError is only set by SetError on the server; the err field is not marshaled to clients.
	if err := res.GetError(); err != nil {
		return nil, err
	}
	texts, types := collectTextAndTypes(res)
	concat := strings.Join(texts, "")
	// IsError is the client-visible soft-fail flag; error text lives in Content.
	if res.IsError {
		if concat == "" {
			return nil, fmt.Errorf("tool error")
		}
		return nil, fmt.Errorf("%s", concat)
	}
	if res.StructuredContent != nil {
		b, err := json.Marshal(res.StructuredContent)
		if err != nil {
			return nil, err
		}
		return json.RawMessage(b), nil
	}
	for _, t := range texts {
		if json.Valid([]byte(t)) {
			return json.RawMessage(t), nil
		}
	}
	if concat == "" {
		return json.RawMessage(`{}`), nil
	}
	if json.Valid([]byte(concat)) {
		return json.RawMessage(concat), nil
	}
	if raw := extractFirstJSON(concat); raw != nil {
		return raw, nil
	}
	structured := "nil"
	if res.StructuredContent != nil {
		structured = "present"
	}
	return nil, fmt.Errorf("non-json tool result (isError=%v types=[%s] structured=%s textLen=%d prefix=%q)",
		res.IsError, strings.Join(types, " "), structured, utf8.RuneCountInString(concat), runePrefix(concat, 80))
}

func collectTextAndTypes(res *mcp.CallToolResult) (texts, types []string) {
	for _, c := range res.Content {
		types = append(types, contentTypeName(c))
		t, ok := c.(*mcp.TextContent)
		if !ok {
			continue
		}
		texts = append(texts, t.Text)
	}
	return texts, types
}

func contentTypeName(c mcp.Content) string {
	switch c.(type) {
	case *mcp.TextContent:
		return "text"
	case *mcp.ImageContent:
		return "image"
	case *mcp.AudioContent:
		return "audio"
	case *mcp.ResourceLink:
		return "resource_link"
	case *mcp.EmbeddedResource:
		return "resource"
	case *mcp.ToolUseContent:
		return "tool_use"
	case *mcp.ToolResultContent:
		return "tool_result"
	default:
		return "unknown"
	}
}

func runePrefix(s string, n int) string {
	if n <= 0 || s == "" {
		return ""
	}
	i := 0
	for j := range s {
		if i == n {
			return s[:j]
		}
		i++
	}
	return s
}

// extractFirstJSON returns the first balanced {...} or [...] substring that is
// valid JSON. Brace matching skips quoted strings and respects escapes.
func extractFirstJSON(s string) json.RawMessage {
	for i := 0; i < len(s); i++ {
		if s[i] != '{' && s[i] != '[' {
			continue
		}
		end := balancedJSONEnd(s, i)
		if end <= i {
			continue
		}
		cand := s[i:end]
		if json.Valid([]byte(cand)) {
			return json.RawMessage(cand)
		}
	}
	return nil
}

func balancedJSONEnd(s string, start int) int {
	depth := 0
	inString := false
	escape := false
	for i := start; i < len(s); i++ {
		c := s[i]
		if inString {
			if escape {
				escape = false
				continue
			}
			if c == '\\' {
				escape = true
				continue
			}
			if c == '"' {
				inString = false
			}
			continue
		}
		switch c {
		case '"':
			inString = true
		case '{', '[':
			depth++
		case '}', ']':
			depth--
			if depth == 0 {
				return i + 1
			}
			if depth < 0 {
				return -1
			}
		}
	}
	return -1
}

type bearerRT struct {
	token string
	base  http.RoundTripper
}

func (b bearerRT) RoundTrip(req *http.Request) (*http.Response, error) {
	r := req.Clone(req.Context())
	r.Header.Set("Authorization", "Bearer "+b.token)
	base := b.base
	if base == nil {
		base = http.DefaultTransport
	}
	return base.RoundTrip(r)
}