Merge pull request #180 from tickernelz/fix-kiro-token-count
refactor(claude-kiro): improve token counting accuracy using API contextUsagePercentage
This commit is contained in:
何夕2077
GitHub
commit
d4e5a25c9d
1 changed files with 359 additions and 281 deletions
|
|
@ -10,6 +10,7 @@ import { getProviderModels } from '../provider-models.js';
|
|||
import { countTokens } from '@anthropic-ai/tokenizer';
|
||||
import { configureAxiosProxy } from '../proxy-utils.js';
|
||||
import { isRetryableNetworkError } from '../common.js';
|
||||
import { CLAUDE_DEFAULT_MAX_TOKENS } from '../converters/utils.js';
|
||||
|
||||
const KIRO_CONSTANTS = {
|
||||
REFRESH_URL: 'https://prod.{{region}}.auth.desktop.kiro.dev/refreshToken',
|
||||
|
|
@ -33,9 +34,9 @@ const KIRO_MODELS = getProviderModels('claude-kiro-oauth');
|
|||
|
||||
// 完整的模型映射表
|
||||
const FULL_MODEL_MAPPING = {
|
||||
"claude-opus-4-5":"claude-opus-4.5",
|
||||
"claude-opus-4-5-20251101":"claude-opus-4.5",
|
||||
"claude-haiku-4-5":"claude-haiku-4.5",
|
||||
"claude-opus-4-5": "claude-opus-4.5",
|
||||
"claude-opus-4-5-20251101": "claude-opus-4.5",
|
||||
"claude-haiku-4-5": "claude-haiku-4.5",
|
||||
"claude-sonnet-4-5": "CLAUDE_SONNET_4_5_20250929_V1_0",
|
||||
"claude-sonnet-4-5-20250929": "CLAUDE_SONNET_4_5_20250929_V1_0",
|
||||
"claude-sonnet-4-20250514": "CLAUDE_SONNET_4_20250514_V1_0",
|
||||
|
|
@ -370,197 +371,197 @@ export class KiroApiService {
|
|||
this.isInitialized = true;
|
||||
}
|
||||
|
||||
async initializeAuth(forceRefresh = false) {
|
||||
if (this.accessToken && !forceRefresh) {
|
||||
console.debug('[Kiro Auth] Access token already available and not forced refresh.');
|
||||
return;
|
||||
}
|
||||
|
||||
// Helper to load credentials from a file
|
||||
const loadCredentialsFromFile = async (filePath) => {
|
||||
try {
|
||||
const fileContent = await fs.readFile(filePath, 'utf8');
|
||||
return JSON.parse(fileContent);
|
||||
} catch (error) {
|
||||
if (error.code === 'ENOENT') {
|
||||
console.debug(`[Kiro Auth] Credential file not found: ${filePath}`);
|
||||
} else if (error instanceof SyntaxError) {
|
||||
console.warn(`[Kiro Auth] Failed to parse JSON from ${filePath}: ${error.message}`);
|
||||
} else {
|
||||
console.warn(`[Kiro Auth] Failed to read credential file ${filePath}: ${error.message}`);
|
||||
}
|
||||
return null;
|
||||
async initializeAuth(forceRefresh = false) {
|
||||
if (this.accessToken && !forceRefresh) {
|
||||
console.debug('[Kiro Auth] Access token already available and not forced refresh.');
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Helper to save credentials to a file
|
||||
const saveCredentialsToFile = async (filePath, newData) => {
|
||||
try {
|
||||
let existingData = {};
|
||||
// Helper to load credentials from a file
|
||||
const loadCredentialsFromFile = async (filePath) => {
|
||||
try {
|
||||
const fileContent = await fs.readFile(filePath, 'utf8');
|
||||
existingData = JSON.parse(fileContent);
|
||||
} catch (readError) {
|
||||
if (readError.code === 'ENOENT') {
|
||||
console.debug(`[Kiro Auth] Token file not found, creating new one: ${filePath}`);
|
||||
return JSON.parse(fileContent);
|
||||
} catch (error) {
|
||||
if (error.code === 'ENOENT') {
|
||||
console.debug(`[Kiro Auth] Credential file not found: ${filePath}`);
|
||||
} else if (error instanceof SyntaxError) {
|
||||
console.warn(`[Kiro Auth] Failed to parse JSON from ${filePath}: ${error.message}`);
|
||||
} else {
|
||||
console.warn(`[Kiro Auth] Could not read existing token file ${filePath}: ${readError.message}`);
|
||||
console.warn(`[Kiro Auth] Failed to read credential file ${filePath}: ${error.message}`);
|
||||
}
|
||||
return null;
|
||||
}
|
||||
const mergedData = { ...existingData, ...newData };
|
||||
await fs.writeFile(filePath, JSON.stringify(mergedData, null, 2), 'utf8');
|
||||
console.info(`[Kiro Auth] Updated token file: ${filePath}`);
|
||||
} catch (error) {
|
||||
console.error(`[Kiro Auth] Failed to write token to file ${filePath}: ${error.message}`);
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
try {
|
||||
let mergedCredentials = {};
|
||||
|
||||
// Priority 1: Load from Base64 credentials if available
|
||||
if (this.base64Creds) {
|
||||
Object.assign(mergedCredentials, this.base64Creds);
|
||||
console.info('[Kiro Auth] Successfully loaded credentials from Base64 (constructor).');
|
||||
// Clear base64Creds after use to prevent re-processing
|
||||
this.base64Creds = null;
|
||||
}
|
||||
|
||||
// Priority 2 & 3 合并: 从指定文件路径或目录加载凭证
|
||||
// 读取指定的 credPath 文件以及目录下的其他 JSON 文件(排除当前文件)
|
||||
const targetFilePath = this.credsFilePath || path.join(this.credPath, KIRO_AUTH_TOKEN_FILE);
|
||||
const dirPath = path.dirname(targetFilePath);
|
||||
const targetFileName = path.basename(targetFilePath);
|
||||
|
||||
console.debug(`[Kiro Auth] Attempting to load credentials from directory: ${dirPath}`);
|
||||
|
||||
try {
|
||||
// 首先尝试读取目标文件
|
||||
const targetCredentials = await loadCredentialsFromFile(targetFilePath);
|
||||
if (targetCredentials) {
|
||||
Object.assign(mergedCredentials, targetCredentials);
|
||||
console.info(`[Kiro Auth] Successfully loaded OAuth credentials from ${targetFilePath}`);
|
||||
}
|
||||
|
||||
// 然后读取目录下的其他 JSON 文件(排除目标文件本身)
|
||||
const files = await fs.readdir(dirPath);
|
||||
for (const file of files) {
|
||||
if (file.endsWith('.json') && file !== targetFileName) {
|
||||
const filePath = path.join(dirPath, file);
|
||||
const credentials = await loadCredentialsFromFile(filePath);
|
||||
if (credentials) {
|
||||
// 保留已有的 expiresAt,避免被覆盖
|
||||
credentials.expiresAt = mergedCredentials.expiresAt;
|
||||
Object.assign(mergedCredentials, credentials);
|
||||
console.debug(`[Kiro Auth] Loaded Client credentials from ${file}`);
|
||||
// Helper to save credentials to a file
|
||||
const saveCredentialsToFile = async (filePath, newData) => {
|
||||
try {
|
||||
let existingData = {};
|
||||
try {
|
||||
const fileContent = await fs.readFile(filePath, 'utf8');
|
||||
existingData = JSON.parse(fileContent);
|
||||
} catch (readError) {
|
||||
if (readError.code === 'ENOENT') {
|
||||
console.debug(`[Kiro Auth] Token file not found, creating new one: ${filePath}`);
|
||||
} else {
|
||||
console.warn(`[Kiro Auth] Could not read existing token file ${filePath}: ${readError.message}`);
|
||||
}
|
||||
}
|
||||
const mergedData = { ...existingData, ...newData };
|
||||
await fs.writeFile(filePath, JSON.stringify(mergedData, null, 2), 'utf8');
|
||||
console.info(`[Kiro Auth] Updated token file: ${filePath}`);
|
||||
} catch (error) {
|
||||
console.error(`[Kiro Auth] Failed to write token to file ${filePath}: ${error.message}`);
|
||||
}
|
||||
} catch (error) {
|
||||
console.warn(`[Kiro Auth] Error loading credentials from directory ${dirPath}: ${error.message}`);
|
||||
}
|
||||
};
|
||||
|
||||
// console.log('[Kiro Auth] Merged credentials:', mergedCredentials);
|
||||
// Apply loaded credentials, prioritizing existing values if they are not null/undefined
|
||||
this.accessToken = this.accessToken || mergedCredentials.accessToken;
|
||||
this.refreshToken = this.refreshToken || mergedCredentials.refreshToken;
|
||||
this.clientId = this.clientId || mergedCredentials.clientId;
|
||||
this.clientSecret = this.clientSecret || mergedCredentials.clientSecret;
|
||||
this.authMethod = this.authMethod || mergedCredentials.authMethod;
|
||||
this.expiresAt = this.expiresAt || mergedCredentials.expiresAt;
|
||||
this.profileArn = this.profileArn || mergedCredentials.profileArn;
|
||||
this.region = this.region || mergedCredentials.region;
|
||||
|
||||
// Ensure region is set before using it in URLs
|
||||
if (!this.region) {
|
||||
console.warn('[Kiro Auth] Region not found in credentials. Using default region us-east-1 for URLs.');
|
||||
this.region = 'us-east-1'; // Set default region
|
||||
}
|
||||
|
||||
this.refreshUrl = (this.config.KIRO_REFRESH_URL || KIRO_CONSTANTS.REFRESH_URL).replace("{{region}}", this.region);
|
||||
this.refreshIDCUrl = (this.config.KIRO_REFRESH_IDC_URL || KIRO_CONSTANTS.REFRESH_IDC_URL).replace("{{region}}", this.region);
|
||||
this.baseUrl = (this.config.KIRO_BASE_URL || KIRO_CONSTANTS.BASE_URL).replace("{{region}}", this.region);
|
||||
this.amazonQUrl = (KIRO_CONSTANTS.AMAZON_Q_URL).replace("{{region}}", this.region);
|
||||
} catch (error) {
|
||||
console.warn(`[Kiro Auth] Error during credential loading: ${error.message}`);
|
||||
}
|
||||
|
||||
// Refresh token if forced or if access token is missing but refresh token is available
|
||||
if (forceRefresh || (!this.accessToken && this.refreshToken)) {
|
||||
if (!this.refreshToken) {
|
||||
throw new Error('No refresh token available to refresh access token.');
|
||||
}
|
||||
try {
|
||||
const requestBody = {
|
||||
refreshToken: this.refreshToken,
|
||||
};
|
||||
let mergedCredentials = {};
|
||||
|
||||
let refreshUrl = this.refreshUrl;
|
||||
if (this.authMethod !== KIRO_CONSTANTS.AUTH_METHOD_SOCIAL) {
|
||||
refreshUrl = this.refreshIDCUrl;
|
||||
requestBody.clientId = this.clientId;
|
||||
requestBody.clientSecret = this.clientSecret;
|
||||
requestBody.grantType = 'refresh_token';
|
||||
// Priority 1: Load from Base64 credentials if available
|
||||
if (this.base64Creds) {
|
||||
Object.assign(mergedCredentials, this.base64Creds);
|
||||
console.info('[Kiro Auth] Successfully loaded credentials from Base64 (constructor).');
|
||||
// Clear base64Creds after use to prevent re-processing
|
||||
this.base64Creds = null;
|
||||
}
|
||||
|
||||
let response = null;
|
||||
if (this.authMethod === KIRO_CONSTANTS.AUTH_METHOD_SOCIAL) {
|
||||
response = await this.axiosSocialRefreshInstance.post(refreshUrl, requestBody);
|
||||
console.log('[Kiro Auth] Token refresh social response: ok');
|
||||
}else{
|
||||
response = await this.axiosInstance.post(refreshUrl, requestBody);
|
||||
console.log('[Kiro Auth] Token refresh idc response: ok');
|
||||
}
|
||||
// Priority 2 & 3 合并: 从指定文件路径或目录加载凭证
|
||||
// 读取指定的 credPath 文件以及目录下的其他 JSON 文件(排除当前文件)
|
||||
const targetFilePath = this.credsFilePath || path.join(this.credPath, KIRO_AUTH_TOKEN_FILE);
|
||||
const dirPath = path.dirname(targetFilePath);
|
||||
const targetFileName = path.basename(targetFilePath);
|
||||
|
||||
if (response.data && response.data.accessToken) {
|
||||
this.accessToken = response.data.accessToken;
|
||||
this.refreshToken = response.data.refreshToken;
|
||||
this.profileArn = response.data.profileArn;
|
||||
const expiresIn = response.data.expiresIn;
|
||||
const expiresAt = new Date(Date.now() + expiresIn * 1000).toISOString();
|
||||
this.expiresAt = expiresAt;
|
||||
console.info('[Kiro Auth] Access token refreshed successfully');
|
||||
console.debug(`[Kiro Auth] Attempting to load credentials from directory: ${dirPath}`);
|
||||
|
||||
// Update the token file - use specified path if configured, otherwise use default
|
||||
const tokenFilePath = this.credsFilePath || path.join(this.credPath, KIRO_AUTH_TOKEN_FILE);
|
||||
const updatedTokenData = {
|
||||
accessToken: this.accessToken,
|
||||
refreshToken: this.refreshToken,
|
||||
expiresAt: expiresAt,
|
||||
};
|
||||
if(this.profileArn){
|
||||
updatedTokenData.profileArn = this.profileArn;
|
||||
try {
|
||||
// 首先尝试读取目标文件
|
||||
const targetCredentials = await loadCredentialsFromFile(targetFilePath);
|
||||
if (targetCredentials) {
|
||||
Object.assign(mergedCredentials, targetCredentials);
|
||||
console.info(`[Kiro Auth] Successfully loaded OAuth credentials from ${targetFilePath}`);
|
||||
}
|
||||
await saveCredentialsToFile(tokenFilePath, updatedTokenData);
|
||||
} else {
|
||||
throw new Error('Invalid refresh response: Missing accessToken');
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('[Kiro Auth] Token refresh failed:', error.message);
|
||||
throw new Error(`Token refresh failed: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
if (!this.accessToken) {
|
||||
throw new Error('No access token available after initialization and refresh attempts.');
|
||||
}
|
||||
// 然后读取目录下的其他 JSON 文件(排除目标文件本身)
|
||||
const files = await fs.readdir(dirPath);
|
||||
for (const file of files) {
|
||||
if (file.endsWith('.json') && file !== targetFileName) {
|
||||
const filePath = path.join(dirPath, file);
|
||||
const credentials = await loadCredentialsFromFile(filePath);
|
||||
if (credentials) {
|
||||
// 保留已有的 expiresAt,避免被覆盖
|
||||
credentials.expiresAt = mergedCredentials.expiresAt;
|
||||
Object.assign(mergedCredentials, credentials);
|
||||
console.debug(`[Kiro Auth] Loaded Client credentials from ${file}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
console.warn(`[Kiro Auth] Error loading credentials from directory ${dirPath}: ${error.message}`);
|
||||
}
|
||||
|
||||
// console.log('[Kiro Auth] Merged credentials:', mergedCredentials);
|
||||
// Apply loaded credentials, prioritizing existing values if they are not null/undefined
|
||||
this.accessToken = this.accessToken || mergedCredentials.accessToken;
|
||||
this.refreshToken = this.refreshToken || mergedCredentials.refreshToken;
|
||||
this.clientId = this.clientId || mergedCredentials.clientId;
|
||||
this.clientSecret = this.clientSecret || mergedCredentials.clientSecret;
|
||||
this.authMethod = this.authMethod || mergedCredentials.authMethod;
|
||||
this.expiresAt = this.expiresAt || mergedCredentials.expiresAt;
|
||||
this.profileArn = this.profileArn || mergedCredentials.profileArn;
|
||||
this.region = this.region || mergedCredentials.region;
|
||||
|
||||
// Ensure region is set before using it in URLs
|
||||
if (!this.region) {
|
||||
console.warn('[Kiro Auth] Region not found in credentials. Using default region us-east-1 for URLs.');
|
||||
this.region = 'us-east-1'; // Set default region
|
||||
}
|
||||
|
||||
this.refreshUrl = (this.config.KIRO_REFRESH_URL || KIRO_CONSTANTS.REFRESH_URL).replace("{{region}}", this.region);
|
||||
this.refreshIDCUrl = (this.config.KIRO_REFRESH_IDC_URL || KIRO_CONSTANTS.REFRESH_IDC_URL).replace("{{region}}", this.region);
|
||||
this.baseUrl = (this.config.KIRO_BASE_URL || KIRO_CONSTANTS.BASE_URL).replace("{{region}}", this.region);
|
||||
this.amazonQUrl = (KIRO_CONSTANTS.AMAZON_Q_URL).replace("{{region}}", this.region);
|
||||
} catch (error) {
|
||||
console.warn(`[Kiro Auth] Error during credential loading: ${error.message}`);
|
||||
}
|
||||
|
||||
// Refresh token if forced or if access token is missing but refresh token is available
|
||||
if (forceRefresh || (!this.accessToken && this.refreshToken)) {
|
||||
if (!this.refreshToken) {
|
||||
throw new Error('No refresh token available to refresh access token.');
|
||||
}
|
||||
try {
|
||||
const requestBody = {
|
||||
refreshToken: this.refreshToken,
|
||||
};
|
||||
|
||||
let refreshUrl = this.refreshUrl;
|
||||
if (this.authMethod !== KIRO_CONSTANTS.AUTH_METHOD_SOCIAL) {
|
||||
refreshUrl = this.refreshIDCUrl;
|
||||
requestBody.clientId = this.clientId;
|
||||
requestBody.clientSecret = this.clientSecret;
|
||||
requestBody.grantType = 'refresh_token';
|
||||
}
|
||||
|
||||
let response = null;
|
||||
if (this.authMethod === KIRO_CONSTANTS.AUTH_METHOD_SOCIAL) {
|
||||
response = await this.axiosSocialRefreshInstance.post(refreshUrl, requestBody);
|
||||
console.log('[Kiro Auth] Token refresh social response: ok');
|
||||
} else {
|
||||
response = await this.axiosInstance.post(refreshUrl, requestBody);
|
||||
console.log('[Kiro Auth] Token refresh idc response: ok');
|
||||
}
|
||||
|
||||
if (response.data && response.data.accessToken) {
|
||||
this.accessToken = response.data.accessToken;
|
||||
this.refreshToken = response.data.refreshToken;
|
||||
this.profileArn = response.data.profileArn;
|
||||
const expiresIn = response.data.expiresIn;
|
||||
const expiresAt = new Date(Date.now() + expiresIn * 1000).toISOString();
|
||||
this.expiresAt = expiresAt;
|
||||
console.info('[Kiro Auth] Access token refreshed successfully');
|
||||
|
||||
// Update the token file - use specified path if configured, otherwise use default
|
||||
const tokenFilePath = this.credsFilePath || path.join(this.credPath, KIRO_AUTH_TOKEN_FILE);
|
||||
const updatedTokenData = {
|
||||
accessToken: this.accessToken,
|
||||
refreshToken: this.refreshToken,
|
||||
expiresAt: expiresAt,
|
||||
};
|
||||
if (this.profileArn) {
|
||||
updatedTokenData.profileArn = this.profileArn;
|
||||
}
|
||||
await saveCredentialsToFile(tokenFilePath, updatedTokenData);
|
||||
} else {
|
||||
throw new Error('Invalid refresh response: Missing accessToken');
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('[Kiro Auth] Token refresh failed:', error.message);
|
||||
throw new Error(`Token refresh failed: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
if (!this.accessToken) {
|
||||
throw new Error('No access token available after initialization and refresh attempts.');
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract text content from OpenAI message format
|
||||
*/
|
||||
getContentText(message) {
|
||||
if(message==null){
|
||||
if (message == null) {
|
||||
return "";
|
||||
}
|
||||
if (Array.isArray(message) ) {
|
||||
if (Array.isArray(message)) {
|
||||
return message
|
||||
.filter(part => part.type === 'text' && part.text)
|
||||
.map(part => part.text)
|
||||
.join('');
|
||||
} else if (typeof message.content === 'string') {
|
||||
return message.content;
|
||||
} else if (Array.isArray(message.content) ) {
|
||||
} else if (Array.isArray(message.content)) {
|
||||
return message.content
|
||||
.filter(part => part.type === 'text' && part.text)
|
||||
.map(part => part.text)
|
||||
|
|
@ -1102,7 +1103,6 @@ async initializeAuth(forceRefresh = false) {
|
|||
|
||||
_processApiResponse(response) {
|
||||
const rawResponseText = Buffer.isBuffer(response.data) ? response.data.toString('utf8') : String(response.data);
|
||||
//console.log(`[Kiro] Raw response length: ${rawResponseText.length}`);
|
||||
if (rawResponseText.includes("[Called")) {
|
||||
console.log("[Kiro] Raw response contains [Called marker.");
|
||||
}
|
||||
|
|
@ -1154,13 +1154,22 @@ async initializeAuth(forceRefresh = false) {
|
|||
const finalModel = MODEL_MAPPING[model] ? model : this.modelName;
|
||||
console.log(`[Kiro] Calling generateContent with model: ${finalModel}`);
|
||||
|
||||
// Estimate input tokens before making the API call
|
||||
const inputTokens = this.estimateInputTokens(requestBody);
|
||||
|
||||
const response = await this.callApi('', finalModel, requestBody);
|
||||
|
||||
try {
|
||||
const { responseText, toolCalls } = this._processApiResponse(response);
|
||||
|
||||
let inputTokens = 0;
|
||||
const rawResponseText = Buffer.isBuffer(response.data)
|
||||
? response.data.toString('utf8')
|
||||
: String(response.data);
|
||||
|
||||
const contextUsageMatch = rawResponseText.match(/"contextUsagePercentage":\s*([\d.]+)/);
|
||||
if (contextUsageMatch) {
|
||||
const percentage = parseFloat(contextUsageMatch[1]);
|
||||
inputTokens = this.calculateInputTokensFromPercentage(percentage);
|
||||
}
|
||||
|
||||
return this.buildClaudeResponse(responseText, false, 'assistant', model, toolCalls, inputTokens);
|
||||
} catch (error) {
|
||||
console.error('[Kiro] Error in generateContent:', error);
|
||||
|
|
@ -1192,9 +1201,10 @@ async initializeAuth(forceRefresh = false) {
|
|||
const followupStart = remaining.indexOf('{"followupPrompt":', searchStart);
|
||||
const inputStart = remaining.indexOf('{"input":', searchStart);
|
||||
const stopStart = remaining.indexOf('{"stop":', searchStart);
|
||||
const contextUsageStart = remaining.indexOf('{"contextUsagePercentage":', searchStart);
|
||||
|
||||
// 找到最早出现的有效 JSON 模式
|
||||
const candidates = [contentStart, nameStart, followupStart, inputStart, stopStart].filter(pos => pos >= 0);
|
||||
const candidates = [contentStart, nameStart, followupStart, inputStart, stopStart, contextUsageStart].filter(pos => pos >= 0);
|
||||
if (candidates.length === 0) break;
|
||||
|
||||
const jsonStart = Math.min(...candidates);
|
||||
|
|
@ -1284,6 +1294,15 @@ async initializeAuth(forceRefresh = false) {
|
|||
}
|
||||
});
|
||||
}
|
||||
// 处理 context usage percentage 事件
|
||||
else if (parsed.contextUsagePercentage !== undefined) {
|
||||
events.push({
|
||||
type: 'contextUsage',
|
||||
data: {
|
||||
percentage: parsed.contextUsagePercentage
|
||||
}
|
||||
});
|
||||
}
|
||||
} catch (e) {
|
||||
// JSON 解析失败，跳过这个位置继续搜索
|
||||
}
|
||||
|
|
@ -1330,7 +1349,7 @@ async initializeAuth(forceRefresh = false) {
|
|||
|
||||
stream = response.data;
|
||||
let buffer = '';
|
||||
let lastContentEvent = null; // 用于检测连续重复的 content 事件
|
||||
let lastContentEvent = null;
|
||||
|
||||
for await (const chunk of stream) {
|
||||
buffer += chunk.toString();
|
||||
|
|
@ -1355,6 +1374,8 @@ async initializeAuth(forceRefresh = false) {
|
|||
yield { type: 'toolUseInput', input: event.data.input };
|
||||
} else if (event.type === 'toolUseStop') {
|
||||
yield { type: 'toolUseStop', stop: event.data.stop };
|
||||
} else if (event.type === 'contextUsage') {
|
||||
yield { type: 'contextUsage', percentage: event.data.percentage };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1438,46 +1459,69 @@ async initializeAuth(forceRefresh = false) {
|
|||
const finalModel = MODEL_MAPPING[model] ? model : this.modelName;
|
||||
console.log(`[Kiro] Calling generateContentStream with model: ${finalModel} (real streaming)`);
|
||||
|
||||
const inputTokens = this.estimateInputTokens(requestBody);
|
||||
let inputTokens = 0;
|
||||
let contextUsagePercentage = null;
|
||||
const messageId = `${uuidv4()}`;
|
||||
|
||||
let messageStartSent = false;
|
||||
const bufferedEvents = [];
|
||||
|
||||
try {
|
||||
// 1. 先发送 message_start 事件
|
||||
yield {
|
||||
type: "message_start",
|
||||
message: {
|
||||
id: messageId,
|
||||
type: "message",
|
||||
role: "assistant",
|
||||
model: model,
|
||||
usage: { input_tokens: inputTokens, output_tokens: 0 },
|
||||
content: []
|
||||
}
|
||||
};
|
||||
|
||||
// 2. 发送 content_block_start 事件
|
||||
yield {
|
||||
type: "content_block_start",
|
||||
index: 0,
|
||||
content_block: { type: "text", text: "" }
|
||||
};
|
||||
|
||||
let totalContent = '';
|
||||
let outputTokens = 0;
|
||||
const toolCalls = [];
|
||||
let currentToolCall = null; // 用于累积结构化工具调用
|
||||
let currentToolCall = null;
|
||||
|
||||
// 3. 流式接收并发送每个 content_block_delta
|
||||
for await (const event of this.streamApiReal('', finalModel, requestBody)) {
|
||||
if (event.type === 'content' && event.content) {
|
||||
totalContent += event.content;
|
||||
// 不再每个 chunk 都计算 token，改为最后统一计算，避免阻塞事件循环
|
||||
if (event.type === 'contextUsage' && event.percentage) {
|
||||
contextUsagePercentage = event.percentage;
|
||||
inputTokens = this.calculateInputTokensFromPercentage(contextUsagePercentage);
|
||||
|
||||
yield {
|
||||
if (!messageStartSent) {
|
||||
yield {
|
||||
type: "message_start",
|
||||
message: {
|
||||
id: messageId,
|
||||
type: "message",
|
||||
role: "assistant",
|
||||
model: model,
|
||||
usage: {
|
||||
input_tokens: inputTokens,
|
||||
output_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0
|
||||
},
|
||||
content: []
|
||||
}
|
||||
};
|
||||
|
||||
yield {
|
||||
type: "content_block_start",
|
||||
index: 0,
|
||||
content_block: { type: "text", text: "" }
|
||||
};
|
||||
|
||||
messageStartSent = true;
|
||||
|
||||
for (const buffered of bufferedEvents) {
|
||||
yield buffered;
|
||||
}
|
||||
bufferedEvents.length = 0;
|
||||
}
|
||||
} else if (event.type === 'content' && event.content) {
|
||||
totalContent += event.content;
|
||||
|
||||
const contentEvent = {
|
||||
type: "content_block_delta",
|
||||
index: 0,
|
||||
delta: { type: "text_delta", text: event.content }
|
||||
};
|
||||
|
||||
if (messageStartSent) {
|
||||
yield contentEvent;
|
||||
} else {
|
||||
bufferedEvents.push(contentEvent);
|
||||
}
|
||||
} else if (event.type === 'toolUse') {
|
||||
const tc = event.toolUse;
|
||||
// 工具调用事件（包含 name 和 toolUseId）
|
||||
|
|
@ -1508,7 +1552,7 @@ async initializeAuth(forceRefresh = false) {
|
|||
if (tc.stop) {
|
||||
try {
|
||||
currentToolCall.input = JSON.parse(currentToolCall.input);
|
||||
} catch (e) {}
|
||||
} catch (e) { }
|
||||
toolCalls.push(currentToolCall);
|
||||
currentToolCall = null;
|
||||
}
|
||||
|
|
@ -1536,11 +1580,17 @@ async initializeAuth(forceRefresh = false) {
|
|||
if (currentToolCall) {
|
||||
try {
|
||||
currentToolCall.input = JSON.parse(currentToolCall.input);
|
||||
} catch (e) {}
|
||||
} catch (e) { }
|
||||
toolCalls.push(currentToolCall);
|
||||
currentToolCall = null;
|
||||
}
|
||||
|
||||
// Fallback: 如果 contextUsagePercentage 没有收到，抛出错误
|
||||
if (!messageStartSent) {
|
||||
console.error('[Kiro Stream] contextUsagePercentage not received from API - cannot calculate accurate input tokens');
|
||||
throw new Error('Failed to receive contextUsagePercentage from Kiro API. Input token calculation requires this data.');
|
||||
}
|
||||
|
||||
// 检查文本内容中的 bracket 格式工具调用
|
||||
const bracketToolCalls = parseBracketToolCalls(totalContent);
|
||||
if (bracketToolCalls && bracketToolCalls.length > 0) {
|
||||
|
|
@ -1595,8 +1645,16 @@ async initializeAuth(forceRefresh = false) {
|
|||
|
||||
yield {
|
||||
type: "message_delta",
|
||||
delta: { stop_reason: toolCalls.length > 0 ? "tool_use" : "end_turn" },
|
||||
usage: { output_tokens: outputTokens }
|
||||
delta: {
|
||||
stop_reason: toolCalls.length > 0 ? "tool_use" : "end_turn",
|
||||
stop_sequence: null
|
||||
},
|
||||
usage: {
|
||||
input_tokens: inputTokens,
|
||||
output_tokens: outputTokens,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0
|
||||
}
|
||||
};
|
||||
|
||||
// 7. 发送 message_stop 事件
|
||||
|
|
@ -1623,9 +1681,27 @@ async initializeAuth(forceRefresh = false) {
|
|||
}
|
||||
|
||||
/**
|
||||
* Convert context usage percentage to actual input tokens
|
||||
* @param {number} percentage - Context usage percentage (0-100)
|
||||
* @returns {number} Actual input tokens
|
||||
*/
|
||||
calculateInputTokensFromPercentage(percentage) {
|
||||
if (!percentage || percentage <= 0) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const contextWindow = CLAUDE_DEFAULT_MAX_TOKENS;
|
||||
const inputTokens = Math.round((percentage / 100) * contextWindow);
|
||||
|
||||
return inputTokens;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated Use contextUsagePercentage from API response instead
|
||||
* Calculate input tokens from request body using Claude's official tokenizer
|
||||
*/
|
||||
estimateInputTokens(requestBody) {
|
||||
console.warn('[Kiro] estimateInputTokens() is deprecated. Use contextUsagePercentage from API response instead.');
|
||||
let totalTokens = 0;
|
||||
|
||||
// Count system prompt tokens
|
||||
|
|
@ -1824,7 +1900,9 @@ async initializeAuth(forceRefresh = false) {
|
|||
stop_sequence: null,
|
||||
usage: {
|
||||
input_tokens: inputTokens,
|
||||
output_tokens: outputTokens
|
||||
output_tokens: outputTokens,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0
|
||||
},
|
||||
content: contentArray
|
||||
};
|
||||
|
|
@ -1950,7 +2028,7 @@ async initializeAuth(forceRefresh = false) {
|
|||
origin: KIRO_CONSTANTS.ORIGIN_AI_EDITOR,
|
||||
resourceType: resourceType
|
||||
});
|
||||
if (this.authMethod === KIRO_CONSTANTS.AUTH_METHOD_SOCIAL && this.profileArn) {
|
||||
if (this.authMethod === KIRO_CONSTANTS.AUTH_METHOD_SOCIAL && this.profileArn) {
|
||||
params.append('profileArn', this.profileArn);
|
||||
}
|
||||
const fullUrl = `${usageLimitsUrl}?${params.toString()}`;
|
||||
|
|
|
|||
Loading…
Reference in a new issue