diff --git a/README.md b/README.md index 29e50e6..52118bf 100644 --- a/README.md +++ b/README.md @@ -79,20 +79,24 @@ or quota response is unavailable. ## Gemini CLI -Register each Google account, then run the official Gemini CLI through Prism: +Register each Google AI Studio account, then run the official Gemini CLI through Prism: ```sh -prism gemini auth login -prism gemini auth list +prism gemini-ai auth add --name personal +prism gemini-ai auth list prism gemini -p 'Reply with exactly GEMINI_OK.' ``` -Prism rotates across registered Gemini accounts unless `--account` selects one: +Prism rotates across registered AI Studio accounts unless `--account` selects one: ```sh prism gemini --account work-admin -p 'Reply with exactly GEMINI_OK.' ``` +Code Assist OAuth accounts remain available through `prism gemini auth login`, +but require a current Code Assist license and are selected only when no AI +Studio account is registered or when `--account` names one explicitly. + The default model is `gemini-3.7-flash`. For harder software-engineering or multi-step tool-use tasks, select `gemini-3.1-pro-preview` explicitly: diff --git a/internal/cli/gemini.go b/internal/cli/gemini.go index 014c355..09a5b63 100644 --- a/internal/cli/gemini.go +++ b/internal/cli/gemini.go @@ -19,6 +19,8 @@ import ( "path/filepath" "strings" "time" + + "github.com/circlesac/prism-cli/internal/api" ) const defaultGeminiModel = "gemini-3.7-flash" @@ -62,26 +64,58 @@ func runGeminiCommand(ctx context.Context, args []string, stdout io.Writer, stde if err != nil { return err } - if account == "" { - accounts, listErr := client.List(ctx, "gemini") - if listErr != nil { - return listErr + aiStudioAccounts, err := client.List(ctx, "gemini-ai") + if err != nil { + return err + } + codeAssistAccounts, err := client.List(ctx, "gemini") + if err != nil { + return err + } + provider, selectedAccount, err := selectGeminiAccount(account, aiStudioAccounts, codeAssistAccounts) + if err != nil { + return err + } + return runGemini(ctx, client.BaseURL, client.Token, provider, selectedAccount, withDefaultGeminiModel(passthrough), os.Stdin, stdout, stderr) +} + +func selectGeminiAccount(selector string, aiStudioAccounts []api.Credential, codeAssistAccounts []api.Credential) (string, string, error) { + if selector != "" { + matches := make([]struct{ provider, id string }, 0, 2) + for _, group := range []struct { + provider string + accounts []api.Credential + }{{"gemini-ai", aiStudioAccounts}, {"gemini", codeAssistAccounts}} { + for _, account := range group.accounts { + if account.ID == selector || account.Name == selector { + matches = append(matches, struct{ provider, id string }{group.provider, account.ID}) + } + } } - if len(accounts) == 0 { - return errors.New("no Gemini accounts are registered; run 'prism gemini auth login'") + if len(matches) == 0 { + return "", "", fmt.Errorf("Gemini account %q is not registered", selector) } - account, err = rotateProviderAccount("gemini", accounts) - if err != nil { - return err + if len(matches) > 1 { + return "", "", fmt.Errorf("Gemini account %q is ambiguous; use its credential ID", selector) } + return matches[0].provider, matches[0].id, nil } - return runGemini(ctx, client.BaseURL, client.Token, account, withDefaultGeminiModel(passthrough), os.Stdin, stdout, stderr) + provider, accounts := "gemini-ai", aiStudioAccounts + if len(accounts) == 0 { + provider, accounts = "gemini", codeAssistAccounts + } + if len(accounts) == 0 { + return "", "", errors.New("no Gemini accounts are registered; run 'prism gemini-ai auth add'") + } + account, err := rotateProviderAccount(provider, accounts) + return provider, account, err } func runGemini( ctx context.Context, prismURL string, prismCredential string, + provider string, account string, args []string, stdin io.Reader, @@ -92,7 +126,7 @@ func runGemini( if err != nil { return err } - bridge, err := startGeminiBridge(prismURL, prismCredential, account, stderr) + bridge, err := startGeminiBridge(prismURL, prismCredential, provider, account, stderr) if err != nil { return err } @@ -157,7 +191,7 @@ func withDefaultGeminiModel(args []string) []string { return append([]string{"--model", defaultGeminiModel}, args...) } -func startGeminiBridge(prismURL string, prismCredential string, account string, stderr io.Writer) (*geminiBridge, error) { +func startGeminiBridge(prismURL string, prismCredential string, provider string, account string, stderr io.Writer) (*geminiBridge, error) { target, err := url.Parse(prismURL) if err != nil || (target.Scheme != "https" && target.Scheme != "http") || target.Host == "" { return nil, errors.New("Prism URL is invalid") @@ -168,6 +202,9 @@ func startGeminiBridge(prismURL string, prismCredential string, account string, if strings.TrimSpace(account) == "" || strings.ContainsAny(account, "\r\n") { return nil, errors.New("Gemini account selector is invalid") } + if provider != "gemini-ai" && provider != "gemini" { + return nil, errors.New("Gemini account provider is invalid") + } credentialBytes := make([]byte, 32) if _, err := rand.Read(credentialBytes); err != nil { return nil, errors.New("could not create a local Gemini credential") @@ -183,8 +220,10 @@ func startGeminiBridge(prismURL string, prismCredential string, account string, request.Header.Del("Authorization") request.Header.Del("X-Goog-Api-Key") request.Header.Del(localHeaderName) + request.Header.Del("X-Prism-Gemini-Provider") request.Header.Set("Authorization", "Bearer "+prismCredential) request.Header.Set("X-Prism-Gemini-Account", "b64:"+base64.RawURLEncoding.EncodeToString([]byte(account))) + request.Header.Set("X-Prism-Gemini-Provider", provider) } proxy.ErrorLog = log.New(stderr, "prism: ", 0) proxy.ErrorHandler = func(response http.ResponseWriter, _ *http.Request, _ error) { @@ -258,7 +297,8 @@ func printGeminiHelp(output io.Writer) { prism gemini [--account ] [Gemini CLI arguments...] Runs the official Gemini CLI through Prism's Vault-backed Google accounts. -Without --account, registered accounts are selected in balanced rotation. +AI Studio accounts are preferred and selected in balanced rotation. A valid +Code Assist account can still be selected explicitly with --account. The default model is gemini-3.7-flash; use --model gemini-3.1-pro-preview for hard software-engineering and multi-step tool-use work. Run 'gemini --help' for Gemini CLI options.`) diff --git a/internal/cli/gemini_test.go b/internal/cli/gemini_test.go index 33a5c35..b008783 100644 --- a/internal/cli/gemini_test.go +++ b/internal/cli/gemini_test.go @@ -12,6 +12,8 @@ import ( "strings" "sync/atomic" "testing" + + "github.com/circlesac/prism-cli/internal/api" ) func TestGeminiDefaultsTo37FlashAndPreservesExplicitModel(t *testing.T) { @@ -30,13 +32,28 @@ func TestGeminiHelpDocumentsOfficialCLIAccountsAndModels(t *testing.T) { if err := runGeminiCommand(context.Background(), []string{"--help"}, &output, io.Discard); err != nil { t.Fatal(err) } - for _, value := range []string{"official Gemini CLI", "--account", "balanced rotation", "gemini-3.7-flash", "gemini-3.1-pro-preview"} { + for _, value := range []string{"official Gemini CLI", "--account", "balanced rotation", "AI Studio", "gemini-3.7-flash", "gemini-3.1-pro-preview"} { if !strings.Contains(output.String(), value) { t.Fatalf("help omitted %q: %s", value, output.String()) } } } +func TestGeminiAccountSelectionPrefersAIStudioAndSupportsExplicitCodeAssist(t *testing.T) { + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + aiStudio := []api.Credential{{ID: "ai-1", Name: "personal"}, {ID: "ai-2", Name: "work-admin"}} + codeAssist := []api.Credential{{ID: "oauth-1", Name: "enterprise"}} + + provider, account, err := selectGeminiAccount("", aiStudio, codeAssist) + if err != nil || provider != "gemini-ai" || account != "ai-1" { + t.Fatalf("default = %q/%q, error = %v", provider, account, err) + } + provider, account, err = selectGeminiAccount("enterprise", aiStudio, codeAssist) + if err != nil || provider != "gemini" || account != "oauth-1" { + t.Fatalf("explicit = %q/%q, error = %v", provider, account, err) + } +} + func TestGeminiBridgeAuthenticatesLocallyAndSelectsAccount(t *testing.T) { var requests atomic.Int32 upstream := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { @@ -50,6 +67,9 @@ func TestGeminiBridgeAuthenticatesLocallyAndSelectsAccount(t *testing.T) { if request.Header.Get("X-Prism-Gemini-Account") != "b64:cGVyc29uQGV4YW1wbGUuY29t" { t.Errorf("account = %q", request.Header.Get("X-Prism-Gemini-Account")) } + if request.Header.Get("X-Prism-Gemini-Provider") != "gemini-ai" { + t.Errorf("provider = %q", request.Header.Get("X-Prism-Gemini-Provider")) + } if request.Header.Get("X-Goog-Api-Key") != "" || request.Header.Get("X-Prism-Gemini-Bridge") != "" { t.Errorf("private headers leaked") } @@ -58,7 +78,7 @@ func TestGeminiBridgeAuthenticatesLocallyAndSelectsAccount(t *testing.T) { })) defer upstream.Close() - bridge, err := startGeminiBridge(upstream.URL, "circles-secret", "person@example.com", io.Discard) + bridge, err := startGeminiBridge(upstream.URL, "circles-secret", "gemini-ai", "person@example.com", io.Discard) if err != nil { t.Fatal(err) } @@ -98,7 +118,7 @@ func TestRunGeminiUsesOfficialCLIWithGatewayEnvironment(t *testing.T) { upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) defer upstream.Close() var output bytes.Buffer - if err := runGemini(context.Background(), upstream.URL, "circles-secret", "person@example.com", withDefaultGeminiModel([]string{"-p", "hello"}), strings.NewReader(""), &output, io.Discard); err != nil { + if err := runGemini(context.Background(), upstream.URL, "circles-secret", "gemini-ai", "person@example.com", withDefaultGeminiModel([]string{"-p", "hello"}), strings.NewReader(""), &output, io.Discard); err != nil { t.Fatal(err) } for _, value := range []string{"--model\ngemini-3.7-flash\n-p\nhello", "base=http://127.0.0.1:", "headers=X-Prism-Gemini-Bridge:", "settings=/", "trust=true", `"selectedType":"gateway"`, `"useExternal":true`} {