diff --git a/pkg/backend/codewhisperer.go b/pkg/backend/codewhisperer.go
index 68cbb2e..bbfed8f 100644
--- a/pkg/backend/codewhisperer.go
+++ b/pkg/backend/codewhisperer.go
@@ -16,6 +16,7 @@ package backend
import (
"context"
+ "crypto/rand"
"encoding/json"
"errors"
"fmt"
@@ -133,7 +134,7 @@ func (b *CodeWhispererBackend) ChatStream(
return nil, fmt.Errorf("kiro profile ARN: %w", err)
}
- req, err := buildKiroRequest(model, messages, tools, profileARN)
+ req, err := buildKiroRequest(model, messages, tools, profileARN, params)
if err != nil {
return nil, err
}
@@ -181,6 +182,16 @@ func (b *CodeWhispererBackend) runStream(ctx context.Context, req *kiroGenerateR
stopped = true
}
},
+ OnReasoningDelta: func(event kiroReasoningContentEvent) {
+ if stopped || (event.Text == "" && event.RedactedContent == "" && event.Signature == "") {
+ return
+ }
+ select {
+ case ch <- StreamEvent{Reasoning: event.Text}:
+ case <-ctx.Done():
+ stopped = true
+ }
+ },
})
if streamErr != nil {
@@ -255,12 +266,15 @@ func kiroSanitizeModel(model string) string {
return model
}
-func buildKiroRequest(model string, messages []Message, tools []Tool, profileARN string) (*kiroGenerateRequest, error) {
+func buildKiroRequest(model string, messages []Message, tools []Tool, profileARN string, params GenerationParams) (*kiroGenerateRequest, error) {
model = kiroSanitizeModel(model)
if len(messages) == 0 {
return nil, fmt.Errorf("at least one message is required")
}
+ // Apply counter-thinking tags to user messages before assistant responses.
+ messages = kiroApplyCounterThinkingTags(messages)
+
encodedTools := encodeKiroToolDefinitions(tools)
lastRole := strings.ToLower(strings.TrimSpace(messages[len(messages)-1].Role))
@@ -275,7 +289,7 @@ func buildKiroRequest(model string, messages []Message, tools []Tool, profileARN
if err != nil {
return nil, err
}
- current = kiroChatMessage{UserInputMessage: encodeKiroUserMessage(messages[len(messages)-1], model, encodedTools, nil)}
+ current = kiroChatMessage{UserInputMessage: encodeKiroUserMessage(messages[len(messages)-1], model, encodedTools, nil, true)}
case "tool":
// Find the last non-tool message.
@@ -288,7 +302,7 @@ func buildKiroRequest(model string, messages []Message, tools []Tool, profileARN
return nil, err
}
toolResults, images := encodeKiroToolResults(messages[cutoff+1:])
- current = kiroChatMessage{UserInputMessage: emptyKiroUserMessage(model, encodedTools, toolResults, images)}
+ current = kiroChatMessage{UserInputMessage: emptyKiroUserMessage(model, encodedTools, toolResults, images, true)}
default:
return nil, fmt.Errorf("last message role %q is not supported (expected user, system, or tool)", lastRole)
@@ -303,7 +317,7 @@ func buildKiroRequest(model string, messages []Message, tools []Tool, profileARN
return nil, err
}
- return &kiroGenerateRequest{
+ req := &kiroGenerateRequest{
ProfileARN: profileARN,
ConversationState: kiroConversationState{
ConversationID: convID,
@@ -313,7 +327,14 @@ func buildKiroRequest(model string, messages []Message, tools []Tool, profileARN
AgentContinuationID: contID,
AgentTaskType: kiroAgentTaskType,
},
- }, nil
+ }
+
+ // Set additionalModelRequestFields for reasoning/thinking support.
+ if fields := kiroModelRequestFieldsFromParams(params); fields != nil {
+ req.AdditionalModelRequestFields = fields
+ }
+
+ return req, nil
}
func encodeKiroHistory(messages []Message) ([]kiroChatMessage, error) {
@@ -322,7 +343,7 @@ func encodeKiroHistory(messages []Message) ([]kiroChatMessage, error) {
m := messages[i]
switch strings.ToLower(strings.TrimSpace(m.Role)) {
case "user", "system":
- out = append(out, kiroChatMessage{UserInputMessage: encodeKiroUserMessage(m, "", nil, nil)})
+ out = append(out, kiroChatMessage{UserInputMessage: encodeKiroUserMessage(m, "", nil, nil, false)})
case "assistant":
out = append(out, kiroChatMessage{AssistantResponseMessage: encodeKiroAssistantMessage(m)})
case "tool":
@@ -332,17 +353,18 @@ func encodeKiroHistory(messages []Message) ([]kiroChatMessage, error) {
i++
}
results, imgs := encodeKiroToolResults(messages[start : i+1])
- out = append(out, kiroChatMessage{UserInputMessage: emptyKiroUserMessage("", nil, results, imgs)})
+ out = append(out, kiroChatMessage{UserInputMessage: emptyKiroUserMessage("", nil, results, imgs, false)})
}
}
return out, nil
}
-func encodeKiroUserMessage(m Message, modelID string, tools []kiroTool, toolResults []kiroToolResult) *kiroUserInputMessage {
+func encodeKiroUserMessage(m Message, modelID string, tools []kiroTool, toolResults []kiroToolResult, includeEnvState bool) *kiroUserInputMessage {
msg := &kiroUserInputMessage{
- Content: kiroEncodeContent(m),
+ Content: kiroExtractTextContent(m),
Origin: kiroDefaultOrigin,
ModelID: modelID,
+ Images: kiroExtractImages(m),
}
ctx := &kiroUserInputContext{}
if len(tools) > 0 {
@@ -351,15 +373,18 @@ func encodeKiroUserMessage(m Message, modelID string, tools []kiroTool, toolResu
if len(toolResults) > 0 {
ctx.ToolResults = toolResults
}
- if len(ctx.Tools) > 0 || len(ctx.ToolResults) > 0 {
+ if includeEnvState {
+ ctx.EnvState = kiroCurrentEnvState()
+ }
+ if len(ctx.Tools) > 0 || len(ctx.ToolResults) > 0 || ctx.EnvState != nil {
msg.UserInputMessageContext = ctx
}
return msg
}
-func emptyKiroUserMessage(modelID string, tools []kiroTool, toolResults []kiroToolResult, images []kiroImageBlock) *kiroUserInputMessage {
+func emptyKiroUserMessage(modelID string, tools []kiroTool, toolResults []kiroToolResult, images []kiroImageBlock, includeEnvState bool) *kiroUserInputMessage {
msg := &kiroUserInputMessage{
- Content: kiroEncodeStringContent(""),
+ Content: "",
Origin: kiroDefaultOrigin,
ModelID: modelID,
Images: images,
@@ -371,7 +396,10 @@ func emptyKiroUserMessage(modelID string, tools []kiroTool, toolResults []kiroTo
if len(toolResults) > 0 {
ctx.ToolResults = toolResults
}
- if len(ctx.Tools) > 0 || len(ctx.ToolResults) > 0 {
+ if includeEnvState {
+ ctx.EnvState = kiroCurrentEnvState()
+ }
+ if len(ctx.Tools) > 0 || len(ctx.ToolResults) > 0 || ctx.EnvState != nil {
msg.UserInputMessageContext = ctx
}
return msg
@@ -383,6 +411,16 @@ func encodeKiroAssistantMessage(m Message) *kiroAssistantResponseMessage {
input := tc.Arguments
if len(input) == 0 || !json.Valid(input) {
input = json.RawMessage("{}")
+ } else {
+ // Normalize the JSON by round-tripping through unmarshal/marshal.
+ // This deduplicates keys (the model sometimes produces duplicate
+ // keys which the upstream API rejects) and ensures well-formed output.
+ var parsed any
+ if err := json.Unmarshal(input, &parsed); err == nil {
+ if normalized, err := json.Marshal(parsed); err == nil {
+ input = normalized
+ }
+ }
}
toolUseID := strings.TrimSpace(tc.ID)
if toolUseID == "" {
@@ -464,63 +502,136 @@ func encodeKiroToolResultBlocks(m Message) ([]kiroToolResultContent, []kiroImage
}
func encodeKiroToolResultContent(content string) kiroToolResultContent {
- content = strings.TrimSpace(content)
- if len(content) > 0 && content[0] == '{' && json.Valid([]byte(content)) {
- return kiroToolResultContent{JSON: json.RawMessage(content)}
- }
- return kiroToolResultContent{Text: content}
+ // Always use the text field. The upstream Amazon Q API rejects bare JSON
+ // primitives (numbers, strings, booleans, null) in the "json" field with
+ // "ValidationException: Improperly formed request." — it expects only
+ // objects or arrays. Using text is always safe; the model sees the same
+ // content either way.
+ return kiroToolResultContent{Text: strings.TrimSpace(content)}
}
-// kiroEncodeContent returns a json.RawMessage that is either a JSON string
-// (plain text) or a JSON array of Bedrock Converse content blocks (when images
-// are present).
-func kiroEncodeContent(m Message) json.RawMessage {
+// kiroExtractTextContent returns the text content of a message. When images are
+// present, they are encoded as separate image blocks on the message; this
+// function only returns the text portion.
+func kiroExtractTextContent(m Message) string {
if len(m.ContentBlocks) == 0 {
- return kiroEncodeStringContent(m.Content)
+ return m.Content
}
- hasImage := false
+ var sb strings.Builder
+ for _, cb := range m.ContentBlocks {
+ if cb.Type == "text" {
+ sb.WriteString(cb.Text)
+ }
+ }
+ return sb.String()
+}
+
+// kiroExtractImages extracts image blocks from a message's ContentBlocks.
+func kiroExtractImages(m Message) []kiroImageBlock {
+ if len(m.ContentBlocks) == 0 {
+ return nil
+ }
+ var images []kiroImageBlock
for _, cb := range m.ContentBlocks {
if cb.Type == "image" && cb.ImageSource != nil {
- hasImage = true
- break
- }
- }
- if !hasImage {
- // Text-only blocks: concatenate and send as string.
- var sb strings.Builder
- for _, cb := range m.ContentBlocks {
- if cb.Type == "text" {
- sb.WriteString(cb.Text)
- }
- }
- return kiroEncodeStringContent(sb.String())
- }
- // Mixed content: encode as Bedrock Converse content block array.
- blocks := make([]kiroContentBlock, 0, len(m.ContentBlocks))
- for _, cb := range m.ContentBlocks {
- switch cb.Type {
- case "text":
- blocks = append(blocks, kiroContentBlock{Text: cb.Text})
- case "image":
- if cb.ImageSource == nil {
- continue
- }
format := strings.TrimPrefix(cb.ImageSource.MediaType, "image/")
- blocks = append(blocks, kiroContentBlock{
- Image: &kiroImageBlock{
- Format: format,
- Source: kiroImageSource{Bytes: cb.ImageSource.Data},
- },
+ images = append(images, kiroImageBlock{
+ Format: format,
+ Source: kiroImageSource{Bytes: cb.ImageSource.Data},
})
}
}
- data, _ := json.Marshal(blocks)
- return json.RawMessage(data)
+ return images
}
-func kiroEncodeStringContent(s string) json.RawMessage {
- data, _ := json.Marshal(s)
- return json.RawMessage(data)
+func kiroCurrentEnvState() *kiroEnvState {
+ return &kiroEnvState{
+ OperatingSystem: runtime.GOOS,
+ CurrentWorkingDirectory: "/",
+ }
+}
+
+// ── Counter-thinking & reasoning support ──────────────────────────────────────
+
+const kiroCounterThinkingInstructions = `Be thorough in your work; don't skip steps or cut corners to save time. Think deeply and step by step, do not skip ahead prematurely.`
+const kiroDefaultReasoningEffort = "95"
+
+// kiroApplyCounterThinkingTags appends corrective tags to user messages that
+// immediately precede assistant responses, countering server-side injections
+// that degrade reasoning quality.
+func kiroApplyCounterThinkingTags(messages []Message) []Message {
+ if len(messages) < 2 {
+ return messages
+ }
+ var mangled []Message
+ for i := 1; i < len(messages); i++ {
+ if !strings.EqualFold(messages[i].Role, "assistant") {
+ continue
+ }
+ prev := messages[i-1]
+ role := strings.ToLower(strings.TrimSpace(prev.Role))
+ if role != "user" && role != "system" {
+ continue
+ }
+ if mangled == nil {
+ mangled = make([]Message, len(messages))
+ copy(mangled, messages)
+ }
+ suffix := "\n\n" + kiroDefaultReasoningEffort + "\n" +
+ "" + kiroDefaultReasoningEffort + "\n" +
+ "" + kiroCounterThinkingInstructions + ""
+ mangled[i-1].Content = prev.Content + suffix
+ }
+ if mangled != nil {
+ return mangled
+ }
+ return messages
+}
+
+// kiroModelRequestFieldsFromParams builds additionalModelRequestFields from
+// GenerationParams. Returns nil if no reasoning/thinking fields are needed.
+func kiroModelRequestFieldsFromParams(params GenerationParams) *kiroModelRequestFields {
+ fields := &kiroModelRequestFields{}
+ hasFields := false
+
+ // Enable summarized thinking by default for reasoning-capable models.
+ if params.ThinkingBudget > 0 || params.ReasoningEffort != "" {
+ fields.Thinking = &kiroThinkingConfig{
+ Type: "adaptive",
+ Display: "summarized",
+ }
+ hasFields = true
+ }
+
+ // Map reasoning effort.
+ if effort := kiroMapEffort(params.ReasoningEffort); effort != "" {
+ fields.OutputConfig = &kiroOutputConfig{Effort: effort}
+ hasFields = true
+ } else if params.ThinkingBudget > 0 {
+ // Default to high effort when thinking is requested but no specific effort.
+ fields.OutputConfig = &kiroOutputConfig{Effort: "high"}
+ hasFields = true
+ }
+
+ if !hasFields {
+ return nil
+ }
+ return fields
+}
+
+func kiroMapEffort(effort string) string {
+ switch strings.ToLower(strings.TrimSpace(effort)) {
+ case "max":
+ return "max"
+ case "high":
+ return "high"
+ case "medium":
+ return "medium"
+ case "low", "minimal":
+ return "low"
+ default:
+ return ""
+ }
}
// ── Auth sources ──────────────────────────────────────────────────────────────
@@ -565,6 +676,10 @@ type sqliteKiroAuth struct {
loaded bool
initMu sync.Mutex
mu sync.Mutex
+ // jitter is computed once at creation time so all checks within this
+ // process use the same refresh threshold. Different processes get
+ // different jitter values, spreading refresh attempts over time.
+ jitter time.Duration
}
type kiroSQLiteState struct {
@@ -609,7 +724,7 @@ func (s *sqliteKiroAuth) AccessToken(ctx context.Context) (string, error) {
if err != nil {
return "", err
}
- if state.AccessToken != "" && !kiroTokenExpiringSoon(state.ExpiresAt) {
+ if state.AccessToken != "" && !s.tokenExpiringSoon(state.ExpiresAt) {
return state.AccessToken, nil
}
if !s.canRefreshState(state) {
@@ -635,7 +750,14 @@ func (s *sqliteKiroAuth) Refresh(ctx context.Context) error {
s.mu.Lock()
defer s.mu.Unlock()
- state, err := s.load(ctx)
+ // Bound the refresh operation to avoid hanging indefinitely.
+ refreshCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+ defer cancel()
+
+ // Always perform a full reload of all sqlite state (including device
+ // registration) before refreshing. Another process (e.g. kiro-cli login)
+ // may have rotated the client credentials since our initial load.
+ state, err := s.fullLoad(refreshCtx)
if err != nil {
return err
}
@@ -644,8 +766,8 @@ func (s *sqliteKiroAuth) Refresh(ctx context.Context) error {
}
// Re-read current token — another process may have already refreshed it.
- currentAccess, currentRefresh, currentExpiry, err := kiroReadSQLiteTokens(ctx, s.path)
- if err == nil && currentAccess != "" && currentAccess != state.AccessToken && !kiroTokenExpiringSoon(currentExpiry) {
+ currentAccess, currentRefresh, currentExpiry, err := kiroReadSQLiteTokens(refreshCtx, s.path)
+ if err == nil && currentAccess != "" && currentAccess != state.AccessToken && !s.tokenExpiringSoon(currentExpiry) {
return nil // already refreshed by another process
}
@@ -655,12 +777,12 @@ func (s *sqliteKiroAuth) Refresh(ctx context.Context) error {
}
if state.IsSocial {
- return s.refreshSocial(ctx, state, refreshToken)
+ return s.refreshSocial(refreshCtx, state, refreshToken)
}
- return s.refreshOIDC(ctx, state, refreshToken)
+ return s.refreshOIDC(refreshCtx, state, refreshToken, currentAccess)
}
-func (s *sqliteKiroAuth) refreshOIDC(ctx context.Context, state kiroSQLiteState, refreshToken string) error {
+func (s *sqliteKiroAuth) refreshOIDC(ctx context.Context, state kiroSQLiteState, refreshToken, currentAccess string) error {
oidcEndpoint := fmt.Sprintf("https://oidc.%s.amazonaws.com/token", strings.TrimSpace(state.Region))
client := newKiroOIDCClient(oidcEndpoint, nil)
@@ -671,6 +793,14 @@ func (s *sqliteKiroAuth) refreshOIDC(ctx context.Context, state kiroSQLiteState,
RefreshToken: refreshToken,
})
if err != nil {
+ // Try SSO cache fallback.
+ if fallbackErr := s.refreshFromSSOCache(ctx, currentAccess); fallbackErr == nil {
+ return nil
+ }
+ // Check if another process refreshed while we were trying.
+ if s.sqliteHasNewUsableAccessToken(ctx, currentAccess) {
+ return nil
+ }
return err
}
@@ -680,7 +810,13 @@ func (s *sqliteKiroAuth) refreshOIDC(ctx context.Context, state kiroSQLiteState,
if resp.ExpiresIn != nil && *resp.ExpiresIn > 0 {
expiresAt = time.Now().Add(time.Duration(*resp.ExpiresIn) * time.Second).UTC().Format(time.RFC3339)
}
- return kiroUpdateSQLiteTokens(ctx, s.path, newAccess, newRefresh, expiresAt)
+ if err := kiroUpdateSQLiteTokens(ctx, s.path, newAccess, newRefresh, expiresAt); err != nil {
+ if s.sqliteHasNewUsableAccessToken(ctx, currentAccess) {
+ return nil
+ }
+ return err
+ }
+ return nil
}
func (s *sqliteKiroAuth) refreshSocial(ctx context.Context, state kiroSQLiteState, refreshToken string) error {
@@ -747,8 +883,117 @@ func (s *sqliteKiroAuth) load(ctx context.Context) (kiroSQLiteState, error) {
return state, nil
}
+// fullLoad performs a complete reload of all sqlite state including device
+// registration (ClientID, ClientSecret). Used in Refresh() to pick up
+// credentials rotated by another process (e.g. kiro-cli login).
+func (s *sqliteKiroAuth) fullLoad(ctx context.Context) (kiroSQLiteState, error) {
+ state, err := kiroLoadSQLiteState(ctx, s.path)
+ if err != nil {
+ return kiroSQLiteState{}, err
+ }
+ s.initMu.Lock()
+ s.state = state
+ s.loaded = true
+ s.initMu.Unlock()
+ return state, nil
+}
+
+// tokenExpiringSoon checks if the token is about to expire, using instance
+// jitter to spread refresh attempts across processes.
+func (s *sqliteKiroAuth) tokenExpiringSoon(expiresAt string) bool {
+ if expiresAt == "" {
+ return false
+ }
+ t, err := time.Parse(time.RFC3339Nano, expiresAt)
+ if err != nil {
+ t, err = time.Parse(time.RFC3339, expiresAt)
+ if err != nil {
+ return false
+ }
+ }
+ return time.Until(t) < kiroTokenExpiryBuffer+s.jitter
+}
+
+// refreshFromSSOCache attempts to refresh using the AWS SSO cache as a
+// fallback when the primary OIDC refresh token is invalid/expired.
+func (s *sqliteKiroAuth) refreshFromSSOCache(ctx context.Context, previousAccess string) error {
+ cacheState, err := kiroLoadSSOCacheAuthState()
+ if err != nil {
+ return err
+ }
+ if strings.TrimSpace(cacheState.RefreshToken) == "" ||
+ strings.TrimSpace(cacheState.Region) == "" ||
+ strings.TrimSpace(cacheState.ClientID) == "" ||
+ strings.TrimSpace(cacheState.ClientSecret) == "" {
+ return fmt.Errorf("SSO cache does not contain enough data to refresh")
+ }
+
+ endpoint := fmt.Sprintf("https://oidc.%s.amazonaws.com/token", strings.TrimSpace(cacheState.Region))
+ client := newKiroOIDCClient(endpoint, nil)
+
+ resp, err := client.RefreshToken(ctx, kiroRefreshTokenRequest{
+ ClientID: cacheState.ClientID,
+ ClientSecret: cacheState.ClientSecret,
+ GrantType: "refresh_token",
+ RefreshToken: cacheState.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.RFC3339)
+ }
+ if err := kiroUpdateSQLiteTokens(ctx, s.path, newAccess, newRefresh, expiresAt); err != nil {
+ if s.sqliteHasNewUsableAccessToken(ctx, previousAccess) {
+ return nil
+ }
+ return err
+ }
+ // Also update SSO cache file.
+ _ = kiroUpdateSSOCacheToken(cacheState.TokenPath, newAccess, newRefresh, expiresAt)
+ return nil
+}
+
+// sqliteHasNewUsableAccessToken retries reading the SQLite token to detect if
+// another process refreshed during our attempt.
+func (s *sqliteKiroAuth) sqliteHasNewUsableAccessToken(ctx context.Context, previousAccess string) bool {
+ const attempts = 3
+ const backoff = 500 * time.Millisecond
+ for i := range attempts {
+ if i > 0 {
+ select {
+ case <-ctx.Done():
+ return false
+ case <-time.After(backoff):
+ }
+ }
+ access, _, expiresAt, err := kiroReadSQLiteTokens(ctx, s.path)
+ if err != nil {
+ continue
+ }
+ if access != "" && access != strings.TrimSpace(previousAccess) && !s.tokenExpiringSoon(expiresAt) {
+ return true
+ }
+ }
+ return false
+}
+
const kiroTokenExpiryBuffer = 2 * time.Minute
+// kiroTokenExpiryJitter is the maximum random offset added to kiroTokenExpiryBuffer.
+const kiroTokenExpiryJitter = 2 * time.Minute
+
+func kiroRandomJitter() time.Duration {
+ b := make([]byte, 4)
+ _, _ = rand.Read(b)
+ n := uint32(b[0]) | uint32(b[1])<<8 | uint32(b[2])<<16 | uint32(b[3])<<24
+ return time.Duration(n%uint32(kiroTokenExpiryJitter/time.Millisecond)) * time.Millisecond
+}
+
func kiroTokenExpiringSoon(expiresAt string) bool {
if expiresAt == "" {
return false
@@ -972,23 +1217,161 @@ func kiroQuerySQLiteOptional(ctx context.Context, path, query string) (string, b
return "", false, err
}
+// sqliteBusyTimeout is the per-connection busy timeout passed to the sqlite3
+// CLI via -cmd. This is session-scoped and does not modify the DB file.
+const kiroSQLiteBusyTimeout = "5000"
+
+// sqliteBusyRetries is the number of times we retry at the Go level when
+// sqlite3 still returns a busy/locked error after its internal timeout.
+const kiroSQLiteBusyRetries = 2
+
+// sqliteBusyBackoff is the delay between Go-level retries.
+const kiroSQLiteBusyBackoff = 200 * time.Millisecond
+
func kiroRunSQLite(ctx context.Context, args ...string) ([]byte, error) {
- cmd := exec.CommandContext(ctx, "sqlite3", args...)
- out, err := cmd.CombinedOutput()
- if err != nil {
- msg := strings.TrimSpace(string(out))
- if msg != "" {
- return nil, fmt.Errorf("%s", msg)
+ // Prepend busy timeout so sqlite3 retries internally on lock contention.
+ fullArgs := make([]string, 0, len(args)+2)
+ fullArgs = append(fullArgs, "-cmd", ".timeout "+kiroSQLiteBusyTimeout)
+ fullArgs = append(fullArgs, args...)
+
+ var out []byte
+ var err error
+ for attempt := range kiroSQLiteBusyRetries + 1 {
+ if attempt > 0 {
+ select {
+ case <-ctx.Done():
+ return out, ctx.Err()
+ case <-time.After(kiroSQLiteBusyBackoff):
+ }
+ }
+ cmd := exec.CommandContext(ctx, "sqlite3", fullArgs...)
+ out, err = cmd.CombinedOutput()
+ if err == nil || !kiroIsSQLiteBusy(out, err) {
+ return out, err
}
- return nil, err
}
- return out, nil
+ return out, err
+}
+
+// kiroIsSQLiteBusy checks whether the sqlite3 invocation failed due to lock
+// contention. The sqlite3 CLI returns exit code 5 (SQLITE_BUSY) or 6
+// (SQLITE_LOCKED).
+func kiroIsSQLiteBusy(output []byte, err error) bool {
+ if exitErr, ok := err.(*exec.ExitError); ok {
+ code := exitErr.ExitCode()
+ if code == 5 || code == 6 {
+ return true
+ }
+ }
+ s := strings.ToLower(string(output))
+ return strings.Contains(s, "database is locked") ||
+ strings.Contains(s, "database table is locked")
}
func kiroSQLiteQuote(s string) string {
return "'" + strings.ReplaceAll(s, "'", "''") + "'"
}
+// ── SSO cache fallback ────────────────────────────────────────────────────────
+
+type kiroSSOCacheState struct {
+ TokenPath string
+ AccessToken string
+ RefreshToken string
+ ExpiresAt string
+ Region string
+ ClientID string
+ ClientSecret string
+}
+
+func kiroLoadSSOCacheAuthState() (kiroSSOCacheState, error) {
+ homeDir, err := os.UserHomeDir()
+ if err != nil {
+ return kiroSSOCacheState{}, fmt.Errorf("cannot determine home directory for SSO cache: %w", err)
+ }
+ cacheDir := filepath.Join(homeDir, ".aws", "sso", "cache")
+ tokenPath := filepath.Join(cacheDir, "kiro-auth-token-cli.json")
+ tokenBytes, err := os.ReadFile(tokenPath)
+ if err != nil {
+ return kiroSSOCacheState{}, fmt.Errorf("read SSO token cache %q: %w", tokenPath, err)
+ }
+ var tokenState struct {
+ AccessToken string `json:"accessToken"`
+ RefreshToken string `json:"refreshToken"`
+ ExpiresAt string `json:"expiresAt"`
+ Region string `json:"region"`
+ ClientIDHash string `json:"clientIdHash"`
+ }
+ if err := json.Unmarshal(tokenBytes, &tokenState); err != nil {
+ return kiroSSOCacheState{}, fmt.Errorf("decode SSO token cache %q: %w", tokenPath, err)
+ }
+ clientIDHash := strings.TrimSpace(tokenState.ClientIDHash)
+ if clientIDHash == "" {
+ return kiroSSOCacheState{}, fmt.Errorf("SSO token cache %q did not contain clientIdHash", tokenPath)
+ }
+ clientPath := filepath.Join(cacheDir, clientIDHash+".json")
+ clientBytes, err := os.ReadFile(clientPath)
+ if err != nil {
+ return kiroSSOCacheState{}, fmt.Errorf("read SSO client cache %q: %w", clientPath, err)
+ }
+ var clientState struct {
+ ClientID string `json:"clientId"`
+ ClientSecret string `json:"clientSecret"`
+ }
+ if err := json.Unmarshal(clientBytes, &clientState); err != nil {
+ return kiroSSOCacheState{}, fmt.Errorf("decode SSO client cache %q: %w", clientPath, err)
+ }
+ return kiroSSOCacheState{
+ TokenPath: tokenPath,
+ AccessToken: strings.TrimSpace(tokenState.AccessToken),
+ RefreshToken: strings.TrimSpace(tokenState.RefreshToken),
+ ExpiresAt: strings.TrimSpace(tokenState.ExpiresAt),
+ Region: strings.TrimSpace(tokenState.Region),
+ ClientID: strings.TrimSpace(clientState.ClientID),
+ ClientSecret: strings.TrimSpace(clientState.ClientSecret),
+ }, nil
+}
+
+func kiroUpdateSSOCacheToken(path, accessToken, refreshToken, expiresAt string) error {
+ accessToken = strings.TrimSpace(accessToken)
+ if accessToken == "" {
+ return fmt.Errorf("access token must not be empty")
+ }
+ body, err := os.ReadFile(path)
+ if err != nil {
+ return fmt.Errorf("read SSO token cache %q: %w", path, err)
+ }
+ var tokenState map[string]any
+ if err := json.Unmarshal(body, &tokenState); err != nil {
+ return fmt.Errorf("decode SSO token cache %q: %w", path, err)
+ }
+ tokenState["accessToken"] = accessToken
+ if strings.TrimSpace(refreshToken) != "" {
+ tokenState["refreshToken"] = strings.TrimSpace(refreshToken)
+ }
+ if strings.TrimSpace(expiresAt) != "" {
+ tokenState["expiresAt"] = strings.TrimSpace(expiresAt)
+ }
+ updated, err := json.MarshalIndent(tokenState, "", " ")
+ if err != nil {
+ return fmt.Errorf("encode SSO token cache %q: %w", path, err)
+ }
+ updated = append(updated, '\n')
+ info, err := os.Stat(path)
+ if err != nil {
+ return fmt.Errorf("stat SSO token cache %q: %w", path, err)
+ }
+ tmpPath := path + ".tmp"
+ if err := os.WriteFile(tmpPath, updated, info.Mode().Perm()); err != nil {
+ return fmt.Errorf("write SSO token cache temp %q: %w", tmpPath, err)
+ }
+ if err := os.Rename(tmpPath, path); err != nil {
+ os.Remove(tmpPath)
+ return fmt.Errorf("rename SSO token cache %q: %w", path, err)
+ }
+ return nil
+}
+
// ── Auth source factory ───────────────────────────────────────────────────────
func newKiroAuthSource(apiKey string) (kiroAuthSource, error) {
@@ -1012,7 +1395,7 @@ func newKiroAuthSource(apiKey string) (kiroAuthSource, error) {
if strings.TrimSpace(path) == "" {
return nil, fmt.Errorf("sqlite URL %q does not contain a database path", apiKey)
}
- return &sqliteKiroAuth{path: path}, nil
+ return &sqliteKiroAuth{path: path, jitter: kiroRandomJitter()}, nil
}
func kiroDefaultSQLiteKey() string {
diff --git a/pkg/backend/codewhisperer_internal.go b/pkg/backend/codewhisperer_internal.go
index 2367933..7d45e5e 100644
--- a/pkg/backend/codewhisperer_internal.go
+++ b/pkg/backend/codewhisperer_internal.go
@@ -21,6 +21,7 @@ import (
"io"
"maps"
"net/http"
+ "net/url"
"strings"
)
@@ -33,11 +34,11 @@ const (
kiroAgentTaskType = "vibe"
kiroAmzSDKRequest = "attempt=1; max=3"
- kiroUserAgent = "aws-sdk-rust/1.3.14 ua/2.1 api/codewhispererstreaming/0.1.16551 os/linux lang/rust/1.92.0 md/appVersion-0.12.333 app/AmazonQ-For-CLI"
- kiroAmzUserAgent = "aws-sdk-rust/1.3.14 ua/2.1 api/codewhispererstreaming/0.1.16551 os/linux lang/rust/1.92.0 m/F app/AmazonQ-For-CLI"
+ kiroUserAgent = "aws-sdk-rust/1.3.15 ua/2.1 api/codewhispererstreaming/0.1.16551 os/linux lang/rust/1.92.0 md/appVersion-2.5.1 app/AmazonQ-For-CLI"
+ kiroAmzUserAgent = "aws-sdk-rust/1.3.15 ua/2.1 api/codewhispererstreaming/0.1.16551 os/linux lang/rust/1.92.0 m/F app/AmazonQ-For-CLI"
- kiroOIDCUserAgent = "aws-sdk-rust/1.3.10 os/linux lang/rust/1.92.0"
- kiroOIDCAmzUserAgent = "aws-sdk-rust/1.3.10 ua/2.1 api/ssooidc/1.92.0 os/linux lang/rust/1.92.0 m/E app/AmazonQ-For-CLI"
+ kiroOIDCUserAgent = "aws-sdk-rust/1.3.15 os/linux lang/rust/1.92.0"
+ kiroOIDCAmzUserAgent = "aws-sdk-rust/1.3.15 ua/2.1 api/ssooidc/1.100.0 os/linux lang/rust/1.92.0 m/E,N app/AmazonQ-For-CLI"
)
type kiroRefreshTokenRequest struct {
@@ -54,8 +55,24 @@ type kiroRefreshTokenResponse struct {
}
type kiroGenerateRequest struct {
- ConversationState kiroConversationState `json:"conversationState"`
- ProfileARN string `json:"profileArn,omitempty"`
+ ConversationState kiroConversationState `json:"conversationState"`
+ ProfileARN string `json:"profileArn,omitempty"`
+ AdditionalModelRequestFields *kiroModelRequestFields `json:"additionalModelRequestFields,omitempty"`
+}
+
+type kiroModelRequestFields struct {
+ Thinking *kiroThinkingConfig `json:"thinking,omitempty"`
+ OutputConfig *kiroOutputConfig `json:"output_config,omitempty"`
+ MaxTokens *int `json:"max_tokens,omitempty"`
+}
+
+type kiroThinkingConfig struct {
+ Type string `json:"type"`
+ Display string `json:"display,omitempty"`
+}
+
+type kiroOutputConfig struct {
+ Effort string `json:"effort,omitempty"`
}
type kiroConversationState struct {
@@ -73,18 +90,13 @@ type kiroChatMessage struct {
}
type kiroUserInputMessage struct {
- Content json.RawMessage `json:"content"`
+ Content string `json:"content"`
UserInputMessageContext *kiroUserInputContext `json:"userInputMessageContext,omitempty"`
Origin string `json:"origin,omitempty"`
Images []kiroImageBlock `json:"images,omitempty"`
ModelID string `json:"modelId,omitempty"`
}
-type kiroContentBlock struct {
- Text string `json:"text,omitempty"`
- Image *kiroImageBlock `json:"image,omitempty"`
-}
-
type kiroImageBlock struct {
Format string `json:"format"`
Source kiroImageSource `json:"source"`
@@ -97,6 +109,12 @@ type kiroImageSource struct {
type kiroUserInputContext struct {
ToolResults []kiroToolResult `json:"toolResults,omitempty"`
Tools []kiroTool `json:"tools,omitempty"`
+ EnvState *kiroEnvState `json:"envState,omitempty"`
+}
+
+type kiroEnvState struct {
+ OperatingSystem string `json:"operatingSystem,omitempty"`
+ CurrentWorkingDirectory string `json:"currentWorkingDirectory,omitempty"`
}
type kiroAssistantResponseMessage struct {
@@ -145,6 +163,12 @@ type kiroAssistantEvent struct {
Content string `json:"content"`
}
+type kiroReasoningContentEvent struct {
+ Text string `json:"text,omitempty"`
+ RedactedContent string `json:"redactedContent,omitempty"`
+ Signature string `json:"signature,omitempty"`
+}
+
type kiroToolUseEvent struct {
ToolUseID string `json:"toolUseId"`
Name string `json:"name"`
@@ -350,11 +374,22 @@ func (e *kiroAPIError) isThrottled() bool {
}
func (e *kiroAPIError) isTransient() bool {
- return e.StatusCode == http.StatusInternalServerError || e.Code == "InternalServerException"
+ switch e.StatusCode {
+ case http.StatusRequestTimeout, // 408
+ http.StatusInternalServerError, // 500
+ http.StatusBadGateway, // 502
+ http.StatusServiceUnavailable: // 503
+ return true
+ }
+ code := strings.ToLower(e.Code)
+ return strings.Contains(code, "insufficientmodelcapacity") ||
+ strings.Contains(code, "serviceunavailable") ||
+ strings.Contains(code, "internalserver")
}
type kiroStreamCallbacks struct {
OnAssistantDelta func(kiroAssistantEvent)
+ OnReasoningDelta func(kiroReasoningContentEvent)
}
func newKiroAPIClient(endpoint, token string, extraHeaders map[string]string, httpClient *http.Client) *kiroAPIClient {
@@ -426,6 +461,14 @@ func (c *kiroAPIClient) Stream(
if callbacks.OnAssistantDelta != nil {
callbacks.OnAssistantDelta(event)
}
+ case "reasoningContentEvent":
+ var event kiroReasoningContentEvent
+ if err := json.Unmarshal(frame.Payload, &event); err != nil {
+ return nil, fmt.Errorf("decode reasoning content event: %w", err)
+ }
+ if callbacks.OnReasoningDelta != nil {
+ callbacks.OnReasoningDelta(event)
+ }
case "toolUseEvent":
var event kiroToolUseEvent
if err := json.Unmarshal(frame.Payload, &event); err != nil {
@@ -483,7 +526,7 @@ func (c *kiroAPIClient) applyHeaders(req *http.Request) error {
req.Header.Set("X-Amz-User-Agent", kiroAmzUserAgent)
req.Header.Set("Amz-Sdk-Request", kiroAmzSDKRequest)
req.Header.Set("Amz-Sdk-Invocation-Id", invocationID)
- req.Header.Set("X-Amzn-Codewhisperer-Optout", "false")
+ req.Header.Set("X-Amzn-Codewhisperer-Optout", "true")
for k, v := range c.extraHeaders {
req.Header.Set(k, v)
}
@@ -553,11 +596,25 @@ func kiroDefaultEndpoint(profileARN string) (string, error) {
if region == "" {
return "", fmt.Errorf("profile ARN %q does not contain a region", profileARN)
}
- return "https://q." + region + ".amazonaws.com/", nil
+ return "https://runtime." + region + ".kiro.dev/", nil
}
func kiroEndpointForRegion(region string) string {
- return "https://q." + strings.TrimSpace(region) + ".amazonaws.com/"
+ return "https://runtime." + strings.TrimSpace(region) + ".kiro.dev/"
+}
+
+func kiroManagementEndpointFor(endpoint string) (string, error) {
+ parsed, err := url.Parse(strings.TrimSpace(endpoint))
+ if err != nil {
+ return "", fmt.Errorf("parse endpoint: %w", err)
+ }
+ if parsed.Host == "" {
+ return "", fmt.Errorf("endpoint %q does not contain a host", endpoint)
+ }
+ if strings.HasPrefix(parsed.Host, "runtime.") && strings.HasSuffix(parsed.Host, ".kiro.dev") {
+ parsed.Host = "management." + strings.TrimPrefix(parsed.Host, "runtime.")
+ }
+ return strings.TrimRight(parsed.String(), "/") + "/", nil
}
func kiroNewUUID() (string, error) {
@@ -638,8 +695,8 @@ func (c *kiroOIDCClient) RefreshToken(ctx context.Context, req kiroRefreshTokenR
const (
kiroListModelsTarget = "AmazonCodeWhispererService.ListAvailableModels"
- kiroListModelsUserAgent = "aws-sdk-rust/1.3.14 ua/2.1 api/codewhispererruntime/0.1.16551 os/linux lang/rust/1.92.0 md/appVersion-0.12.333 app/AmazonQ-For-CLI"
- kiroListModelsAmzUserAgent = "aws-sdk-rust/1.3.14 ua/2.1 api/codewhispererruntime/0.1.16551 os/linux lang/rust/1.92.0 m/F,C app/AmazonQ-For-CLI"
+ kiroListModelsUserAgent = "aws-sdk-rust/1.3.15 ua/2.1 api/codewhispererruntime/0.1.16551 os/linux lang/rust/1.92.0 md/appVersion-2.5.1 app/AmazonQ-For-CLI"
+ kiroListModelsAmzUserAgent = "aws-sdk-rust/1.3.15 ua/2.1 api/codewhispererruntime/0.1.16551 os/linux lang/rust/1.92.0 m/F,C app/AmazonQ-For-CLI"
)
type kiroListModelsRequest struct {
@@ -670,7 +727,24 @@ func (c *kiroAPIClient) ListModels(ctx context.Context, profileARN string) (*kir
if err != nil {
return nil, err
}
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint, bytes.NewReader(body))
+ // ListModels uses the management endpoint (management.*.kiro.dev) rather
+ // than the streaming endpoint (runtime.*.kiro.dev).
+ mgmtEndpoint, err := kiroManagementEndpointFor(c.endpoint)
+ if err != nil {
+ mgmtEndpoint = c.endpoint // fall back to streaming endpoint
+ }
+ targetURL, err := url.Parse(mgmtEndpoint)
+ if err != nil {
+ return nil, fmt.Errorf("parse management endpoint: %w", err)
+ }
+ query := targetURL.Query()
+ query.Set("origin", kiroDefaultOrigin)
+ if profileARN != "" {
+ query.Set("profileArn", profileARN)
+ }
+ targetURL.RawQuery = query.Encode()
+
+ req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL.String(), bytes.NewReader(body))
if err != nil {
return nil, err
}