kiro-proxy 0.1.15

package/server.js

@@ -0,0 +1,377 @@
+#!/usr/bin/env node
+import express from 'express';
+import crypto from 'crypto';
+import { getAccessToken } from './token-reader.js';
+import { createClient, chat, chatStream, listAvailableModels } from './q-client.js';
+import { c, log, logSummary, reqId, tagError } from './logger.js';
+import { countMessages, countContent } from './token-counter.js';
+import { recordUsage, queryUsage, todaySummary } from './usage-tracker.js';
+
+const app = express();
+app.use(express.json({ limit: '10mb' }));
+
+const PORT = process.env.PORT || 3456;
+
+let cachedClient = null;
+let cachedToken = null;
+
+async function getClient() {
+  const tokenData = await getAccessToken();
+  if (!cachedClient || cachedToken !== tokenData.accessToken) {
+    cachedClient = createClient(tokenData.accessToken, {
+      authMethod: tokenData.authMethod,
+      profileArn: tokenData.profileArn,
+      provider: tokenData.provider,
+    });
+    cachedToken = tokenData.accessToken;
+  }
+  return { client: cachedClient, tokenData };
+}
+
+function msgId() {
+  return `msg_${crypto.randomUUID().replace(/-/g, '').slice(0, 20)}`;
+}
+
+// ============================================================
+// POST /v1/messages — Anthropic Messages API (with tool support)
+// ============================================================
+app.post('/v1/messages', async (req, res) => {
+  try {
+    const { model, messages, system, tools, stream, max_tokens, tool_choice } = req.body;
+    if (!messages?.length) {
+      return res.status(400).json({ type: 'error', error: { type: 'invalid_request_error', message: 'messages required' } });
+    }
+
+    const { client, tokenData } = await getClient();
+    const opts = { messages, system, tools, profileArn: tokenData.profileArn, modelId: model };
+    const rid = reqId();
+    const start = Date.now();
+
+    log('POST', '/v1/messages', rid, {
+      model: model || 'default',
+      stream: !!stream,
+      messages: messages.length,
+      tools: tools?.length || 0,
+    });
+
+    if (stream) {
+      res.setHeader('Content-Type', 'text/event-stream');
+      res.setHeader('Cache-Control', 'no-cache');
+      res.setHeader('Connection', 'keep-alive');
+
+      const id = msgId();
+      const usedModel = model || 'q-developer';
+      let blockIndex = 0;
+      let hasTextBlock = false;
+      const inputTokens = countMessages(messages, system);
+
+      // message_start
+      const send = (event, data) => { res.write(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`); };
+
+      send('message_start', {
+        type: 'message_start',
+        message: {
+          id, type: 'message', role: 'assistant', content: [],
+          model: usedModel, stop_reason: null, stop_sequence: null,
+          usage: { input_tokens: inputTokens, output_tokens: 0 },
+        },
+      });
+      send('ping', { type: 'ping' });
+
+      try {
+        let hasToolUse = false;
+        let hasThinkingBlock = false;
+        let summary;
+        const outputParts = [];
+
+        for await (const chunk of chatStream(client, opts)) {
+          if (chunk.type === 'thinking') {
+            // 开启 thinking 块（如果还没有）
+            if (!hasThinkingBlock) {
+              send('content_block_start', {
+                type: 'content_block_start', index: blockIndex,
+                content_block: { type: 'thinking', thinking: '' },
+              });
+              hasThinkingBlock = true;
+            }
+            outputParts.push(chunk.text);
+            send('content_block_delta', {
+              type: 'content_block_delta', index: blockIndex,
+              delta: { type: 'thinking_delta', thinking: chunk.text },
+            });
+          } else if (chunk.type === 'thinking_signature') {
+            // 关闭 thinking 块，附带 signature
+            if (hasThinkingBlock) {
+              send('content_block_delta', {
+                type: 'content_block_delta', index: blockIndex,
+                delta: { type: 'signature_delta', signature: chunk.signature },
+              });
+              send('content_block_stop', { type: 'content_block_stop', index: blockIndex });
+              blockIndex++;
+              hasThinkingBlock = false;
+            }
+          } else if (chunk.type === 'content') {
+            // 关闭未关闭的 thinking 块
+            if (hasThinkingBlock) {
+              send('content_block_stop', { type: 'content_block_stop', index: blockIndex });
+              blockIndex++;
+              hasThinkingBlock = false;
+            }
+            // 开启文本块（如果还没有）
+            if (!hasTextBlock) {
+              send('content_block_start', {
+                type: 'content_block_start', index: blockIndex,
+                content_block: { type: 'text', text: '' },
+              });
+              hasTextBlock = true;
+            }
+            outputParts.push(chunk.content);
+            send('content_block_delta', {
+              type: 'content_block_delta', index: blockIndex,
+              delta: { type: 'text_delta', text: chunk.content },
+            });
+          } else if (chunk.type === 'tool_use_start') {
+            // 关闭之前的 thinking 块
+            if (hasThinkingBlock) {
+              send('content_block_stop', { type: 'content_block_stop', index: blockIndex });
+              blockIndex++;
+              hasThinkingBlock = false;
+            }
+            // 关闭之前的文本块
+            if (hasTextBlock) {
+              send('content_block_stop', { type: 'content_block_stop', index: blockIndex });
+              blockIndex++;
+              hasTextBlock = false;
+            }
+          } else if (chunk.type === 'tool_use_end') {
+            hasToolUse = true;
+            outputParts.push(JSON.stringify(chunk.input));
+            // 发送完整的 tool_use content block
+            send('content_block_start', {
+              type: 'content_block_start', index: blockIndex,
+              content_block: { type: 'tool_use', id: chunk.toolUseId, name: chunk.name, input: {} },
+            });
+            // 发送 input_json_delta（完整 JSON 一次性发送）
+            send('content_block_delta', {
+              type: 'content_block_delta', index: blockIndex,
+              delta: { type: 'input_json_delta', partial_json: JSON.stringify(chunk.input) },
+            });
+            send('content_block_stop', { type: 'content_block_stop', index: blockIndex });
+            blockIndex++;
+          } else if (chunk.type === 'summary') {
+            summary = chunk.stats;
+            if (typeof chunk.meteringUsage === 'number') recordUsage(chunk.meteringUsage, model);
+          }
+        }
+
+        // 关闭最后的 thinking 块
+        if (hasThinkingBlock) {
+          send('content_block_stop', { type: 'content_block_stop', index: blockIndex });
+        }
+
+        // 关闭最后的文本块
+        if (hasTextBlock) {
+          send('content_block_stop', { type: 'content_block_stop', index: blockIndex });
+        }
+
+        const stopReason = hasToolUse ? 'tool_use' : 'end_turn';
+        const outputTokens = countContent(outputParts.join(''));
+        send('message_delta', {
+          type: 'message_delta',
+          delta: { stop_reason: stopReason, stop_sequence: null },
+          usage: { output_tokens: outputTokens },
+        });
+        send('message_stop', { type: 'message_stop' });
+        res.end();
+        const s = summary || {};
+        s.estTokens = `~tokens: in=${inputTokens} out=${outputTokens}`;
+        logSummary(rid, Date.now() - start, s);
+      } catch (err) {
+        tagError('stream', err.message);
+        res.write(`event: error\ndata: ${JSON.stringify({ type: 'error', error: { type: 'api_error', message: err.message } })}\n\n`);
+        res.end();
+      }
+    } else {
+      // 非流式
+      const result = await chat(client, opts);
+      if (typeof result.meteringUsage === 'number') recordUsage(result.meteringUsage, model);
+      const inputTokens = countMessages(messages, system);
+      const outputTokens = countContent(result.content);
+      const s = result.stats || {};
+      s.estTokens = `~tokens: in=${inputTokens} out=${outputTokens}`;
+      logSummary(rid, Date.now() - start, s);
+      res.json({
+        id: msgId(), type: 'message', role: 'assistant',
+        content: result.content,
+        model: model || 'q-developer',
+        stop_reason: result.stopReason,
+        stop_sequence: null,
+        usage: { input_tokens: inputTokens, output_tokens: outputTokens },
+      });
+    }
+  } catch (err) {
+    tagError('anthropic', err.message || err);
+    const status = err.message?.includes('expired') ? 401 : 500;
+    res.status(status).json({ type: 'error', error: { type: status === 401 ? 'authentication_error' : 'api_error', message: err.message } });
+  }
+});
+
+// ============================================================
+// POST /v1/chat/completions — OpenAI compatible
+// ============================================================
+app.post('/v1/chat/completions', async (req, res) => {
+  try {
+    const { messages: rawMsgs, model, stream } = req.body;
+    if (!rawMsgs?.length) return res.status(400).json({ error: 'messages required' });
+
+    // 简单转换 OpenAI → Anthropic 格式
+    let system;
+    const messages = [];
+    for (const m of rawMsgs) {
+      if (m.role === 'system') { system = m.content; continue; }
+      messages.push({ role: m.role, content: m.content });
+    }
+
+    const { client, tokenData } = await getClient();
+    const opts = { messages, system, profileArn: tokenData.profileArn, modelId: model };
+    const rid = reqId();
+    const start = Date.now();
+
+    log('POST', '/v1/chat/completions', rid, {
+      model: model || 'default',
+      stream: !!stream,
+      messages: messages.length,
+    });
+
+    if (stream) {
+      res.setHeader('Content-Type', 'text/event-stream');
+      res.setHeader('Cache-Control', 'no-cache');
+      res.setHeader('Connection', 'keep-alive');
+      const responseId = `chatcmpl-${crypto.randomUUID()}`;
+      const created = Math.floor(Date.now() / 1000);
+      const inputTokens = countMessages(messages, system);
+      let summary;
+      const outputParts = [];
+
+      for await (const chunk of chatStream(client, opts)) {
+        if (chunk.type === 'content') {
+          outputParts.push(chunk.content);
+          res.write(`data: ${JSON.stringify({
+            id: responseId, object: 'chat.completion.chunk', created,
+            model: model || 'q-developer',
+            choices: [{ index: 0, delta: { content: chunk.content }, finish_reason: null }],
+          })}\n\n`);
+        } else if (chunk.type === 'summary') {
+          summary = chunk.stats;
+          if (typeof chunk.meteringUsage === 'number') recordUsage(chunk.meteringUsage, model);
+        }
+      }
+      res.write(`data: ${JSON.stringify({
+        id: responseId, object: 'chat.completion.chunk', created,
+        model: model || 'q-developer',
+        choices: [{ index: 0, delta: {}, finish_reason: 'stop' }],
+      })}\n\n`);
+      res.write('data: [DONE]\n\n');
+      res.end();
+      const outputTokens = countContent(outputParts.join(''));
+      const s = summary || {};
+      s.estTokens = `~tokens: in=${inputTokens} out=${outputTokens}`;
+      logSummary(rid, Date.now() - start, s);
+    } else {
+      const result = await chat(client, opts);
+      if (typeof result.meteringUsage === 'number') recordUsage(result.meteringUsage, model);
+      const text = result.content.filter(b => b.type === 'text').map(b => b.text).join('');
+      const promptTokens = countMessages(messages, system);
+      const completionTokens = countContent(text);
+      const s = result.stats || {};
+      s.estTokens = `~tokens: in=${promptTokens} out=${completionTokens}`;
+      logSummary(rid, Date.now() - start, s);
+      res.json({
+        id: `chatcmpl-${crypto.randomUUID()}`, object: 'chat.completion',
+        created: Math.floor(Date.now() / 1000), model: model || 'q-developer',
+        choices: [{ index: 0, message: { role: 'assistant', content: text }, finish_reason: 'stop' }],
+        usage: { prompt_tokens: promptTokens, completion_tokens: completionTokens, total_tokens: promptTokens + completionTokens },
+      });
+    }
+  } catch (err) {
+    tagError('openai', err.message || err);
+    res.status(500).json({ error: { message: err.message } });
+  }
+});
+
+// ============================================================
+// GET /v1/models
+// ============================================================
+app.get('/v1/models', async (_req, res) => {
+  try {
+    const tokenData = await getAccessToken();
+    const { models, defaultModel } = await listAvailableModels(tokenData.accessToken, {
+      profileArn: tokenData.profileArn, authMethod: tokenData.authMethod, provider: tokenData.provider,
+    });
+    res.json({
+      object: 'list',
+      data: models.map(m => ({
+        id: m.modelId, object: 'model', created: Math.floor(Date.now() / 1000), owned_by: 'amazon',
+        name: m.modelName || m.modelId, description: m.description,
+        is_default: defaultModel?.modelId === m.modelId,
+      })),
+    });
+  } catch (err) {
+    res.status(500).json({ error: { message: err.message } });
+  }
+});
+
+app.get('/q/models', async (_req, res) => {
+  try {
+    const tokenData = await getAccessToken();
+    const result = await listAvailableModels(tokenData.accessToken, {
+      profileArn: tokenData.profileArn, authMethod: tokenData.authMethod, provider: tokenData.provider,
+    });
+    res.json(result);
+  } catch (err) {
+    res.status(500).json({ error: { message: err.message } });
+  }
+});
+
+app.get('/health', async (_req, res) => {
+  try {
+    const tokenData = await getAccessToken();
+    const expired = tokenData.expiresAt && new Date(tokenData.expiresAt) < new Date();
+    res.json({ status: expired ? 'token_expired' : 'ok', provider: tokenData.provider || 'unknown', expiresAt: tokenData.expiresAt });
+  } catch (err) {
+    res.status(503).json({ status: 'error', message: err.message });
+  }
+});
+
+// ============================================================
+// GET /credits — Usage statistics
+// ============================================================
+app.get('/credits', (_req, res) => {
+  const period = _req.query.period || 'today';
+  res.json(queryUsage(period));
+});
+
+app.listen(PORT, async () => {
+  console.log(`${c.cyan}Kiro Proxy${c.reset} running on ${c.green}http://localhost:${PORT}${c.reset}`);
+  console.log(`  ${c.gray}Anthropic:${c.reset} http://localhost:${PORT}/v1/messages`);
+  console.log(`  ${c.gray}OpenAI:   ${c.reset} http://localhost:${PORT}/v1/chat/completions`);
+  console.log(`  ${c.gray}Models:   ${c.reset} http://localhost:${PORT}/v1/models`);
+  console.log(`  ${c.gray}Credits: ${c.reset} http://localhost:${PORT}/credits`);
+  try {
+    const t = await getAccessToken();
+    console.log(`  ${c.gray}Provider: ${c.yellow}${t.provider || 'unknown'}${c.reset}, Expires: ${c.dim}${t.expiresAt || 'unknown'}${c.reset}`);
+  } catch (err) {
+    console.warn(`  ${c.yellow}Warning:${c.reset} ${err.message}`);
+  }
+});
+
+function shutdown() {
+  const today = todaySummary();
+  if (today.requests > 0) {
+    console.log(`\n${c.cyan}Today:${c.reset} ${c.yellow}${today.credits.toFixed(4)} credits${c.reset} (${today.requests} requests)`);
+  }
+  process.exit(0);
+}
+
+process.on('SIGINT', shutdown);
+process.on('SIGTERM', shutdown);

package/token-counter.js

@@ -0,0 +1,33 @@
+function countText(text) {
+  if (!text) return 0;
+  return Math.round(text.length / 4);
+}
+
+export function countMessages(messages, system) {
+  let tokens = 0;
+  if (system) tokens += countText(typeof system === 'string' ? system : JSON.stringify(system));
+  for (const msg of messages || []) {
+    tokens += 4;
+    if (typeof msg.content === 'string') {
+      tokens += countText(msg.content);
+    } else if (Array.isArray(msg.content)) {
+      for (const block of msg.content) {
+        if (block.type === 'text') tokens += countText(block.text);
+        else if (block.type === 'tool_result') tokens += countText(typeof block.content === 'string' ? block.content : JSON.stringify(block.content));
+        else if (block.type === 'tool_use') tokens += countText(JSON.stringify(block.input));
+      }
+    }
+  }
+  return tokens;
+}
+
+export function countContent(content) {
+  if (typeof content === 'string') return countText(content);
+  let tokens = 0;
+  for (const block of content || []) {
+    if (block.type === 'text') tokens += countText(block.text);
+    else if (block.type === 'tool_use') tokens += countText(JSON.stringify(block.input));
+    else if (block.type === 'thinking') tokens += countText(block.thinking);
+  }
+  return tokens;
+}

package/token-reader.js

@@ -0,0 +1,246 @@
+import fs from 'fs';
+import path from 'path';
+import os from 'os';
+import { tagLog, tagWarn, tagError } from './logger.js';
+
+const SSO_CACHE_DIR = path.join(os.homedir(), '.aws', 'sso', 'cache');
+const KIRO_TOKEN_FILE = 'kiro-auth-token.json';
+
+// Social 刷新 URL
+const SOCIAL_REFRESH_URL = 'https://prod.us-east-1.auth.desktop.kiro.dev/refreshToken';
+
+// token 提前刷新的缓冲时间（5 分钟）
+const REFRESH_BUFFER_MS = 5 * 60 * 1000;
+
+// Kiro profile 缓存路径
+const KIRO_PROFILE_PATHS = [
+  path.join(os.homedir(), 'Library', 'Application Support', 'Kiro', 'User', 'globalStorage', 'kiro.kiroagent', 'profile.json'),
+  path.join(os.homedir(), '.config', 'Kiro', 'User', 'globalStorage', 'kiro.kiroagent', 'profile.json'),
+  path.join(os.homedir(), 'AppData', 'Roaming', 'Kiro', 'User', 'globalStorage', 'kiro.kiroagent', 'profile.json'),
+];
+
+// 内存缓存
+let cachedToken = null;
+let refreshPromise = null;
+
+/**
+ * 读取 ~/.aws/sso/cache/kiro-auth-token.json
+ */
+function readKiroToken() {
+  const tokenPath = path.join(SSO_CACHE_DIR, KIRO_TOKEN_FILE);
+  if (!fs.existsSync(tokenPath)) return null;
+  try {
+    return JSON.parse(fs.readFileSync(tokenPath, 'utf8'));
+  } catch { return null; }
+}
+
+/**
+ * 写回 token 到磁盘（让 Kiro 也能用刷新后的 token）
+ */
+function writeKiroToken(tokenData) {
+  try {
+    const tokenPath = path.join(SSO_CACHE_DIR, KIRO_TOKEN_FILE);
+    fs.mkdirSync(SSO_CACHE_DIR, { recursive: true });
+    fs.writeFileSync(tokenPath, JSON.stringify(tokenData, null, 2));
+  } catch (err) {
+    tagWarn('token', 'Failed to write token to disk:', err.message);
+  }
+}
+
+/**
+ * 读取 Kiro profile 缓存
+ */
+function readKiroProfile() {
+  for (const p of KIRO_PROFILE_PATHS) {
+    try {
+      if (fs.existsSync(p)) return JSON.parse(fs.readFileSync(p, 'utf8'));
+    } catch { /* skip */ }
+  }
+  return null;
+}
+
+/**
+ * 读取 IdC 的 client registration（从 ~/.aws/sso/cache/{hash}.json）
+ */
+function readClientRegistration(clientIdHash) {
+  if (!clientIdHash) return null;
+  const filePath = path.join(SSO_CACHE_DIR, `${clientIdHash}.json`);
+  try {
+    if (fs.existsSync(filePath)) return JSON.parse(fs.readFileSync(filePath, 'utf8'));
+  } catch { /* skip */ }
+  return null;
+}
+
+/**
+ * 判断 token 是否过期或即将过期
+ */
+function isTokenExpired(tokenData) {
+  if (!tokenData?.expiresAt) return true;
+  return new Date(tokenData.expiresAt).getTime() < Date.now() + REFRESH_BUFFER_MS;
+}
+
+// ============================================================
+// Social token 刷新（Google / Github 登录）
+// POST https://prod.us-east-1.auth.desktop.kiro.dev/refreshToken
+// ============================================================
+async function refreshSocialToken(tokenData) {
+  tagLog('token', 'Refreshing Social token...');
+  const res = await fetch(SOCIAL_REFRESH_URL, {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify({ refreshToken: tokenData.refreshToken }),
+  });
+
+  if (!res.ok) {
+    const body = await res.text();
+    throw new Error(`Social token refresh failed (${res.status}): ${body}`);
+  }
+
+  const data = await res.json();
+  // 响应: { accessToken, refreshToken?, expiresIn, profileArn? }
+  const now = new Date();
+  const expiresAt = new Date(now.getTime() + (data.expiresIn || 3600) * 1000).toISOString();
+
+  return {
+    ...tokenData,
+    accessToken: data.accessToken,
+    // refreshToken 可能会更新
+    ...(data.refreshToken && { refreshToken: data.refreshToken }),
+    ...(data.profileArn && { profileArn: data.profileArn }),
+    expiresAt,
+  };
+}
+
+// ============================================================
+// IdC token 刷新（Enterprise / BuilderId）
+// POST https://oidc.us-east-1.amazonaws.com/token
+// ============================================================
+async function refreshIdCToken(tokenData) {
+  tagLog('token', 'Refreshing IdC token...');
+
+  // 读取 client registration
+  const clientReg = readClientRegistration(tokenData.clientIdHash);
+  if (!clientReg?.clientId || !clientReg?.clientSecret) {
+    throw new Error('IdC refresh failed: no valid client registration found. Please re-login in Kiro.');
+  }
+
+  const region = tokenData.region || 'us-east-1';
+  const endpoint = `https://oidc.${region}.amazonaws.com/token`;
+
+  const res = await fetch(endpoint, {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify({
+      clientId: clientReg.clientId,
+      clientSecret: clientReg.clientSecret,
+      grantType: 'refresh_token',
+      refreshToken: tokenData.refreshToken,
+    }),
+  });
+
+  if (!res.ok) {
+    const body = await res.text();
+    throw new Error(`IdC token refresh failed (${res.status}): ${body}`);
+  }
+
+  const data = await res.json();
+  const now = new Date();
+  const expiresAt = new Date(now.getTime() + (data.expiresIn || 3600) * 1000).toISOString();
+
+  return {
+    ...tokenData,
+    accessToken: data.accessToken,
+    ...(data.refreshToken && { refreshToken: data.refreshToken }),
+    expiresAt,
+  };
+}
+
+// ============================================================
+// 刷新调度
+// ============================================================
+async function refreshToken(tokenData) {
+  const method = tokenData.authMethod;
+  if (method === 'social' || method === 'Social') {
+    return refreshSocialToken(tokenData);
+  }
+  if (method === 'IdC' || method === 'idc') {
+    return refreshIdCToken(tokenData);
+  }
+  throw new Error(`Unknown auth method: ${method}. Cannot refresh token.`);
+}
+
+/**
+ * 获取可用的 access token
+ * - 优先使用内存缓存
+ * - 过期时自动刷新（带去重，避免并发刷新）
+ * - 刷新后写回磁盘
+ * - 如果没有 profileArn，从 Kiro profile 缓存补充
+ */
+export async function getAccessToken() {
+  // 1. 内存缓存未过期，直接返回
+  if (cachedToken && !isTokenExpired(cachedToken)) {
+    return cachedToken;
+  }
+
+  // 2. 从磁盘读取
+  let tokenData = readKiroToken();
+  if (!tokenData?.accessToken) {
+    throw new Error('No token found in ~/.aws/sso/cache/kiro-auth-token.json. Please login in Kiro first.');
+  }
+
+  // 3. 如果未过期，缓存并返回
+  if (!isTokenExpired(tokenData)) {
+    tokenData = enrichWithProfile(tokenData);
+    cachedToken = tokenData;
+    return tokenData;
+  }
+
+  // 4. 过期了，需要刷新
+  if (!tokenData.refreshToken) {
+    throw new Error('Token expired and no refreshToken available. Please re-login in Kiro.');
+  }
+
+  // 去重：如果已经有刷新在进行，等待它完成
+  if (refreshPromise) {
+    tagLog('token', 'Waiting for ongoing refresh...');
+    return refreshPromise;
+  }
+
+  refreshPromise = (async () => {
+    try {
+      tagLog('token', `Token expired (${tokenData.expiresAt}), refreshing...`);
+      const newToken = await refreshToken(tokenData);
+      const enriched = enrichWithProfile(newToken);
+
+      // 写回磁盘
+      writeKiroToken(enriched);
+      cachedToken = enriched;
+
+      tagLog('token', `Token refreshed, new expiry: ${enriched.expiresAt}`);
+      return enriched;
+    } catch (err) {
+      tagError('token', 'Refresh failed:', err.message);
+      // 刷新失败，如果旧 token 还没完全过期（只是在缓冲期内），仍然可以用
+      if (tokenData.expiresAt && new Date(tokenData.expiresAt) > new Date()) {
+        tagWarn('token', 'Using existing token despite refresh failure');
+        cachedToken = enrichWithProfile(tokenData);
+        return cachedToken;
+      }
+      throw err;
+    } finally {
+      refreshPromise = null;
+    }
+  })();
+
+  return refreshPromise;
+}
+
+function enrichWithProfile(tokenData) {
+  if (!tokenData.profileArn) {
+    const profile = readKiroProfile();
+    if (profile?.arn) {
+      tokenData.profileArn = profile.arn;
+    }
+  }
+  return tokenData;
+}