kiro: implement social token refresh via desktop.kiro.dev

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 {

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
+}