diff --git a/internal/claudebridge/mcp.go b/internal/claudebridge/mcp.go index 3c0860be..da8dd4a6 100644 --- a/internal/claudebridge/mcp.go +++ b/internal/claudebridge/mcp.go @@ -42,8 +42,13 @@ type callbackResponse struct { } // RunMCP runs the hidden stdio MCP subprocess. args are callback URL and tools -// JSON path. It intentionally implements only the MCP methods Claude Code needs. +// JSON path, or, for an agent whose MCP servers are set once for every run +// (Grok), MAGPIE_MCP_CALLBACK and MAGPIE_MCP_TOOLS in the environment it +// hands down. It intentionally implements only the MCP methods Claude Code needs. func RunMCP(args []string) error { + if len(args) == 0 && os.Getenv("MAGPIE_MCP_CALLBACK") != "" { + args = []string{os.Getenv("MAGPIE_MCP_CALLBACK"), os.Getenv("MAGPIE_MCP_TOOLS")} + } if len(args) != 2 { return errors.New("claude MCP helper expects callback URL and tools file") } diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index e1131bb3..1aef02ef 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -213,7 +213,7 @@ func (s *Server) countTokens(w http.ResponseWriter, r *http.Request) { p, model, ok := provider.Resolve(modelOf(body)) // Claude Subscription generations run through the Claude Code binary. Its // OAuth token must not take a direct HTTP side path just for token counting. - if ok && p.Account != nil && (p.Account.Agent == "claude" || p.Account.Agent == "cursor") { + if ok && p.Account != nil && (p.Account.Agent == "claude" || p.Account.Agent == "cursor" || p.Account.Agent == "grok") { req, err := parseAnthropic(body) if err != nil { writeError(w, provider.Anthropic, 400, err.Error()) @@ -410,6 +410,14 @@ func (s *Server) attempt(w http.ResponseWriter, r *http.Request, from provider.P } return s.serveSubscription(w, r, from, "Cursor", model, body, &call.Usage, start) } + // and so does Grok's + if p.Account != nil && p.Account.Agent == "grok" { + call.To = from + start := func(ctx context.Context, req *Request) (*subscriptionRun, <-chan Event, error) { + return s.subscription.startGrok(ctx, req, model) + } + return s.serveSubscription(w, r, from, "Grok", model, body, &call.Usage, start) + } // a backend that only streams gets a non-streaming request translated // (the provider is always streamed on that path) rather than relayed relay := p.Base(from) != "" && (p.Account == nil || !p.Account.Stream || streamOf(body)) diff --git a/internal/gateway/grok_subscription.go b/internal/gateway/grok_subscription.go new file mode 100644 index 00000000..9fecad9a --- /dev/null +++ b/internal/gateway/grok_subscription.go @@ -0,0 +1,332 @@ +package gateway + +// A Grok subscription runs through xAI's Grok Build CLI, the way a Cursor one +// runs through cursor-agent (cursor_subscription.go): the caller's tools are +// an MCP server whose calls park until the caller sends results, and each +// tool_use goes out when the helper calls back — grok reaches MCP tools +// through its own search_tool and use_tool, so its stream never names them. +// +// It runs in a home of magpie's own, where the user's hooks, skills, rules +// and MCP servers (grok's and those it borrows from Claude and Cursor) are +// not. The sign-in is borrowed without being shared: grok there takes its +// token from `magpie grok-token`, which reads the user's own grok home and +// leaves refreshing it to the user's grok. + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "runtime" + "strconv" + "strings" + "sync" + "time" + + "github.com/yetone/magpie/internal/provider" +) + +func (b *subscriptionBridge) startGrok(ctx context.Context, req *Request, model string) (*subscriptionRun, <-chan Event, error) { + binary := provider.GrokExecutable() + if binary == "" { + return nil, nil, errors.New("the Grok CLI is not installed; install it with `curl -fsSL https://x.ai/cli/install.sh | bash` and run `grok login`") + } + exe, err := os.Executable() + if err != nil { + return nil, nil, err + } + auth := grokAuthCommand(exe, provider.GrokHome(), binary) + home, err := grokHome(exe, binary, auth) + if err != nil { + return nil, nil, err + } + tmp, err := os.MkdirTemp("", "magpie-grok-") + if err != nil { + return nil, nil, err + } + cleanup := func() { _ = os.RemoveAll(tmp) } + + tools := bridgeTools(req) + toolsPath := filepath.Join(tmp, "tools.json") + toolBytes, _ := json.Marshal(tools) + if err := os.WriteFile(toolsPath, toolBytes, 0o600); err != nil { + cleanup() + return nil, nil, err + } + promptPath := filepath.Join(tmp, "prompt.txt") + if err := os.WriteFile(promptPath, []byte(renderGrokPrompt(req, tools)), 0o600); err != nil { + cleanup() + return nil, nil, err + } + ws := filepath.Join(tmp, "workspace") + if err := os.MkdirAll(ws, 0o700); err != nil { + cleanup() + return nil, nil, err + } + token := randomToken() + callback := callbackBaseURL() + "/_magpie/claude-mcp/" + token + + // only the ways to the magpie server's tools: no shell, files or web + args := []string{"--prompt-file", promptPath, "--output-format", "streaming-messages-json", "--include-partial-messages", + "-m", model, "--tools", "search_tool,use_tool", "--disable-web-search", "--no-subagents", "--no-plan"} + cmd := exec.CommandContext(context.Background(), binary, args...) + cmd.Dir = ws + cmd.Env = append(grokEnv(os.Environ(), home, auth), "MAGPIE_MCP_CALLBACK="+callback, "MAGPIE_MCP_TOOLS="+toolsPath) + stdout, err := cmd.StdoutPipe() + if err != nil { + cleanup() + return nil, nil, err + } + stderr, err := cmd.StderrPipe() + if err != nil { + cleanup() + return nil, nil, err + } + + run := &subscriptionRun{bridge: b, token: token, model: model, cmd: cmd, tmp: tmp, pending: map[string]chan mcpToolResult{}} + c := &cursorTurn{run: run} + run.onCall = c.called + run.resume = c.resumed + run.begin = func() Event { return Event{Kind: KStart, MsgID: "msg_" + randomToken()[:24], Model: model} } + run.timer = time.AfterFunc(30*time.Minute, run.abort) + segment := run.attach() + run.emit(run.begin()) + b.mu.Lock() + b.runs[token] = run + b.mu.Unlock() + + if err := cmd.Start(); err != nil { + b.removeRun(run) + return nil, nil, err + } + go func() { + _, _ = io.Copy(&lockedWriter{run: run}, io.LimitReader(stderr, 1<<20)) + }() + go c.readGrok(stdout) + go func() { + _ = cmd.Wait() + run.finish() + }() + return run, segment, nil +} + +// grokAuthCommand is how grok in magpie's home asks for the user's token. +func grokAuthCommand(exe, userHome, binary string) string { + return shellQuote(exe) + " grok-token " + shellQuote(userHome) + " " + shellQuote(binary) +} + +func shellQuote(s string) string { + return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'" +} + +var grokHomeMu sync.Mutex + +// grokHome is the home grok runs in: magpie's, kept between runs so its +// caches are, with a config of magpie's and a sign-in through auth. +func grokHome(exe, binary, auth string) (string, error) { + grokHomeMu.Lock() + defer grokHomeMu.Unlock() + base, err := os.UserCacheDir() + if err != nil { + base = os.TempDir() + } + home := filepath.Join(base, "magpie", "grok-home") + dir := filepath.Join(home, ".grok") + if err := os.MkdirAll(dir, 0o700); err != nil { + return "", err + } + cfg := grokConfig(exe) + path := filepath.Join(dir, "config.toml") + if cur, err := os.ReadFile(path); err != nil || !strings.HasPrefix(string(cur), cfg) { + if err := os.WriteFile(path, []byte(cfg), 0o600); err != nil { + return "", err + } + } + // grok signs in to an external provider once; the account it keeps + // must be the one the user's grok is signed in to + user, ok := provider.GrokUser(provider.GrokHome()) + if !ok { + return "", errors.New("Grok is not signed in; run `grok login`") + } + if cur, ok := provider.GrokUser(dir); ok && strings.EqualFold(cur, user) { + return home, nil + } + _ = os.Remove(filepath.Join(dir, "auth.json")) + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + cmd := exec.CommandContext(ctx, binary, "login") + cmd.Dir = home + cmd.Env = grokEnv(os.Environ(), home, auth) + out, err := cmd.CombinedOutput() + if _, ok := provider.GrokUser(dir); !ok { + msg := strings.TrimSpace(string(out)) + if i := strings.LastIndex(msg, "\n"); i >= 0 { + msg = msg[i+1:] + } + if msg == "" && err != nil { + msg = err.Error() + } + return "", fmt.Errorf("grok couldn't sign in: %s", msg) + } + return home, nil +} + +// grokConfig turns off all grok would load that is the user's rather than +// the caller's, and has the magpie server's calls wait as long as the +// caller takes. +func grokConfig(exe string) string { + var b strings.Builder + b.WriteString("# written by magpie for the Grok runs behind its gateway\n") + b.WriteString("[cli]\nauto_update = false\n\n") + b.WriteString("[marketplace]\ndefault_skills_installs_purged = true\nofficial_marketplace_auto_installed = true\n\n") + for _, vendor := range []string{"claude", "cursor"} { + fmt.Fprintf(&b, "[compat.%s]\nmcps = false\nskills = false\nhooks = false\nrules = false\nagents = false\n\n", vendor) + } + b.WriteString("[compat.codex]\nskills = false\n\n") + fmt.Fprintf(&b, "[mcp_servers.magpie]\ncommand = %s\nargs = [\"claude-mcp-helper\"]\nstartup_timeout_sec = 30\ntool_timeout_sec = 86400\n\n", strconv.Quote(exe)) + b.WriteString("[mcp]\nmax_output_bytes = 4194304\n\n") + b.WriteString("[permission]\nallow = [\"MCPTool(magpie__*)\"]\n") + return b.String() +} + +func grokEnv(env []string, home, auth string) []string { + out := make([]string, 0, len(env)+4) + for _, e := range env { + k, _, _ := strings.Cut(e, "=") + k = strings.ToUpper(k) + if k == "HOME" || k == "USERPROFILE" || k == "XAI_API_KEY" || strings.HasPrefix(k, "GROK_") || strings.HasPrefix(k, "MAGPIE_MCP_") { + continue + } + out = append(out, e) + } + out = append(out, "HOME="+home, "GROK_AUTH_PROVIDER_COMMAND="+auth, "GROK_AUTH_EARLY_INVALIDATION_SECS=60") + if runtime.GOOS == "windows" { + out = append(out, "USERPROFILE="+home) + } + return out +} + +// renderGrokPrompt is the conversation as one text, with the caller's +// tools named the way grok calls them. +func renderGrokPrompt(req *Request, tools []bridgeTool) string { + blocks, _ := renderClaudePrompt(req) + var b strings.Builder + if len(tools) > 0 { + names := make([]string, len(tools)) + for i, t := range tools { + names[i] = "magpie__" + t.Name + } + b.WriteString("\nThe only tools in this session are those of the magpie MCP server; Grok's own tools (shell, reading, editing and searching files, the web) are turned off, and this workspace is empty. Call them with use_tool, by these names: " + strings.Join(names, ", ") + ". Call one whenever the conversation needs it.\n\n\n") + } + for _, bl := range blocks { + if t, ok := bl["text"].(string); ok { + b.WriteString(t) + } + } + return b.String() +} + +// readGrok follows grok's stream: Claude's stream events, then the whole +// message again, then the result. +func (c *cursorTurn) readGrok(rd io.Reader) { + s := bufio.NewScanner(rd) + s.Buffer(make([]byte, 64<<10), 64<<20) + var thought, said bool // this message's thinking and text came in pieces + for s.Scan() { + var e struct { + Type string `json:"type"` + Parent *string `json:"parent_tool_use_id"` + IsError bool `json:"is_error"` + Result string `json:"result"` + Errors []string `json:"errors"` + Event json.RawMessage `json:"event"` + Message struct { + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + Thinking string `json:"thinking"` + } `json:"content"` + } `json:"message"` + Usage grokUsage `json:"usage"` + } + if json.Unmarshal(s.Bytes(), &e) != nil || (e.Parent != nil && *e.Parent != "") { + continue + } + switch e.Type { + case "stream_event": + var ev struct { + Type string `json:"type"` + Delta struct { + Type string `json:"type"` + Text string `json:"text"` + Thinking string `json:"thinking"` + } `json:"delta"` + } + if json.Unmarshal(e.Event, &ev) != nil { + continue + } + switch { + case ev.Type == "message_start": + thought, said = false, false + case ev.Delta.Type == "thinking_delta" && ev.Delta.Thinking != "": + thought = true + c.run.emit(Event{Kind: KThink, Text: ev.Delta.Thinking}) + case ev.Delta.Type == "text_delta" && ev.Delta.Text != "": + said = true + c.run.emit(Event{Kind: KText, Text: ev.Delta.Text}) + } + case "assistant": + // what came whole and not in pieces + for _, p := range e.Message.Content { + if p.Type == "thinking" && p.Thinking != "" && !thought { + c.run.emit(Event{Kind: KThink, Text: p.Thinking}) + } + if p.Type == "text" && p.Text != "" && !said { + c.run.emit(Event{Kind: KText, Text: p.Text}) + } + } + case "result": + c.mu.Lock() + c.done = true + c.mu.Unlock() + if e.IsError { + msg := strings.Join(e.Errors, "; ") + if msg == "" { + msg = e.Result + } + if msg == "" { + msg = c.run.lastError("grok ended with an error") + } + c.run.emit(Event{Kind: KError, Text: msg}) + } else { + c.run.emit(Event{Kind: KUsage, Usage: e.Usage.usage()}) + c.run.emit(Event{Kind: KStop, Stop: "stop"}) + } + c.run.endSegment() + time.AfterFunc(2*time.Second, c.run.abort) + } + } + c.mu.Lock() + done := c.done + c.mu.Unlock() + if !done { + c.run.emit(Event{Kind: KError, Text: c.run.lastError("grok ended without an answer")}) + } +} + +// grokUsage is the Messages API's usage, whose input leaves out the cache. +type grokUsage struct { + Input int `json:"input_tokens"` + Output int `json:"output_tokens"` + CacheRead int `json:"cache_read_input_tokens"` + CacheWrite int `json:"cache_creation_input_tokens"` +} + +func (u grokUsage) usage() Usage { + return Usage{Input: u.Input, Output: u.Output, CacheRead: u.CacheRead, CacheWrite: u.CacheWrite} +} diff --git a/internal/gateway/grok_subscription_test.go b/internal/gateway/grok_subscription_test.go new file mode 100644 index 00000000..21e8a3f7 --- /dev/null +++ b/internal/gateway/grok_subscription_test.go @@ -0,0 +1,115 @@ +package gateway + +import ( + "strings" + "testing" + "time" +) + +func TestGrokEnvRunsInMagpiesHome(t *testing.T) { + got := strings.Join(grokEnv([]string{"PATH=/bin", "HOME=/Users/me", "GROK_HOME=/x", "XAI_API_KEY=k", "GROK_API_KEY=k", "MAGPIE_MCP_CALLBACK=old", "KEEP=1"}, "/cache/home", "'magpie' grok-token"), "\n") + for _, bad := range []string{"HOME=/Users/me", "GROK_HOME=", "XAI_API_KEY=", "GROK_API_KEY=", "MAGPIE_MCP_CALLBACK="} { + if strings.Contains(got, bad) { + t.Fatalf("kept %s in %q", bad, got) + } + } + for _, want := range []string{"PATH=/bin", "KEEP=1", "HOME=/cache/home", "GROK_AUTH_PROVIDER_COMMAND='magpie' grok-token"} { + if !strings.Contains(got, want) { + t.Fatalf("missing %s in %q", want, got) + } + } +} + +func TestGrokConfigKeepsOnlyTheCallersTools(t *testing.T) { + cfg := grokConfig(`/Applications/magpie.app/Contents/MacOS/magpie`) + for _, want := range []string{ + "[mcp_servers.magpie]\ncommand = \"/Applications/magpie.app/Contents/MacOS/magpie\"\nargs = [\"claude-mcp-helper\"]", + "tool_timeout_sec = 86400", + "[compat.claude]\nmcps = false\nskills = false\nhooks = false", + "[compat.cursor]\nmcps = false", + "auto_update = false", + `allow = ["MCPTool(magpie__*)"]`, + } { + if !strings.Contains(cfg, want) { + t.Fatalf("config lacks %q:\n%s", want, cfg) + } + } +} + +func TestGrokAuthCommandQuotes(t *testing.T) { + got := grokAuthCommand("/Apps/it's magpie", "/Users/me/.grok", "/Users/me/.grok/bin/grok") + if got != `'/Apps/it'\''s magpie' grok-token '/Users/me/.grok' '/Users/me/.grok/bin/grok'` { + t.Fatalf("got %s", got) + } +} + +func TestRenderGrokPromptNamesTools(t *testing.T) { + req := &Request{Messages: []Message{{Role: "user", Parts: []Part{{Kind: Text, Text: "weather?"}}}}} + got := renderGrokPrompt(req, []bridgeTool{{Name: "get_weather"}, {Name: "read"}}) + if !strings.Contains(got, "magpie__get_weather, magpie__read") || !strings.Contains(got, "use_tool") || !strings.Contains(got, "weather?") { + t.Fatalf("prompt = %q", got) + } + if strings.Contains(renderGrokPrompt(req, nil), "external_system_instructions") { + t.Fatal("no tools, yet told of them") + } +} + +func TestReadGrokStream(t *testing.T) { + stream := strings.Join([]string{ + `{"type":"system","subtype":"init","tools":["search_tool","use_tool"]}`, + `{"type":"stream_event","event":{"type":"message_start","message":{"id":"msg_0"}}}`, + `{"type":"stream_event","event":{"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"hmm"}}}`, + `{"type":"stream_event","event":{"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"po"}}}`, + `{"type":"stream_event","event":{"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"ng"}}}`, + `{"type":"stream_event","event":{"type":"content_block_start","index":2,"content_block":{"type":"tool_use","name":"use_tool","input":{}}}}`, + `{"type":"assistant","message":{"content":[{"type":"thinking","thinking":"hmm"},{"type":"text","text":"pong"}]}}`, + `{"type":"stream_event","event":{"type":"message_start","message":{"id":"msg_1"}}}`, + `{"type":"assistant","message":{"content":[{"type":"thinking","thinking":"whole"}]}}`, + `{"type":"stream_event","parent_tool_use_id":"x","event":{"type":"content_block_delta","delta":{"type":"text_delta","text":"sub"}}}`, + `{"type":"result","subtype":"success","is_error":false,"result":"pong","usage":{"input_tokens":262,"output_tokens":27,"cache_read_input_tokens":28416,"cache_creation_input_tokens":0}}`, + }, "\n") + run := &subscriptionRun{pending: map[string]chan mcpToolResult{}} + run.timer = time.AfterFunc(time.Hour, func() {}) + c := &cursorTurn{run: run} + seg := run.attach() + go c.readGrok(strings.NewReader(stream)) + var think, text string + var usage Usage + var stop string + for ev := range seg { + switch ev.Kind { + case KThink: + think += ev.Text + "|" + case KText: + text += ev.Text + case KUsage: + usage = ev.Usage + case KStop: + stop = ev.Stop + case KError: + t.Fatalf("error %q", ev.Text) + } + } + if think != "hmm|whole|" || text != "pong" || stop != "stop" { + t.Fatalf("think=%q text=%q stop=%q", think, text, stop) + } + if usage.Input != 262 || usage.CacheRead != 28416 || usage.Output != 27 { + t.Fatalf("usage = %+v", usage) + } +} + +func TestReadGrokError(t *testing.T) { + run := &subscriptionRun{pending: map[string]chan mcpToolResult{}} + c := &cursorTurn{run: run} + seg := run.attach() + go c.readGrok(strings.NewReader(`{"type":"result","subtype":"error_during_execution","is_error":true,"errors":["rate limited"]}`)) + var got string + for ev := range seg { + if ev.Kind == KError { + got = ev.Text + } + } + if got != "rate limited" { + t.Fatalf("error = %q", got) + } +} diff --git a/internal/gui/assets/app.js b/internal/gui/assets/app.js index ed78c8a3..1a8c4c38 100644 --- a/internal/gui/assets/app.js +++ b/internal/gui/assets/app.js @@ -671,6 +671,7 @@ function accountPlan(a) { if (a.agent === "copilot") return "GitHub"; if (a.agent === "claude") return "Claude" + (a.plan ? " " + a.plan[0].toUpperCase() + a.plan.slice(1) : ""); if (a.agent === "cursor") return "Cursor" + (a.plan ? " " + a.plan[0].toUpperCase() + a.plan.slice(1) : ""); + if (a.agent === "grok") return "Grok"; return t("signed in"); } @@ -1890,6 +1891,8 @@ const SUBS = [ { agent: "codex", name: "ChatGPT", icon: "openai", plans: "Plus · Pro · Business" }, // cursor-agent keeps one account; signing in again replaces it { agent: "cursor", name: "Cursor", icon: "cursor", plans: "Pro · Ultra · Teams", single: true }, + // so does the Grok CLI + { agent: "grok", name: "Grok", icon: "xai", plans: "SuperGrok · X Premium+", single: true }, ]; const subOf = (agent) => SUBS.find((x) => x.agent === agent); let signing = null; // the sign-in under way: { id, agent, url, state, error } diff --git a/internal/gui/providers.go b/internal/gui/providers.go index b4fe6c7f..c66a1967 100644 --- a/internal/gui/providers.go +++ b/internal/gui/providers.go @@ -147,6 +147,8 @@ func providerInfo(p provider.Provider, agents []*agent.Agent) providerJSON { } else if a.Agent == "cursor" { // cursor-agent runs behind the gateway, not as an agent magpie configures out.Account.Name, out.Account.Icon = "Cursor CLI", "cursor" + } else if a.Agent == "grok" { + out.Account.Name, out.Account.Icon = "Grok CLI", "xai" } out.Account.Logins = provider.Logins(a.Agent) } diff --git a/internal/provider/account.go b/internal/provider/account.go index b102b6bc..711353a1 100644 --- a/internal/provider/account.go +++ b/internal/provider/account.go @@ -545,6 +545,9 @@ func Accounts() []Provider { if p, ok := cursorAccount(); ok { out = append(out, p) } + if p, ok := grokAccount(); ok { + out = append(out, p) + } return out } diff --git a/internal/provider/cursor.go b/internal/provider/cursor.go index 826ab6e5..33afdd46 100644 --- a/internal/provider/cursor.go +++ b/internal/provider/cursor.go @@ -159,65 +159,8 @@ func startCursorSignIn(s *signInFlow) error { if path == "" { return errorf("install Cursor's CLI first: curl https://cursor.com/install -fsS | bash") } - ctx, cancel := context.WithCancel(context.Background()) - cmd := exec.CommandContext(ctx, path, "login") - cmd.Env = append(os.Environ(), "NO_OPEN_BROWSER=1") - out, err := cmd.StdoutPipe() - if err != nil { - cancel() - return err - } - cmd.Stderr = cmd.Stdout - if err := cmd.Start(); err != nil { - cancel() - return err - } - s.stop = cancel - got := make(chan string, 1) - go func() { - sc := bufio.NewScanner(out) - sent := false - var tail []string - for sc.Scan() { - line := ansi.ReplaceAllString(sc.Text(), "") - if u := cursorLoginURL.FindString(line); u != "" && !sent { - sent = true - got <- u - } - if strings.TrimSpace(line) != "" { - tail = append(tail, strings.TrimSpace(line)) - } - } - if !sent { - close(got) - } - err := cmd.Wait() - cancel() + return runCLISignIn(s, "cursor-agent login", append(os.Environ(), "NO_OPEN_BROWSER=1"), func() (string, string, bool) { forgetCursorStatus() - forgetAccountCaches() - if err == nil { - if user, plan, ok := askCursorIdentity(); ok { - s.finish(SignInState{State: "done", User: user, Plan: plan, Using: true}) - return - } - } - msg := "cursor-agent login didn't finish" - if n := len(tail); n > 0 { - msg = tail[n-1] - } - s.finish(SignInState{State: "failed", Error: msg}) - }() - select { - case u, ok := <-got: - if !ok { - return errorf("cursor-agent login gave no link to open") - } - s.mu.Lock() - s.st.URL = u - s.mu.Unlock() - return nil - case <-time.After(30 * time.Second): - cancel() - return errorf("cursor-agent login gave no link to open") - } + return askCursorIdentity() + }, path, "login") } diff --git a/internal/provider/grok.go b/internal/provider/grok.go new file mode 100644 index 00000000..a6061735 --- /dev/null +++ b/internal/provider/grok.go @@ -0,0 +1,276 @@ +package provider + +// A Grok subscription (SuperGrok, X Premium+) is used the way a Cursor one +// is: through xAI's own Grok Build CLI, which magpie runs behind the gateway +// (see gateway/grok_subscription.go). Here is who the CLI is signed in to, +// the models it offers, the sign-in, which is the CLI's own `login`, and the +// token the gateway's runs borrow from it. + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "regexp" + "sort" + "strings" + "sync" + "time" + + "github.com/yetone/magpie/internal/catalog" +) + +// GrokExecutable finds the Grok Build CLI; a var so tests can fake it. +var GrokExecutable = func() string { + home, _ := os.UserHomeDir() + if p := filepath.Join(GrokHome(), "bin", "grok"); isFile(p) { + return p + } + if p, err := exec.LookPath("grok"); err == nil && isGrokBuild(p) { + return p + } + if p := filepath.Join(home, ".local", "bin", "grok"); isFile(p) && isGrokBuild(p) { + return p + } + return "" +} + +func isFile(p string) bool { + st, err := os.Stat(p) + return err == nil && !st.IsDir() +} + +// isGrokBuild tells xAI's grok from any other program by that name: its +// installer keeps it under the grok home. +func isGrokBuild(path string) bool { + real, err := filepath.EvalSymlinks(path) + return err == nil && strings.Contains(filepath.ToSlash(real), "/.grok/") +} + +// GrokHome is where the CLI keeps its sign-in and settings. +func GrokHome() string { + if h := os.Getenv("GROK_HOME"); h != "" { + return h + } + home, _ := os.UserHomeDir() + return filepath.Join(home, ".grok") +} + +// grokCredential is the CLI's sign-in, as its auth.json keeps it. +type grokCredential struct { + Key string `json:"key"` + Email string `json:"email"` + ExpiresAt time.Time `json:"expires_at"` + Issuer string `json:"oidc_issuer"` +} + +// readGrokCredential reads the sign-in the CLI keeps in its home. magpie +// only reads it: the CLI refreshes it with a token that changes each time. +func readGrokCredential(home string) (grokCredential, bool) { + var all map[string]grokCredential + if !readJSON(filepath.Join(home, "auth.json"), &all) { + return grokCredential{}, false + } + keys := make([]string, 0, len(all)) + for k := range all { + keys = append(keys, k) + } + sort.Strings(keys) + for _, k := range keys { + if c := all[k]; c.Key != "" { + return c, true + } + } + return grokCredential{}, false +} + +func grokAccount() (Provider, bool) { + if GrokExecutable() == "" { + return Provider{}, false + } + c, ok := readGrokCredential(GrokHome()) + if !ok { + return Provider{}, false + } + acct := &Account{Agent: "grok", User: c.Email} + acct.models = func() []catalog.Model { return []catalog.Model{{ID: "grok-4.7", Name: "grok-4.7"}} } + acct.fetch = func(ctx context.Context) ([]catalog.Model, error) { + ms, err := grokModels(ctx) + if err != nil { + return nil, err + } + return ms, catalog.SaveLive("grok", "", ms) + } + return Provider{ID: "grok", Name: "Grok", Icon: "xai", Website: "https://x.ai/cli", Account: acct}, true +} + +var grokModelL = regexp.MustCompile(`^[*-]\s+([A-Za-z0-9][\w.:-]*)`) + +// grokModels lists what the account can use, as `grok models` prints it. +func grokModels(ctx context.Context) ([]catalog.Model, error) { + path := GrokExecutable() + if path == "" { + return nil, errorf("the Grok CLI is not installed") + } + ctx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + cmd := exec.CommandContext(ctx, path, "models") + cmd.Dir, _ = os.UserHomeDir() + out, err := cmd.Output() + if err != nil { + return nil, errorf("grok models: %v", err) + } + return parseGrokModels(string(out)), nil +} + +func parseGrokModels(out string) []catalog.Model { + var ms []catalog.Model + s := bufio.NewScanner(strings.NewReader(ansi.ReplaceAllString(out, ""))) + for s.Scan() { + if m := grokModelL.FindStringSubmatch(strings.TrimSpace(s.Text())); m != nil { + ms = append(ms, catalog.Model{ID: m[1], Name: m[1]}) + } + } + return ms +} + +// grokRefreshMargin is how close to its expiry a borrowed token has the CLI +// refresh it first. +const grokRefreshMargin = 5 * time.Minute + +var grokRefresh sync.Mutex + +// GrokToken writes the CLI's access token for a Grok run in magpie's home, +// as an external auth provider answers (GROK_AUTH_PROVIDER_COMMAND). A +// token about to expire is first refreshed by the CLI in its own home: a +// `grok models` there renews it the way any use of the CLI does; so is one +// the run has found expired. +func GrokToken(w io.Writer, home, binary string, expired bool) error { + c, ok := readGrokCredential(home) + if ok && (expired || time.Until(c.ExpiresAt) < grokRefreshMargin) && binary != "" { + grokRefresh.Lock() + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + cmd := exec.CommandContext(ctx, binary, "models") + cmd.Dir = filepath.Dir(home) + cmd.Env = grokOwnEnv(os.Environ(), home) + _ = cmd.Run() + cancel() + grokRefresh.Unlock() + c, ok = readGrokCredential(home) + } + if !ok { + return errorf("Grok is not signed in; run `grok login`") + } + left := int(time.Until(c.ExpiresAt).Seconds()) + if c.ExpiresAt.IsZero() { + left = 3600 + } + if left <= 0 { + return errorf("Grok's sign-in has expired; run `grok login`") + } + return json.NewEncoder(w).Encode(map[string]any{"access_token": c.Key, "expires_in": left, "issuer": c.Issuer}) +} + +// grokOwnEnv runs the CLI as the user runs it, in its own home. +func grokOwnEnv(env []string, home string) []string { + out := make([]string, 0, len(env)+1) + for _, e := range env { + k, _, _ := strings.Cut(e, "=") + switch strings.ToUpper(k) { + case "GROK_HOME", "GROK_AUTH_PROVIDER_COMMAND", "GROK_AUTH_EXPIRED": + continue + } + out = append(out, e) + } + return append(out, "GROK_HOME="+home) +} + +// startGrokSignIn runs `grok login` with its device code, hands its link to +// the window, and finishes when the CLI says the account is in. +func startGrokSignIn(s *signInFlow) error { + path := GrokExecutable() + if path == "" { + return errorf("install the Grok CLI first: curl -fsSL https://x.ai/cli/install.sh | bash") + } + return runCLISignIn(s, "grok login", nil, func() (string, string, bool) { + c, ok := readGrokCredential(GrokHome()) + return c.Email, "", ok + }, path, "login", "--device-auth") +} + +// runCLISignIn runs an agent's own login command, hands the first link it +// prints to the window, and finishes when the command does and identity +// says who is signed in. +func runCLISignIn(s *signInFlow, what string, env []string, identity func() (user, plan string, ok bool), path string, args ...string) error { + ctx, cancel := context.WithCancel(context.Background()) + cmd := exec.CommandContext(ctx, path, args...) + cmd.Dir, _ = os.UserHomeDir() + cmd.Env = env + out, err := cmd.StdoutPipe() + if err != nil { + cancel() + return err + } + cmd.Stderr = cmd.Stdout + if err := cmd.Start(); err != nil { + cancel() + return err + } + s.stop = cancel + got := make(chan string, 1) + go func() { + sc := bufio.NewScanner(out) + sent := false + var tail []string + for sc.Scan() { + line := ansi.ReplaceAllString(sc.Text(), "") + if u := cursorLoginURL.FindString(line); u != "" && !sent { + sent = true + got <- u + } + if strings.TrimSpace(line) != "" { + tail = append(tail, strings.TrimSpace(line)) + } + } + if !sent { + close(got) + } + err := cmd.Wait() + cancel() + forgetAccountCaches() + if err == nil { + if user, plan, ok := identity(); ok { + s.finish(SignInState{State: "done", User: user, Plan: plan, Using: true}) + return + } + } + msg := what + " didn't finish" + if n := len(tail); n > 0 { + msg = tail[n-1] + } + s.finish(SignInState{State: "failed", Error: msg}) + }() + select { + case u, ok := <-got: + if !ok { + return fmt.Errorf("%s gave no link to open", what) + } + s.mu.Lock() + s.st.URL = u + s.mu.Unlock() + return nil + case <-time.After(30 * time.Second): + cancel() + return fmt.Errorf("%s gave no link to open", what) + } +} + +// GrokUser is who the grok with this home is signed in to. +func GrokUser(home string) (string, bool) { + c, ok := readGrokCredential(home) + return c.Email, ok +} diff --git a/internal/provider/grok_test.go b/internal/provider/grok_test.go new file mode 100644 index 00000000..3967da16 --- /dev/null +++ b/internal/provider/grok_test.go @@ -0,0 +1,52 @@ +package provider + +import ( + "bytes" + "encoding/json" + "os" + "path/filepath" + "testing" + "time" +) + +func TestParseGrokModels(t *testing.T) { + out := "You are logged in with grok.com.\n\nDefault model: grok-4.7\n\nAvailable models:\n * grok-4.7 (default)\n - grok-4.7-build-fast\n - grok-4.6\n" + ms := parseGrokModels(out) + if len(ms) != 3 || ms[0].ID != "grok-4.7" || ms[1].ID != "grok-4.7-build-fast" || ms[2].ID != "grok-4.6" { + t.Fatalf("models = %+v", ms) + } +} + +func TestGrokTokenReadsTheCLIsSignIn(t *testing.T) { + home := t.TempDir() + exp := time.Now().Add(2 * time.Hour).UTC() + auth := map[string]any{"https://auth.x.ai::u1": map[string]any{ + "key": "tok", "email": "me@example.com", "auth_mode": "oidc", "refresh_token": "r", + "expires_at": exp.Format(time.RFC3339Nano), "oidc_issuer": "https://auth.x.ai"}} + b, _ := json.Marshal(auth) + if err := os.WriteFile(filepath.Join(home, "auth.json"), b, 0o600); err != nil { + t.Fatal(err) + } + if u, ok := GrokUser(home); !ok || u != "me@example.com" { + t.Fatalf("user = %q %v", u, ok) + } + var out bytes.Buffer + if err := GrokToken(&out, home, "", false); err != nil { + t.Fatal(err) + } + var got struct { + Token string `json:"access_token"` + Expires int `json:"expires_in"` + Issuer string `json:"issuer"` + Refresh string `json:"refresh_token"` + } + if err := json.Unmarshal(out.Bytes(), &got); err != nil { + t.Fatal(err) + } + if got.Token != "tok" || got.Issuer != "https://auth.x.ai" || got.Refresh != "" || got.Expires < 7000 || got.Expires > 7200 { + t.Fatalf("token answer = %+v", got) + } + if err := GrokToken(&out, t.TempDir(), "", false); err == nil { + t.Fatal("no sign-in, yet a token") + } +} diff --git a/internal/provider/signin.go b/internal/provider/signin.go index 771971ec..01ba734b 100644 --- a/internal/provider/signin.go +++ b/internal/provider/signin.go @@ -125,6 +125,11 @@ func StartSignIn(agent string) (SignInState, error) { if err := startCursorSignIn(s); err != nil { return SignInState{}, err } + case "grok": + // so is Grok: its CLI signs in with a device code + if err := startGrokSignIn(s); err != nil { + return SignInState{}, err + } default: return SignInState{}, fmt.Errorf("magpie can't sign in to %s accounts", agent) } diff --git a/main.go b/main.go index f929fc0d..a99c9458 100644 --- a/main.go +++ b/main.go @@ -16,6 +16,7 @@ import ( "github.com/yetone/magpie/internal/gateway" "github.com/yetone/magpie/internal/netproxy" "github.com/yetone/magpie/internal/profile" + "github.com/yetone/magpie/internal/provider" "github.com/yetone/magpie/internal/settings" "github.com/yetone/magpie/internal/tui" "github.com/yetone/magpie/internal/update" @@ -78,6 +79,11 @@ func main() { } func run(args []string) error { + // internal: the auth provider of a Grok run behind the gateway, asked + // for a token often; nothing else of magpie's needs to start for it + if len(args) == 3 && args[0] == "grok-token" { + return provider.GrokToken(os.Stdout, args[1], args[2], os.Getenv("GROK_AUTH_EXPIRED") == "1") + } settings.Migrate() agent.RenameLegacy() if len(args) == 0 {