-
-
Notifications
You must be signed in to change notification settings - Fork 40
feat(kiro): 新增授权码登录流程,优化邮箱获取与官方 Thinking 模式解析 预支持 #42
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
Summary of ChangesHello @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
Using Gemini Code AssistThe 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
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 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
|
There was a problem hiding this 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.
| } else if errObj, ok := event["error"].(map[string]interface{}); ok { | ||
| if msg, ok := errObj["message"].(string); ok { | ||
| errMsg = msg | ||
| } | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
}
}| 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) | ||
| } | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
}Amp-Thread-ID: https://ampcode.com/threads/T-019b2ecc-fb2d-713f-b30d-1196c7dce3e2 Co-authored-by: Amp <[email protected]>
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
d687ee27which implemented officialreasoningContentEventsupport 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
d687ee27772ac3541048c9ad60bee834aaac7356reasoningContentEventfor thinking modemessageMetadataEvent/metadataEventparsing for accurate token usageinterleavedtoenabledfor 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-authcodeflaginternal/cmd/kiro_login.go- AddedDoKiroAWSAuthCodeLogin()functionsdk/auth/kiro.go- AddedLoginWithAuthCode()methodinternal/auth/kiro/oauth.go- AddedLoginWithBuilderIDAuthCode()methodinternal/auth/kiro/sso_oidc.go- Core implementationKey Implementation Details:
How it works:
19877authorization_codegrant typeAdvantages over Device Code Flow:
2. CodeWhisperer API Client for User Email Retrieval
New File:
internal/auth/kiro/codewhisperer_client.goA dedicated client to fetch user information via the official CodeWhisperer API.
API Endpoint:
https://codewhisperer.us-east-1.amazonaws.com/getUsageLimits?isEmailRequired=true&origin=AI_EDITOR&resourceType=AGENTIC_REQUESTKey Features:
3. Enhanced User Email Retrieval with Fallback
Function:
FetchUserEmailWithFallback(ctx, cfg, accessToken)Implements a priority-based email retrieval strategy:
🔧 Improvements
1. Official Kiro IDE Client Registration
Updated client registration to match official Kiro IDE behavior:
2. User-Agent Header Consistency
Added consistent
KiroIDEUser-Agent across all SSO OIDC requests:Affected Methods:
RegisterClient()RegisterClientForAuthCode()StartDeviceAuthorization()CreateToken()CreateTokenWithAuthCode()RefreshToken()3. UserInfo Endpoint Integration
Added SSO OIDC userinfo endpoint support for email retrieval:
📊 Changes from Commit d687ee2
Thinking Mode Improvements
Before (interleaved mode with prompt injection):
After (official enabled mode):
New Event Handling
Added support for official
reasoningContentEvent:Enhanced Metadata Parsing
Improved
messageMetadataEvent/metadataEventhandling for accurate token usage:📁 Files Changed Summary
cmd/server/main.go--kiro-aws-authcodeflaginternal/auth/kiro/oauth.goLoginWithBuilderIDAuthCode()internal/auth/kiro/sso_oidc.gointernal/auth/kiro/codewhisperer_client.gointernal/cmd/kiro_login.goDoKiroAWSAuthCodeLogin()sdk/auth/kiro.goLoginWithAuthCode()Total: +695 lines, -7 lines
🧪 Testing
Manual Testing Steps
Authorization Code Flow:
Device Code Flow (existing):
Email Retrieval Fallback:
⚙️ Configuration
No new configuration required. The new
--kiro-aws-authcodeflag is opt-in.Callback Server
19877/oauth/callback🔒 Security Considerations
127.0.0.1📝 Usage Examples
New Authorization Code Flow
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/uuidcrypto/rand,crypto/sha256,net/http✅ Checklist