diff --git a/pkg/backend/codewhisperer.go b/pkg/backend/codewhisperer.go index 1f86891..4156dca 100644 --- a/pkg/backend/codewhisperer.go +++ b/pkg/backend/codewhisperer.go @@ -540,13 +540,21 @@ func (s *sqliteKiroAuth) Refresh(ctx context.Context) error { return nil // already refreshed by another process } - oidcEndpoint := fmt.Sprintf("https://oidc.%s.amazonaws.com/token", strings.TrimSpace(state.Region)) - client := newKiroOIDCClient(oidcEndpoint, nil) - refreshToken := currentRefresh if refreshToken == "" { refreshToken = state.RefreshToken } + + if state.IsSocial { + return s.refreshSocial(ctx, state, refreshToken) + } + return s.refreshOIDC(ctx, state, refreshToken) +} + +func (s *sqliteKiroAuth) refreshOIDC(ctx context.Context, state kiroSQLiteState, refreshToken string) error { + oidcEndpoint := fmt.Sprintf("https://oidc.%s.amazonaws.com/token", strings.TrimSpace(state.Region)) + client := newKiroOIDCClient(oidcEndpoint, nil) + resp, err := client.RefreshToken(ctx, kiroRefreshTokenRequest{ ClientID: state.ClientID, ClientSecret: state.ClientSecret, @@ -566,14 +574,39 @@ func (s *sqliteKiroAuth) Refresh(ctx context.Context) error { return kiroUpdateSQLiteTokens(ctx, s.path, newAccess, newRefresh, expiresAt) } +func (s *sqliteKiroAuth) refreshSocial(ctx context.Context, state kiroSQLiteState, refreshToken string) error { + client := newKiroSocialAuthClient(nil) + resp, err := client.RefreshToken(ctx, refreshToken) + if err != nil { + return err + } + + newAccess := strings.TrimSpace(resp.AccessToken) + newRefresh := strings.TrimSpace(resp.RefreshToken) + var expiresAt string + if resp.ExpiresIn != nil && *resp.ExpiresIn > 0 { + expiresAt = time.Now().Add(time.Duration(*resp.ExpiresIn) * time.Second).UTC().Format(time.RFC3339Nano) + } + profileARN := strings.TrimSpace(resp.ProfileARN) + if profileARN == "" { + profileARN = state.ProfileARN + } + return kiroUpdateSQLiteSocialTokens(ctx, s.path, newAccess, newRefresh, expiresAt, profileARN) +} + func (s *sqliteKiroAuth) CanRefresh() bool { state, err := s.load(context.Background()) return err == nil && s.canRefreshState(state) } func (s *sqliteKiroAuth) canRefreshState(state kiroSQLiteState) bool { - return strings.TrimSpace(state.RefreshToken) != "" && - strings.TrimSpace(state.Region) != "" && + if strings.TrimSpace(state.RefreshToken) == "" { + return false + } + if state.IsSocial { + return true + } + return strings.TrimSpace(state.Region) != "" && strings.TrimSpace(state.ClientID) != "" && strings.TrimSpace(state.ClientSecret) != "" } @@ -781,6 +814,32 @@ func kiroUpdateSQLiteTokens(ctx context.Context, path, access, refresh, expiresA return err } +func kiroUpdateSQLiteSocialTokens(ctx context.Context, path, access, refresh, expiresAt, profileARN string) error { + access = strings.TrimSpace(access) + if access == "" { + return fmt.Errorf("access token must not be empty") + } + var sb strings.Builder + sb.WriteString("BEGIN IMMEDIATE;\n") + sb.WriteString("UPDATE auth_kv SET value = json_set(value, '$.access_token', ") + sb.WriteString(kiroSQLiteQuote(access)) + if r := strings.TrimSpace(refresh); r != "" { + sb.WriteString(", '$.refresh_token', ") + sb.WriteString(kiroSQLiteQuote(r)) + } + if e := strings.TrimSpace(expiresAt); e != "" { + sb.WriteString(", '$.expires_at', ") + sb.WriteString(kiroSQLiteQuote(e)) + } + if p := strings.TrimSpace(profileARN); p != "" { + sb.WriteString(", '$.profile_arn', ") + sb.WriteString(kiroSQLiteQuote(p)) + } + sb.WriteString(") WHERE key = 'kirocli:social:token';\nCOMMIT;\n") + _, err := kiroRunSQLite(ctx, "-batch", path, sb.String()) + return err +} + func kiroQuerySQLite(ctx context.Context, path, query string) (string, error) { out, err := kiroRunSQLite(ctx, "-readonly", "-batch", "-noheader", path, query) if err != nil { diff --git a/pkg/backend/codewhisperer_internal.go b/pkg/backend/codewhisperer_internal.go index 7927cb0..566406d 100644 --- a/pkg/backend/codewhisperer_internal.go +++ b/pkg/backend/codewhisperer_internal.go @@ -677,3 +677,56 @@ func (c *kiroAPIClient) ListModels(ctx context.Context, profileARN string) (*kir } return &result, nil } + +// ── Social token refresh client ─────────────────────────────────────────────── + +const kiroSocialAuthEndpoint = "https://prod.us-east-1.auth.desktop.kiro.dev" + +type kiroSocialAuthClient struct { + endpoint string + httpClient *http.Client +} + +type kiroSocialRefreshResponse struct { + AccessToken string `json:"accessToken"` + RefreshToken string `json:"refreshToken,omitempty"` + ExpiresIn *int `json:"expiresIn,omitempty"` + ProfileARN string `json:"profileArn,omitempty"` +} + +func newKiroSocialAuthClient(httpClient *http.Client) *kiroSocialAuthClient { + if httpClient == nil { + httpClient = &http.Client{} + } + return &kiroSocialAuthClient{endpoint: kiroSocialAuthEndpoint, httpClient: httpClient} +} + +func (c *kiroSocialAuthClient) RefreshToken(ctx context.Context, refreshToken string) (*kiroSocialRefreshResponse, error) { + body, err := json.Marshal(map[string]string{"refreshToken": refreshToken}) + if err != nil { + return nil, fmt.Errorf("marshal social refresh request: %w", err) + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint+"/refreshToken", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("build social refresh request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("send social refresh request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, kiroDecodeHTTPError(resp) + } + var result kiroSocialRefreshResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("decode social refresh response: %w", err) + } + if strings.TrimSpace(result.AccessToken) == "" { + return nil, fmt.Errorf("social refresh response did not contain accessToken") + } + return &result, nil +}