Skip to content

Conversation

@Ravens2121
Copy link

@Ravens2121 Ravens2121 commented Dec 17, 2025

Pull Request: Kiro Authentication Enhancements & API Improvements

📋 Summary

This PR introduces significant improvements to the Kiro authentication system and API integration, building upon the recent changes in commit d687ee27 which implemented official reasoningContentEvent support and improved metadata handling. The new changes focus on enhancing the authentication experience with AWS Builder ID Authorization Code Flow and improving user email retrieval mechanisms.

🔗 Related Commits

  • Base Commit: d687ee27772ac3541048c9ad60bee834aaac7356
    • Implemented official reasoningContentEvent for thinking mode
    • Improved messageMetadataEvent / metadataEvent parsing for accurate token usage
    • Simplified thinking tag parsing logic
    • Updated thinking mode from interleaved to enabled for official API support

✨ New Features

1. AWS Builder ID Authorization Code Flow (--kiro-aws-authcode)

A new authentication method that provides a superior user experience compared to device code flow.

Files Changed:

  • cmd/server/main.go - Added --kiro-aws-authcode flag
  • internal/cmd/kiro_login.go - Added DoKiroAWSAuthCodeLogin() function
  • sdk/auth/kiro.go - Added LoginWithAuthCode() method
  • internal/auth/kiro/oauth.go - Added LoginWithBuilderIDAuthCode() method
  • internal/auth/kiro/sso_oidc.go - Core implementation

Key Implementation Details:

// New flag added to main.go
flag.BoolVar(&kiroAWSAuthCode, "kiro-aws-authcode", false, 
    "Login to Kiro using AWS Builder ID (authorization code flow, better UX)")

How it works:

  1. Generates PKCE code verifier and challenge
  2. Starts a local HTTP callback server on port 19877
  3. Registers OIDC client with authorization_code grant type
  4. Opens browser with authorization URL
  5. Receives callback with authorization code
  6. Exchanges code for tokens
  7. Fetches user profile and email

Advantages over Device Code Flow:

  • No manual code entry required
  • Automatic browser callback
  • Better UX with instant authentication completion
  • Browser auto-closes after success

2. CodeWhisperer API Client for User Email Retrieval

New File: internal/auth/kiro/codewhisperer_client.go

A dedicated client to fetch user information via the official CodeWhisperer API.

type CodeWhispererClient struct {
    httpClient *http.Client
    machineID  string
}

API Endpoint: https://codewhisperer.us-east-1.amazonaws.com/getUsageLimits?isEmailRequired=true&origin=AI_EDITOR&resourceType=AGENTIC_REQUEST

Key Features:

  • Uses official Kiro IDE User-Agent for compatibility
  • Returns user email, subscription info, and usage limits
  • Reliable email retrieval compared to JWT parsing

3. Enhanced User Email Retrieval with Fallback

Function: FetchUserEmailWithFallback(ctx, cfg, accessToken)

Implements a priority-based email retrieval strategy:

Priority Method Description
1 CodeWhisperer API Most reliable, official API
2 SSO OIDC UserInfo Standard OIDC userinfo endpoint
3 JWT Parsing Fallback to token claim extraction

🔧 Improvements

1. Official Kiro IDE Client Registration

Updated client registration to match official Kiro IDE behavior:

// Before
payload := map[string]interface{}{
    "clientName": fmt.Sprintf("CLI-Proxy-API-%d", time.Now().UnixNano()),
    "clientType": "public",
    "scopes":     []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations"},
}

// After
payload := map[string]interface{}{
    "clientName": "Kiro IDE",
    "clientType": "public",
    "scopes":     []string{
        "codewhisperer:completions", 
        "codewhisperer:analysis", 
        "codewhisperer:conversations", 
        "codewhisperer:transformations",  // NEW
        "codewhisperer:taskassist",        // NEW
    },
    "grantTypes": []string{"urn:ietf:params:oauth:grant-type:device_code", "refresh_token"},
}

2. User-Agent Header Consistency

Added consistent KiroIDE User-Agent across all SSO OIDC requests:

const kiroUserAgent = "KiroIDE"

req.Header.Set("User-Agent", kiroUserAgent)

Affected Methods:

  • RegisterClient()
  • RegisterClientForAuthCode()
  • StartDeviceAuthorization()
  • CreateToken()
  • CreateTokenWithAuthCode()
  • RefreshToken()

3. UserInfo Endpoint Integration

Added SSO OIDC userinfo endpoint support for email retrieval:

func (c *SSOOIDCClient) FetchUserEmail(ctx context.Context, accessToken string) string {
    // Method 1: Try userinfo endpoint (standard OIDC)
    email := c.tryUserInfoEndpoint(ctx, accessToken)
    if email != "" {
        return email
    }
    // Method 2: Fallback to JWT parsing
    return ExtractEmailFromJWT(accessToken)
}

📊 Changes from Commit d687ee2

Thinking Mode Improvements

Before (interleaved mode with prompt injection):

thinkingHint := `<thinking_mode>interleaved</thinking_mode>
<max_thinking_length>200000</max_thinking_length>

IMPORTANT: You MUST use <thinking>...</thinking> tags to show your reasoning...`

After (official enabled mode):

thinkingHint := `<thinking_mode>enabled</thinking_mode>
<max_thinking_length>200000</max_thinking_length>`

New Event Handling

Added support for official reasoningContentEvent:

case "reasoningContentEvent":
    // Handle official reasoningContentEvent from Kiro API
    // Format: { text: string, signature?: string, redactedContent?: base64 }

Enhanced Metadata Parsing

Improved messageMetadataEvent / metadataEvent handling for accurate token usage:

if tokenUsage, ok := metadata["tokenUsage"].(map[string]interface{}); ok {
    // outputTokens - precise output token count
    // totalTokens - precise total token count
    // uncachedInputTokens - input tokens not from cache
    // cacheReadInputTokens - tokens read from cache
    // contextUsagePercentage - context usage percentage
}

📁 Files Changed Summary

File Changes Description
cmd/server/main.go +6 Added --kiro-aws-authcode flag
internal/auth/kiro/oauth.go +7 Added LoginWithBuilderIDAuthCode()
internal/auth/kiro/sso_oidc.go +403 Auth code flow, PKCE, callback server, userinfo
internal/auth/kiro/codewhisperer_client.go +166 (new) CodeWhisperer API client
internal/cmd/kiro_login.go +48 Added DoKiroAWSAuthCodeLogin()
sdk/auth/kiro.go +65 Added LoginWithAuthCode()

Total: +695 lines, -7 lines


🧪 Testing

Manual Testing Steps

  1. Authorization Code Flow:

    ./cliproxy --kiro-aws-authcode
    • Verify browser opens automatically
    • Complete AWS Builder ID login
    • Verify callback received
    • Verify email displayed correctly
  2. Device Code Flow (existing):

    ./cliproxy --kiro-aws-login
    • Verify enhanced email retrieval works
  3. Email Retrieval Fallback:

    • Test with valid token: CodeWhisperer API should work
    • Test with restricted token: UserInfo fallback
    • Test with JWT-only: JWT parsing fallback

⚙️ Configuration

No new configuration required. The new --kiro-aws-authcode flag is opt-in.

Callback Server

  • Default Port: 19877
  • Fallback: Dynamic port if default is busy
  • Callback Path: /oauth/callback
  • Timeout: 10 minutes

🔒 Security Considerations

  1. PKCE Implementation: Uses SHA-256 code challenge method
  2. State Parameter: Random 16-byte state for CSRF protection
  3. Local Callback: Only listens on 127.0.0.1
  4. Token Storage: Follows existing secure storage patterns

📝 Usage Examples

New Authorization Code Flow

# Better UX - automatic browser callback
./cliproxy --kiro-aws-authcode

# With incognito mode disabled (use existing browser session)
./cliproxy --kiro-aws-authcode --no-incognito

Existing Device Code Flow

# Device code flow (manual code entry)
./cliproxy --kiro-aws-login

🎯 Breaking Changes

None. All changes are backward compatible.


📌 Dependencies

No new dependencies added. Uses existing:

  • github.com/google/uuid
  • Standard library crypto/rand, crypto/sha256, net/http

✅ Checklist

  • Code follows project conventions
  • Added proper error handling
  • Added debug logging
  • User-Agent matches official Kiro IDE
  • PKCE security implemented
  • Fallback mechanisms implemented
  • Documentation updated

@gemini-code-assist
Copy link

Summary of Changes

Hello @Ravens2121, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly upgrades the Kiro executor's interaction with the Kiro API. It introduces support for the official reasoning mode, moving away from heuristic tag parsing to event-driven processing for thinking content. Additionally, it refines the extraction of token usage from metadata events, incorporates handling for metering events to track billing, and enhances the overall robustness of error handling for various API stream errors.

Highlights

  • Official Kiro API Reasoning Mode Support: The executor now supports the official Kiro API's reasoning mode by injecting the <thinking_mode>enabled</thinking_mode> tag, which leads to reasoningContentEvents being returned instead of relying on heuristic tag-based parsing.
  • Enhanced Metadata Event Handling: Improved parsing for messageMetadataEvent and metadataEvent to accurately extract token usage information, including support for nested tokenUsage objects and a fallback mechanism for direct fields.
  • Metering Event Processing: Added specific handling for meteringEvent to capture and log usage billing information from the Kiro API stream.
  • Improved Error Event Management: The system now gracefully manages various error event types from the Kiro API stream, including error, exception, internalServerException, and invalidStateEvent, with appropriate logging and error propagation.
  • Refactored Thinking Block Detection: The previous heuristic logic for detecting and parsing <thinking> blocks within content has been removed, replaced by the more robust and official reasoningContentEvent handling.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request significantly refactors the handling of streaming events from the Kiro API, primarily by adopting official event types for token usage, metering, and reasoning content. The changes introduce a new parsing mechanism for messageMetadataEvent and metadataEvent to support a nested tokenUsage object, while retaining fallback for legacy formats. Dedicated meteringEvent and various error events (error, exception, internalServerException, invalidStateEvent) are now explicitly handled in both parseEventStream and streamToChannel functions. Crucially, the complex heuristic-based parsing of <thinking> tags within the response content has been removed and replaced with handling an official reasoningContentEvent, which provides structured reasoning text. Correspondingly, the system prompt injection for enabling thinking mode has been updated to use <thinking_mode>enabled</thinking_mode>, removing the instruction for the model to use inline <thinking> tags. Review comments noted an inconsistency in error type extraction for event["error"] in streamToChannel, where the errType was not updated, and pointed out duplicated and inconsistent logging logic for meteringEvent across parseEventStream and streamToChannel, suggesting a helper function and consistent logging checks.

Comment on lines +2056 to +2060
} else if errObj, ok := event["error"].(map[string]interface{}); ok {
if msg, ok := errObj["message"].(string); ok {
errMsg = msg
}
}

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

There's an inconsistency in error type extraction. When handling an error from event["error"], the errType is not updated from the nested type field. This differs from the logic for event[eventType] and the implementation in parseEventStream, which does update errType. This could lead to less specific error logging and reporting.

} else if errObj, ok := event["error"].(map[string]interface{}); ok {
				if msg, ok := errObj["message"].(string); ok {
					errMsg = msg
				}
				if t, ok := errObj["type"].(string); ok {
					errType = t
				}
			}

Comment on lines +1408 to +1436
case "meteringEvent":
// Handle metering events from Kiro API (usage billing information)
// Official format: { unit: string, unitPlural: string, usage: number }
if metering, ok := event["meteringEvent"].(map[string]interface{}); ok {
unit := ""
if u, ok := metering["unit"].(string); ok {
unit = u
}
usageVal := 0.0
if u, ok := metering["usage"].(float64); ok {
usageVal = u
}
log.Infof("kiro: parseEventStream received meteringEvent: usage=%.2f %s", usageVal, unit)
// Store metering info for potential billing/statistics purposes
// Note: This is separate from token counts - it's AWS billing units
} else {
// Try direct fields
unit := ""
if u, ok := event["unit"].(string); ok {
unit = u
}
usageVal := 0.0
if u, ok := event["usage"].(float64); ok {
usageVal = u
}
if unit != "" || usageVal > 0 {
log.Infof("kiro: parseEventStream received meteringEvent (direct): usage=%.2f %s", usageVal, unit)
}
}

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The logic for handling meteringEvent is duplicated in parseEventStream and streamToChannel (lines 2015-2039). This can be simplified by extracting the parsing logic into a helper function. Additionally, the logging for the nested meteringEvent case occurs even if no unit or usage value is found, which can create log noise. The direct field case correctly checks if values were found before logging. This should be consistent across both implementations.

case "meteringEvent":
			// Handle metering events from Kiro API (usage billing information)
			// Official format: { unit: string, unitPlural: string, usage: number }
			var meteringData map[string]interface{}
			logMsgPrefix := "kiro: parseEventStream received meteringEvent"

			if metering, ok := event["meteringEvent"].(map[string]interface{}); ok {
				meteringData = metering
			} else {
				meteringData = event
				logMsgPrefix += " (direct)"
			}

			unit := ""
			if u, ok := meteringData["unit"].(string); ok {
				unit = u
			}
			usageVal := 0.0
			if u, ok := meteringData["usage"].(float64); ok {
				usageVal = u
			}

			if unit != "" || usageVal > 0 {
				log.Infof("%s: usage=%.2f %s", logMsgPrefix, usageVal, unit)
				// Store metering info for potential billing/statistics purposes
				// Note: This is separate from token counts - it's AWS billing units
			}

@Ravens2121 Ravens2121 changed the title 添加了对于 计量事件 的处理,改进了错误事件的处理。 添加了KIRO供应商 对于 计量事件 的处理,改进了错误事件的处理。 Dec 17, 2025
@Ravens2121 Ravens2121 changed the title 添加了KIRO供应商 对于 计量事件 的处理,改进了错误事件的处理。 feat(kiro): 新增授权码登录流程,优化邮箱获取与官方 Thinking 模式支持 Dec 18, 2025
@Ravens2121 Ravens2121 changed the title feat(kiro): 新增授权码登录流程,优化邮箱获取与官方 Thinking 模式支持 feat(kiro): 新增授权码登录流程,优化邮箱获取与官方 Thinking 模式解析 预支持 Dec 18, 2025
@luispater luispater merged commit ca99323 into router-for-me:main Dec 18, 2025
1 of 2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants