From 853e023605a88a41efa5a7d5491135c563e0faa9 Mon Sep 17 00:00:00 2001 From: Baller Date: Mon, 2 Mar 2026 23:28:34 -0500 Subject: [PATCH] feat(auth): update related functionality --- .env.example | 17 ++++ cmd/picoclaw/internal/auth/helpers.go | 57 ++++++------- cmd/picoclaw/internal/auth/login.go | 6 +- pkg/auth/{usage.go => anthropic_usage.go} | 27 +++++-- pkg/auth/anthropic_usage_test.go | 98 +++++++++++++++++++++++ pkg/auth/token.go | 3 - pkg/auth/token_test.go | 61 ++++++++++++++ pkg/providers/anthropic/provider.go | 7 +- pkg/providers/anthropic/provider_test.go | 59 ++++++++++++++ 9 files changed, 293 insertions(+), 42 deletions(-) create mode 100644 .env.example rename pkg/auth/{usage.go => anthropic_usage.go} (57%) create mode 100644 pkg/auth/anthropic_usage_test.go create mode 100644 pkg/auth/token_test.go diff --git a/.env.example b/.env.example new file mode 100644 index 000000000..8c2cfbbdd --- /dev/null +++ b/.env.example @@ -0,0 +1,17 @@ +# ── LLM Provider ────────────────────────── +# Uncomment and set the API key for your provider +# OPENROUTER_API_KEY=sk-or-v1-xxx +# ZHIPU_API_KEY=xxx +# ANTHROPIC_API_KEY=sk-ant-xxx +# OPENAI_API_KEY=sk-xxx +# GEMINI_API_KEY=xxx +# CLAUDE_CODE_OAUTH=xxx +# ── Chat Channel ────────────────────────── +# TELEGRAM_BOT_TOKEN=123456:ABC... +# DISCORD_BOT_TOKEN=xxx + +# ── Web Search (optional) ──────────────── +# BRAVE_SEARCH_API_KEY=BSA... + +# ── Timezone ────────────────────────────── +TZ=Asia/Tokyo diff --git a/cmd/picoclaw/internal/auth/helpers.go b/cmd/picoclaw/internal/auth/helpers.go index f2236ce3d..e3db7f333 100644 --- a/cmd/picoclaw/internal/auth/helpers.go +++ b/cmd/picoclaw/internal/auth/helpers.go @@ -17,13 +17,14 @@ import ( ) const supportedProvidersMsg = "supported providers: openai, anthropic, google-antigravity" +const defaultAnthropicModel = "claude-sonnet-4.6" -func authLoginCmd(provider string, useDeviceCode bool, setupToken bool) error { +func authLoginCmd(provider string, useDeviceCode bool, useOauth bool) error { switch provider { case "openai": return authLoginOpenAI(useDeviceCode) case "anthropic": - return authLoginAnthropic(setupToken) + return authLoginAnthropic(useOauth) case "google-antigravity", "antigravity": return authLoginGoogleAntigravity() default: @@ -164,32 +165,34 @@ func authLoginGoogleAntigravity() error { return nil } -func authLoginAnthropic(setupToken bool) error { - if setupToken { +func authLoginAnthropic(useOauth bool) error { + if useOauth { return authLoginAnthropicSetupToken() } fmt.Println("Anthropic login method:") fmt.Println(" 1) Setup token (from `claude setup-token`) (Recommended)") fmt.Println(" 2) API key (from console.anthropic.com)") - fmt.Print("Choose [1]: ") scanner := bufio.NewScanner(os.Stdin) - choice := "1" - if scanner.Scan() { - text := strings.TrimSpace(scanner.Text()) - if text != "" { - choice = text + for { + fmt.Print("Choose [1]: ") + choice := "1" + if scanner.Scan() { + text := strings.TrimSpace(scanner.Text()) + if text != "" { + choice = text + } } - } - switch choice { - case "1": - return authLoginAnthropicSetupToken() - case "2": - return authLoginPasteToken("anthropic") - default: - return fmt.Errorf("invalid choice: %s", choice) + switch choice { + case "1": + return authLoginAnthropicSetupToken() + case "2": + return authLoginPasteToken("anthropic") + default: + fmt.Printf("Invalid choice: %s. Please enter 1 or 2.\n", choice) + } } } @@ -217,21 +220,22 @@ func authLoginAnthropicSetupToken() error { } if !found { appCfg.ModelList = append(appCfg.ModelList, config.ModelConfig{ - ModelName: "claude-sonnet-4.6", - Model: "anthropic/claude-sonnet-4.6", + ModelName: defaultAnthropicModel, + Model: "anthropic/" + defaultAnthropicModel, AuthMethod: "oauth", }) + // Only set default model if user has no default configured yet + if appCfg.Agents.Defaults.GetModelName() == "" { + appCfg.Agents.Defaults.ModelName = defaultAnthropicModel + } } - appCfg.Agents.Defaults.ModelName = "claude-sonnet-4.6" - if err := config.SaveConfig(internal.GetConfigPath(), appCfg); err != nil { return fmt.Errorf("could not update config: %w", err) } } fmt.Println("Setup token saved for Anthropic!") - fmt.Println("Default model set to: claude-sonnet-4.6") return nil } @@ -290,13 +294,12 @@ func authLoginPasteToken(provider string) error { } if !found { appCfg.ModelList = append(appCfg.ModelList, config.ModelConfig{ - ModelName: "claude-sonnet-4.6", - Model: "anthropic/claude-sonnet-4.6", + ModelName: defaultAnthropicModel, + Model: "anthropic/" + defaultAnthropicModel, AuthMethod: "token", }) + appCfg.Agents.Defaults.ModelName = defaultAnthropicModel } - // Update default model - appCfg.Agents.Defaults.ModelName = "claude-sonnet-4.6" case "openai": appCfg.Providers.OpenAI.AuthMethod = "token" // Update ModelList diff --git a/cmd/picoclaw/internal/auth/login.go b/cmd/picoclaw/internal/auth/login.go index 1cc238896..223823489 100644 --- a/cmd/picoclaw/internal/auth/login.go +++ b/cmd/picoclaw/internal/auth/login.go @@ -6,7 +6,7 @@ func newLoginCommand() *cobra.Command { var ( provider string useDeviceCode bool - setupToken bool + useOauth bool ) cmd := &cobra.Command{ @@ -14,13 +14,13 @@ func newLoginCommand() *cobra.Command { Short: "Login via OAuth or paste token", Args: cobra.NoArgs, RunE: func(cmd *cobra.Command, _ []string) error { - return authLoginCmd(provider, useDeviceCode, setupToken) + return authLoginCmd(provider, useDeviceCode, useOauth) }, } cmd.Flags().StringVarP(&provider, "provider", "p", "", "Provider to login with (openai, anthropic)") cmd.Flags().BoolVar(&useDeviceCode, "device-code", false, "Use device code flow (for headless environments)") - cmd.Flags().BoolVar(&setupToken, "setup-token", false, "Use setup-token flow for Anthropic (from `claude setup-token`)") + cmd.Flags().BoolVar(&useOauth, "setup-token", false, "Use setup-token flow for Anthropic (from `claude setup-token`)") _ = cmd.MarkFlagRequired("provider") return cmd diff --git a/pkg/auth/usage.go b/pkg/auth/anthropic_usage.go similarity index 57% rename from pkg/auth/usage.go rename to pkg/auth/anthropic_usage.go index ded3d9723..3fffc3aa8 100644 --- a/pkg/auth/usage.go +++ b/pkg/auth/anthropic_usage.go @@ -8,19 +8,30 @@ import ( "time" ) +const ( + anthropicBetaHeader = "oauth-2025-04-20" + anthropicAPIVersion = "2023-06-01" +) + +// anthropicUsageURL is the endpoint for fetching OAuth usage stats. +// It is a var (not const) to allow overriding in tests. +var anthropicUsageURL = "https://api.anthropic.com/api/oauth/usage" + +func setAnthropicUsageURL(url string) { anthropicUsageURL = url } + type AnthropicUsage struct { FiveHourUtilization float64 SevenDayUtilization float64 } func FetchAnthropicUsage(token string) (*AnthropicUsage, error) { - req, err := http.NewRequest("GET", "https://api.anthropic.com/api/oauth/usage", nil) + req, err := http.NewRequest("GET", anthropicUsageURL, nil) if err != nil { return nil, err } req.Header.Set("Authorization", "Bearer "+token) - req.Header.Set("anthropic-version", "2023-06-01") - req.Header.Set("anthropic-beta", "oauth-2025-04-20") + req.Header.Set("anthropic-version", anthropicAPIVersion) + req.Header.Set("anthropic-beta", anthropicBetaHeader) client := &http.Client{Timeout: 10 * time.Second} resp, err := client.Do(req) @@ -29,13 +40,15 @@ func FetchAnthropicUsage(token string) (*AnthropicUsage, error) { } defer resp.Body.Close() - body, _ := io.ReadAll(resp.Body) - - if resp.StatusCode == http.StatusForbidden { - return nil, fmt.Errorf("insufficient scope: usage endpoint requires oauth scope") + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("reading usage response: %w", err) } if resp.StatusCode != http.StatusOK { + if resp.StatusCode == http.StatusForbidden { + return nil, fmt.Errorf("insufficient scope: usage endpoint requires oauth scope") + } return nil, fmt.Errorf("usage request failed (%d): %s", resp.StatusCode, string(body)) } diff --git a/pkg/auth/anthropic_usage_test.go b/pkg/auth/anthropic_usage_test.go new file mode 100644 index 000000000..2c11ac05a --- /dev/null +++ b/pkg/auth/anthropic_usage_test.go @@ -0,0 +1,98 @@ +package auth + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestFetchAnthropicUsage_Success(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "Bearer test-token" { + t.Errorf("Authorization = %q, want %q", got, "Bearer test-token") + } + if got := r.Header.Get("anthropic-beta"); got != anthropicBetaHeader { + t.Errorf("anthropic-beta = %q, want %q", got, anthropicBetaHeader) + } + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"five_hour":{"utilization":0.42},"seven_day":{"utilization":0.85}}`)) + })) + defer srv.Close() + + // Temporarily override the URL by using the test server + origURL := anthropicUsageURL + defer func() { setAnthropicUsageURL(origURL) }() + setAnthropicUsageURL(srv.URL) + + usage, err := FetchAnthropicUsage("test-token") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if usage.FiveHourUtilization != 0.42 { + t.Errorf("FiveHourUtilization = %v, want 0.42", usage.FiveHourUtilization) + } + if usage.SevenDayUtilization != 0.85 { + t.Errorf("SevenDayUtilization = %v, want 0.85", usage.SevenDayUtilization) + } +} + +func TestFetchAnthropicUsage_Forbidden(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + w.Write([]byte(`{"error":"forbidden"}`)) + })) + defer srv.Close() + + origURL := anthropicUsageURL + defer func() { setAnthropicUsageURL(origURL) }() + setAnthropicUsageURL(srv.URL) + + _, err := FetchAnthropicUsage("test-token") + if err == nil { + t.Fatal("expected error for 403, got nil") + } + if !strings.Contains(err.Error(), "insufficient scope") { + t.Errorf("expected 'insufficient scope' error, got %q", err.Error()) + } +} + +func TestFetchAnthropicUsage_ServerError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`internal error`)) + })) + defer srv.Close() + + origURL := anthropicUsageURL + defer func() { setAnthropicUsageURL(origURL) }() + setAnthropicUsageURL(srv.URL) + + _, err := FetchAnthropicUsage("test-token") + if err == nil { + t.Fatal("expected error for 500, got nil") + } + if !strings.Contains(err.Error(), "500") { + t.Errorf("expected error containing '500', got %q", err.Error()) + } +} + +func TestFetchAnthropicUsage_MalformedJSON(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`not json`)) + })) + defer srv.Close() + + origURL := anthropicUsageURL + defer func() { setAnthropicUsageURL(origURL) }() + setAnthropicUsageURL(srv.URL) + + _, err := FetchAnthropicUsage("test-token") + if err == nil { + t.Fatal("expected error for malformed JSON, got nil") + } + if !strings.Contains(err.Error(), "parsing usage response") { + t.Errorf("expected 'parsing usage response' error, got %q", err.Error()) + } +} diff --git a/pkg/auth/token.go b/pkg/auth/token.go index 3280f629b..0e69e60ac 100644 --- a/pkg/auth/token.go +++ b/pkg/auth/token.go @@ -44,9 +44,6 @@ func LoginSetupToken(r io.Reader) (*AuthCredential, error) { } token := strings.TrimSpace(scanner.Text()) - if token == "" { - return nil, fmt.Errorf("token cannot be empty") - } if !strings.HasPrefix(token, "sk-ant-oat01-") { return nil, fmt.Errorf("invalid setup token: expected prefix sk-ant-oat01-") diff --git a/pkg/auth/token_test.go b/pkg/auth/token_test.go new file mode 100644 index 000000000..673cd9d5d --- /dev/null +++ b/pkg/auth/token_test.go @@ -0,0 +1,61 @@ +package auth + +import ( + "strings" + "testing" +) + +func TestLoginSetupToken(t *testing.T) { + // A valid token: correct prefix + at least 80 chars + validToken := "sk-ant-oat01-" + strings.Repeat("a", 80) + + tests := []struct { + name string + input string + wantErr string + }{ + {"valid token", validToken, ""}, + {"empty input", "", "expected prefix sk-ant-oat01-"}, + {"wrong prefix", "sk-ant-api-" + strings.Repeat("a", 80), "expected prefix sk-ant-oat01-"}, + {"too short", "sk-ant-oat01-short", "too short"}, + {"whitespace only", " ", "expected prefix sk-ant-oat01-"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + r := strings.NewReader(tt.input + "\n") + cred, err := LoginSetupToken(r) + + if tt.wantErr != "" { + if err == nil { + t.Fatalf("expected error containing %q, got nil", tt.wantErr) + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("expected error containing %q, got %q", tt.wantErr, err.Error()) + } + return + } + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cred.AccessToken != validToken { + t.Errorf("AccessToken = %q, want %q", cred.AccessToken, validToken) + } + if cred.Provider != "anthropic" { + t.Errorf("Provider = %q, want %q", cred.Provider, "anthropic") + } + if cred.AuthMethod != "oauth" { + t.Errorf("AuthMethod = %q, want %q", cred.AuthMethod, "oauth") + } + }) + } +} + +func TestLoginSetupToken_EmptyReader(t *testing.T) { + r := strings.NewReader("") + _, err := LoginSetupToken(r) + if err == nil { + t.Fatal("expected error for empty reader, got nil") + } +} diff --git a/pkg/providers/anthropic/provider.go b/pkg/providers/anthropic/provider.go index c68a34219..0d761f3a5 100644 --- a/pkg/providers/anthropic/provider.go +++ b/pkg/providers/anthropic/provider.go @@ -23,7 +23,10 @@ type ( ToolFunctionDefinition = protocoltypes.ToolFunctionDefinition ) -const defaultBaseURL = "https://api.anthropic.com" +const ( + defaultBaseURL = "https://api.anthropic.com" + anthropicBetaHeader = "oauth-2025-04-20" +) type Provider struct { client *anthropic.Client @@ -79,7 +82,7 @@ func (p *Provider) Chat( } opts = append(opts, option.WithAuthToken(tok), - option.WithHeader("anthropic-beta", "oauth-2025-04-20"), + option.WithHeader("anthropic-beta", anthropicBetaHeader), ) } diff --git a/pkg/providers/anthropic/provider_test.go b/pkg/providers/anthropic/provider_test.go index d10fde81c..3542447dd 100644 --- a/pkg/providers/anthropic/provider_test.go +++ b/pkg/providers/anthropic/provider_test.go @@ -262,6 +262,65 @@ func TestProvider_ChatUsesTokenSource(t *testing.T) { } } +func TestProvider_ChatStreamingRoundTrip(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/messages" { + http.Error(w, "not found", http.StatusNotFound) + return + } + if got := r.Header.Get("Authorization"); got != "Bearer refreshed-token" { + t.Errorf("Authorization = %q, want %q", got, "Bearer refreshed-token") + } + if got := r.Header.Get("anthropic-beta"); got != anthropicBetaHeader { + t.Errorf("anthropic-beta = %q, want %q", got, anthropicBetaHeader) + } + + w.Header().Set("Content-Type", "text/event-stream") + flusher, _ := w.(http.Flusher) + + events := []string{ + "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_stream\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-sonnet-4-6\",\"stop_reason\":null,\"usage\":{\"input_tokens\":12,\"output_tokens\":0}}}\n\n", + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\" world\"}}\n\n", + "event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\n", + "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":5}}\n\n", + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + } + for _, e := range events { + w.Write([]byte(e)) + if flusher != nil { + flusher.Flush() + } + } + })) + defer server.Close() + + p := NewProviderWithTokenSourceAndBaseURL("stale-token", func() (string, error) { + return "refreshed-token", nil + }, server.URL) + + resp, err := p.Chat( + t.Context(), + []Message{{Role: "user", Content: "Hello"}}, + nil, + "claude-sonnet-4.6", + map[string]any{}, + ) + if err != nil { + t.Fatalf("Chat() error: %v", err) + } + if resp.Content != "Hello world" { + t.Errorf("Content = %q, want %q", resp.Content, "Hello world") + } + if resp.FinishReason != "stop" { + t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop") + } + if resp.Usage.CompletionTokens != 5 { + t.Errorf("CompletionTokens = %d, want 5", resp.Usage.CompletionTokens) + } +} + func createAnthropicTestClient(baseURL, token string) *anthropic.Client { c := anthropic.NewClient( anthropicoption.WithAuthToken(token),