diff --git a/auth/login.go b/auth/login.go index 404e098..b1b1024 100644 --- a/auth/login.go +++ b/auth/login.go @@ -23,5 +23,5 @@ func Login(ctx context.Context, cfg Config) (accountID string, err error) { } return "", nil } - return "", fmt.Errorf("login: no ROBINHOOD_ACCESS_TOKEN and oauth not wired") + return loginOAuth(ctx, cfg) } diff --git a/auth/oauth.go b/auth/oauth.go new file mode 100644 index 0000000..d833252 --- /dev/null +++ b/auth/oauth.go @@ -0,0 +1,223 @@ +package auth + +import ( + "context" + "encoding/json" + "fmt" + "net" + "net/http" + "os" + "os/exec" + "time" + + mcpauth "github.com/modelcontextprotocol/go-sdk/auth" + mcp "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" + + "s1d3sw1ped/robinhood-agentic-mcp/client" + "s1d3sw1ped/robinhood-agentic-mcp/internal/wire" +) + +func oauthIdentity(cfg Config) (name, version string) { + cfg = cfg.WithDefaults() + return cfg.Name, cfg.Version +} + +func loginOAuth(ctx context.Context, cfg Config) (string, error) { + cfg = cfg.WithDefaults() + name, version := oauthIdentity(cfg.WithDefaults()) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return "", fmt.Errorf("listen: %w", err) + } + defer ln.Close() + port := ln.Addr().(*net.TCPAddr).Port + redirect := fmt.Sprintf("http://127.0.0.1:%d/callback", port) + + codeCh := make(chan mcpauth.AuthorizationResult, 1) + mux := http.NewServeMux() + mux.HandleFunc("/callback", func(w http.ResponseWriter, r *http.Request) { + q := r.URL.Query() + if q.Get("error") != "" { + http.Error(w, q.Get("error_description"), http.StatusBadRequest) + select { + case codeCh <- mcpauth.AuthorizationResult{}: + default: + } + return + } + res := mcpauth.AuthorizationResult{Code: q.Get("code"), State: q.Get("state"), Iss: q.Get("iss")} + _, _ = w.Write([]byte(name + " is signed in. You can close this tab.")) + select { + case codeCh <- res: + default: + } + }) + srv := &http.Server{Handler: mux} + go func() { _ = srv.Serve(ln) }() + defer func() { + shctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _ = srv.Shutdown(shctx) + }() + + var saved TokenSet + oauthCfg := &mcpauth.AuthorizationCodeHandlerConfig{ + RedirectURL: redirect, + DynamicClientRegistrationConfig: &mcpauth.DynamicClientRegistrationConfig{ + Metadata: &oauthex.ClientRegistrationMetadata{ + RedirectURIs: []string{redirect}, + ClientName: name, + ApplicationType: "native", + TokenEndpointAuthMethod: "none", + GrantTypes: []string{"authorization_code", "refresh_token"}, + ResponseTypes: []string{"code"}, + }, + }, + RequestRefreshToken: true, + AuthorizationCodeFetcher: func(ctx context.Context, args *mcpauth.AuthorizationArgs) (*mcpauth.AuthorizationResult, error) { + fmt.Fprintf(os.Stderr, "Open this URL to authorize %s with Robinhood:\n\n %s\n\n", name, args.URL) + _ = openBrowser(args.URL) + select { + case <-ctx.Done(): + return nil, ctx.Err() + case res := <-codeCh: + if res.Code == "" { + return nil, fmt.Errorf("oauth callback missing code") + } + return &res, nil + } + }, + NewTokenSource: func(ctx context.Context, oc *oauth2.Config, tok *oauth2.Token) (oauth2.TokenSource, error) { + saved = tokenSetFrom(oc, tok, redirect, saved.AccountID) + _ = WriteTokenSet(cfg.TokenFile, saved) + return oc.TokenSource(ctx, tok), nil + }, + } + handler, err := mcpauth.NewAuthorizationCodeHandler(oauthCfg) + if err != nil { + return "", fmt.Errorf("oauth handler: %w", err) + } + transport := &mcp.StreamableClientTransport{ + Endpoint: cfg.URL, + OAuthHandler: handler, + DisableStandaloneSSE: true, + } + mcpClient := mcp.NewClient(&mcp.Implementation{Name: name, Version: version}, nil) + ctx, cancel := context.WithTimeout(ctx, 5*time.Minute) + defer cancel() + sess, err := mcpClient.Connect(ctx, transport, nil) + if err != nil { + return "", fmt.Errorf("connect robinhood mcp: %w", err) + } + defer sess.Close() + + c := &client.Client{URL: cfg.URL, Name: name, Version: version} + c.AttachSession(sess) + raw, err := c.Call(ctx, "get_accounts", map[string]any{}) + if err != nil { + return "", fmt.Errorf("get_accounts: %w", err) + } + id, err := accountIDFrom(raw) + if err != nil { + return "", err + } + saved.AccountID = id + if ts, err := handler.TokenSource(ctx); err == nil && ts != nil { + if tok, err := ts.Token(); err == nil { + saved.AccessToken = tok.AccessToken + saved.RefreshToken = tok.RefreshToken + saved.Expiry = tok.Expiry + saved.TokenType = tok.TokenType + } + } + if err := WriteTokenSet(cfg.TokenFile, saved); err != nil { + return "", err + } + return id, nil +} + +func tokenSetFrom(oc *oauth2.Config, tok *oauth2.Token, redirect, accountID string) TokenSet { + if oc == nil { + oc = &oauth2.Config{} + } + if tok == nil { + tok = &oauth2.Token{} + } + return TokenSet{ + AccessToken: tok.AccessToken, + RefreshToken: tok.RefreshToken, + TokenType: tok.TokenType, + Expiry: tok.Expiry, + ClientID: oc.ClientID, + ClientSecret: oc.ClientSecret, + AuthURL: oc.Endpoint.AuthURL, + TokenURL: oc.Endpoint.TokenURL, + RedirectURL: redirect, + AccountID: accountID, + } +} + +func openBrowser(rawURL string) error { + cmds := [][]string{ + {"xdg-open", rawURL}, + {"gio", "open", rawURL}, + {"open", rawURL}, + } + for _, c := range cmds { + if err := exec.Command(c[0], c[1:]...).Start(); err == nil { + return nil + } + } + return fmt.Errorf("open browser") +} + +func accountIDFrom(raw json.RawMessage) (string, error) { + data := wire.Unwrap(raw) + type row struct { + ID string `json:"id"` + AccountNumber string `json:"account_number"` + AgenticAllowed bool `json:"agentic_allowed"` + Agentic bool `json:"agentic"` + } + var wrap struct { + Accounts []row `json:"accounts"` + } + if err := json.Unmarshal(data, &wrap); err != nil || wrap.Accounts == nil { + var one row + if err2 := json.Unmarshal(data, &one); err2 != nil { + if err != nil { + return "", fmt.Errorf("parse accounts: %w", err) + } + return "", fmt.Errorf("parse accounts") + } + wrap.Accounts = []row{one} + } + for _, a := range wrap.Accounts { + id := a.AccountNumber + if id == "" { + id = a.ID + } + if a.AgenticAllowed || a.Agentic { + return id, nil + } + } + return "", fmt.Errorf("no agentic account") +} + +// ClientToken maps a persisted TokenSet to the fields ConnectSession needs. +func ClientToken(t TokenSet) client.Token { + return client.Token{ + AccessToken: t.AccessToken, + RefreshToken: t.RefreshToken, + TokenType: t.TokenType, + Expiry: t.Expiry, + ClientID: t.ClientID, + ClientSecret: t.ClientSecret, + AuthURL: t.AuthURL, + TokenURL: t.TokenURL, + RedirectURL: t.RedirectURL, + } +} diff --git a/auth/oauth_test.go b/auth/oauth_test.go new file mode 100644 index 0000000..c381afe --- /dev/null +++ b/auth/oauth_test.go @@ -0,0 +1,14 @@ +package auth + +import "testing" + +func TestOAuthIdentity(t *testing.T) { + gotN, gotV := oauthIdentity(Config{Name: "tradey", Version: "0.9"}) + if gotN != "tradey" || gotV != "0.9" { + t.Fatalf("%s %s", gotN, gotV) + } + gotN, gotV = oauthIdentity(Config{}) + if gotN != DefaultName || gotV != DefaultVersion { + t.Fatalf("%s %s", gotN, gotV) + } +} diff --git a/client/client.go b/client/client.go index 51feb03..3da93b9 100644 --- a/client/client.go +++ b/client/client.go @@ -4,6 +4,8 @@ import ( "context" "encoding/json" "net/http" + + mcp "github.com/modelcontextprotocol/go-sdk/mcp" ) // Client is a Robinhood Agentic MCP client. @@ -11,7 +13,7 @@ type Client struct { URL, Token, Name, Version string HTTP *http.Client Hook Caller - session any + session *mcp.ClientSession } // Call invokes an MCP tool by name and returns its JSON result. @@ -19,5 +21,8 @@ func (c *Client) Call(ctx context.Context, name string, args map[string]any) (js if c.Hook != nil { return c.Hook.Call(ctx, name, args) } + if c.session != nil { + return c.sessionCall(ctx, name, args) + } return c.rpcCall(ctx, name, args) } diff --git a/client/session.go b/client/session.go new file mode 100644 index 0000000..d95a577 --- /dev/null +++ b/client/session.go @@ -0,0 +1,134 @@ +package client + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "time" + + "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") + } + if err := res.GetError(); err != nil { + return nil, err + } + var b []byte + for _, c := range res.Content { + t, ok := c.(*mcp.TextContent) + if !ok { + continue + } + b = append(b, t.Text...) + } + if len(b) == 0 { + return json.RawMessage(`{}`), nil + } + if !json.Valid(b) { + return nil, fmt.Errorf("non-json tool result") + } + return json.RawMessage(b), nil +} + +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) +} diff --git a/go.mod b/go.mod index f0d342d..72cc7f7 100644 --- a/go.mod +++ b/go.mod @@ -2,11 +2,19 @@ module s1d3sw1ped/robinhood-agentic-mcp go 1.25.0 -require github.com/alpacahq/alpacadecimal v0.0.9 +require ( + github.com/alpacahq/alpacadecimal v0.0.9 + github.com/modelcontextprotocol/go-sdk v1.7.0 + golang.org/x/oauth2 v0.35.0 +) require ( - github.com/google/go-cmp v0.7.0 // indirect - github.com/modelcontextprotocol/go-sdk v1.7.0 // indirect + github.com/google/jsonschema-go v0.4.3 // indirect + github.com/segmentio/asm v1.1.3 // indirect + github.com/segmentio/encoding v0.5.4 // indirect github.com/shopspring/decimal v1.4.0 // indirect - golang.org/x/oauth2 v0.35.0 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + golang.org/x/sync v0.20.0 // indirect + golang.org/x/sys v0.41.0 // indirect + golang.org/x/time v0.15.0 // indirect ) diff --git a/go.sum b/go.sum index 9dc417f..a4fb8a5 100644 --- a/go.sum +++ b/go.sum @@ -2,17 +2,35 @@ github.com/alpacahq/alpacadecimal v0.0.9 h1:geeT3ZMyfgBV1mrqRDJWQ/4u+FaHep9MmQNF github.com/alpacahq/alpacadecimal v0.0.9/go.mod h1:DmR0Qs+sFJ7nyhfYD0/UUzE+9tCbEF1Wa4iz1lZnlwg= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY= +github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/jsonschema-go v0.4.3 h1:/DBOLZTfDow7pe2GmaJNhltueGTtDKICi8V8p+DQPd0= +github.com/google/jsonschema-go v0.4.3/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= github.com/modelcontextprotocol/go-sdk v1.7.0 h1:yqjY2dsbKAC0LSuWZVBMrHgiG8ukXv6NRo0JiALay44= github.com/modelcontextprotocol/go-sdk v1.7.0/go.mod h1:dL7u98E/zjJTGzEq+j30jQ8K2k1mb6LeAH4inEcSGts= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/segmentio/asm v1.1.3 h1:WM03sfUOENvvKexOLp+pCqgb/WDjsi7EK8gIsICtzhc= +github.com/segmentio/asm v1.1.3/go.mod h1:Ld3L4ZXGNcSLRg4JBsZ3//1+f/TjYl0Mzen/DQy1EJg= +github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfvNt0= +github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0= github.com/shopspring/decimal v1.4.0 h1:bxl37RwXBklmTi0C79JfXCEBD1cqqHt0bbgBAGFp81k= github.com/shopspring/decimal v1.4.0/go.mod h1:gawqmDU56v4yIKSwfBSFip1HdCCXN8/+DMd9qYNcwME= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= golang.org/x/oauth2 v0.35.0 h1:Mv2mzuHuZuY2+bkyWXIHMfhNdJAdwW3FuWeCPYN5GVQ= golang.org/x/oauth2 v0.35.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k= +golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= +golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k= +golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=