diff --git a/src/__tests__/accounts.test.ts b/src/__tests__/accounts.test.ts new file mode 100644 index 0000000..60d7376 --- /dev/null +++ b/src/__tests__/accounts.test.ts @@ -0,0 +1,460 @@ +import { describe, expect, mock, test } from 'bun:test' +import { AccountManager, createDeterministicAccountId } from '../plugin/accounts.js' +import type { ManagedAccount } from '../plugin/types.js' + +// Mock DB and external dependencies +mock.module('../plugin/storage/sqlite.js', () => ({ + kiroDb: { + getAccounts: () => [], + upsertAccount: () => Promise.resolve(), + deleteAccount: () => Promise.resolve(), + batchUpsertAccounts: () => Promise.resolve() + } +})) +mock.module('../plugin/sync/kiro-cli.js', () => ({ + syncFromKiroCli: () => Promise.resolve(), + writeToKiroCli: () => Promise.resolve() +})) +mock.module('../kiro/auth.js', () => ({ + decodeRefreshToken: (t: string) => ({ refreshToken: t }), + encodeRefreshToken: (p: any) => p.refreshToken, + accessTokenExpired: () => false +})) + +function makeAccount(overrides: Partial = {}): ManagedAccount { + return { + id: 'test-id', + email: 'test@example.com', + authMethod: 'idc', + region: 'eu-central-1', + refreshToken: 'refresh', + accessToken: 'access', + expiresAt: Date.now() + 3600000, + rateLimitResetTime: 0, + isHealthy: true, + failCount: 0, + lastUsed: 0, + usedCount: 0, + limitCount: 0, + ...overrides + } +} + +// ── createDeterministicAccountId ────────────────────────────────────────────── + +describe('createDeterministicAccountId', () => { + test('IDC uses email + method + profileArn, ignores clientId', () => { + const id1 = createDeterministicAccountId('a@b.com', 'idc', 'client-1', 'arn:aws:123') + const id2 = createDeterministicAccountId('a@b.com', 'idc', 'client-2', 'arn:aws:123') + expect(id1).toBe(id2) + }) + + test('non-IDC uses email + method + clientId + profileArn', () => { + const id1 = createDeterministicAccountId('a@b.com', 'builderid', 'client-1') + const id2 = createDeterministicAccountId('a@b.com', 'builderid', 'client-2') + expect(id1).not.toBe(id2) + }) + + test('different emails produce different IDs', () => { + const id1 = createDeterministicAccountId('a@b.com', 'idc', 'c', 'arn') + const id2 = createDeterministicAccountId('x@b.com', 'idc', 'c', 'arn') + expect(id1).not.toBe(id2) + }) + + test('returns 64-char hex string', () => { + const id = createDeterministicAccountId('a@b.com', 'idc', 'c', 'arn') + expect(id).toMatch(/^[a-f0-9]{64}$/) + }) +}) + +// ── AccountManager ──────────────────────────────────────────────────────────── + +describe('AccountManager.getCurrentOrNext', () => { + test('returns healthy account', () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + expect(mgr.getCurrentOrNext()).not.toBeNull() + }) + + test('returns null when no accounts', () => { + expect(new AccountManager([]).getCurrentOrNext()).toBeNull() + }) + + test('skips permanently unhealthy accounts', () => { + const acc = makeAccount({ isHealthy: false, unhealthyReason: 'HTTP_403', failCount: 10 }) + expect(new AccountManager([acc]).getCurrentOrNext()).toBeNull() + }) + + test('skips rate-limited accounts', () => { + const acc = makeAccount({ rateLimitResetTime: Date.now() + 60000 }) + expect(new AccountManager([acc]).getCurrentOrNext()).toBeNull() + }) + + test('recovers unhealthy account past recoveryTime', () => { + const acc = makeAccount({ + isHealthy: false, + unhealthyReason: 'temporary', + failCount: 3, + recoveryTime: Date.now() - 1000 + }) + const mgr = new AccountManager([acc]) + const selected = mgr.getCurrentOrNext() + expect(selected).not.toBeNull() + expect(selected!.isHealthy).toBe(true) + }) + + test('does NOT recover permanently unhealthy account past recoveryTime', () => { + const acc = makeAccount({ + isHealthy: false, + unhealthyReason: 'HTTP_403', + failCount: 10, + recoveryTime: Date.now() - 1000 + }) + expect(new AccountManager([acc]).getCurrentOrNext()).toBeNull() + }) + + test('increments usedCount and sets lastUsed on selection', () => { + const acc = makeAccount({ usedCount: 5 }) + const mgr = new AccountManager([acc]) + mgr.getCurrentOrNext() + expect(acc.usedCount).toBe(6) + expect(acc.lastUsed).toBeGreaterThan(0) + }) + + test('round-robin cycles through multiple accounts', () => { + const a = makeAccount({ id: 'a', email: 'a@x.com' }) + const b = makeAccount({ id: 'b', email: 'b@x.com' }) + const mgr = new AccountManager([a, b], 'round-robin') + const first = mgr.getCurrentOrNext() + const second = mgr.getCurrentOrNext() + expect(first!.id).not.toBe(second!.id) + }) + + test('round-robin skips rate-limited account and resumes correct position', () => { + const a = makeAccount({ id: 'a', email: 'a@x.com' }) + const b = makeAccount({ id: 'b', email: 'b@x.com' }) + const c = makeAccount({ id: 'c', email: 'c@x.com' }) + const mgr = new AccountManager([a, b, c], 'round-robin') + // a→b→c normal cycle + expect(mgr.getCurrentOrNext()!.id).toBe('a') + expect(mgr.getCurrentOrNext()!.id).toBe('b') + // Now rate-limit b before the next call + b.rateLimitResetTime = Date.now() + 60_000 + // cursor is at c — c is still available + expect(mgr.getCurrentOrNext()!.id).toBe('c') + // cursor wraps: a is next (b skipped) + expect(mgr.getCurrentOrNext()!.id).toBe('a') + // b's rate limit expires + b.rateLimitResetTime = 0 + // cursor is at b — b is available again + expect(mgr.getCurrentOrNext()!.id).toBe('b') + }) + + test('lowest-usage picks account with fewer usedCount', () => { + const a = makeAccount({ id: 'a', email: 'a@x.com', usedCount: 10 }) + const b = makeAccount({ id: 'b', email: 'b@x.com', usedCount: 2 }) + const mgr = new AccountManager([a, b], 'lowest-usage') + expect(mgr.getCurrentOrNext()!.id).toBe('b') + }) +}) + +describe('AccountManager.markUnhealthy', () => { + test('permanent error sets isHealthy=false, failCount=10, no recoveryTime', () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + mgr.markUnhealthy(acc, 'ExpiredTokenException') + expect(acc.isHealthy).toBe(false) + expect(acc.failCount).toBe(10) + expect(acc.recoveryTime).toBeUndefined() + }) + + test('non-permanent error increments failCount', () => { + const acc = makeAccount({ failCount: 2 }) + const mgr = new AccountManager([acc]) + mgr.markUnhealthy(acc, 'Server Error') + expect(acc.failCount).toBe(3) + expect(acc.isHealthy).toBe(true) + }) + + test('non-permanent error sets isHealthy=false after 10 failures', () => { + const acc = makeAccount({ failCount: 9 }) + const mgr = new AccountManager([acc]) + mgr.markUnhealthy(acc, 'Server Error') + expect(acc.failCount).toBe(10) + expect(acc.isHealthy).toBe(false) + expect(acc.recoveryTime).toBeGreaterThan(Date.now()) + }) + + test('expired token exception is treated as permanent', () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + mgr.markUnhealthy(acc, 'ExpiredTokenException') + expect(acc.isHealthy).toBe(false) + expect(acc.failCount).toBe(10) + expect(acc.recoveryTime).toBeUndefined() + }) + + test('does nothing for unknown account id', () => { + const acc = makeAccount({ id: 'known' }) + const mgr = new AccountManager([acc]) + const unknown = makeAccount({ id: 'unknown' }) + mgr.markUnhealthy(unknown, 'HTTP_403') + expect(acc.isHealthy).toBe(true) // unchanged + }) +}) + +describe('AccountManager.removeAccount', () => { + test('removes account from list', () => { + const a = makeAccount({ id: 'a' }) + const b = makeAccount({ id: 'b' }) + const mgr = new AccountManager([a, b]) + mgr.removeAccount(a) + expect(mgr.getAccountCount()).toBe(1) + expect(mgr.getAccounts()[0]!.id).toBe('b') + }) + + test('cursor resets to 0 when list becomes empty', () => { + const a = makeAccount() + const mgr = new AccountManager([a]) + mgr.removeAccount(a) + expect(mgr.getAccountCount()).toBe(0) + expect(mgr.getCurrentOrNext()).toBeNull() + }) + + test('ignores unknown account', () => { + const a = makeAccount({ id: 'a' }) + const mgr = new AccountManager([a]) + mgr.removeAccount(makeAccount({ id: 'unknown' })) + expect(mgr.getAccountCount()).toBe(1) + }) +}) + +describe('AccountManager.addAccount', () => { + test('adds new account', () => { + const mgr = new AccountManager([]) + mgr.addAccount(makeAccount({ id: 'new' })) + expect(mgr.getAccountCount()).toBe(1) + }) + + test('replaces existing account with same id', () => { + const original = makeAccount({ id: 'x', email: 'old@x.com' }) + const updated = makeAccount({ id: 'x', email: 'new@x.com' }) + const mgr = new AccountManager([original]) + mgr.addAccount(updated) + expect(mgr.getAccountCount()).toBe(1) + expect(mgr.getAccounts()[0]!.email).toBe('new@x.com') + }) +}) + +describe('AccountManager.getMinWaitTime', () => { + test('returns 0 with no rate-limited accounts', () => { + const mgr = new AccountManager([makeAccount()]) + expect(mgr.getMinWaitTime()).toBe(0) + }) + + test('returns minimum wait across rate-limited accounts', () => { + const a = makeAccount({ id: 'a', rateLimitResetTime: Date.now() + 5000 }) + const b = makeAccount({ id: 'b', rateLimitResetTime: Date.now() + 10000 }) + const mgr = new AccountManager([a, b]) + expect(mgr.getMinWaitTime()).toBeGreaterThan(0) + expect(mgr.getMinWaitTime()).toBeLessThanOrEqual(5000) + }) +}) + +// ── updateUsage ─────────────────────────────────────────────────────────────── + +describe('AccountManager.updateUsage', () => { + test('updates usedCount and limitCount', () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + mgr.updateUsage(acc.id, { usedCount: 50, limitCount: 500 }) + expect(acc.usedCount).toBe(50) + expect(acc.limitCount).toBe(500) + }) + + test('updates email when provided', () => { + const acc = makeAccount({ email: 'old@example.com' }) + const mgr = new AccountManager([acc]) + mgr.updateUsage(acc.id, { usedCount: 0, limitCount: 0, email: 'new@example.com' }) + expect(acc.email).toBe('new@example.com') + }) + + test('resets failCount and marks healthy for non-permanent error', () => { + const acc = makeAccount({ failCount: 5, isHealthy: false, unhealthyReason: 'transient' }) + const mgr = new AccountManager([acc]) + mgr.updateUsage(acc.id, { usedCount: 0, limitCount: 0 }) + expect(acc.failCount).toBe(0) + expect(acc.isHealthy).toBe(true) + expect(acc.unhealthyReason).toBeUndefined() + }) + + test('does not reset health for permanent error', () => { + const acc = makeAccount({ + failCount: 10, + isHealthy: false, + unhealthyReason: 'ExpiredTokenException' + }) + const mgr = new AccountManager([acc]) + mgr.updateUsage(acc.id, { usedCount: 0, limitCount: 0 }) + expect(acc.isHealthy).toBe(false) + }) + + test('no-ops on unknown id', () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + expect(() => mgr.updateUsage('unknown', { usedCount: 99, limitCount: 99 })).not.toThrow() + expect(acc.usedCount).toBe(0) // unchanged + }) +}) + +// ── addAccount / removeAccount ───────────────────────────────────────────────── + +describe('AccountManager.addAccount / removeAccount', () => { + test('addAccount appends new account', () => { + const mgr = new AccountManager([]) + const acc = makeAccount() + mgr.addAccount(acc) + expect(mgr.getAccountCount()).toBe(1) + }) + + test('addAccount replaces existing account with same id', () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + mgr.addAccount({ ...acc, email: 'updated@example.com' }) + expect(mgr.getAccountCount()).toBe(1) + expect(mgr.getAccounts()[0]!.email).toBe('updated@example.com') + }) + + test('removeAccount removes the account', () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + mgr.removeAccount(acc) + expect(mgr.getAccountCount()).toBe(0) + }) + + test('removeAccount is no-op for unknown account', () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + mgr.removeAccount(makeAccount({ id: 'unknown' })) + expect(mgr.getAccountCount()).toBe(1) + }) + + test('cursor adjusts after removeAccount', () => { + const a = makeAccount({ id: 'a' }) + const b = makeAccount({ id: 'b', email: 'b@example.com' }) + const mgr = new AccountManager([a, b]) + mgr.removeAccount(a) + expect(mgr.getAccountCount()).toBe(1) + // Should not throw selecting next account + expect(mgr.getCurrentOrNext()).toBeDefined() + }) +}) + +// ── lowest-usage strategy ───────────────────────────────────────────────────── + +describe('AccountManager: lowest-usage strategy', () => { + test('selects account with lowest usedCount', () => { + const a = makeAccount({ id: 'a', usedCount: 100 }) + const b = makeAccount({ id: 'b', email: 'b@example.com', usedCount: 10 }) + const mgr = new AccountManager([a, b], 'lowest-usage') + const selected = mgr.getCurrentOrNext() + expect(selected?.id).toBe('b') + }) + + test('breaks ties by lastUsed', () => { + const a = makeAccount({ id: 'a', usedCount: 5, lastUsed: 1000 }) + const b = makeAccount({ id: 'b', email: 'b@example.com', usedCount: 5, lastUsed: 500 }) + const mgr = new AccountManager([a, b], 'lowest-usage') + const selected = mgr.getCurrentOrNext() + expect(selected?.id).toBe('b') // lower lastUsed = used less recently + }) +}) + +// ── markRateLimited ─────────────────────────────────────────────────────────── + +describe('AccountManager.markRateLimited', () => { + test('sets rateLimitResetTime in the future', () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + mgr.markRateLimited(acc, 30000) + expect(acc.rateLimitResetTime).toBeGreaterThan(Date.now()) + }) + + test('rate-limited account is excluded from getCurrentOrNext', () => { + const a = makeAccount({ id: 'a' }) + const mgr = new AccountManager([a]) + mgr.markRateLimited(a, 60000) + expect(mgr.getCurrentOrNext()).toBeNull() + }) +}) + +// ── shouldShowToast / shouldShowUsageToast ───────────────────────────────────── + +describe('AccountManager.shouldShowToast', () => { + test('returns true on first call', () => { + const mgr = new AccountManager([makeAccount()]) + expect(mgr.shouldShowToast()).toBe(true) + }) + + test('returns false within debounce window', () => { + const mgr = new AccountManager([makeAccount()]) + mgr.shouldShowToast(5000) + expect(mgr.shouldShowToast(5000)).toBe(false) + }) + + test('shouldShowUsageToast returns true on first call', () => { + const mgr = new AccountManager([makeAccount()]) + expect(mgr.shouldShowUsageToast()).toBe(true) + }) +}) + +// ── recovery after unhealthy ────────────────────────────────────────────────── + +describe('AccountManager: recovery from temporary unhealthy', () => { + test('account with past recoveryTime becomes available again', () => { + const acc = makeAccount({ + isHealthy: false, + failCount: 3, + recoveryTime: Date.now() - 1000 // past + }) + const mgr = new AccountManager([acc]) + const selected = mgr.getCurrentOrNext() + expect(selected).not.toBeNull() + expect(acc.isHealthy).toBe(true) + }) + + test('account with future recoveryTime is not returned (respects the wait)', () => { + const acc = makeAccount({ + isHealthy: false, + failCount: 3, + recoveryTime: Date.now() + 3_600_000 + }) + const mgr = new AccountManager([acc]) + expect(mgr.getCurrentOrNext()).toBeNull() + }) + + test('fallback returns unhealthy account with no recoveryTime (limbo state)', () => { + // An account can end up isHealthy=false with no recoveryTime due to incremental + // failCount increases that haven't yet hit the threshold. The fallback rescues it. + const acc = makeAccount({ + isHealthy: false, + failCount: 3 + // no recoveryTime + }) + const mgr = new AccountManager([acc]) + const selected = mgr.getCurrentOrNext() + expect(selected).not.toBeNull() + expect(acc.isHealthy).toBe(true) + expect(acc.recoveryTime).toBeUndefined() + }) + + test('permanently unhealthy account is never returned', () => { + const acc = makeAccount({ + isHealthy: false, + failCount: 10, + unhealthyReason: 'bearer token included in the request is invalid' + }) + const mgr = new AccountManager([acc]) + expect(mgr.getCurrentOrNext()).toBeNull() + }) +}) diff --git a/src/__tests__/auth-handler.test.ts b/src/__tests__/auth-handler.test.ts new file mode 100644 index 0000000..b5d686b --- /dev/null +++ b/src/__tests__/auth-handler.test.ts @@ -0,0 +1,135 @@ +import { describe, expect, mock, test } from 'bun:test' + +mock.module('../plugin/sync/kiro-cli.js', () => ({ + syncFromKiroCli: () => Promise.resolve(), + writeToKiroCli: () => Promise.resolve() +})) +mock.module('../kiro/auth.js', () => ({ + decodeRefreshToken: (t: string) => ({ refreshToken: t }), + encodeRefreshToken: (p: any) => p.refreshToken, + accessTokenExpired: () => false +})) + +import { AuthHandler } from '../core/auth/auth-handler.js' +import type { KiroAuthDetails, ManagedAccount } from '../plugin/types.js' + +function makeAccount(overrides: Partial = {}): ManagedAccount { + return { + id: 'acc-1', + email: 'test@example.com', + authMethod: 'idc', + region: 'eu-central-1', + refreshToken: 'r', + accessToken: 'a', + expiresAt: Date.now() + 3600000, + rateLimitResetTime: 0, + isHealthy: true, + failCount: 0, + lastUsed: 0, + usedCount: 0, + limitCount: 0, + ...overrides + } +} + +function makeAuth(): KiroAuthDetails { + return { + refresh: 'refresh-token', + access: 'access-token', + expires: Date.now() + 3600000, // not expired -> no refresh attempted + authMethod: 'idc', + region: 'eu-central-1', + profileArn: 'arn:aws:codewhisperer:eu-central-1:000000:profile/ABC' + } +} + +function makeManager(acc: ManagedAccount) { + return { + getAccounts: () => [acc], + toAuthDetails: () => makeAuth(), + updateUsage: () => {} + } +} + +const fakeRepo: any = { + batchSave: async () => {}, + invalidateCache: () => {}, + findAll: async () => [] +} + +const CREDIT_RESPONSE = JSON.stringify({ + usageBreakdownList: [ + { + freeTrialInfo: null, + currentUsage: 70, + currentUsageWithPrecision: 70.45, + usageLimit: 10000, + usageLimitWithPrecision: 10000, + displayNamePlural: 'Credits', + resourceType: 'CREDIT' + } + ], + userInfo: { email: 'test@example.com' } +}) + +describe('AuthHandler.refreshUsageFromApi', () => { + test('fetches live usage and updates the account with dashboard credits', async () => { + const acc = makeAccount({ usedCount: 4292, limitCount: 10000 }) // stale prior-period value + const handler = new AuthHandler( + { usage_tracking_enabled: true, token_expiry_buffer_ms: 300000, auto_sync_kiro_cli: false }, + fakeRepo + ) + handler.setAccountManager(makeManager(acc)) + + const original = globalThis.fetch + globalThis.fetch = mock(async () => new Response(CREDIT_RESPONSE, { status: 200 })) as any + try { + await handler.refreshUsageFromApi() + expect(acc.usedCount).toBe(70.45) + expect(acc.limitCount).toBe(10000) + } finally { + globalThis.fetch = original + } + }) + + test('keeps stored value when the live fetch fails', async () => { + const acc = makeAccount({ usedCount: 70.45, limitCount: 10000 }) + const handler = new AuthHandler( + { usage_tracking_enabled: true, token_expiry_buffer_ms: 300000, auto_sync_kiro_cli: false }, + fakeRepo + ) + handler.setAccountManager(makeManager(acc)) + + const original = globalThis.fetch + globalThis.fetch = mock(async () => new Response('boom', { status: 500 })) as any + try { + await handler.refreshUsageFromApi() + expect(acc.usedCount).toBe(70.45) // unchanged + } finally { + globalThis.fetch = original + } + }) + + test('is a one-time guard (skips the second call)', async () => { + const acc = makeAccount() + const handler = new AuthHandler( + { usage_tracking_enabled: true, token_expiry_buffer_ms: 300000, auto_sync_kiro_cli: false }, + fakeRepo + ) + handler.setAccountManager(makeManager(acc)) + + let calls = 0 + const original = globalThis.fetch + globalThis.fetch = mock(async () => { + calls++ + return new Response(CREDIT_RESPONSE, { status: 200 }) + }) as any + try { + await handler.refreshUsageFromApi() + await handler.refreshUsageFromApi() + expect(calls).toBe(1) + } finally { + globalThis.fetch = original + } + }) +}) diff --git a/src/__tests__/health.test.ts b/src/__tests__/health.test.ts new file mode 100644 index 0000000..86604b2 --- /dev/null +++ b/src/__tests__/health.test.ts @@ -0,0 +1,71 @@ +import { describe, expect, test } from 'bun:test' +import { isPermanentError } from '../plugin/health.js' + +describe('isPermanentError', () => { + test('returns false for undefined', () => { + expect(isPermanentError(undefined)).toBe(false) + }) + + test('returns false for empty string', () => { + expect(isPermanentError('')).toBe(false) + }) + + test('returns false for generic error', () => { + expect(isPermanentError('Internal Server Error')).toBe(false) + expect(isPermanentError('Rate limited')).toBe(false) + expect(isPermanentError('Network timeout')).toBe(false) + }) + + test('detects Invalid refresh token', () => { + expect(isPermanentError('Invalid refresh token')).toBe(true) + expect(isPermanentError('Error: Invalid refresh token provided')).toBe(true) + }) + + test('detects Invalid grant provided', () => { + expect(isPermanentError('Invalid grant provided')).toBe(true) + }) + + test('detects invalid_grant', () => { + expect(isPermanentError('invalid_grant')).toBe(true) + expect(isPermanentError('error: invalid_grant')).toBe(true) + }) + + test('detects ExpiredTokenException', () => { + expect(isPermanentError('ExpiredTokenException')).toBe(true) + expect(isPermanentError('AWS: ExpiredTokenException: token expired')).toBe(true) + }) + + test('detects InvalidTokenException', () => { + expect(isPermanentError('InvalidTokenException')).toBe(true) + }) + + test('detects ExpiredClientException', () => { + expect(isPermanentError('ExpiredClientException')).toBe(true) + }) + + test('detects Client is expired', () => { + expect(isPermanentError('Client is expired')).toBe(true) + }) + + test('detects HTTP_401', () => { + expect(isPermanentError('HTTP_401')).toBe(true) + expect(isPermanentError('error HTTP_401 Unauthorized')).toBe(true) + }) + + test('does not treat HTTP_403 as permanent (token expiry — should refresh, not reauth)', () => { + // HTTP_403 from Kiro means the access token expired mid-request. + // This is recoverable via token refresh, not a permanent error. + expect(isPermanentError('HTTP_403')).toBe(false) + expect(isPermanentError('error HTTP_403 Forbidden')).toBe(false) + }) + + test('does not treat bearer token invalid as permanent (handled in error-handler with refresh)', () => { + // bearer token invalid triggers a forced token refresh in ErrorHandler, not permanent unhealthy. + expect(isPermanentError('The bearer token included in the request is invalid')).toBe(false) + expect(isPermanentError('bearer token included in the request is invalid')).toBe(false) + }) + + test('detects Account Suspended', () => { + expect(isPermanentError('Account Suspended')).toBe(true) + }) +}) diff --git a/src/__tests__/kiro-cli-parser.test.ts b/src/__tests__/kiro-cli-parser.test.ts new file mode 100644 index 0000000..36aecf5 --- /dev/null +++ b/src/__tests__/kiro-cli-parser.test.ts @@ -0,0 +1,131 @@ +import { describe, expect, test } from 'bun:test' +import { + findClientCredsRecursive, + getCliDbPath, + makePlaceholderEmail, + normalizeExpiresAt, + safeJsonParse +} from '../plugin/sync/kiro-cli-parser.js' + +// ── getCliDbPath ────────────────────────────────────────────────────────────── + +describe('getCliDbPath', () => { + test('respects KIROCLI_DB_PATH override', () => { + process.env.KIROCLI_DB_PATH = '/custom/path.db' + expect(getCliDbPath()).toBe('/custom/path.db') + delete process.env.KIROCLI_DB_PATH + }) + + test('returns a string path without override', () => { + delete process.env.KIROCLI_DB_PATH + const path = getCliDbPath() + expect(typeof path).toBe('string') + expect(path.length).toBeGreaterThan(0) + }) +}) + +// ── safeJsonParse ───────────────────────────────────────────────────────────── + +describe('safeJsonParse', () => { + test('parses valid JSON string', () => { + expect(safeJsonParse('{"key":"value"}')).toEqual({ key: 'value' }) + }) + + test('returns null for invalid JSON', () => { + expect(safeJsonParse('not json')).toBeNull() + expect(safeJsonParse('{')).toBeNull() + }) + + test('returns null for non-string input', () => { + expect(safeJsonParse(42)).toBeNull() + expect(safeJsonParse(null)).toBeNull() + expect(safeJsonParse(undefined)).toBeNull() + expect(safeJsonParse({})).toBeNull() + }) +}) + +// ── normalizeExpiresAt ──────────────────────────────────────────────────────── + +describe('normalizeExpiresAt', () => { + test('ms timestamp stays as-is', () => { + const ms = 1700000000000 + expect(normalizeExpiresAt(ms)).toBe(ms) + }) + + test('seconds timestamp is converted to ms', () => { + const sec = 1700000000 // < 10_000_000_000 + expect(normalizeExpiresAt(sec)).toBe(sec * 1000) + }) + + test('ISO date string is converted to ms', () => { + const iso = '2024-01-01T00:00:00.000Z' + const expected = new Date(iso).getTime() + expect(normalizeExpiresAt(iso)).toBe(expected) + }) + + test('numeric string is converted', () => { + expect(normalizeExpiresAt('1700000000')).toBe(1700000000 * 1000) + }) + + test('returns 0 for invalid input', () => { + expect(normalizeExpiresAt(null)).toBe(0) + expect(normalizeExpiresAt('')).toBe(0) + expect(normalizeExpiresAt('not-a-date')).toBe(0) + }) +}) + +// ── findClientCredsRecursive ────────────────────────────────────────────────── + +describe('findClientCredsRecursive', () => { + test('finds flat clientId/clientSecret', () => { + const result = findClientCredsRecursive({ client_id: 'cid', client_secret: 'csec' }) + expect(result).toEqual({ clientId: 'cid', clientSecret: 'csec' }) + }) + + test('finds camelCase variant', () => { + const result = findClientCredsRecursive({ clientId: 'cid', clientSecret: 'csec' }) + expect(result).toEqual({ clientId: 'cid', clientSecret: 'csec' }) + }) + + test('finds nested credentials', () => { + const result = findClientCredsRecursive({ + nested: { deeper: { client_id: 'n-id', client_secret: 'n-sec' } } + }) + expect(result).toEqual({ clientId: 'n-id', clientSecret: 'n-sec' }) + }) + + test('finds credentials inside array', () => { + const result = findClientCredsRecursive([ + { unrelated: true }, + { client_id: 'arr-id', client_secret: 'arr-sec' } + ]) + expect(result).toEqual({ clientId: 'arr-id', clientSecret: 'arr-sec' }) + }) + + test('returns empty object when not found', () => { + expect(findClientCredsRecursive({})).toEqual({}) + expect(findClientCredsRecursive(null)).toEqual({}) + expect(findClientCredsRecursive('string')).toEqual({}) + }) +}) + +// ── makePlaceholderEmail ────────────────────────────────────────────────────── + +describe('makePlaceholderEmail', () => { + test('returns a valid placeholder email', () => { + const email = makePlaceholderEmail('idc', 'eu-central-1', 'cid', 'arn') + expect(email).toMatch(/^idc-placeholder\+[a-f0-9]+@awsapps\.local$/) + }) + + test('same inputs produce same email (deterministic)', () => { + const a = makePlaceholderEmail('idc', 'us-east-1', 'c1', 'arn1') + const b = makePlaceholderEmail('idc', 'us-east-1', 'c1', 'arn1') + expect(a).toBe(b) + }) + + test('different inputs produce different emails', () => { + const a = makePlaceholderEmail('idc', 'us-east-1', 'c1', 'arn1') + const b = makePlaceholderEmail('idc', 'eu-central-1', 'c1', 'arn1') + expect(a).not.toBe(b) + }) +}) diff --git a/src/__tests__/kiro-cli-profile.test.ts b/src/__tests__/kiro-cli-profile.test.ts new file mode 100644 index 0000000..bbebd72 --- /dev/null +++ b/src/__tests__/kiro-cli-profile.test.ts @@ -0,0 +1,68 @@ +import { Database } from 'bun:sqlite' +import { afterEach, beforeEach, describe, expect, test } from 'bun:test' +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' + +let dir: string +let dbPath: string + +beforeEach(() => { + dir = mkdtempSync(join(tmpdir(), 'kiro-profile-test-')) + dbPath = join(dir, 'data.sqlite3') +}) + +afterEach(() => { + rmSync(dir, { recursive: true, force: true }) +}) + +describe('readActiveProfileArnFromKiroCli', () => { + test('returns undefined when DB file does not exist', async () => { + process.env.KIROCLI_DB_PATH = join(dir, 'nonexistent.sqlite3') + const { readActiveProfileArnFromKiroCli } = await import('../plugin/sync/kiro-cli-profile.js') + expect(readActiveProfileArnFromKiroCli()).toBeUndefined() + delete process.env.KIROCLI_DB_PATH + }) + + test('returns profileArn from state table', async () => { + process.env.KIROCLI_DB_PATH = dbPath + const db = new Database(dbPath) + db.run('CREATE TABLE state (key TEXT PRIMARY KEY, value TEXT)') + db.run('INSERT INTO state (key, value) VALUES (?, ?)', [ + 'api.codewhisperer.profile', + JSON.stringify({ arn: 'arn:aws:codewhisperer:eu-central-1:123:profile/ABC' }) + ]) + db.close() + + const { readActiveProfileArnFromKiroCli } = await import('../plugin/sync/kiro-cli-profile.js') + const result = readActiveProfileArnFromKiroCli() + expect(result).toBe('arn:aws:codewhisperer:eu-central-1:123:profile/ABC') + delete process.env.KIROCLI_DB_PATH + }) + + test('returns undefined when row is missing', async () => { + process.env.KIROCLI_DB_PATH = dbPath + const db = new Database(dbPath) + db.run('CREATE TABLE state (key TEXT PRIMARY KEY, value TEXT)') + db.close() + + const { readActiveProfileArnFromKiroCli } = await import('../plugin/sync/kiro-cli-profile.js') + expect(readActiveProfileArnFromKiroCli()).toBeUndefined() + delete process.env.KIROCLI_DB_PATH + }) + + test('returns undefined when JSON has no arn field', async () => { + process.env.KIROCLI_DB_PATH = dbPath + const db = new Database(dbPath) + db.run('CREATE TABLE state (key TEXT PRIMARY KEY, value TEXT)') + db.run('INSERT INTO state (key, value) VALUES (?, ?)', [ + 'api.codewhisperer.profile', + JSON.stringify({ other: 'field' }) + ]) + db.close() + + const { readActiveProfileArnFromKiroCli } = await import('../plugin/sync/kiro-cli-profile.js') + expect(readActiveProfileArnFromKiroCli()).toBeUndefined() + delete process.env.KIROCLI_DB_PATH + }) +}) diff --git a/src/__tests__/usage.test.ts b/src/__tests__/usage.test.ts new file mode 100644 index 0000000..5951794 --- /dev/null +++ b/src/__tests__/usage.test.ts @@ -0,0 +1,247 @@ +import { describe, expect, mock, test } from 'bun:test' +import type { KiroAuthDetails, ManagedAccount } from '../plugin/types.js' +import { fetchUsageLimits, updateAccountQuota } from '../plugin/usage.js' + +function makeAuth(overrides: Partial = {}): KiroAuthDetails { + return { + refresh: 'refresh-token', + access: 'access-token', + expires: Date.now() + 3600000, + authMethod: 'idc', + region: 'eu-central-1', + profileArn: 'arn:aws:codewhisperer:eu-central-1:000000:profile/ABC', + ...overrides + } +} + +function makeAccount(overrides: Partial = {}): ManagedAccount { + return { + id: 'acc-1', + email: 'test@example.com', + authMethod: 'idc', + region: 'eu-central-1', + refreshToken: 'r', + accessToken: 'a', + expiresAt: Date.now() + 3600000, + rateLimitResetTime: 0, + isHealthy: true, + failCount: 0, + lastUsed: 0, + usedCount: 0, + limitCount: 0, + ...overrides + } +} + +// ── updateAccountQuota ──────────────────────────────────────────────────────── + +describe('updateAccountQuota', () => { + test('updates usedCount and limitCount on account', () => { + const acc = makeAccount() + updateAccountQuota(acc, { usedCount: 150, limitCount: 2000 }) + expect(acc.usedCount).toBe(150) + expect(acc.limitCount).toBe(2000) + }) + + test('updates email when provided', () => { + const acc = makeAccount({ email: 'old@example.com' }) + updateAccountQuota(acc, { usedCount: 0, limitCount: 0, email: 'new@example.com' }) + expect(acc.email).toBe('new@example.com') + }) + + test('does not update email when not provided', () => { + const acc = makeAccount({ email: 'keep@example.com' }) + updateAccountQuota(acc, { usedCount: 5, limitCount: 100 }) + expect(acc.email).toBe('keep@example.com') + }) + + test('calls accountManager.updateUsage when provided', () => { + const acc = makeAccount() + const calls: any[] = [] + const mgr = { updateUsage: (id: string, meta: any) => calls.push({ id, meta }) } + updateAccountQuota(acc, { usedCount: 10, limitCount: 50 }, mgr) + expect(calls).toHaveLength(1) + expect(calls[0].id).toBe('acc-1') + expect(calls[0].meta.usedCount).toBe(10) + expect(calls[0].meta.limitCount).toBe(50) + }) + + test('handles missing usedCount/limitCount gracefully', () => { + const acc = makeAccount() + updateAccountQuota(acc, {}) + expect(acc.usedCount).toBe(0) + expect(acc.limitCount).toBe(0) + }) +}) + +// ── fetchUsageLimits ────────────────────────────────────────────────────────── + +describe('fetchUsageLimits', () => { + test('returns usedCount and limitCount from usageBreakdownList', async () => { + const mockFetch = mock( + async () => + new Response( + JSON.stringify({ + usageBreakdownList: [ + { + freeTrialInfo: { currentUsage: 100, usageLimit: 1000 }, + currentUsage: 50, + usageLimit: 500 + } + ], + userInfo: { email: 'test@example.com' } + }), + { status: 200 } + ) + ) + const original = globalThis.fetch + globalThis.fetch = mockFetch as any + try { + const result = await fetchUsageLimits(makeAuth()) + expect(result.usedCount).toBe(150) // 100 + 50 + expect(result.limitCount).toBe(1500) // 1000 + 500 + expect(result.email).toBe('test@example.com') + } finally { + globalThis.fetch = original + } + }) + + test('prefers WithPrecision fields (matches Kiro dashboard credits)', async () => { + // Mirrors the real Kiro Power getUsageLimits response: the integer + // currentUsage is rounded, the dashboard shows currentUsageWithPrecision. + const mockFetch = mock( + async () => + new Response( + JSON.stringify({ + usageBreakdownList: [ + { + freeTrialInfo: null, + currentUsage: 70, + currentUsageWithPrecision: 70.45, + usageLimit: 10000, + usageLimitWithPrecision: 10000, + displayNamePlural: 'Credits', + resourceType: 'CREDIT' + } + ], + userInfo: { email: 'test@example.com' } + }), + { status: 200 } + ) + ) + const original = globalThis.fetch + globalThis.fetch = mockFetch as any + try { + const result = await fetchUsageLimits(makeAuth()) + expect(result.usedCount).toBe(70.45) + expect(result.limitCount).toBe(10000) + } finally { + globalThis.fetch = original + } + }) + + test('falls back to integer fields when precision absent', async () => { + const mockFetch = mock( + async () => + new Response( + JSON.stringify({ + usageBreakdownList: [{ currentUsage: 50, usageLimit: 500 }], + userInfo: { email: 'test@example.com' } + }), + { status: 200 } + ) + ) + const original = globalThis.fetch + globalThis.fetch = mockFetch as any + try { + const result = await fetchUsageLimits(makeAuth()) + expect(result.usedCount).toBe(50) + expect(result.limitCount).toBe(500) + } finally { + globalThis.fetch = original + } + }) + + test('retries on FEATURE_NOT_SUPPORTED and succeeds on later attempt', async () => { + let callCount = 0 + const mockFetch = mock(async () => { + callCount++ + if (callCount < 3) { + return new Response('FEATURE_NOT_SUPPORTED', { status: 400 }) + } + return new Response(JSON.stringify({ usageBreakdownList: [], userInfo: {} }), { status: 200 }) + }) + const original = globalThis.fetch + globalThis.fetch = mockFetch as any + try { + const result = await fetchUsageLimits(makeAuth()) + expect(callCount).toBeGreaterThanOrEqual(3) + expect(result.usedCount).toBe(0) + } finally { + globalThis.fetch = original + } + }) + + test('throws when all attempts fail', async () => { + const mockFetch = mock(async () => new Response('Server Error', { status: 500 })) + const original = globalThis.fetch + globalThis.fetch = mockFetch as any + try { + await expect(fetchUsageLimits(makeAuth())).rejects.toThrow() + } finally { + globalThis.fetch = original + } + }) + + test('does NOT chain to next param combo on 429 (rate limit)', async () => { + let callCount = 0 + const mockFetch = mock(async () => { + callCount++ + return new Response(JSON.stringify({ message: 'rate limited' }), { + status: 429, + headers: { 'x-amzn-errortype': 'ThrottlingException' } + }) + }) + const original = globalThis.fetch + globalThis.fetch = mockFetch as any + try { + await expect(fetchUsageLimits(makeAuth())).rejects.toThrow(/429|Throttling/i) + // Old behaviour would call all 4 attempts. New behaviour stops at 1. + expect(callCount).toBe(1) + } finally { + globalThis.fetch = original + } + }) + + test('does NOT chain to next param combo on 401', async () => { + let callCount = 0 + const mockFetch = mock(async () => { + callCount++ + return new Response(JSON.stringify({ message: 'unauthorized' }), { status: 401 }) + }) + const original = globalThis.fetch + globalThis.fetch = mockFetch as any + try { + await expect(fetchUsageLimits(makeAuth())).rejects.toThrow(/401/) + expect(callCount).toBe(1) + } finally { + globalThis.fetch = original + } + }) + + test('does NOT retry on network error across all combos', async () => { + let callCount = 0 + const mockFetch = mock(async () => { + callCount++ + throw new Error('fetch failed: ECONNRESET') + }) + const original = globalThis.fetch + globalThis.fetch = mockFetch as any + try { + await expect(fetchUsageLimits(makeAuth())).rejects.toThrow(/ECONNRESET/) + expect(callCount).toBe(1) + } finally { + globalThis.fetch = original + } + }) +}) diff --git a/src/core/account/account-selector.ts b/src/core/account/account-selector.ts index fa26969..f26e6d0 100644 --- a/src/core/account/account-selector.ts +++ b/src/core/account/account-selector.ts @@ -1,6 +1,7 @@ import type { AccountRepository } from '../../infrastructure/database/account-repository' import type { AccountManager } from '../../plugin/accounts' import type { ManagedAccount } from '../../plugin/types' +import { summarizeUsage } from '../../plugin/usage' type ToastFunction = (message: string, variant: 'info' | 'warning' | 'success' | 'error') => void @@ -72,11 +73,8 @@ export class AccountSelector { } private formatUsageMessage(usedCount: number, limitCount: number, email: string): string { - if (limitCount > 0) { - const percentage = Math.round((usedCount / limitCount) * 100) - return `Usage (${email}): ${usedCount}/${limitCount} (${percentage}%)` - } - return `Usage (${email}): ${usedCount}` + const { used, limit, pct } = summarizeUsage(usedCount, limitCount) + return limit > 0 ? `Usage (${email}): ${used}/${limit} (${pct}%)` : `Usage (${email}): ${used}` } private checkCircuitBreaker(): void { diff --git a/src/core/account/usage-tracker.ts b/src/core/account/usage-tracker.ts index 2ea1a74..cf55498 100644 --- a/src/core/account/usage-tracker.ts +++ b/src/core/account/usage-tracker.ts @@ -37,32 +37,50 @@ export class UsageTracker { }) } + // Fetch usage once and persist it, bypassing the cooldown/retry loop. Used by + // the startup refresh, where the caller handles token refresh and fallback. + async syncNow(account: ManagedAccount, auth: KiroAuthDetails): Promise { + const u = await fetchUsageLimits(auth) + updateAccountQuota(account, u, this.accountManager) + await this.repository.batchSave(this.accountManager.getAccounts()) + } + private async syncWithRetry( account: ManagedAccount, auth: KiroAuthDetails, attempt: number ): Promise { try { - const u = await fetchUsageLimits(auth) - updateAccountQuota(account, u, this.accountManager) - await this.repository.batchSave(this.accountManager.getAccounts()) + await this.syncNow(account, auth) } catch (e: any) { - if (attempt < this.config.usage_sync_max_retries) { + const msg = e?.message || '' + + // Don't retry rate-limit errors — that just amplifies the problem. + const isRateLimit = + msg.includes('429') || + msg.includes('ThrottlingException') || + msg.includes('TooManyRequests') + + if (!isRateLimit && attempt < this.config.usage_sync_max_retries) { await this.sleep(1000 * Math.pow(2, attempt)) return this.syncWithRetry(account, auth, attempt + 1) } - if (e.message?.includes('FEATURE_NOT_SUPPORTED')) { - // Some IDC profiles don't support getUsageLimits; don't penalize the account. + if (msg.includes('FEATURE_NOT_SUPPORTED')) { + // Some IDC profiles don't expose getUsageLimits — not an error. + return + } + + if (isRateLimit) { + // Don't penalize the account; the request flow has its own 429 handler. + logger.warn('Usage sync rate-limited; skipping until next cooldown', { + accountId: account.id + }) return } - if ( - e.message?.includes('403') || - e.message?.includes('invalid') || - e.message?.includes('bearer token') - ) { - this.accountManager.markUnhealthy(account, e.message) + if (msg.includes('403') || msg.includes('invalid') || msg.includes('bearer token')) { + this.accountManager.markUnhealthy(account, msg) this.repository.save(account).catch(() => {}) } diff --git a/src/core/auth/auth-handler.ts b/src/core/auth/auth-handler.ts index ab00448..a3e1147 100644 --- a/src/core/auth/auth-handler.ts +++ b/src/core/auth/auth-handler.ts @@ -2,12 +2,16 @@ import type { AuthHook } from '@opencode-ai/plugin' import type { AccountRepository } from '../../infrastructure/database/account-repository.js' import { RegionSchema } from '../../plugin/config/schema.js' import * as logger from '../../plugin/logger.js' +import { summarizeUsage } from '../../plugin/usage.js' +import { UsageTracker } from '../account/usage-tracker.js' import { IdcAuthMethod } from './idc-auth-method.js' +import { TokenRefresher } from './token-refresher.js' type ToastFunction = (message: string, variant: 'info' | 'warning' | 'success' | 'error') => void export class AuthHandler { private accountManager?: any + private startupUsageFetched = false constructor( private config: any, @@ -29,7 +33,53 @@ export class AuthHandler { logger.log('Kiro CLI sync: done', { importedAccounts: accounts.length }) } - this.logUsageSummary(showToast) + // Refresh usage before the summary toast: the persisted value is stale after + // the monthly reset until the first request syncs. Backgrounded so it never + // delays the auth loader, and falls back to the stored value on error. + void (async () => { + try { + await this.refreshUsageFromApi(showToast) + } catch (e) { + logger.warn('Startup usage refresh failed', { + error: e instanceof Error ? e.message : String(e) + }) + } + this.logUsageSummary(showToast) + })() + } + + async refreshUsageFromApi(showToast?: ToastFunction): Promise { + if (!this.accountManager || this.config.usage_tracking_enabled === false) return + if (this.startupUsageFetched) return + this.startupUsageFetched = true + + const { syncFromKiroCli } = await import('../../plugin/sync/kiro-cli.js') + const tokenRefresher = new TokenRefresher( + this.config, + this.accountManager, + syncFromKiroCli, + this.repository + ) + const usageTracker = new UsageTracker(this.config, this.accountManager, this.repository) + const toast: ToastFunction = showToast ?? (() => {}) + + for (const acc of this.accountManager.getAccounts()) { + if (!acc.isHealthy) continue + try { + const { account: usable } = await tokenRefresher.refreshIfNeeded( + acc, + this.accountManager.toAuthDetails(acc), + toast + ) + if (!usable.isHealthy) continue + await usageTracker.syncNow(usable, this.accountManager.toAuthDetails(usable)) + } catch (e) { + logger.warn('Startup usage fetch failed; keeping stored value', { + email: acc.email, + error: e instanceof Error ? e.message : String(e) + }) + } + } } private logUsageSummary(showToast?: ToastFunction): void { @@ -38,10 +88,8 @@ export class AuthHandler { if (!accounts.length) return for (const acc of accounts) { - const used = acc.usedCount ?? 0 - const limit = acc.limitCount ?? 0 + const { used, limit, pct } = summarizeUsage(acc.usedCount ?? 0, acc.limitCount ?? 0) if (limit > 0) { - const pct = Math.round((used / limit) * 100) const msg = `Kiro usage (${acc.email}): ${used}/${limit} (${pct}%)` logger.log(msg) if (showToast) { diff --git a/src/core/auth/idc-auth-method.ts b/src/core/auth/idc-auth-method.ts index 10e56a8..9b16da4 100644 --- a/src/core/auth/idc-auth-method.ts +++ b/src/core/auth/idc-auth-method.ts @@ -1,5 +1,5 @@ import type { AuthOuathResult } from '@opencode-ai/plugin' -import { exec } from 'node:child_process' +import { execFile } from 'node:child_process' import { extractRegionFromArn, normalizeRegion } from '../../constants.js' import type { AccountRepository } from '../../infrastructure/database/account-repository.js' import { authorizeKiroIDC, pollKiroIDCToken } from '../../kiro/oauth-idc.js' @@ -11,15 +11,14 @@ import type { KiroRegion, ManagedAccount } from '../../plugin/types.js' import { fetchUsageLimits } from '../../plugin/usage.js' const openBrowser = (url: string) => { - const escapedUrl = url.replace(/"/g, '\\"') const platform = process.platform - const cmd = + const [bin, ...args] = platform === 'win32' - ? `cmd /c start "" "${escapedUrl}"` + ? ['cmd', '/c', 'start', '', url] : platform === 'darwin' - ? `open "${escapedUrl}"` - : `xdg-open "${escapedUrl}"` - exec(cmd, (error) => { + ? ['open', url] + : ['xdg-open', url] + execFile(bin!, args, (error) => { if (error) logger.warn(`Browser error: ${error.message}`) }) } @@ -62,8 +61,14 @@ export class IdcAuthMethod { const invokedWithoutPrompts = !inputs || Object.keys(inputs).length === 0 const startUrl = normalizeStartUrl(inputs?.start_url || this.config.idc_start_url) || undefined - const oidcRegion: KiroRegion = normalizeRegion(inputs?.idc_region || this.config.idc_region) + // For the OIDC device-code flow, prefer explicit idc_region, then fall back to + // the region from a pre-configured profileArn, then default_region. This ensures + // accounts with a eu-central-1 profileArn don't hit oidc.us-east-1.amazonaws.com. const configuredProfileArn = this.config.idc_profile_arn + const arnRegion = extractRegionFromArn(configuredProfileArn) + const oidcRegion: KiroRegion = normalizeRegion( + inputs?.idc_region || this.config.idc_region || arnRegion || configuredServiceRegion + ) logger.log('IDC authorize: resolved defaults', { hasInputs: !!inputs && Object.keys(inputs).length > 0, invokedWithoutPrompts, @@ -161,7 +166,7 @@ export class IdcAuthMethod { email, authMethod: 'idc', region: serviceRegion, - oidcRegion, + oidcRegion: oidcRegion, clientId: token.clientId, clientSecret: token.clientSecret, profileArn, diff --git a/src/core/auth/token-refresher.ts b/src/core/auth/token-refresher.ts index a1b927c..7eb6672 100644 --- a/src/core/auth/token-refresher.ts +++ b/src/core/auth/token-refresher.ts @@ -34,7 +34,9 @@ export class TokenRefresher { try { const newAuth = await refreshAccessToken(auth) this.accountManager.updateFromAuth(account, newAuth) - await this.repository.batchSave(this.accountManager.getAccounts()) + // Persist only the updated account instead of all accounts — avoids + // invalidating the whole AccountCache on every token refresh. + await this.repository.save(account) return { account, shouldContinue: false } } catch (e: any) { return await this.handleRefreshError(e, account, showToast) @@ -81,7 +83,7 @@ export class TokenRefresher { error.message.includes('Invalid grant provided') || error.message.includes('Client is expired')) ) { - this.accountManager.markUnhealthy(account, error.message) + this.accountManager.markUnhealthy(account, error.code || error.message) await this.repository.batchSave(this.accountManager.getAccounts()) return { account, shouldContinue: true } } diff --git a/src/plugin/accounts.ts b/src/plugin/accounts.ts index d0a623e..0fd7ca1 100644 --- a/src/plugin/accounts.ts +++ b/src/plugin/accounts.ts @@ -17,9 +17,12 @@ export function createDeterministicAccountId( clientId?: string, profileArn?: string ): string { - return createHash('sha256') - .update(`${email}:${method}:${clientId || ''}:${profileArn || ''}`) - .digest('hex') + // IDC clientId rotates on re-auth; profileArn + email is the stable identity. + const key = + method === 'idc' + ? `${email}:${method}:${profileArn || ''}` + : `${email}:${method}:${clientId || ''}:${profileArn || ''}` + return createHash('sha256').update(key).digest('hex') } export class AccountManager { @@ -34,6 +37,19 @@ export class AccountManager { this.strategy = strategy } static async loadFromDisk(strategy?: AccountSelectionStrategy): Promise { + // Sweep stale rows (test placeholders, never-upgraded placeholders, + // long-expired permanently-unhealthy accounts) before loading so the + // account pool reflects only real, usable credentials. + try { + const removed = await kiroDb.cleanupTestAndStaleAccounts() + if (removed > 0) { + logger.log(`Accounts: swept ${removed} stale row(s) from DB`) + } + } catch (e) { + logger.debug( + `Accounts: cleanup sweep failed (non-fatal): ${e instanceof Error ? e.message : String(e)}` + ) + } const rows = kiroDb.getAccounts() const accounts: ManagedAccount[] = rows.map((r: any) => ({ id: r.id, @@ -102,8 +118,16 @@ export class AccountManager { if (this.strategy === 'sticky') { selected = available.find((_, i) => i === this.cursor) || available[0] } else if (this.strategy === 'round-robin') { - selected = available[this.cursor % available.length] - this.cursor = (this.cursor + 1) % available.length + // Cursor anchored to this.accounts, not the filtered `available` list + const n = this.accounts.length + for (let i = 0; i < n; i++) { + const candidate = this.accounts[(this.cursor + i) % n] + if (candidate && available.includes(candidate)) { + selected = candidate + this.cursor = (this.accounts.indexOf(candidate) + 1) % n + break + } + } } else if (this.strategy === 'lowest-usage') { selected = [...available].sort( (a, b) => (a.usedCount || 0) - (b.usedCount || 0) || (a.lastUsed || 0) - (b.lastUsed || 0) @@ -111,22 +135,30 @@ export class AccountManager { } } if (!selected) { + // Fallback: unhealthy accounts without a scheduled recoveryTime const fallback = this.accounts - .filter((a) => !a.isHealthy && a.failCount < 10 && !isPermanentError(a.unhealthyReason)) + .filter( + (a) => + !a.isHealthy && + a.failCount < 10 && + !isPermanentError(a.unhealthyReason) && + !a.recoveryTime + ) .sort( (a, b) => (a.usedCount || 0) - (b.usedCount || 0) || (a.lastUsed || 0) - (b.lastUsed || 0) )[0] if (fallback) { fallback.isHealthy = true delete fallback.unhealthyReason - delete fallback.recoveryTime selected = fallback } } if (selected) { selected.lastUsed = now selected.usedCount = (selected.usedCount || 0) + 1 - this.cursor = this.accounts.indexOf(selected) + if (this.strategy !== 'round-robin') { + this.cursor = this.accounts.indexOf(selected) + } return selected } return null diff --git a/src/plugin/health.ts b/src/plugin/health.ts index 10eb464..6916b69 100644 --- a/src/plugin/health.ts +++ b/src/plugin/health.ts @@ -9,6 +9,6 @@ export function isPermanentError(reason?: string): boolean { reason.includes('ExpiredClientException') || reason.includes('Client is expired') || reason.includes('HTTP_401') || - reason.includes('HTTP_403') + reason.includes('Account Suspended') ) } diff --git a/src/plugin/storage/sqlite.ts b/src/plugin/storage/sqlite.ts index 53cedfc..6e2e753 100644 --- a/src/plugin/storage/sqlite.ts +++ b/src/plugin/storage/sqlite.ts @@ -171,6 +171,54 @@ export class KiroDatabase { }) } + async cleanupTestAndStaleAccounts(staleDays = 30): Promise { + const cutoffMs = Date.now() - staleDays * 24 * 60 * 60 * 1000 + return withDatabaseLock(this.path, async () => { + const before = (this.db.prepare('SELECT COUNT(*) AS n FROM accounts').get() as { n: number }) + .n + this.db + .prepare( + `DELETE FROM accounts + WHERE email = 'test@example.com' + OR email LIKE 'placeholder-%@awsapps.local' + OR (is_healthy = 0 + AND unhealthy_reason IN ('Account Suspended', 'ExpiredTokenException') + AND (recovery_time IS NULL OR recovery_time < ?))` + ) + .run(cutoffMs) + const after = (this.db.prepare('SELECT COUNT(*) AS n FROM accounts').get() as { n: number }).n + return before - after + }) + } + + async deleteStaleIdcDuplicates( + canonicalId: string, + email: string, + profileArn: string + ): Promise { + await withDatabaseLock(this.path, async () => { + this.db + .prepare( + `DELETE FROM accounts + WHERE auth_method = 'idc' + AND email = ? + AND profile_arn = ? + AND id != ?` + ) + .run(email, profileArn, canonicalId) + // Also clean up placeholder rows for the same profileArn. + this.db + .prepare( + `DELETE FROM accounts + WHERE auth_method = 'idc' + AND profile_arn = ? + AND email LIKE 'placeholder-%' + AND id != ?` + ) + .run(profileArn, canonicalId) + }) + } + private rowToAccount(row: any): ManagedAccount { return { id: row.id, diff --git a/src/plugin/sync/kiro-cli.ts b/src/plugin/sync/kiro-cli.ts index 1344363..d5d33f9 100644 --- a/src/plugin/sync/kiro-cli.ts +++ b/src/plugin/sync/kiro-cli.ts @@ -52,10 +52,12 @@ export async function syncFromKiroCli() { const isIdc = row.key.includes('odic') const authMethod = isIdc ? 'idc' : 'desktop' - const oidcRegion = normalizeRegion(data.region) let profileArn: string | undefined = data.profile_arn || data.profileArn if (!profileArn && isIdc) profileArn = activeProfileArn || readActiveProfileArnFromKiroCli() - const serviceRegion = extractRegionFromArn(profileArn) || oidcRegion + // serviceRegion wins over data.region: kiro-cli stores data.region as the + // OIDC region (often us-east-1) regardless of where the account actually lives. + const serviceRegion = extractRegionFromArn(profileArn) || normalizeRegion(data.region) + const oidcRegion = serviceRegion const startUrl: string | undefined = typeof data.start_url === 'string' ? data.start_url @@ -128,8 +130,33 @@ export async function syncFromKiroCli() { } } - const resolvedEmail = + // Reuse known email for this profileArn to avoid duplicate placeholder rows + let resolvedEmail: string = email || makePlaceholderEmail(authMethod, serviceRegion, clientId, profileArn) + if (resolvedEmail.startsWith('placeholder-')) { + let existingReal: any | undefined + if (profileArn) { + existingReal = all.find( + (a) => + a.auth_method === authMethod && + a.profile_arn === profileArn && + a.email && + !a.email.startsWith('placeholder-') + ) + } + if (!existingReal && authMethod === 'idc' && clientId) { + existingReal = all.find( + (a) => + a.auth_method === 'idc' && + a.client_id === clientId && + a.email && + !a.email.startsWith('placeholder-') + ) + } + if (existingReal) { + resolvedEmail = existingReal.email + } + } const id = createDeterministicAccountId(resolvedEmail, authMethod, clientId, profileArn) const existingById = all.find((a) => a.id === id) @@ -210,6 +237,10 @@ export async function syncFromKiroCli() { clientId, profileArn }) + + if (authMethod === 'idc' && profileArn) { + await kiroDb.deleteStaleIdcDuplicates(id, resolvedEmail, profileArn) + } } } diff --git a/src/plugin/usage.ts b/src/plugin/usage.ts index 41e1284..6c058fc 100644 --- a/src/plugin/usage.ts +++ b/src/plugin/usage.ts @@ -39,20 +39,23 @@ export async function fetchUsageLimits(auth: KiroAuthDetails): Promise { const errType = res.headers.get('x-amzn-errortype') || res.headers.get('x-amzn-error-type') || '' - if (body.includes('FEATURE_NOT_SUPPORTED') && index < attempts.length - 1) { - continue - } - const msg = body && body.length > 0 ? `${body.slice(0, 2000)}${body.length > 2000 ? '…' : ''}` : `HTTP ${res.status}` - lastError = new Error( - `Status: ${res.status}${errType ? ` (${errType})` : ''}${ - requestId ? ` [${requestId}]` : '' - }: ${msg}` - ) - continue + const errorMessage = `Status: ${res.status}${errType ? ` (${errType})` : ''}${ + requestId ? ` [${requestId}]` : '' + }: ${msg}` + + // Only chain to the next param combo for FEATURE_NOT_SUPPORTED. + // Other failures (429, 401, 5xx, network) bubble up so the caller's + // retry/backoff handles them, instead of hitting the API 4x per call. + if (body.includes('FEATURE_NOT_SUPPORTED') && index < attempts.length - 1) { + lastError = new Error(errorMessage) + continue + } + + throw new Error(errorMessage) } const data: any = await res.json() @@ -60,24 +63,40 @@ export async function fetchUsageLimits(auth: KiroAuthDetails): Promise { limitCount = 0 if (Array.isArray(data.usageBreakdownList)) { for (const s of data.usageBreakdownList) { + // Kiro reports a rounded integer (currentUsage) plus the exact value + // (currentUsageWithPrecision) — the latter is what the Kiro dashboard + // shows (e.g. 70.45 credits). Prefer it; fall back to the integer. if (s.freeTrialInfo) { - usedCount += s.freeTrialInfo.currentUsage || 0 - limitCount += s.freeTrialInfo.usageLimit || 0 + usedCount += + s.freeTrialInfo.currentUsageWithPrecision ?? s.freeTrialInfo.currentUsage ?? 0 + limitCount += s.freeTrialInfo.usageLimitWithPrecision ?? s.freeTrialInfo.usageLimit ?? 0 } - usedCount += s.currentUsage || 0 - limitCount += s.usageLimit || 0 + usedCount += s.currentUsageWithPrecision ?? s.currentUsage ?? 0 + limitCount += s.usageLimitWithPrecision ?? s.usageLimit ?? 0 } } return { usedCount, limitCount, email: data.userInfo?.email } } catch (e) { - lastError = e instanceof Error ? e : new Error(String(e)) - if (index < attempts.length - 1) continue + // Network errors bubble up — don't try the next param combo. + throw e instanceof Error ? e : new Error(String(e)) } } throw lastError || new Error('All getUsageLimits attempts failed') } +// Credits come back fractional (e.g. 70.45); round for display and derive the +// percentage in one place so the startup summary and the high-usage warning agree. +export function summarizeUsage( + usedCount: number, + limitCount: number +): { used: number; limit: number; pct: number } { + const used = Number(usedCount.toFixed(2)) + const limit = Number(limitCount.toFixed(2)) + const pct = limit > 0 ? Math.round((used / limit) * 100) : 0 + return { used, limit, pct } +} + export function updateAccountQuota( account: ManagedAccount, usage: any,