diff --git a/README.md b/README.md index 4e646a4..5ceb710 100644 --- a/README.md +++ b/README.md @@ -30,165 +30,35 @@ Add the plugin to your `opencode.json` or `opencode.jsonc`: ```json { - "plugin": ["@zhafron/opencode-kiro-auth"], - "provider": { - "kiro": { - "models": { - "claude-sonnet-4-5": { - "name": "Claude Sonnet 4.5", - "limit": { "context": 200000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] } - }, - "claude-sonnet-4-5-thinking": { - "name": "Claude Sonnet 4.5 Thinking", - "limit": { "context": 200000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] }, - "variants": { - "low": { "thinkingConfig": { "thinkingBudget": 8192 } }, - "medium": { "thinkingConfig": { "thinkingBudget": 16384 } }, - "max": { "thinkingConfig": { "thinkingBudget": 32768 } } - } - }, - "claude-sonnet-4-6": { - "name": "Claude Sonnet 4.6", - "limit": { "context": 1000000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] } - }, - "claude-sonnet-4-6-thinking": { - "name": "Claude Sonnet 4.6 Thinking", - "limit": { "context": 1000000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] }, - "variants": { - "low": { "thinkingConfig": { "thinkingBudget": 8192 } }, - "medium": { "thinkingConfig": { "thinkingBudget": 16384 } }, - "max": { "thinkingConfig": { "thinkingBudget": 32768 } } - } - }, - "claude-haiku-4-5": { - "name": "Claude Haiku 4.5", - "limit": { "context": 200000, "output": 64000 }, - "modalities": { "input": ["text", "image"], "output": ["text"] } - }, - "claude-opus-4-5": { - "name": "Claude Opus 4.5", - "limit": { "context": 200000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] } - }, - "claude-opus-4-5-thinking": { - "name": "Claude Opus 4.5 Thinking", - "limit": { "context": 200000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] }, - "variants": { - "low": { "thinkingConfig": { "thinkingBudget": 8192 } }, - "medium": { "thinkingConfig": { "thinkingBudget": 16384 } }, - "max": { "thinkingConfig": { "thinkingBudget": 32768 } } - } - }, - "claude-opus-4-6": { - "name": "Claude Opus 4.6", - "limit": { "context": 1000000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] } - }, - "claude-opus-4-6-thinking": { - "name": "Claude Opus 4.6 Thinking", - "limit": { "context": 1000000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] }, - "variants": { - "low": { "thinkingConfig": { "thinkingBudget": 8192 } }, - "medium": { "thinkingConfig": { "thinkingBudget": 16384 } }, - "max": { "thinkingConfig": { "thinkingBudget": 32768 } } - } - }, - "claude-opus-4-6-1m": { - "name": "Claude Opus 4.6 (1M Context)", - "limit": { "context": 1000000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] } - }, - "claude-opus-4-6-1m-thinking": { - "name": "Claude Opus 4.6 (1M Context) Thinking", - "limit": { "context": 1000000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] }, - "variants": { - "low": { "thinkingConfig": { "thinkingBudget": 8192 } }, - "medium": { "thinkingConfig": { "thinkingBudget": 16384 } }, - "max": { "thinkingConfig": { "thinkingBudget": 32768 } } - } - }, - "claude-opus-4-7": { - "name": "Claude Opus 4.7", - "limit": { "context": 1000000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] } - }, - "claude-opus-4-7-thinking": { - "name": "Claude Opus 4.7 Thinking", - "limit": { "context": 1000000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] }, - "variants": { - "low": { "thinkingConfig": { "thinkingBudget": 8192 } }, - "medium": { "thinkingConfig": { "thinkingBudget": 16384 } }, - "max": { "thinkingConfig": { "thinkingBudget": 32768 } } - } - }, - "claude-sonnet-4-5-1m": { - "name": "Claude Sonnet 4.5 (1M Context)", - "limit": { "context": 1000000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] } - }, - "claude-sonnet-4-6-1m": { - "name": "Claude Sonnet 4.6 (1M Context)", - "limit": { "context": 1000000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] } - }, - "claude-sonnet-4-6-1m-thinking": { - "name": "Claude Sonnet 4.6 (1M Context) Thinking", - "limit": { "context": 1000000, "output": 64000 }, - "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] }, - "variants": { - "low": { "thinkingConfig": { "thinkingBudget": 8192 } }, - "medium": { "thinkingConfig": { "thinkingBudget": 16384 } }, - "max": { "thinkingConfig": { "thinkingBudget": 32768 } } - } - }, - "auto": { "name": "Auto (1.0x)" }, - "claude-sonnet-4": { - "name": "Claude Sonnet 4.0 (1.3x)", - "limit": { "context": 200000, "output": 64000 } - }, - "deepseek-3.2": { - "name": "DeepSeek 3.2 (0.25x)", - "limit": { "context": 128000, "output": 64000 } - }, - "glm-5": { "name": "GLM-5 (0.5x)", "limit": { "context": 200000, "output": 64000 } }, - "minimax-m2.5": { - "name": "MiniMax 2.5 (0.25x)", - "limit": { "context": 200000, "output": 64000 } - }, - "minimax-m2.1": { - "name": "MiniMax 2.1 (0.15x)", - "limit": { "context": 200000, "output": 64000 } - }, - "qwen3-coder-next": { - "name": "Qwen3 Coder Next (0.05x)", - "limit": { "context": 256000, "output": 64000 } - } - } - } - } + "plugin": ["@zhafron/opencode-kiro-auth"] } ``` +That's all you need — the plugin registers the `kiro-auth` provider and its +models automatically. + +**Provider id — `kiro-auth` (with a `kiro` fallback):** use `kiro-auth` for new +setups. If you previously configured this plugin under the old `kiro` id, it +keeps working — the plugin also registers `kiro` as a fallback alias, routing it +through the same backend. `kiro-auth` is recommended going forward because +OpenCode is expected to ship a built-in `kiro` provider that would otherwise +clash with this plugin. + +You only need a `provider` block if you want to override the defaults — for +example a custom model list or your own thinking budgets. + ## Setup 1. **Authentication via Kiro CLI (Recommended)**: - Perform login directly in your terminal using `kiro-cli login`. - - The plugin automatically bootstraps a minimal `kiro` placeholder in + - The plugin automatically bootstraps a minimal `kiro-auth` placeholder in OpenCode's `auth.json` when it detects the Kiro CLI database, then imports and synchronizes your active session on startup. - For AWS IAM Identity Center (SSO/IDC), the plugin imports both the token and device registration (OIDC client credentials) from the `kiro-cli` database. 2. **Direct Authentication**: - Run `opencode auth login`. - - Select `Other`, type `kiro`, and press enter. + - Select `Other`, type `kiro-auth`, and press enter. - You'll be prompted for your **IAM Identity Center Start URL** and **IAM Identity Center region** (`sso_region`). - Leave it blank to sign in with **AWS Builder ID**. @@ -262,7 +132,7 @@ It will replace it with the real email once usage/email lookup succeeds. ### Kiro CLI (Google/GitHub OAuth) users: plugin sync does not start If you authenticated via `kiro-cli login` using Google or GitHub OAuth (not AWS Builder -ID or IAM Identity Center), OpenCode still needs a stored `kiro` auth entry before it +ID or IAM Identity Center), OpenCode still needs a stored `kiro-auth` auth entry before it will call the plugin loader. The plugin now creates that minimal placeholder automatically when it detects the local diff --git a/src/__tests__/accounts.test.ts b/src/__tests__/accounts.test.ts new file mode 100644 index 0000000..625d90a --- /dev/null +++ b/src/__tests__/accounts.test.ts @@ -0,0 +1,458 @@ +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', () => ({ + writeToKiroCli: () => Promise.resolve() +})) +mock.module('../kiro/auth.js', () => ({ + decodeRefreshToken: (t: string) => ({ refreshToken: t }), + encodeRefreshToken: (p: any) => p.refreshToken +})) + +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..3d0ea12 --- /dev/null +++ b/src/__tests__/auth-handler.test.ts @@ -0,0 +1,124 @@ +import { describe, expect, mock, test } from 'bun:test' +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__/edge-cases.test.ts b/src/__tests__/edge-cases.test.ts new file mode 100644 index 0000000..1c52c41 --- /dev/null +++ b/src/__tests__/edge-cases.test.ts @@ -0,0 +1,692 @@ +import { describe, expect, mock, test } from 'bun:test' + +mock.module('../plugin/storage/sqlite.js', () => ({ + kiroDb: { + getConversationId: () => undefined, + setConversationId: () => {}, + deleteConversationId: () => {}, + getAccounts: () => [], + upsertAccount: () => Promise.resolve(), + deleteAccount: () => Promise.resolve(), + batchUpsertAccounts: () => Promise.resolve() + } +})) + +mock.module('../plugin/sync/kiro-cli.js', () => ({ writeToKiroCli: () => Promise.resolve() })) +mock.module('../kiro/auth.js', () => ({ + decodeRefreshToken: (t: string) => ({ refreshToken: t }), + encodeRefreshToken: (p: any) => p.refreshToken +})) + +import { + buildHistory, + collapseAgenticLoops +} from '../infrastructure/transformers/history-builder.js' +import { mergeAdjacentMessages } from '../infrastructure/transformers/message-transformer.js' +import { + buildToolNameMaps, + convertToolsToCodeWhisperer, + deduplicateToolCallsByContent, + shortenToolName +} from '../infrastructure/transformers/tool-transformer.js' +import { transformToSdkRequest } from '../plugin/request.js' +import type { KiroAuthDetails } from '../plugin/types.js' + +const auth: KiroAuthDetails = { + refresh: 'refresh', + access: 'token', + expires: Date.now() + 3600000, + authMethod: 'idc', + region: 'us-east-1', + email: 'test@test.com', + profileArn: 'arn:aws:codewhisperer:us-east-1:123:profile/ABC' +} + +describe('shortenToolName edge cases', () => { + test('null input', () => { + expect(shortenToolName(null as any)).toBeFalsy() + }) + + test('undefined input', () => { + expect(shortenToolName(undefined as any)).toBeFalsy() + }) + + test('empty string', () => { + expect(shortenToolName('')).toBe('') + }) + + test('exactly 64 chars', () => { + const name = 'a'.repeat(64) + expect(shortenToolName(name)).toBe(name) + }) + + test('65 chars', () => { + const name = 'a'.repeat(65) + const result = shortenToolName(name) + expect(result.length).toBeLessThanOrEqual(64) + }) + + test('200 chars', () => { + const name = 'a'.repeat(200) + const result = shortenToolName(name) + expect(result.length).toBeLessThanOrEqual(64) + }) + + test('special chars', () => { + const name = '🎉'.repeat(40) + '\n\t' + 'ñ'.repeat(30) + const result = shortenToolName(name) + expect(result.length).toBeLessThanOrEqual(64) + }) + + test('does not split surrogate pair', () => { + // 🎉 is U+1F389 = 2 UTF-16 code units. 40 of them = 80 UTF-16 chars. + // The naive slice would land between high+low surrogate. + const name = '🎉'.repeat(40) + const result = shortenToolName(name) + expect(result.length).toBeLessThanOrEqual(64) + // No unpaired surrogate left in the prefix + expect(/[\uD800-\uDBFF](?![\uDC00-\uDFFF])/.test(result)).toBe(false) + expect(/(? { + const a = 'tool_a_'.repeat(20) + const b = 'tool_b_'.repeat(20) + expect(shortenToolName(a)).not.toBe(shortenToolName(b)) + }) + + test('same long name is deterministic', () => { + const name = 'x'.repeat(100) + expect(shortenToolName(name)).toBe(shortenToolName(name)) + }) +}) + +describe('sanitizeSchema via convertToolsToCodeWhisperer', () => { + test('deeply nested schema strips additionalProperties', () => { + const tools = [ + { + name: 'deep', + description: 'test', + input_schema: { + type: 'object', + additionalProperties: false, + required: [], + properties: { + nested: { + type: 'object', + additionalProperties: true, + required: [], + properties: { + items: { + type: 'array', + items: { + type: 'object', + additionalProperties: false, + required: [], + properties: { + value: { + anyOf: [ + { type: 'string', additionalProperties: false, required: [] }, + { type: 'number', additionalProperties: true } + ] + } + } + } + } + } + } + } + } + } + ] + const result = convertToolsToCodeWhisperer(tools) + const json = JSON.stringify(result) + expect(json).not.toContain('additionalProperties') + expect(json).not.toContain('"required":[]') + }) + + test('schema with empty required at multiple levels', () => { + const tools = [ + { + name: 'req', + description: 'test', + input_schema: { + type: 'object', + required: [], + properties: { + a: { type: 'object', required: [], properties: { b: { type: 'string', required: [] } } } + } + } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(JSON.stringify(result)).not.toContain('"required":[]') + }) + + test('schema with ONLY additionalProperties and required:[]', () => { + const tools = [ + { + name: 'minimal', + description: 'test', + input_schema: { additionalProperties: false, required: [] } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(result[0].toolSpecification.inputSchema.json).toEqual({}) + }) + + test('null schema', () => { + const tools = [{ name: 'n', description: 'test', input_schema: null }] + expect(() => convertToolsToCodeWhisperer(tools)).not.toThrow() + }) + + test('undefined schema', () => { + const tools = [{ name: 'u', description: 'test' }] + expect(() => convertToolsToCodeWhisperer(tools)).not.toThrow() + }) + + test('strips additionalProperties from "not" subschema', () => { + const tools = [ + { + name: 'n', + description: 'test', + input_schema: { type: 'object', not: { additionalProperties: false, required: [] } } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(JSON.stringify(result)).not.toContain('additionalProperties') + expect(JSON.stringify(result)).not.toContain('"required":[]') + }) + + test('strips additionalProperties from patternProperties', () => { + const tools = [ + { + name: 'p', + description: 'test', + input_schema: { + type: 'object', + patternProperties: { '^x': { additionalProperties: false, required: [] } } + } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(JSON.stringify(result)).not.toContain('additionalProperties') + }) + + test('strips additionalProperties from $defs', () => { + const tools = [ + { + name: 'd', + description: 'test', + input_schema: { + type: 'object', + $defs: { foo: { additionalProperties: false, required: [] } } + } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(JSON.stringify(result)).not.toContain('additionalProperties') + }) + + test('strips additionalProperties from prefixItems', () => { + const tools = [ + { + name: 'pi', + description: 'test', + input_schema: { + type: 'array', + prefixItems: [{ additionalProperties: false }, { additionalProperties: true }] + } + } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(JSON.stringify(result)).not.toContain('additionalProperties') + }) + + test('handles circular schema reference without stack overflow', () => { + const obj: any = { type: 'object', additionalProperties: false, properties: {} } + obj.properties.self = obj + const tools = [{ name: 'c', description: 'test', input_schema: obj }] + expect(() => convertToolsToCodeWhisperer(tools)).not.toThrow() + }) + + test('does not mutate input schema', () => { + const original = { + type: 'object', + additionalProperties: false, + required: [], + properties: { x: { type: 'string', additionalProperties: false } } + } + const before = JSON.stringify(original) + convertToolsToCodeWhisperer([{ name: 't', description: 't', input_schema: original }]) + expect(JSON.stringify(original)).toBe(before) + }) +}) + +describe('payload trim preserves valid structure', () => { + test('large payload trimmed to ≤600KB', () => { + const messages: any[] = [] + for (let i = 0; i < 80; i++) { + messages.push({ role: 'user', content: 'x'.repeat(10000) }) + messages.push({ + role: 'assistant', + content: [ + { type: 'text', text: 'y'.repeat(5000) }, + { type: 'tool_use', id: `tool-${i}`, name: 'myTool', input: { q: 'test' } } + ] + }) + } + messages.push({ + role: 'user', + content: [{ type: 'tool_result', tool_use_id: 'tool-79', content: 'result' }] + }) + messages.push({ role: 'user', content: 'final question' }) + + const body = { + messages, + tools: [{ name: 'myTool', description: 'a tool', input_schema: { type: 'object' } }] + } + const result = transformToSdkRequest(body, 'auto', auth) + const payload = JSON.stringify(result.conversationState) + expect(payload.length).toBeLessThanOrEqual(600_000) + }) + + test('history starts with userInputMessage after trim', () => { + const messages: any[] = [] + for (let i = 0; i < 80; i++) { + messages.push({ role: 'user', content: 'x'.repeat(10000) }) + messages.push({ role: 'assistant', content: 'y'.repeat(10000) }) + } + messages.push({ role: 'user', content: 'final' }) + + const body = { messages } + const result = transformToSdkRequest(body, 'auto', auth) + const history = result.conversationState.history + if (history && history.length > 0) { + expect(history[0]!.userInputMessage).toBeDefined() + } + }) + + test('toolResults have matching toolUseIds in history after trim', () => { + const messages: any[] = [] + for (let i = 0; i < 40; i++) { + messages.push({ role: 'user', content: 'x'.repeat(20000) }) + messages.push({ + role: 'assistant', + content: [ + { type: 'text', text: 'y'.repeat(5000) }, + { type: 'tool_use', id: `t-${i}`, name: 'myTool', input: { q: 'test' } } + ] + }) + messages.push({ + role: 'user', + content: [{ type: 'tool_result', tool_use_id: `t-${i}`, content: 'result '.repeat(1000) }] + }) + } + messages.push({ role: 'user', content: 'final' }) + + const body = { + messages, + tools: [{ name: 'myTool', description: 'd', input_schema: { type: 'object' } }] + } + const result = transformToSdkRequest(body, 'auto', auth) + const history = result.conversationState.history || [] + + const allToolUseIds = new Set() + for (const h of history) { + if (h.assistantResponseMessage?.toolUses) { + for (const tu of h.assistantResponseMessage.toolUses) allToolUseIds.add(tu.toolUseId) + } + } + for (const h of history) { + if (h.userInputMessage?.userInputMessageContext?.toolResults) { + for (const tr of h.userInputMessage.userInputMessageContext.toolResults) { + expect(allToolUseIds.has(tr.toolUseId)).toBe(true) + } + } + } + }) +}) + +describe('mergeAdjacentMessages edge cases', () => { + test('empty array', () => { + expect(mergeAdjacentMessages([])).toEqual([]) + }) + + test('assistant string + assistant array with tool_use', () => { + const msgs = [ + { role: 'assistant', content: 'hello' }, + { role: 'assistant', content: [{ type: 'tool_use', id: '1', name: 'x', input: {} }] } + ] + const result = mergeAdjacentMessages(msgs) + expect(result).toHaveLength(1) + expect(result[0].content).toContainEqual({ type: 'tool_use', id: '1', name: 'x', input: {} }) + }) + + test('assistant with tool_calls + assistant with tool_calls', () => { + const msgs = [ + { + role: 'assistant', + content: 'a', + tool_calls: [{ id: '1', function: { name: 'x', arguments: '{}' } }] + }, + { + role: 'assistant', + content: 'b', + tool_calls: [{ id: '2', function: { name: 'y', arguments: '{}' } }] + } + ] + const result = mergeAdjacentMessages(msgs) + expect(result).toHaveLength(1) + expect(result[0].tool_calls).toHaveLength(2) + }) + + test('user string + user array', () => { + const msgs = [ + { role: 'user', content: 'hello' }, + { role: 'user', content: [{ type: 'text', text: 'world' }] } + ] + const result = mergeAdjacentMessages(msgs) + expect(result).toHaveLength(1) + expect(Array.isArray(result[0].content)).toBe(true) + }) + + test('3 consecutive same-role messages', () => { + const msgs = [ + { role: 'assistant', content: 'a' }, + { role: 'assistant', content: 'b' }, + { role: 'assistant', content: 'c' } + ] + const result = mergeAdjacentMessages(msgs) + expect(result).toHaveLength(1) + expect(result[0].content).toBe('a\nb\nc') + }) +}) + +describe('buildHistory with null/undefined content', () => { + test('assistant with content: null', () => { + const msgs = [ + { role: 'user', content: 'hi' }, + { role: 'assistant', content: null }, + { role: 'user', content: 'bye' } + ] + expect(() => buildHistory(msgs, 'model')).not.toThrow() + }) + + test('assistant with content: undefined', () => { + const msgs = [ + { role: 'user', content: 'hi' }, + { role: 'assistant', content: undefined }, + { role: 'user', content: 'bye' } + ] + expect(() => buildHistory(msgs, 'model')).not.toThrow() + }) + + test('user with content: null', () => { + const msgs = [ + { role: 'user', content: null }, + { role: 'user', content: 'bye' } + ] + expect(() => buildHistory(msgs, 'model')).not.toThrow() + }) + + test('assistant with empty tool_calls array', () => { + const msgs = [ + { role: 'user', content: 'hi' }, + { role: 'assistant', content: 'text', tool_calls: [] }, + { role: 'user', content: 'bye' } + ] + const history = buildHistory(msgs, 'model') + const asst = history.find((h) => h.assistantResponseMessage) + expect(asst?.assistantResponseMessage?.toolUses).toBeUndefined() + }) + + test('assistant with tool_use where name is undefined', () => { + const msgs = [ + { role: 'user', content: 'hi' }, + { role: 'assistant', content: [{ type: 'tool_use', id: '1', name: undefined, input: {} }] }, + { role: 'user', content: 'bye' } + ] + expect(() => buildHistory(msgs, 'model')).not.toThrow() + }) +}) + +describe('collapseAgenticLoops edge cases', () => { + test('single assistant+user pair passes through unchanged', () => { + const history = [ + { + assistantResponseMessage: { + content: 'text', + toolUses: [{ name: 'x', toolUseId: '1', input: {} }] + } + }, + { + userInputMessage: { + content: 'result', + userInputMessageContext: { + toolResults: [{ toolUseId: '1', content: [{ text: 'r' }], status: 'success' }] + } + } + } + ] + const result = collapseAgenticLoops(history as any) + expect(result).toHaveLength(2) + expect(result[0]!.assistantResponseMessage?.content).toBe('text') + }) + + test('assistant with toolUses without following user passes through', () => { + const history = [ + { userInputMessage: { content: 'hi' } }, + { + assistantResponseMessage: { + content: 'text', + toolUses: [{ name: 'x', toolUseId: '1', input: {} }] + } + } + ] + const result = collapseAgenticLoops(history as any) + expect(result).toHaveLength(2) + }) + + test('3 pairs collapse correctly', () => { + const history = [ + { + assistantResponseMessage: { + content: 'first', + toolUses: [{ name: 'a', toolUseId: '1', input: {} }] + } + }, + { + userInputMessage: { + content: 'r1', + userInputMessageContext: { + toolResults: [{ toolUseId: '1', content: [{ text: 'r' }], status: 'success' }] + } + } + }, + { + assistantResponseMessage: { + content: 'second', + toolUses: [{ name: 'b', toolUseId: '2', input: {} }] + } + }, + { + userInputMessage: { + content: 'r2', + userInputMessageContext: { + toolResults: [{ toolUseId: '2', content: [{ text: 'r' }], status: 'success' }] + } + } + }, + { + assistantResponseMessage: { + content: 'third', + toolUses: [{ name: 'c', toolUseId: '3', input: {} }] + } + }, + { + userInputMessage: { + content: 'r3', + userInputMessageContext: { + toolResults: [{ toolUseId: '3', content: [{ text: 'r' }], status: 'success' }] + } + } + } + ] + const result = collapseAgenticLoops(history as any) + expect(result[0]!.assistantResponseMessage?.content).toBe('first') + expect(result[2]!.assistantResponseMessage?.content).toBe('[system: tool calling continues]') + expect(result[4]!.assistantResponseMessage?.content).toBe('[system: tool calling continues]') + }) +}) + +describe('deduplicateToolCallsByContent edge cases', () => { + test('same name different string inputs kept', () => { + const calls = [ + { name: 'tool', input: 'input1' }, + { name: 'tool', input: 'input2' } + ] + expect(deduplicateToolCallsByContent(calls)).toHaveLength(2) + }) + + test('same name same string input deduplicated', () => { + const calls = [ + { name: 'tool', input: 'same' }, + { name: 'tool', input: 'same' } + ] + expect(deduplicateToolCallsByContent(calls)).toHaveLength(1) + }) + + test('empty array', () => { + expect(deduplicateToolCallsByContent([])).toEqual([]) + }) + + test('separator collision: name="a-b" input="c" vs name="a" input="b-c"', () => { + // Old impl used `${name}-${input}` so these would collapse to "a-b-c". + const calls = [ + { name: 'a-b', input: 'c' }, + { name: 'a', input: 'b-c' } + ] + expect(deduplicateToolCallsByContent(calls)).toHaveLength(2) + }) + + test('object input is stringified before compare', () => { + const calls = [ + { name: 'tool', input: { x: 1 } }, + { name: 'tool', input: { x: 1 } }, + { name: 'tool', input: { x: 2 } } + ] + expect(deduplicateToolCallsByContent(calls)).toHaveLength(2) + }) +}) + +describe('convertToolsToCodeWhisperer edge cases', () => { + test('empty tools array', () => { + expect(convertToolsToCodeWhisperer([])).toEqual([]) + }) + + test('tool with undefined name', () => { + const tools = [{ description: 'test', input_schema: { type: 'object' } }] + expect(() => convertToolsToCodeWhisperer(tools)).not.toThrow() + }) + + test('tool with description >9216 chars truncated', () => { + const tools = [ + { name: 'big', description: 'x'.repeat(10000), input_schema: { type: 'object' } } + ] + const result = convertToolsToCodeWhisperer(tools) + expect(result[0].toolSpecification.description.length).toBe(9216) + }) + + test('tool with empty input_schema', () => { + const tools = [{ name: 'empty', description: 'test', input_schema: {} }] + const result = convertToolsToCodeWhisperer(tools) + expect(result[0].toolSpecification.inputSchema.json).toEqual({}) + }) +}) + +describe('buildToolNameMaps edge cases', () => { + test('empty tools array', () => { + const maps = buildToolNameMaps([]) + expect(maps.toKiroName('anything')).toBeDefined() + expect(maps.fromKiroName('anything')).toBe('anything') + }) + + test('two tools with same name', () => { + const tools = [{ name: 'dup' }, { name: 'dup' }] + expect(() => buildToolNameMaps(tools)).not.toThrow() + const maps = buildToolNameMaps(tools) + expect(maps.toKiroName('dup')).toBe('dup') + }) + + test('tool with undefined name skipped', () => { + const tools = [{ name: undefined }, { name: 'valid' }] + const maps = buildToolNameMaps(tools) + expect(maps.toKiroName('valid')).toBe('valid') + }) +}) + +describe('history alternation after trim', () => { + test('after splice first entry is userInputMessage', () => { + const messages: any[] = [] + for (let i = 0; i < 100; i++) { + messages.push({ role: 'user', content: 'u'.repeat(8000) }) + messages.push({ role: 'assistant', content: 'a'.repeat(8000) }) + } + messages.push({ role: 'user', content: 'final' }) + + const body = { messages } + const result = transformToSdkRequest(body, 'auto', auth) + const history = result.conversationState.history + if (history && history.length > 0) { + expect(history[0]!.userInputMessage).toBeDefined() + expect(history[0]!.assistantResponseMessage).toBeUndefined() + } + }) + + test('remaining history alternates after trim', () => { + const messages: any[] = [] + for (let i = 0; i < 100; i++) { + messages.push({ role: 'user', content: 'u'.repeat(8000) }) + messages.push({ role: 'assistant', content: 'a'.repeat(8000) }) + } + messages.push({ role: 'user', content: 'final' }) + + const body = { messages } + const result = transformToSdkRequest(body, 'auto', auth) + const history = result.conversationState.history || [] + for (let i = 0; i < history.length - 1; i++) { + const curr = history[i]! + const next = history[i + 1]! + if (curr.userInputMessage) { + expect(next.assistantResponseMessage).toBeDefined() + } + if (curr.assistantResponseMessage) { + expect(next.userInputMessage).toBeDefined() + } + } + }) +}) + +describe('payload trim performance', () => { + test('trims very large history in <50ms', () => { + // 200 user/asst pairs of 10KB each = ~4MB raw — would be O(N²) without + // the incremental size accounting. + const messages: any[] = [] + for (let i = 0; i < 200; i++) { + messages.push({ role: 'user', content: 'u'.repeat(10000) }) + messages.push({ role: 'assistant', content: 'a'.repeat(10000) }) + } + messages.push({ role: 'user', content: 'final' }) + + const start = performance.now() + const result = transformToSdkRequest({ messages }, 'auto', auth) + const elapsed = performance.now() - start + + // Result must still be under the limit. + const size = JSON.stringify(result.conversationState).length + expect(size).toBeLessThanOrEqual(600_000) + // Allow 100ms on slow CI; fail if we regress to seconds (O(N²)). + expect(elapsed).toBeLessThan(100) + }) +}) diff --git a/src/__tests__/error-handler.test.ts b/src/__tests__/error-handler.test.ts new file mode 100644 index 0000000..f6fd133 --- /dev/null +++ b/src/__tests__/error-handler.test.ts @@ -0,0 +1,227 @@ +import { describe, expect, mock, test } from 'bun:test' +import { ErrorHandler } from '../core/request/error-handler.js' +import { AccountManager } from '../plugin/accounts.js' +import type { ManagedAccount } from '../plugin/types.js' + +mock.module('../plugin/storage/sqlite.js', () => ({ + kiroDb: { + getAccounts: () => [], + upsertAccount: () => Promise.resolve(), + deleteAccount: () => Promise.resolve(), + batchUpsertAccounts: () => Promise.resolve() + } +})) +mock.module('../plugin/sync/kiro-cli.js', () => ({ writeToKiroCli: () => Promise.resolve() })) +mock.module('../kiro/auth.js', () => ({ + decodeRefreshToken: (t: string) => ({ refreshToken: t }), + encodeRefreshToken: (p: any) => p.refreshToken +})) + +const defaultConfig = { rate_limit_max_retries: 3, rate_limit_retry_delay_ms: 100 } +const noToast = () => {} + +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 makeRepo(accounts: ManagedAccount[]) { + return { + findAll: async () => accounts, + batchSave: async () => {}, + save: async () => {}, + invalidateCache: () => {} + } as any +} + +function makeResponse(status: number, body: any, headers: Record = {}): Response { + return new Response(JSON.stringify(body), { + status, + headers: { 'Content-Type': 'application/json', ...headers } + }) +} + +// ── 400 ─────────────────────────────────────────────────────────────────────── + +describe('ErrorHandler: 400', () => { + test('returns shouldRetry=false', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(400, { message: 'Bad Request' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(false) + }) +}) + +// ── 401 ─────────────────────────────────────────────────────────────────────── + +describe('ErrorHandler: 401', () => { + test('retries when under max retries', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(401, { message: 'Unauthorized' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(result.newContext?.retry).toBe(1) + }) + + test('stops retrying at max retries', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(401, { message: 'Unauthorized' }) + const result = await handler.handle(null, res, acc, { retry: 3 }, noToast) + expect(result.shouldRetry).toBe(false) + }) +}) + +// ── 403 single account ──────────────────────────────────────────────────────── + +describe('ErrorHandler: 403 single account', () => { + test('bearer token invalid 403 forces token refresh (sets expiresAt=0) and retries', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(403, { + message: 'The bearer token included in the request is invalid' + }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + // Should retry so the token refresher can get a fresh token + expect(result.shouldRetry).toBe(true) + // Account stays healthy — refresh will handle it + expect(acc.isHealthy).toBe(true) + // expiresAt zeroed so refreshIfNeeded triggers on next iteration + expect(acc.expiresAt).toBe(0) + }) + + test('TEMPORARILY_SUSPENDED marks account unhealthy', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(403, { reason: 'TEMPORARILY_SUSPENDED', message: 'Suspended' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(false) + expect(acc.isHealthy).toBe(false) + }) + + test('non-permanent 403 retries with backoff', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(403, { message: 'Forbidden' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(result.newContext?.retry).toBe(1) + }) + + test('INVALID_MODEL_ID throws immediately', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(403, { reason: 'INVALID_MODEL_ID', message: 'bad model' }) + await expect(handler.handle(null, res, acc, { retry: 0 }, noToast)).rejects.toThrow( + 'Invalid model: bad model' + ) + }) +}) + +// ── 403 multi account ───────────────────────────────────────────────────────── + +describe('ErrorHandler: 403 multi account', () => { + test('switches account on any 403, increments failCount', async () => { + const a = makeAccount({ id: 'a' }) + const b = makeAccount({ id: 'b', email: 'b@example.com' }) + const mgr = new AccountManager([a, b]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([a, b])) + const res = makeResponse(403, { message: 'Forbidden' }) + const result = await handler.handle(null, res, a, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(result.switchAccount).toBe(true) + // non-permanent 403: failCount incremented but account still available + expect(a.failCount).toBe(1) + expect(a.isHealthy).toBe(true) + }) +}) + +// ── 429 ─────────────────────────────────────────────────────────────────────── + +describe('ErrorHandler: 429', () => { + test('marks account rate-limited', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = new Response('', { + status: 429, + headers: { 'retry-after': '1' } // 1s not 30s + }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(acc.rateLimitResetTime).toBeGreaterThan(Date.now() - 100) + }) + + test('with single account, sleep is excluded from request timeout budget', async () => { + const acc = makeAccount() + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = new Response('', { + status: 429, + headers: { 'retry-after': '1' } // 1s sleep + }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(result.newContext?.excludedMs).toBeGreaterThanOrEqual(1000) + }) + + test('with multiple accounts, switches without sleeping', async () => { + const acc1 = makeAccount({ id: 'a' }) + const acc2 = makeAccount({ id: 'b' }) + const mgr = new AccountManager([acc1, acc2]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc1, acc2])) + const res = new Response('', { status: 429, headers: { 'retry-after': '60' } }) + const start = Date.now() + const result = await handler.handle(null, res, acc1, { retry: 0 }, noToast) + const elapsed = Date.now() - start + expect(result.switchAccount).toBe(true) + expect(elapsed).toBeLessThan(500) // didn't sleep + }) +}) + +// ── 500 ─────────────────────────────────────────────────────────────────────── + +describe('ErrorHandler: 500', () => { + test('retries with backoff on first failure', async () => { + const acc = makeAccount({ failCount: 0 }) + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(500, { message: 'Internal Server Error' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.shouldRetry).toBe(true) + expect(acc.failCount).toBe(1) + }) + + test('marks unhealthy after 5 failures', async () => { + const acc = makeAccount({ failCount: 4 }) + const mgr = new AccountManager([acc]) + const handler = new ErrorHandler(defaultConfig, mgr, makeRepo([acc])) + const res = makeResponse(500, { message: 'Internal Server Error' }) + const result = await handler.handle(null, res, acc, { retry: 0 }, noToast) + expect(result.switchAccount).toBe(true) + expect(acc.isHealthy).toBe(false) + }) +}) diff --git a/src/__tests__/event-stream-parser.test.ts b/src/__tests__/event-stream-parser.test.ts new file mode 100644 index 0000000..eabc5ae --- /dev/null +++ b/src/__tests__/event-stream-parser.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, test } from 'bun:test' +import { + parseAwsEventStreamBuffer, + parseEventLine +} from '../infrastructure/transformers/event-stream-parser.js' + +// ── parseEventLine ──────────────────────────────────────────────────────────── + +describe('parseEventLine', () => { + test('parses valid JSON', () => { + expect(parseEventLine('{"content":"hello"}')).toEqual({ content: 'hello' }) + }) + + test('returns null on invalid JSON', () => { + expect(parseEventLine('not json')).toBeNull() + expect(parseEventLine('{')).toBeNull() + }) +}) + +// ── parseAwsEventStreamBuffer ───────────────────────────────────────────────── + +describe('parseAwsEventStreamBuffer', () => { + test('empty buffer returns empty array', () => { + expect(parseAwsEventStreamBuffer('')).toEqual([]) + }) + + test('parses content event', () => { + const result = parseAwsEventStreamBuffer('{"content":"Hello"}') + expect(result).toHaveLength(1) + expect(result[0]!.type).toBe('content') + expect(result[0]!.data).toBe('Hello') + }) + + test('skips followupPrompt as not content', () => { + const result = parseAwsEventStreamBuffer('{"content":"x","followupPrompt":"y"}') + expect(result).toHaveLength(0) + }) + + test('parses toolUse event', () => { + const result = parseAwsEventStreamBuffer( + '{"name":"bash","toolUseId":"t-1","input":"ls","stop":false}' + ) + expect(result).toHaveLength(1) + expect(result[0]!.type).toBe('toolUse') + expect(result[0]!.data.name).toBe('bash') + expect(result[0]!.data.toolUseId).toBe('t-1') + }) + + test('parses toolUseInput event', () => { + const result = parseAwsEventStreamBuffer('{"input":"-la"}') + expect(result).toHaveLength(1) + expect(result[0]!.type).toBe('toolUseInput') + expect(result[0]!.data.input).toBe('-la') + }) + + test('parses toolUseStop event', () => { + const result = parseAwsEventStreamBuffer('{"stop":true}') + expect(result).toHaveLength(1) + expect(result[0]!.type).toBe('toolUseStop') + expect(result[0]!.data.stop).toBe(true) + }) + + test('parses contextUsage event', () => { + const result = parseAwsEventStreamBuffer('{"contextUsagePercentage":42}') + expect(result).toHaveLength(1) + expect(result[0]!.type).toBe('contextUsage') + expect(result[0]!.data.contextUsagePercentage).toBe(42) + }) + + test('parses multiple events from a single buffer', () => { + const buffer = '{"content":"Hello"}\n{"content":" world"}\n{"contextUsagePercentage":25}' + const result = parseAwsEventStreamBuffer(buffer) + expect(result).toHaveLength(3) + expect(result[0]!.data).toBe('Hello') + expect(result[1]!.data).toBe(' world') + expect(result[2]!.type).toBe('contextUsage') + }) + + test('handles incomplete JSON at end (no jsonEnd)', () => { + const result = parseAwsEventStreamBuffer('{"content":"truncated') + expect(result).toHaveLength(0) + }) + + test('handles escaped strings inside JSON values', () => { + const result = parseAwsEventStreamBuffer('{"content":"say \\"hello\\""}') + expect(result).toHaveLength(1) + expect(result[0]!.data).toBe('say "hello"') + }) +}) 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__/model-resolution.test.ts b/src/__tests__/model-resolution.test.ts index 9b54c9d..02f0c5f 100644 --- a/src/__tests__/model-resolution.test.ts +++ b/src/__tests__/model-resolution.test.ts @@ -16,6 +16,11 @@ describe('resolveKiroModel', () => { expect(resolveKiroModel('claude-sonnet-4')).toBe('claude-sonnet-4') }) + test('resolves Claude Opus 4.8 (and its thinking variant)', () => { + expect(resolveKiroModel('claude-opus-4-8')).toBe('claude-opus-4.8') + expect(resolveKiroModel('claude-opus-4-8-thinking')).toBe('claude-opus-4.8') + }) + test('rejects removed qwen3-coder-480b slug', () => { expect(() => resolveKiroModel('qwen3-coder-480b')).toThrow( 'Unsupported model: qwen3-coder-480b' diff --git a/src/__tests__/plugin-module.test.ts b/src/__tests__/plugin-module.test.ts index 8ad0f95..4cab8b6 100644 --- a/src/__tests__/plugin-module.test.ts +++ b/src/__tests__/plugin-module.test.ts @@ -2,7 +2,9 @@ import { describe, expect, test } from 'bun:test' import pluginModule from '../index.js' describe('package plugin module', () => { - test('uses the kiro provider id in the default export', () => { - expect(pluginModule.id).toBe('kiro') + test('uses the kiro-auth provider id in the default export', () => { + // kiro-auth is the primary id (avoids clashing with a future built-in `kiro`); + // `kiro` remains registered as a back-compat alias by the config hook. + expect(pluginModule.id).toBe('kiro-auth') }) }) diff --git a/src/__tests__/response.test.ts b/src/__tests__/response.test.ts new file mode 100644 index 0000000..baca78f --- /dev/null +++ b/src/__tests__/response.test.ts @@ -0,0 +1,69 @@ +import { describe, expect, test } from 'bun:test' +import { estimateTokens, parseEventStream } from '../plugin/response.js' + +// ── estimateTokens ──────────────────────────────────────────────────────────── + +describe('estimateTokens', () => { + test('empty string returns 0', () => { + expect(estimateTokens('')).toBe(0) + }) + + test('4-char string returns 1 token', () => { + expect(estimateTokens('abcd')).toBe(1) + }) + + test('5-char string returns 2 tokens (ceil)', () => { + expect(estimateTokens('abcde')).toBe(2) + }) + + test('100-char string returns 25 tokens', () => { + expect(estimateTokens('a'.repeat(100))).toBe(25) + }) +}) + +// ── parseEventStream ────────────────────────────────────────────────────────── + +describe('parseEventStream', () => { + test('parses plain text content', () => { + const result = parseEventStream('{"content":"Hello world"}') + expect(result.content).toBe('Hello world') + expect(result.toolCalls).toHaveLength(0) + expect(result.stopReason).toBe('end_turn') + }) + + test('parses multiple content chunks', () => { + const result = parseEventStream('{"content":"Hello "}{"content":"world"}') + expect(result.content).toBe('Hello world') + }) + + test('parses tool use', () => { + const raw = [ + '{"name":"bash","toolUseId":"t-1","input":"","stop":false}', + '{"input":"{\\\"command\\\":\\\"ls\\\"}"}', + '{"stop":true}' + ].join('\n') + const result = parseEventStream(raw) + expect(result.toolCalls).toHaveLength(1) + expect(result.toolCalls[0]!.name).toBe('bash') + expect(result.stopReason).toBe('tool_use') + }) + + test('parses context usage into token counts', () => { + const raw = '{"content":"hi"}{"contextUsagePercentage":50}' + const result = parseEventStream(raw, 'CLAUDE_SONNET_4_5') + // With a known context window, inputTokens should be set + expect(result.outputTokens).toBeDefined() + expect(result.inputTokens).toBeDefined() + }) + + test('returns end_turn when no tool calls', () => { + const result = parseEventStream('{"content":"answer"}') + expect(result.stopReason).toBe('end_turn') + }) + + test('empty input returns empty content', () => { + const result = parseEventStream('') + expect(result.content).toBe('') + expect(result.toolCalls).toHaveLength(0) + }) +}) diff --git a/src/__tests__/sqlite.test.ts b/src/__tests__/sqlite.test.ts new file mode 100644 index 0000000..ec06644 --- /dev/null +++ b/src/__tests__/sqlite.test.ts @@ -0,0 +1,246 @@ +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' +import { KiroDatabase } from '../plugin/storage/sqlite.js' +import type { 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 + } +} + +let dir: string +let dbPath: string +let db: KiroDatabase + +beforeEach(() => { + dir = mkdtempSync(join(tmpdir(), 'kiro-test-')) + dbPath = join(dir, 'test.db') + db = new KiroDatabase(dbPath) +}) + +afterEach(() => { + db.close() + rmSync(dir, { recursive: true, force: true }) +}) + +// ── accounts CRUD ───────────────────────────────────────────────────────────── + +describe('KiroDatabase: accounts', () => { + test('starts empty', () => { + expect(db.getAccounts()).toHaveLength(0) + }) + + test('upsertAccount stores and retrieves account', async () => { + const acc = makeAccount() + await db.upsertAccount(acc) + const rows = db.getAccounts() + expect(rows).toHaveLength(1) + expect(rows[0].email).toBe('test@example.com') + expect(rows[0].is_healthy).toBe(1) + }) + + test('upsertAccount updates token fields on existing account', async () => { + const acc = makeAccount() + await db.upsertAccount(acc) + await db.upsertAccount({ ...acc, accessToken: 'new-token' }) + const rows = db.getAccounts() + expect(rows).toHaveLength(1) + expect(rows[0].access_token).toBe('new-token') + }) + + test('upsertAccount with permanent error sets isHealthy=false', async () => { + const acc = makeAccount() + await db.upsertAccount(acc) + await db.upsertAccount({ + ...acc, + isHealthy: false, + unhealthyReason: 'ExpiredTokenException' + }) + const rows = db.getAccounts() + expect(rows).toHaveLength(1) + expect(rows[0].is_healthy).toBe(0) + }) + + test('batchUpsertAccounts stores multiple accounts', async () => { + const a = makeAccount({ id: 'a', email: 'a@example.com' }) + const b = makeAccount({ id: 'b', email: 'b@example.com' }) + await db.batchUpsertAccounts([a, b]) + expect(db.getAccounts()).toHaveLength(2) + }) + + test('deleteAccount removes account', async () => { + const acc = makeAccount() + await db.upsertAccount(acc) + await db.deleteAccount('acc-1') + expect(db.getAccounts()).toHaveLength(0) + }) + + test('deleteAccount on non-existent id is a no-op', async () => { + await expect(db.deleteAccount('does-not-exist')).resolves.toBeUndefined() + }) +}) + +// ── reauth lock ─────────────────────────────────────────────────────────────── + +describe('KiroDatabase: reauth lock', () => { + test('acquireReauthLock returns true when no lock held', () => { + expect(db.acquireReauthLock()).toBe(true) + }) + + test('acquireReauthLock returns false when lock already held by this process', () => { + db.acquireReauthLock() + // Same process — process.kill(pid, 0) succeeds, so it's not dead + expect(db.acquireReauthLock()).toBe(false) + }) + + test('isReauthLockHeld returns false when no lock', () => { + expect(db.isReauthLockHeld()).toBe(false) + }) + + test('isReauthLockHeld returns true after acquire', () => { + db.acquireReauthLock() + expect(db.isReauthLockHeld()).toBe(true) + }) + + test('releaseReauthLock clears the lock', () => { + db.acquireReauthLock() + db.releaseReauthLock() + expect(db.isReauthLockHeld()).toBe(false) + }) + + test('after release, lock can be acquired again', () => { + db.acquireReauthLock() + db.releaseReauthLock() + expect(db.acquireReauthLock()).toBe(true) + }) + + test('stale lock (dead pid) is evicted on acquire', () => { + // Insert a lock row with a pid that no process uses (high number) + const { Database } = require('bun:sqlite') + const rawDb = new Database(dbPath) + rawDb.prepare('INSERT INTO reauth_lock (id, pid, acquired_at) VALUES (1, 9999999, ?)').run( + Date.now() - 1000 // recent but dead pid + ) + rawDb.close() + // Re-open our db instance + db.close() + db = new KiroDatabase(dbPath) + // Should evict dead-pid lock and acquire + expect(db.acquireReauthLock()).toBe(true) + }) + + test('expired lock is evicted on acquire', () => { + const { Database } = require('bun:sqlite') + const rawDb = new Database(dbPath) + rawDb.prepare('INSERT INTO reauth_lock (id, pid, acquired_at) VALUES (1, ?, ?)').run( + process.pid, + Date.now() - 200_000 // 200s ago, well past 120s TTL + ) + rawDb.close() + db.close() + db = new KiroDatabase(dbPath) + expect(db.acquireReauthLock()).toBe(true) + }) + + test('race-safe: row-replacement when prior dead-pid row exists', () => { + // Simulate a row that the SELECT picked up as "expired" but is actually + // the same one INSERT will try to write. INSERT OR REPLACE handles this. + const { Database } = require('bun:sqlite') + const rawDb = new Database(dbPath) + rawDb + .prepare('INSERT INTO reauth_lock (id, pid, acquired_at) VALUES (1, 9999998, ?)') + .run(Date.now() - 200_000) + rawDb.close() + db.close() + db = new KiroDatabase(dbPath) + // Should not throw on PRIMARY KEY conflict + expect(db.acquireReauthLock()).toBe(true) + // The row now belongs to this process + expect(db.isReauthLockHeld()).toBe(true) + }) +}) + +// ── conversations ───────────────────────────────────────────────────────────── + +describe('KiroDatabase: conversations', () => { + test('getConversationId returns undefined when not set', () => { + expect(db.getConversationId('ws', 'fp')).toBeUndefined() + }) + + test('setConversationId and getConversationId round-trip', () => { + db.setConversationId('ws', 'fp', 'conv-123', 'agent-123') + expect(db.getConversationId('ws', 'fp')).toEqual({ + convId: 'conv-123', + agentContinuationId: 'agent-123' + }) + }) + + test('setConversationId updates existing entry', () => { + db.setConversationId('ws', 'fp', 'conv-1', 'agent-1') + db.setConversationId('ws', 'fp', 'conv-2', 'agent-2') + expect(db.getConversationId('ws', 'fp')).toEqual({ + convId: 'conv-2', + agentContinuationId: 'agent-2' + }) + }) + + test('different fingerprints are independent', () => { + db.setConversationId('ws', 'fp1', 'conv-A', 'agent-A') + db.setConversationId('ws', 'fp2', 'conv-B', 'agent-B') + expect(db.getConversationId('ws', 'fp1')).toEqual({ + convId: 'conv-A', + agentContinuationId: 'agent-A' + }) + expect(db.getConversationId('ws', 'fp2')).toEqual({ + convId: 'conv-B', + agentContinuationId: 'agent-B' + }) + }) + + test('TTL cleanup removes old entries', () => { + const { Database } = require('bun:sqlite') + const rawDb = new Database(dbPath) + rawDb + .prepare( + 'INSERT INTO conversations (workspace, fingerprint, conv_id, agent_continuation_id, last_used) VALUES (?, ?, ?, ?, ?)' + ) + .run('ws', 'old', 'conv-old', 'agent-old', Date.now() - 10 * 24 * 3600000) // 10 days old + rawDb.close() + db.close() + db = new KiroDatabase(dbPath) + // Trigger cleanup by setting a new one (ttlDays=7) + db.setConversationId('ws', 'new', 'conv-new', 'agent-new', 7) + expect(db.getConversationId('ws', 'old')).toBeUndefined() + expect(db.getConversationId('ws', 'new')).toEqual({ + convId: 'conv-new', + agentContinuationId: 'agent-new' + }) + }) + + test('deleteConversationId removes entry so next lookup returns undefined', () => { + db.setConversationId('ws', 'fp', 'conv-del', 'agent-del') + expect(db.getConversationId('ws', 'fp')).toBeDefined() + db.deleteConversationId('ws', 'fp') + expect(db.getConversationId('ws', 'fp')).toBeUndefined() + }) + + test('deleteConversationId is a no-op for non-existent entry', () => { + expect(() => db.deleteConversationId('ws', 'missing')).not.toThrow() + }) +}) diff --git a/src/__tests__/stream-transformer.test.ts b/src/__tests__/stream-transformer.test.ts new file mode 100644 index 0000000..75b0110 --- /dev/null +++ b/src/__tests__/stream-transformer.test.ts @@ -0,0 +1,325 @@ +import { describe, expect, test } from 'bun:test' +import { transformKiroStream } from '../plugin/streaming/stream-transformer.js' + +// Helper: create a Response with a ReadableStream from text chunks +function makeStreamResponse(chunks: string[]): Response { + const stream = new ReadableStream({ + start(controller) { + for (const chunk of chunks) { + controller.enqueue(new TextEncoder().encode(chunk)) + } + controller.close() + } + }) + return new Response(stream) +} + +async function collect(gen: AsyncGenerator): Promise { + const result: any[] = [] + for await (const item of gen) result.push(item) + return result +} + +const MODEL = 'CLAUDE_SONNET_4_5' +const CONV = 'conv-test' + +// ── basic text content ──────────────────────────────────────────────────────── + +describe('transformKiroStream: text', () => { + test('plain text content produces text delta events', async () => { + const response = makeStreamResponse(['{"content":"Hello world"}']) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + // Should contain at least a content_block_delta and message_stop + const deltas = events.filter( + (e) => e.choices?.[0]?.delta?.content !== undefined || e.choices?.[0]?.delta?.type === 'text' + ) + expect(deltas.length).toBeGreaterThan(0) + const stop = events.find((e) => e.object === 'chat.completion.chunk') + expect(stop).toBeDefined() + }) + + test('multiple content chunks are concatenated', async () => { + const response = makeStreamResponse(['{"content":"Hello "}', '{"content":"world"}']) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + const allText = events + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('Hello') + expect(allText).toContain('world') + }) + + test('empty body throws', async () => { + const response = new Response(null) + await expect(collect(transformKiroStream(response, MODEL, CONV))).rejects.toThrow( + 'Response body is null' + ) + }) +}) + +// ── thinking tags ───────────────────────────────────────────────────────────── + +describe('transformKiroStream: thinking tags', () => { + test('thinking content produces thinking delta events', async () => { + const response = makeStreamResponse(['{"content":"I think...Answer"}']) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + // Thinking events have type 'thinking' in the delta + const thinkingEvents = events.filter( + (e) => e.choices?.[0]?.delta?.reasoning_content !== undefined + ) + expect(thinkingEvents.length).toBeGreaterThan(0) + }) + + test('text after thinking tags is emitted as normal text', async () => { + const response = makeStreamResponse(['{"content":"ThinkFinal answer"}']) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + const textContent = events + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e) => e.choices[0].delta.content) + .join('') + expect(textContent).toContain('Final answer') + }) + + test('thinking tag inside code block is not treated as thinking', async () => { + const response = makeStreamResponse([ + '{"content":"```\\nnot thinking\\n```"}' + ]) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + const thinkingEvents = events.filter((e) => e.choices?.[0]?.delta?.thinking) + expect(thinkingEvents.length).toBe(0) + }) +}) + +// ── tool use ────────────────────────────────────────────────────────────────── + +describe('transformKiroStream: tool use', () => { + test('toolUse events produce tool_call chunks', async () => { + const chunks = [ + '{"name":"bash","toolUseId":"t-1","input":"","stop":false}', + '{"input":"{\\\"command\\\":\\\"ls\\\"}"}', + '{"stop":true}' + ] + const response = makeStreamResponse(chunks) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + // Should have tool_calls in at least one event + const toolEvents = events.filter( + (e) => e.choices?.[0]?.delta?.tool_calls || e.choices?.[0]?.delta?.type === 'tool_use' + ) + expect(toolEvents.length).toBeGreaterThan(0) + }) +}) + +// ── context usage ───────────────────────────────────────────────────────────── + +describe('transformKiroStream: context usage', () => { + test('contextUsagePercentage is reflected in final usage', async () => { + const response = makeStreamResponse(['{"content":"hi"}', '{"contextUsagePercentage":50}']) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + const delta = events.find((e) => e.choices?.[0]?.delta?.stop_reason || e.usage) + expect(delta).toBeDefined() + }) +}) + +// ── TextDecoder flush ───────────────────────────────────────────────────────── + +describe('transformKiroStream: TextDecoder flush', () => { + test('multi-byte UTF-8 split across chunks is decoded correctly', async () => { + // € = E2 82 AC (3 bytes), split: [E2 82] + [AC ...] + const euro = new TextEncoder().encode('{"content":"€"}') + const part1 = euro.slice(0, euro.length - 5) // cut before '€' completes + const part2 = euro.slice(euro.length - 5) + + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(part1) + controller.enqueue(part2) + controller.close() + } + }) + const response = new Response(stream) + const events = await collect(transformKiroStream(response, MODEL, CONV)) + const allText = events + .filter((e) => e.choices?.[0]?.delta?.content) + .map((e) => e.choices[0].delta.content) + .join('') + expect(allText).toContain('€') + }) +}) + +// ── SDK stream: tool call chunking ─────────────────────────────────────────── + +import { transformSdkStream } from '../plugin/streaming/sdk-stream-transformer.js' + +function makeSdkResponse(events: any[]) { + return { + generateAssistantResponseResponse: (async function* () { + for (const e of events) yield e + })() + } +} + +async function collectSdk(gen: AsyncGenerator): Promise { + const result: any[] = [] + for await (const item of gen) result.push(item) + return result +} + +describe('transformSdkStream: tool call streaming', () => { + test('tool call input chunks without name are concatenated', async () => { + const events = [ + { toolUseEvent: { toolUseId: 't-1', name: 'write', input: '{"file' } }, + { toolUseEvent: { toolUseId: 't-1', input: 'Path":"/tmp/x.txt","content":"hello"}' } }, + { toolUseEvent: { toolUseId: 't-1', stop: true } } + ] + const sdk = makeSdkResponse(events) + const result = await collectSdk(transformSdkStream(sdk, 'auto', 'conv-1')) + const toolBlock = result.find( + (e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.arguments + ) + expect(toolBlock).toBeDefined() + const args = JSON.parse(toolBlock.choices[0].delta.tool_calls[0].function.arguments) + expect(args.filePath).toBe('/tmp/x.txt') + expect(args.content).toBe('hello') + }) + + test('multiple tool calls in one response are parsed correctly', async () => { + const events = [ + { assistantResponseEvent: { content: 'I will run two tools.' } }, + { toolUseEvent: { toolUseId: 't-1', name: 'read', input: '{"path":"/a"}', stop: true } }, + { toolUseEvent: { toolUseId: 't-2', name: 'read', input: '{"path":"/b"}', stop: true } } + ] + const sdk = makeSdkResponse(events) + const result = await collectSdk(transformSdkStream(sdk, 'auto', 'conv-1')) + const toolBlocks = result.filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name) + expect(toolBlocks.length).toBe(2) + }) +}) + +// ── SDK stream: real token usage ───────────────────────────────────────────── + +describe('transformSdkStream: token usage', () => { + test('real tokenUsage from metadata wins over the context% estimate', async () => { + const events = [ + { assistantResponseEvent: { content: 'hello' } }, + { + metadataEvent: { + contextUsagePercentage: 80, + tokenUsage: { inputTokens: 1234, outputTokens: 56 } + } + } + ] + const result = await collectSdk(transformSdkStream(makeSdkResponse(events), 'auto', 'conv-1')) + const usageEvent = result.find((e) => e.usage) + expect(usageEvent.usage.prompt_tokens).toBe(1234) + expect(usageEvent.usage.completion_tokens).toBe(56) + }) +}) + +// ── isNewThread after merge ────────────────────────────────────────────────── + +import { buildHistory } from '../infrastructure/transformers/history-builder.js' +import { mergeAdjacentMessages } from '../infrastructure/transformers/message-transformer.js' + +describe('isNewThread detection after merge', () => { + test('consecutive user messages merge to one — treated as new thread', () => { + const msgs = [ + { role: 'user', content: 'msg1' }, + { role: 'user', content: 'msg2' }, + { role: 'user', content: 'msg3' } + ] + const merged = mergeAdjacentMessages([...msgs]) + expect(merged.length).toBe(1) + const history = buildHistory(merged, 'auto') + expect(history.length).toBe(0) + }) + + test('user+assistant+user produces history', () => { + const msgs = [ + { role: 'user', content: 'hello' }, + { role: 'assistant', content: 'hi' }, + { role: 'user', content: 'next' } + ] + const merged = mergeAdjacentMessages([...msgs]) + expect(merged.length).toBe(3) + const history = buildHistory(merged, 'auto') + expect(history.length).toBe(2) + }) + + test('single user message produces empty history', () => { + const msgs = [{ role: 'user', content: 'only one' }] + const merged = mergeAdjacentMessages([...msgs]) + expect(merged.length).toBe(1) + const history = buildHistory(merged, 'auto') + expect(history.length).toBe(0) + }) + + test('consecutive assistant messages are merged', () => { + const msgs = [ + { role: 'user', content: 'start' }, + { role: 'assistant', content: 'part1' }, + { role: 'assistant', content: 'part2' }, + { role: 'user', content: 'next' } + ] + const merged = mergeAdjacentMessages([...msgs]) + expect(merged.length).toBe(3) + expect(merged[1].content).toContain('part1') + expect(merged[1].content).toContain('part2') + }) +}) + +// ── Tool name consistency (>64 chars must be shortened everywhere) ─────────── + +import { shortenToolName } from '../infrastructure/transformers/tool-transformer.js' + +describe('tool name >64 chars is shortened consistently in history', () => { + const longName = + 'a_very_long_tool_name_that_definitely_exceeds_the_sixty_four_character_limit_imposed_by_kiro' + + test('buildHistory shortens tool_use names', () => { + const msgs = [ + { role: 'user', content: 'do it' }, + { + role: 'assistant', + content: [ + { type: 'text', text: 'calling tool' }, + { type: 'tool_use', id: 'tu-1', name: longName, input: { x: 1 } } + ] + }, + { + role: 'user', + content: [{ type: 'tool_result', tool_use_id: 'tu-1', content: 'done' }] + } + ] + const merged = mergeAdjacentMessages([...msgs]) + const history = buildHistory(merged, 'auto') + const toolUses = history.flatMap((h) => h.assistantResponseMessage?.toolUses || []) + expect(toolUses.length).toBeGreaterThan(0) + for (const tu of toolUses) { + expect(tu.name.length).toBeLessThanOrEqual(64) + expect(tu.name).toBe(shortenToolName(longName)) + } + }) + + test('buildHistory shortens tool_calls names', () => { + const msgs = [ + { role: 'user', content: 'do it' }, + { + role: 'assistant', + content: 'ok', + tool_calls: [{ id: 'tc-1', function: { name: longName, arguments: '{"a":1}' } }] + }, + { + role: 'user', + content: [{ type: 'tool_result', tool_use_id: 'tc-1', content: 'result' }] + } + ] + const merged = mergeAdjacentMessages([...msgs]) + const history = buildHistory(merged, 'auto') + const toolUses = history.flatMap((h) => h.assistantResponseMessage?.toolUses || []) + expect(toolUses.length).toBeGreaterThan(0) + for (const tu of toolUses) { + expect(tu.name.length).toBeLessThanOrEqual(64) + expect(tu.name).toBe(shortenToolName(longName)) + } + }) +}) diff --git a/src/__tests__/tool-call-parser.test.ts b/src/__tests__/tool-call-parser.test.ts new file mode 100644 index 0000000..a22a4a6 --- /dev/null +++ b/src/__tests__/tool-call-parser.test.ts @@ -0,0 +1,88 @@ +import { describe, expect, test } from 'bun:test' +import { + cleanToolCallsFromText, + deduplicateToolCalls, + parseBracketToolCalls +} from '../infrastructure/transformers/tool-call-parser.js' + +describe('parseBracketToolCalls', () => { + test('returns empty array for plain text', () => { + expect(parseBracketToolCalls('Hello world')).toHaveLength(0) + }) + + test('parses single bracket tool call', () => { + const text = '[Called bash with args: {"command":"ls"}]' + const result = parseBracketToolCalls(text) + expect(result).toHaveLength(1) + expect(result[0]!.name).toBe('bash') + expect(result[0]!.input).toEqual({ command: 'ls' }) + }) + + test('parses multiple bracket tool calls', () => { + const text = + '[Called bash with args: {"command":"ls"}] [Called read with args: {"path":"/tmp"}]' + const result = parseBracketToolCalls(text) + expect(result).toHaveLength(2) + expect(result[0]!.name).toBe('bash') + expect(result[1]!.name).toBe('read') + }) + + test('skips malformed JSON args', () => { + const text = '[Called bash with args: {not valid json}]' + const result = parseBracketToolCalls(text) + expect(result).toHaveLength(0) + }) + + test('assigns unique toolUseIds', () => { + const text = + '[Called bash with args: {"command":"ls"}][Called bash with args: {"command":"pwd"}]' + const result = parseBracketToolCalls(text) + expect(result[0]!.toolUseId).not.toBe(result[1]!.toolUseId) + }) +}) + +describe('deduplicateToolCalls', () => { + test('returns empty array for empty input', () => { + expect(deduplicateToolCalls([])).toHaveLength(0) + }) + + test('keeps all unique tool calls', () => { + const calls = [ + { toolUseId: 'a', name: 'bash', input: {} }, + { toolUseId: 'b', name: 'read', input: {} } + ] + expect(deduplicateToolCalls(calls)).toHaveLength(2) + }) + + test('removes duplicates by toolUseId', () => { + const calls = [ + { toolUseId: 'a', name: 'bash', input: {} }, + { toolUseId: 'a', name: 'bash', input: { x: 1 } } + ] + const result = deduplicateToolCalls(calls) + expect(result).toHaveLength(1) + expect(result[0]!.input).toEqual({}) // first one kept + }) +}) + +describe('cleanToolCallsFromText', () => { + test('removes bracket tool call from text', () => { + const text = 'Here is the result [Called bash with args: {"command":"ls"}] done.' + const toolCalls = [{ toolUseId: 'x', name: 'bash', input: {} }] + const result = cleanToolCallsFromText(text, toolCalls) + expect(result).not.toContain('[Called bash') + expect(result).toContain('done.') + }) + + test('leaves text unchanged when no tool calls match', () => { + const text = 'Hello world' + const result = cleanToolCallsFromText(text, []) + expect(result).toBe('Hello world') + }) + + test('trims and collapses whitespace', () => { + const text = ' hello world ' + const result = cleanToolCallsFromText(text, []) + expect(result).toBe('hello world') + }) +}) 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/constants.ts b/src/constants.ts index b816f8d..43bb22f 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -72,6 +72,8 @@ export const MODEL_MAPPING: Record = { 'claude-opus-4-6-1m-thinking': 'claude-opus-4.6-1m', 'claude-opus-4-7': 'claude-opus-4.7', 'claude-opus-4-7-thinking': 'claude-opus-4.7', + 'claude-opus-4-8': 'claude-opus-4.8', + 'claude-opus-4-8-thinking': 'claude-opus-4.8', // Auto auto: 'auto', // Open weight models 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/token-refresher.ts b/src/core/auth/token-refresher.ts index a1b927c..b1b8a06 100644 --- a/src/core/auth/token-refresher.ts +++ b/src/core/auth/token-refresher.ts @@ -81,7 +81,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/core/request/error-handler.ts b/src/core/request/error-handler.ts index c9f968f..5446f0b 100644 --- a/src/core/request/error-handler.ts +++ b/src/core/request/error-handler.ts @@ -1,11 +1,15 @@ import type { AccountRepository } from '../../infrastructure/database/account-repository' import type { AccountManager } from '../../plugin/accounts' +import * as logger from '../../plugin/logger' import type { ManagedAccount } from '../../plugin/types' type ToastFunction = (message: string, variant: 'info' | 'warning' | 'success' | 'error') => void interface RequestContext { retry: number + // Time spent in rate-limit sleeps, propagated up so the retry strategy can + // exclude it from the request timeout budget. + excludedMs?: number } interface ErrorHandlerConfig { @@ -38,12 +42,16 @@ export class ErrorHandler { if (response.status === 400) { const reason = await readBody() + logger.warn(`HTTP 400 on ${account.email}: ${reason || 'unknown'}`) showToast(`400: ${reason || 'unknown'}`, 'error') return { shouldRetry: false } } if (response.status === 401 && context.retry < this.config.rate_limit_max_retries) { const reason = await readBody() + logger.warn( + `HTTP 401 on ${account.email} (retry ${context.retry}): ${reason || 'Unauthorized'}` + ) showToast(`401: ${reason || 'Unauthorized'}. Retrying...`, 'warning') return { shouldRetry: true, @@ -64,15 +72,18 @@ export class ErrorHandler { } } catch (e) {} + logger.warn(`HTTP 500 on ${account.email} (failCount ${account.failCount}): ${errorMessage}`) if (account.failCount < 5) { const delay = 1000 * Math.pow(2, account.failCount - 1) showToast(`500: ${errorMessage}. Retrying in ${Math.ceil(delay / 1000)}s...`, 'warning') await this.sleep(delay) return { shouldRetry: true } } else { + account.failCount = 9 // markUnhealthy will increment to 10 and set isHealthy=false this.accountManager.markUnhealthy( account, - `Server Error (500) after 5 attempts: ${errorMessage}` + `Server error (500) after 5 attempts: ${errorMessage}`, + Date.now() + 1800000 // 30 min recovery ) await this.repository.batchSave(this.accountManager.getAccounts()) showToast(`500: ${errorMessage}. Marking account as unhealthy and switching...`, 'warning') @@ -82,6 +93,7 @@ export class ErrorHandler { if (response.status === 429) { const w = parseInt(response.headers.get('retry-after') || '60') * 1000 + logger.warn(`HTTP 429 on ${account.email}: rate limited, retry-after=${Math.ceil(w / 1000)}s`) this.accountManager.markRateLimited(account, w) await this.repository.batchSave(this.accountManager.getAccounts()) const count = this.accountManager.getAccountCount() @@ -90,7 +102,11 @@ export class ErrorHandler { } showToast(`429: Rate limited. Waiting ${Math.ceil(w / 1000)}s...`, 'warning') await this.sleep(w) - return { shouldRetry: true } + // The wait is not request runtime — exclude it from the timeout budget. + return { + shouldRetry: true, + newContext: { ...context, excludedMs: (context.excludedMs ?? 0) + w } + } } if (response.status === 402 || response.status === 403) { @@ -114,16 +130,18 @@ export class ErrorHandler { errorReason = 'Account Suspended' isPermanent = true } - if ( - errorReason.includes('bearer token included in the request is invalid') || - errorReason.includes('The bearer token included in the request is invalid') - ) { - isPermanent = true - } - if (isPermanent) { - account.failCount = 10 + if (errorReason.includes('bearer token included in the request is invalid')) { + // Force token refresh on next retry + account.expiresAt = 0 + return { shouldRetry: true } } + logger.warn(`HTTP ${response.status} on ${account.email}: ${errorReason}`, { + isPermanent, + retry: context.retry, + reason: errorData?.reason + }) + if (this.accountManager.getAccountCount() > 1) { showToast(`${response.status}: ${errorReason}. Switching account...`, 'warning') this.accountManager.markUnhealthy(account, errorReason) @@ -131,6 +149,12 @@ export class ErrorHandler { return { shouldRetry: true, switchAccount: true } } + if (isPermanent) { + this.accountManager.markUnhealthy(account, errorReason) + await this.repository.batchSave(this.accountManager.getAccounts()) + return { shouldRetry: false } + } + if ( response.status === 403 && !isPermanent && @@ -150,6 +174,7 @@ export class ErrorHandler { } const reason = await readBody() + logger.warn(`HTTP ${response.status} on ${account.email}: ${reason || response.statusText}`) showToast(`${response.status}: ${reason || response.statusText}`, 'error') return { shouldRetry: false } } diff --git a/src/core/request/request-handler.ts b/src/core/request/request-handler.ts index 2f05af7..7009d30 100644 --- a/src/core/request/request-handler.ts +++ b/src/core/request/request-handler.ts @@ -6,6 +6,7 @@ import { isPermanentError } from '../../plugin/health' import * as logger from '../../plugin/logger' import { transformToSdkRequest } from '../../plugin/request' import { createSdkClient } from '../../plugin/sdk-client' +import { kiroDb } from '../../plugin/storage/sqlite' import { syncFromKiroCli } from '../../plugin/sync/kiro-cli' import type { KiroAuthDetails, ManagedAccount, SdkPreparedRequest } from '../../plugin/types' import { AccountSelector } from '../account/account-selector' @@ -19,6 +20,7 @@ type ToastFunction = (message: string, variant: 'info' | 'warning' | 'success' | const KIRO_API_PATTERN = /^(https?:\/\/)?q\.[a-z0-9-]+\.amazonaws\.com/ const REAUTH_FAILURE_COOLDOWN_MS = 60000 +const REAUTH_TIMEOUT_MS = 90_000 export class RequestHandler { private accountSelector: AccountSelector @@ -34,7 +36,8 @@ export class RequestHandler { private accountManager: AccountManager, private config: KiroConfig, private repository: AccountRepository, - private client?: any + private client?: any, + private workspace = '' ) { this.accountSelector = new AccountSelector(accountManager, config, syncFromKiroCli, repository) this.tokenRefresher = new TokenRefresher(config, accountManager, syncFromKiroCli, repository) @@ -66,6 +69,7 @@ export class RequestHandler { let retry = 0 let consecutiveNullAccounts = 0 + let forceNewConversation = false const retryContext = this.retryStrategy.createContext() while (true) { @@ -108,7 +112,13 @@ export class RequestHandler { continue } - const sdkPrep = this.prepareSdkRequest(init?.body, model, auth, think, budget, showToast) + const sdkPrep = this.prepareSdkRequest(body, model, auth, think, budget, showToast) + + const histLen = (sdkPrep.conversationState as any).history?.length || 0 + const agentContId = (sdkPrep.conversationState as any).agentContinuationId || 'none' + logger.debug( + `[REQ] convId=${sdkPrep.conversationId} history=${histLen} agentCont=${agentContId} model=${model}` + ) const apiTimestamp = this.config.enable_log_api_request ? logger.getTimestamp() : null if (apiTimestamp) { @@ -131,13 +141,19 @@ export class RequestHandler { this.handleSuccessfulRequest(acc) this.usageTracker.syncUsage(acc, auth) - return await this.responseHandler.handleSdkSuccess( + const result = await this.responseHandler.handleSdkSuccess( sdkResponse, model, sdkPrep.conversationId, - sdkPrep.streaming + sdkPrep.streaming, + sdkPrep.toolNameMapper ) + logger.debug(`[REQ] done convId=${sdkPrep.conversationId}`) + return result } catch (e: any) { + logger.warn( + `[REQ] error convId=${sdkPrep.conversationId}: ${e?.name || ''} ${e?.message?.slice(0, 200) || String(e).slice(0, 200)}` + ) const httpStatus = e?.$metadata?.httpStatusCode if (httpStatus) { @@ -158,13 +174,15 @@ export class RequestHandler { e, mockResponse, acc, - { retry }, + { retry, excludedMs: retryContext.excludedMs }, showToast ) if (errorResult.shouldRetry) { if (errorResult.newContext) { retry = errorResult.newContext.retry + const sleptMs = (errorResult.newContext.excludedMs ?? 0) - retryContext.excludedMs + if (sleptMs > 0) this.retryStrategy.markSleep(retryContext, sleptMs) } if (errorResult.switchAccount) { continue @@ -172,6 +190,21 @@ export class RequestHandler { continue } + if (httpStatus === 400 && e?.name === 'ValidationException' && !forceNewConversation) { + const { workspace, fingerprint } = sdkPrep.conversationKey + kiroDb.deleteConversationId(workspace, fingerprint) + logger.warn( + `[REQ] stale conversationId reset, retrying convId=${sdkPrep.conversationId}` + ) + forceNewConversation = true + continue + } + + if (this.allAccountsPermanentlyUnhealthy()) { + const reauthed = await this.triggerReauth(showToast) + if (reauthed) continue + } + throw new Error(`Kiro Error: ${httpStatus}`) } @@ -201,7 +234,7 @@ export class RequestHandler { budget: number, showToast?: (message: string, variant: 'info' | 'warning' | 'success' | 'error') => void ): SdkPreparedRequest { - return transformToSdkRequest(body, model, auth, think, budget, showToast) + return transformToSdkRequest(body, model, auth, think, budget, showToast, this.workspace) } private handleSuccessfulRequest(acc: ManagedAccount): void { @@ -302,9 +335,26 @@ export class RequestHandler { return this.reauthInFlight } + if (!kiroDb.acquireReauthLock()) { + logger.warn('Reauth lock held by another instance — polling for completion') + showToast('Another session is re-authenticating. Please wait...', 'info') + const deadline = Date.now() + 10_000 + while (Date.now() < deadline) { + await this.sleep(1000) + if (kiroDb.isReauthLockHeld()) continue + this.repository.invalidateCache() + const accounts = await this.repository.findAll() + for (const acc of accounts) this.accountManager.addAccount(acc) + return this.hasUsableAccount(accounts) + } + showToast('Re-authentication timed out. Please try again.', 'error') + return false + } + this.reauthInFlight = this.performReauth(showToast) const success = await this.reauthInFlight.finally(() => { this.reauthInFlight = null + kiroDb.releaseReauthLock() }) if (!success) this.lastFailedReauthAt = Date.now() return success @@ -313,15 +363,31 @@ export class RequestHandler { private async performReauth(showToast: ToastFunction): Promise { try { showToast('Session expired. Re-authenticating...', 'warning') - await this.client.provider.oauth.authorize({ - path: { id: 'kiro' }, - body: { method: 0 } - }) + logger.warn('Reauth: starting oauth flow') + + const withTimeout = (promise: Promise, label: string): Promise => { + let timer: ReturnType | undefined + return Promise.race([ + promise.finally(() => clearTimeout(timer)), + new Promise( + (_, reject) => + (timer = setTimeout( + () => reject(new Error(`Reauth timed out waiting for ${label}`)), + REAUTH_TIMEOUT_MS + )) + ) + ]) + } - await this.client.provider.oauth.callback({ - path: { id: 'kiro' }, - body: { method: 0 } - }) + await withTimeout( + this.client.provider.oauth.authorize({ path: { id: 'kiro' }, body: { method: 0 } }), + 'oauth.authorize' + ) + + await withTimeout( + this.client.provider.oauth.callback({ path: { id: 'kiro' }, body: { method: 0 } }), + 'oauth.callback' + ) this.repository.invalidateCache() const accounts = await this.repository.findAll() @@ -339,6 +405,12 @@ export class RequestHandler { return true } catch (e) { logger.error('Re-auth failed', e instanceof Error ? e : new Error(String(e))) + showToast( + e instanceof Error && e.message.includes('timed out') + ? 'Re-authentication timed out. Please try again.' + : 'Re-authentication failed. Please try again.', + 'error' + ) return false } } diff --git a/src/core/request/response-handler.ts b/src/core/request/response-handler.ts index e95c78f..3036e39 100644 --- a/src/core/request/response-handler.ts +++ b/src/core/request/response-handler.ts @@ -19,10 +19,11 @@ export class ResponseHandler { sdkResponse: any, model: string, conversationId: string, - streaming: boolean + streaming: boolean, + toolNameMapper?: (name: string) => string ): Promise { if (streaming) { - return this.handleSdkStreaming(sdkResponse, model, conversationId) + return this.handleSdkStreaming(sdkResponse, model, conversationId, toolNameMapper) } return this.handleSdkNonStreaming(sdkResponse, model, conversationId) } @@ -33,12 +34,13 @@ export class ResponseHandler { conversationId: string ): Promise { const s = transformKiroStream(response, model, conversationId) + const enc = new TextEncoder() return new Response( new ReadableStream({ async start(c) { try { for await (const e of s) { - c.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(e)}\n\n`)) + c.enqueue(enc.encode(`data: ${JSON.stringify(e)}\n\n`)) } c.close() } catch (err) { @@ -53,15 +55,17 @@ export class ResponseHandler { private async handleSdkStreaming( sdkResponse: any, model: string, - conversationId: string + conversationId: string, + toolNameMapper?: (name: string) => string ): Promise { - const s = transformSdkStream(sdkResponse, model, conversationId) + const s = transformSdkStream(sdkResponse, model, conversationId, toolNameMapper) + const enc = new TextEncoder() return new Response( new ReadableStream({ async start(c) { try { for await (const e of s) { - c.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(e)}\n\n`)) + c.enqueue(enc.encode(`data: ${JSON.stringify(e)}\n\n`)) } c.close() } catch (err) { diff --git a/src/core/request/retry-strategy.ts b/src/core/request/retry-strategy.ts index 4acf5db..b72616e 100644 --- a/src/core/request/retry-strategy.ts +++ b/src/core/request/retry-strategy.ts @@ -6,6 +6,8 @@ interface RetryConfig { interface RetryContext { iterations: number startTime: number + // Time spent in rate-limit waits, excluded from the request timeout budget. + excludedMs: number } export class RetryStrategy { @@ -21,7 +23,8 @@ export class RetryStrategy { } } - if (Date.now() - context.startTime > this.config.request_timeout_ms) { + const elapsed = Date.now() - context.startTime - context.excludedMs + if (elapsed > this.config.request_timeout_ms) { return { canContinue: false, error: 'Request timeout' @@ -31,10 +34,15 @@ export class RetryStrategy { return { canContinue: true } } + markSleep(context: RetryContext, ms: number): void { + context.excludedMs += Math.max(0, ms) + } + createContext(): RetryContext { return { iterations: 0, - startTime: Date.now() + startTime: Date.now(), + excludedMs: 0 } } } diff --git a/src/index.ts b/src/index.ts index 9174053..ffa4183 100644 --- a/src/index.ts +++ b/src/index.ts @@ -3,4 +3,4 @@ export { KiroOAuthPlugin } from './plugin.js' export type { KiroConfig } from './plugin/config/index.js' export type { KiroAuthMethod, KiroRegion, ManagedAccount } from './plugin/types.js' -export default { id: 'kiro', server: (await import('./plugin.js')).KiroOAuthPlugin } +export default { id: 'kiro-auth', server: (await import('./plugin.js')).KiroOAuthPlugin } diff --git a/src/infrastructure/transformers/history-builder.ts b/src/infrastructure/transformers/history-builder.ts index 951b914..0456898 100644 --- a/src/infrastructure/transformers/history-builder.ts +++ b/src/infrastructure/transformers/history-builder.ts @@ -6,7 +6,7 @@ import { } from '../../plugin/image-handler.js' import type { CodeWhispererMessage } from '../../plugin/types' import { getContentText } from './message-transformer.js' -import { deduplicateToolResults } from './tool-transformer.js' +import { deduplicateToolResults, shortenToolName } from './tool-transformer.js' /** * Collapse agentic loop sequences in the built history. @@ -154,7 +154,7 @@ export function buildHistory(msgs: any[], resolved: string): CodeWhispererMessag if (p.type === 'text') arm.content += p.text || '' else if (p.type === 'thinking') th += p.thinking || p.text || '' else if (p.type === 'tool_use') - tus.push({ input: p.input, name: p.name, toolUseId: p.id }) + tus.push({ input: p.input, name: shortenToolName(p.name), toolUseId: p.id }) } } else arm.content = getContentText(m) if (m.tool_calls && Array.isArray(m.tool_calls)) { @@ -164,7 +164,7 @@ export function buildHistory(msgs: any[], resolved: string): CodeWhispererMessag typeof tc.function?.arguments === 'string' ? JSON.parse(tc.function.arguments) : tc.function?.arguments, - name: tc.function?.name, + name: shortenToolName(tc.function?.name), toolUseId: tc.id }) } diff --git a/src/infrastructure/transformers/tool-transformer.ts b/src/infrastructure/transformers/tool-transformer.ts index ef785d9..aa60b92 100644 --- a/src/infrastructure/transformers/tool-transformer.ts +++ b/src/infrastructure/transformers/tool-transformer.ts @@ -1,9 +1,103 @@ +import * as crypto from 'crypto' + +const MAX_TOOL_NAME_LENGTH = 64 + +// Slice without breaking surrogate pairs. +function safeSlice(s: string, end: number): string { + if (end <= 0) return '' + if (end >= s.length) return s + const code = s.charCodeAt(end - 1) + if (code >= 0xd800 && code <= 0xdbff) end -= 1 + return s.slice(0, end) +} + +export function shortenToolName(name: string): string { + if (!name || name.length <= MAX_TOOL_NAME_LENGTH) return name + const hash = crypto.createHash('sha256').update(name).digest('hex').slice(0, 12) + const prefix = safeSlice(name, MAX_TOOL_NAME_LENGTH - hash.length - 1) + return `${prefix}_${hash}` +} + +export function buildToolNameMaps(tools: any[]): { + toKiroName: (name: string) => string + fromKiroName: (name: string) => string +} { + const originalToAlias = new Map() + const aliasToOriginal = new Map() + + for (const t of tools) { + const name = t.name || t.function?.name + if (!name) continue + const alias = shortenToolName(name) + originalToAlias.set(name, alias) + if (alias !== name) aliasToOriginal.set(alias, name) + } + + return { + toKiroName: (name: string) => originalToAlias.get(name) || shortenToolName(name), + fromKiroName: (name: string) => aliasToOriginal.get(name) || name + } +} + +function sanitizeToolInput(input: any): any { + if (!input || typeof input !== 'object' || Array.isArray(input)) return input + const result: Record = {} + for (const [key, value] of Object.entries(input)) { + if (key === '') continue + result[key] = value + } + return result +} + +function sanitizeSchema(schema: any, seen: WeakSet = new WeakSet()): any { + if (!schema || typeof schema !== 'object' || Array.isArray(schema)) return schema + if (seen.has(schema)) return {} + seen.add(schema) + + const result: Record = {} + for (const [key, value] of Object.entries(schema)) { + if (key === 'additionalProperties') continue + if (key === 'required' && Array.isArray(value) && value.length === 0) continue + + if ( + (key === 'properties' || + key === 'patternProperties' || + key === '$defs' || + key === 'definitions') && + typeof value === 'object' && + value !== null && + !Array.isArray(value) + ) { + const props: Record = {} + for (const [pk, pv] of Object.entries(value)) { + props[pk] = sanitizeSchema(pv, seen) + } + result[key] = props + } else if ( + (key === 'anyOf' || key === 'oneOf' || key === 'allOf' || key === 'prefixItems') && + Array.isArray(value) + ) { + result[key] = value.map((v) => sanitizeSchema(v, seen)) + } else if ( + (key === 'items' || key === 'not' || key === 'contains') && + typeof value === 'object' + ) { + result[key] = sanitizeSchema(value, seen) + } else { + result[key] = value + } + } + return result +} + export function convertToolsToCodeWhisperer(tools: any[]): any[] { return tools.map((t) => ({ toolSpecification: { - name: t.name || t.function?.name, + name: shortenToolName(t.name || t.function?.name || ''), description: (t.description || t.function?.description || '').substring(0, 9216), - inputSchema: { json: t.input_schema || t.function?.parameters || {} } + inputSchema: { + json: sanitizeSchema(sanitizeToolInput(t.input_schema || t.function?.parameters || {})) + } } })) } @@ -19,3 +113,19 @@ export function deduplicateToolResults(trs: any[]): any[] { } return u } + +export function deduplicateToolCallsByContent(toolCalls: any[]): any[] { + const seen = new Set() + const unique: any[] = [] + for (const tc of toolCalls) { + // \x00 as separator (can't appear in a tool name) + const name = tc.name || tc.function?.name || '' + const input = tc.input || tc.function?.arguments || '' + const key = `${name}\x00${typeof input === 'string' ? input : JSON.stringify(input)}` + if (!seen.has(key)) { + seen.add(key) + unique.push(tc) + } + } + return unique +} diff --git a/src/plugin.ts b/src/plugin.ts index f44e8b5..541ba82 100644 --- a/src/plugin.ts +++ b/src/plugin.ts @@ -9,7 +9,89 @@ import { loadConfig } from './plugin/config/index.js' type ToastFunction = (message: string, variant: string) => void -const KIRO_PROVIDER_ID = 'kiro' +// `kiro-auth` is the recommended provider id. OpenCode is expected to ship a +// built-in `kiro` provider, which would clash with our default — so new installs +// use `kiro-auth`. We still register `kiro` as a back-compat alias so existing +// installs configured against `kiro` keep working. +const KIRO_PROVIDER_ID = 'kiro-auth' +const KIRO_LEGACY_PROVIDER_ID = 'kiro' + +const DEFAULT_MODELS: Record = { + auto: { + name: 'Auto (1.0x)', + limit: { context: 200000, output: 64000 }, + modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } + }, + // Claude Sonnet + 'claude-sonnet-4': { + name: 'Claude Sonnet 4.0 (1.3x)', + limit: { context: 200000, output: 64000 }, + modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } + }, + 'claude-sonnet-4-5': { + name: 'Claude Sonnet 4.5 (1.3x)', + limit: { context: 200000, output: 64000 }, + modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } + }, + 'claude-sonnet-4-6': { + name: 'Claude Sonnet 4.6 (1.3x)', + limit: { context: 1000000, output: 64000 }, + modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } + }, + // Claude Haiku + 'claude-haiku-4-5': { + name: 'Claude Haiku 4.5 (0.4x)', + limit: { context: 200000, output: 64000 }, + modalities: { input: ['text', 'image'], output: ['text'] } + }, + // Claude Opus + 'claude-opus-4-5': { + name: 'Claude Opus 4.5 (2.2x)', + limit: { context: 200000, output: 64000 }, + modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } + }, + 'claude-opus-4-6': { + name: 'Claude Opus 4.6 (2.2x)', + limit: { context: 1000000, output: 64000 }, + modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } + }, + 'claude-opus-4-7': { + name: 'Claude Opus 4.7 (2.2x)', + limit: { context: 1000000, output: 64000 }, + modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } + }, + 'claude-opus-4-8': { + name: 'Claude Opus 4.8 (2.2x)', + limit: { context: 1000000, output: 64000 }, + modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } + }, + // Open weight models + 'deepseek-3.2': { + name: 'DeepSeek 3.2 (0.25x)', + limit: { context: 128000, output: 64000 }, + modalities: { input: ['text'], output: ['text'] } + }, + 'glm-5': { + name: 'GLM-5 (0.5x)', + limit: { context: 200000, output: 64000 }, + modalities: { input: ['text'], output: ['text'] } + }, + 'minimax-m2.5': { + name: 'MiniMax M2.5 (0.25x)', + limit: { context: 200000, output: 64000 }, + modalities: { input: ['text'], output: ['text'] } + }, + 'minimax-m2.1': { + name: 'MiniMax M2.1 (0.15x)', + limit: { context: 200000, output: 64000 }, + modalities: { input: ['text'], output: ['text'] } + }, + 'qwen3-coder-next': { + name: 'Qwen3 Coder Next (0.05x)', + limit: { context: 256000, output: 64000 }, + modalities: { input: ['text'], output: ['text'] } + } +} export const createKiroPlugin = (id: string) => @@ -27,7 +109,7 @@ export const createKiroPlugin = const accountManager = await AccountManager.loadFromDisk(config.account_selection_strategy) authHandler.setAccountManager(accountManager) - const requestHandler = new RequestHandler(accountManager, config, repository, client) + const requestHandler = new RequestHandler(accountManager, config, repository, client, directory) // Compute the base URL once so both the config hook and auth loader use the same value const baseURL = KIRO_CONSTANTS.BASE_URL.replace('/generateAssistantResponse', '').replace( @@ -35,6 +117,24 @@ export const createKiroPlugin = config.default_region || 'us-east-1' ) + // The custom fetch self-identifies Kiro requests by URL, so a single instance + // serves any provider id. OpenCode binds auth.loader to one id only, so we + // attach this fetch via provider.options in the config hook — resolveSDK reads + // options.fetch per provider, which is how both `kiro-auth` and `kiro` route + // through us. + const kiroFetch = (input: any, init?: any) => requestHandler.handle(input, init, showToast) + + const registerProvider = (input: any, providerId: string) => { + if (!input.provider[providerId]) input.provider[providerId] = {} + const p = input.provider[providerId] + p.npm = '@ai-sdk/openai-compatible' + // OpenCode resolves model.api.url / model.api.npm from these provider-level + // fields, so the models don't need per-model api entries. + if (!p.api) p.api = baseURL + p.options = { ...(p.options ?? {}), fetch: kiroFetch } + if (!p.models) p.models = { ...DEFAULT_MODELS } + } + return { config: async (input: any) => { // Ensure there's an auth entry so OpenCode calls the loader on startup. @@ -42,89 +142,10 @@ export const createKiroPlugin = bootstrapAuthIfNeeded(id) if (!input.provider) input.provider = {} - if (!input.provider[id]) input.provider[id] = {} - // Always set npm and api — these must be present regardless of whether - // the user has already defined the provider in their opencode.json. - input.provider[id].npm = '@ai-sdk/openai-compatible' - // Set the base URL at the provider level. OpenCode reads provider.api as - // model.api.url, which resolveSDK() uses to construct the endpoint URL. - // Only set if not already overridden by the user. - if (!input.provider[id].api) { - input.provider[id].api = baseURL - } - if (!input.provider[id].models) { - input.provider[id].models = { - auto: { - name: 'Auto (1.0x)', - limit: { context: 200000, output: 64000 }, - modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } - }, - // Claude Sonnet - 'claude-sonnet-4': { - name: 'Claude Sonnet 4.0 (1.3x)', - limit: { context: 200000, output: 64000 }, - modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } - }, - 'claude-sonnet-4-5': { - name: 'Claude Sonnet 4.5 (1.3x)', - limit: { context: 200000, output: 64000 }, - modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } - }, - 'claude-sonnet-4-6': { - name: 'Claude Sonnet 4.6 (1.3x)', - limit: { context: 1000000, output: 64000 }, - modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } - }, - // Claude Haiku - 'claude-haiku-4-5': { - name: 'Claude Haiku 4.5 (0.4x)', - limit: { context: 200000, output: 64000 }, - modalities: { input: ['text', 'image'], output: ['text'] } - }, - // Claude Opus - 'claude-opus-4-5': { - name: 'Claude Opus 4.5 (2.2x)', - limit: { context: 200000, output: 64000 }, - modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } - }, - 'claude-opus-4-6': { - name: 'Claude Opus 4.6 (2.2x)', - limit: { context: 1000000, output: 64000 }, - modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } - }, - 'claude-opus-4-7': { - name: 'Claude Opus 4.7 (2.2x)', - limit: { context: 1000000, output: 64000 }, - modalities: { input: ['text', 'image', 'pdf'], output: ['text'] } - }, - // Open weight models - 'deepseek-3.2': { - name: 'DeepSeek 3.2 (0.25x)', - limit: { context: 128000, output: 64000 }, - modalities: { input: ['text'], output: ['text'] } - }, - 'glm-5': { - name: 'GLM-5 (0.5x)', - limit: { context: 200000, output: 64000 }, - modalities: { input: ['text'], output: ['text'] } - }, - 'minimax-m2.5': { - name: 'MiniMax M2.5 (0.25x)', - limit: { context: 200000, output: 64000 }, - modalities: { input: ['text'], output: ['text'] } - }, - 'minimax-m2.1': { - name: 'MiniMax M2.1 (0.15x)', - limit: { context: 200000, output: 64000 }, - modalities: { input: ['text'], output: ['text'] } - }, - 'qwen3-coder-next': { - name: 'Qwen3 Coder Next (0.05x)', - limit: { context: 256000, output: 64000 }, - modalities: { input: ['text'], output: ['text'] } - } - } - } + // Primary id (kiro-auth) plus the back-compat alias (kiro), both sharing + // the same custom fetch. + registerProvider(input, id) + registerProvider(input, KIRO_LEGACY_PROVIDER_ID) }, auth: { provider: id, @@ -135,10 +156,10 @@ export const createKiroPlugin = return { apiKey: '', // Provide baseURL explicitly so the @ai-sdk/openai-compatible provider - // always has a valid URL. The custom fetch below intercepts all Kiro - // API calls, so this value is only used for URL construction. + // always has a valid URL. The custom fetch intercepts all Kiro API + // calls, so this value is only used for URL construction. baseURL, - fetch: (input: any, init?: any) => requestHandler.handle(input, init, showToast) + fetch: kiroFetch } }, methods: authHandler.getMethods() diff --git a/src/plugin/accounts.ts b/src/plugin/accounts.ts index d0a623e..819fb4e 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 { @@ -102,8 +105,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 +122,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/logger.ts b/src/plugin/logger.ts index b0f85fe..1938668 100644 --- a/src/plugin/logger.ts +++ b/src/plugin/logger.ts @@ -72,7 +72,7 @@ export function warn(message: string, ...args: unknown[]): void { } export function debug(message: string, ...args: unknown[]): void { - if (process.env.DEBUG) { + if (process.env.DEBUG || process.env.OPENCODE_LOG_LEVEL === 'debug') { writeToFile('DEBUG', message, ...args) } } diff --git a/src/plugin/request.ts b/src/plugin/request.ts index 71542f7..ad7b548 100644 --- a/src/plugin/request.ts +++ b/src/plugin/request.ts @@ -13,8 +13,10 @@ import { mergeAdjacentMessages } from '../infrastructure/transformers/message-transformer.js' import { + buildToolNameMaps, convertToolsToCodeWhisperer, - deduplicateToolResults + deduplicateToolResults, + shortenToolName } from '../infrastructure/transformers/tool-transformer.js' import { convertImagesToKiroFormat, @@ -22,6 +24,7 @@ import { extractTextFromParts } from './image-handler.js' import { resolveKiroModel } from './models.js' +import { kiroDb } from './storage/sqlite.js' import type { CodeWhispererRequest, KiroAuthDetails, @@ -29,10 +32,47 @@ import type { SdkPreparedRequest } from './types' +// Stable conversationId + agentContinuationId per thread, persisted in SQLite. +function deriveConversationIds( + workspace: string, + firstUserContent: string, + isNewThread: boolean +): { convId: string; agentContinuationId: string; fingerprint: string } { + const fingerprint = crypto + .createHash('sha256') + .update(workspace + '\0' + (firstUserContent || '_empty_')) + .digest('hex') + .slice(0, 32) + + if (!isNewThread) { + const existing = kiroDb.getConversationId(workspace, fingerprint) + if (existing) { + if (!existing.agentContinuationId) { + existing.agentContinuationId = crypto.randomUUID() + kiroDb.setConversationId( + workspace, + fingerprint, + existing.convId, + existing.agentContinuationId + ) + } + return { ...existing, fingerprint } + } + } + + const convId = crypto.randomUUID() + const agentContinuationId = crypto.randomUUID() + kiroDb.setConversationId(workspace, fingerprint, convId, agentContinuationId) + return { convId, agentContinuationId, fingerprint } +} + interface TransformResult { request: CodeWhispererRequest resolved: string convId: string + agentContinuationId: string + fingerprint: string + toolNameMapper?: (name: string) => string } type ToastFunction = (message: string, variant: 'info' | 'warning' | 'success' | 'error') => void @@ -43,13 +83,13 @@ function buildCodeWhispererRequest( auth: KiroAuthDetails, think = false, budget = 20000, - showToast?: ToastFunction + showToast?: ToastFunction, + workspace = '' ): TransformResult { const req = typeof body === 'string' ? JSON.parse(body) : body const { messages, tools, system } = req - const convId = crypto.randomUUID() if (!messages || messages.length === 0) throw new Error('No messages') - const resolved = resolveKiroModel(model) + const systemMsgs = messages.filter((m: any) => m.role === 'system') const otherMsgs = messages.filter((m: any) => m.role !== 'system') let sys = system || '' @@ -64,6 +104,22 @@ function buildCodeWhispererRequest( const msgs = mergeAdjacentMessages([...otherMsgs]) const lastMsg = msgs[msgs.length - 1] if (lastMsg && lastMsg.role === 'assistant' && getContentText(lastMsg) === '{') msgs.pop() + + // isNewThread after merge — consecutive same-role messages collapse into one + const isNewThread = msgs.length <= 1 + const firstUserMsg = msgs.find((m: any) => m.role === 'user') + const firstUserContent = firstUserMsg + ? typeof firstUserMsg.content === 'string' + ? firstUserMsg.content + : JSON.stringify(firstUserMsg.content) + : '' + const { convId, agentContinuationId, fingerprint } = deriveConversationIds( + workspace, + firstUserContent, + isNewThread + ) + const resolved = resolveKiroModel(model) + const toolMaps = tools ? buildToolNameMaps(tools) : undefined const cwTools = tools ? convertToolsToCodeWhisperer(tools) : [] let history = buildHistory(msgs, resolved) @@ -107,7 +163,7 @@ function buildCodeWhispererRequest( else if (p.type === 'thinking') th += p.thinking || p.text || '' else if (p.type === 'tool_use') { if (!arm.toolUses) arm.toolUses = [] - arm.toolUses.push({ input: p.input, name: p.name, toolUseId: p.id }) + arm.toolUses.push({ input: p.input, name: shortenToolName(p.name), toolUseId: p.id }) } } } else arm.content = getContentText(curMsg) @@ -119,7 +175,7 @@ function buildCodeWhispererRequest( typeof tc.function?.arguments === 'string' ? JSON.parse(tc.function.arguments) : tc.function?.arguments, - name: tc.function?.name, + name: shortenToolName(tc.function?.name), toolUseId: tc.id }) } @@ -179,6 +235,8 @@ function buildCodeWhispererRequest( } const request: CodeWhispererRequest = { conversationState: { + agentContinuationId, + agentTaskType: 'vibe', chatTriggerType: KIRO_CONSTANTS.CHAT_TRIGGER_TYPE_MANUAL, conversationId: convId, currentMessage: { @@ -202,7 +260,7 @@ function buildCodeWhispererRequest( if (originalCall) { orphanedTrs.push({ call: { - name: originalCall.name || originalCall.function?.name || 'tool', + name: shortenToolName(originalCall.name || originalCall.function?.name || 'tool'), toolUseId: tr.toolUseId, input: originalCall.input || @@ -270,7 +328,60 @@ function buildCodeWhispererRequest( } } - return { request, resolved, convId } + // Strip empty toolUses arrays from history (Kiro quirk) + for (const h of history) { + if (h.assistantResponseMessage?.toolUses && h.assistantResponseMessage.toolUses.length === 0) { + delete h.assistantResponseMessage.toolUses + } + } + + // Trim history if payload exceeds Kiro's ~615KB limit. + // Compute per-entry sizes once and shrink incrementally to avoid O(N²) + // re-stringifying the full request on every iteration. + const MAX_PAYLOAD_BYTES = 600_000 + if (history.length > 2) { + const sizes = history.map((h) => JSON.stringify(h).length + 1) + const baseRequest: any = { ...request, conversationState: { ...request.conversationState } } + delete baseRequest.conversationState.history + let totalSize = JSON.stringify(baseRequest).length + 2 // for `"history":[]` + for (const s of sizes) totalSize += s + + while (history.length > 2 && totalSize > MAX_PAYLOAD_BYTES) { + // Drop the two oldest entries (typically user + assistant pair). + totalSize -= (sizes.shift() || 0) + (sizes.shift() || 0) + history.splice(0, 2) + + // Strip leading orphans: assistantResponseMessage can't start the history. + while (history.length > 0 && history[0]?.assistantResponseMessage) { + totalSize -= sizes.shift() || 0 + history.shift() + } + + // Strip leading toolResult-only userInputMessages whose toolUseIds are gone. + const toolUseIds = new Set( + history.flatMap( + (h) => h.assistantResponseMessage?.toolUses?.map((tu: any) => tu.toolUseId) ?? [] + ) + ) + while (history.length > 0) { + const trs = history[0]?.userInputMessage?.userInputMessageContext?.toolResults + if (!trs) break + const allMatched = trs.every((tr: any) => toolUseIds.has(tr.toolUseId)) + if (allMatched) break + totalSize -= sizes.shift() || 0 + history.shift() + } + } + } + + return { + request, + resolved, + convId, + agentContinuationId, + fingerprint, + toolNameMapper: toolMaps?.fromKiroName + } } export function transformToCodeWhisperer( @@ -279,9 +390,18 @@ export function transformToCodeWhisperer( model: string, auth: KiroAuthDetails, think = false, - budget = 20000 + budget = 20000, + workspace = '' ): PreparedRequest { - const { request, resolved, convId } = buildCodeWhispererRequest(body, model, auth, think, budget) + const { request, resolved, convId } = buildCodeWhispererRequest( + body, + model, + auth, + think, + budget, + undefined, + workspace + ) const osP = os.platform(), osR = os.release(), nodeV = process.version.replace('v', '') @@ -317,15 +437,17 @@ export function transformToSdkRequest( auth: KiroAuthDetails, think = false, budget = 20000, - showToast?: ToastFunction + showToast?: ToastFunction, + workspace = '' ): SdkPreparedRequest { - const { request, resolved, convId } = buildCodeWhispererRequest( + const { request, resolved, convId, fingerprint, toolNameMapper } = buildCodeWhispererRequest( body, model, auth, think, budget, - showToast + showToast, + workspace ) return { conversationState: request.conversationState, @@ -333,6 +455,8 @@ export function transformToSdkRequest( streaming: true, effectiveModel: resolved, conversationId: convId, - region: extractRegionFromArn(auth.profileArn) ?? auth.region + conversationKey: { workspace, fingerprint }, + region: extractRegionFromArn(auth.profileArn) ?? auth.region, + toolNameMapper } } diff --git a/src/plugin/sdk-client.ts b/src/plugin/sdk-client.ts index ba29311..948e218 100644 --- a/src/plugin/sdk-client.ts +++ b/src/plugin/sdk-client.ts @@ -1,7 +1,15 @@ import { CodeWhispererStreamingClient } from '@aws/codewhisperer-streaming-client' +import * as crypto from 'crypto' import { KIRO_CONSTANTS } from '../constants.js' import type { KiroAuthDetails } from './types' +const KIRO_VERSION = '0.11.63' + +function getMachineId(auth: KiroAuthDetails): string { + const key = auth.profileArn || auth.email || 'default' + return crypto.createHash('sha256').update(key).digest('hex') +} + const clientCache = new Map() export function createSdkClient( @@ -15,18 +23,32 @@ export function createSdkClient( return cached.client } + // Token rotated (refresh) — tear down the stale client so its sockets/agent + // don't leak before we replace the cache entry. + if (cached) { + try { + cached.client.destroy() + } catch {} + } + + const machineId = getMachineId(auth) const token = auth.access const client = new CodeWhispererStreamingClient({ region, endpoint: `https://q.${region}.amazonaws.com`, token: () => Promise.resolve({ token }), maxAttempts: 1, - customUserAgent: [[KIRO_CONSTANTS.USER_AGENT]] + customUserAgent: [[`${KIRO_CONSTANTS.USER_AGENT}-${KIRO_VERSION}-${machineId}`]], + requestHandler: { + connectionTimeout: 10000, + requestTimeout: 120000 + } }) client.middlewareStack.add( (next: any) => async (args: any) => { args.request.headers['x-amzn-kiro-agent-mode'] = 'vibe' + args.request.headers['x-amzn-codewhisperer-optout'] = 'true' return next(args) }, { step: 'build', name: 'addKiroHeaders' } diff --git a/src/plugin/storage/migrations.ts b/src/plugin/storage/migrations.ts index b892b64..d6e96c2 100644 --- a/src/plugin/storage/migrations.ts +++ b/src/plugin/storage/migrations.ts @@ -7,6 +7,37 @@ export function runMigrations(db: Database): void { migrateStartUrlColumn(db) migrateOidcRegionColumn(db) migrateDropRefreshTokenUniqueIndex(db) + migrateConversationsTable(db) + migrateReauthLockTable(db) + migrateConversationsAgentContinuationId(db) +} + +function migrateReauthLockTable(db: Database): void { + db.run(` + CREATE TABLE IF NOT EXISTS reauth_lock ( + id INTEGER PRIMARY KEY CHECK (id = 1), + pid INTEGER NOT NULL, + acquired_at INTEGER NOT NULL + ) + `) +} + +function migrateConversationsTable(db: Database): void { + const hasTable = db + .prepare("SELECT name FROM sqlite_master WHERE type='table' AND name='conversations'") + .get() + if (hasTable) return + + db.run(` + CREATE TABLE conversations ( + workspace TEXT NOT NULL, + fingerprint TEXT NOT NULL, + conv_id TEXT NOT NULL, + last_used INTEGER NOT NULL, + PRIMARY KEY (workspace, fingerprint) + ) + `) + db.run('CREATE INDEX idx_conversations_last_used ON conversations(last_used)') } function migrateToUniqueRefreshToken(db: Database): void { @@ -181,3 +212,11 @@ function migrateDropRefreshTokenUniqueIndex(db: Database): void { } } } + +function migrateConversationsAgentContinuationId(db: Database): void { + const columns = db.prepare('PRAGMA table_info(conversations)').all() as any[] + const names = new Set(columns.map((c: any) => c.name)) + if (!names.has('agent_continuation_id')) { + db.run('ALTER TABLE conversations ADD COLUMN agent_continuation_id TEXT') + } +} diff --git a/src/plugin/storage/sqlite.ts b/src/plugin/storage/sqlite.ts index eefc69a..0673371 100644 --- a/src/plugin/storage/sqlite.ts +++ b/src/plugin/storage/sqlite.ts @@ -137,6 +137,34 @@ export class KiroDatabase { }) } + 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, @@ -163,9 +191,132 @@ export class KiroDatabase { } } + private static readonly REAUTH_LOCK_TTL_MS = 120_000 + + acquireReauthLock(): boolean { + const now = Date.now() + try { + this.db.run('BEGIN IMMEDIATE') + } catch { + // Another write transaction is active — treat as lock held + return false + } + try { + const existing = this.db + .prepare('SELECT pid, acquired_at FROM reauth_lock WHERE id = 1') + .get() as { pid: number; acquired_at: number } | undefined + + if (existing) { + const expired = now - existing.acquired_at >= KiroDatabase.REAUTH_LOCK_TTL_MS + const dead = (() => { + try { + process.kill(existing.pid, 0) + return false + } catch { + return true + } + })() + if (expired || dead) { + this.db.prepare('DELETE FROM reauth_lock WHERE id = 1').run() + } else { + this.db.run('ROLLBACK') + return false + } + } + + // INSERT OR REPLACE handles a race where two instances both saw the + // same dead/expired lock and both reach this branch. + this.db + .prepare('INSERT OR REPLACE INTO reauth_lock (id, pid, acquired_at) VALUES (1, ?, ?)') + .run(process.pid, now) + this.db.run('COMMIT') + return true + } catch { + try { + this.db.run('ROLLBACK') + } catch { + // already rolled back + } + return false + } + } + + isReauthLockHeld(): boolean { + const row = this.db.prepare('SELECT pid FROM reauth_lock WHERE id = 1').get() as + | { pid: number } + | undefined + if (!row) return false + try { + process.kill(row.pid, 0) + return true + } catch { + return false + } + } + + releaseReauthLock(): void { + this.db.prepare('DELETE FROM reauth_lock WHERE id = 1 AND pid = ?').run(process.pid) + } + close() { this.db.close() } + + getConversationId( + workspace: string, + fingerprint: string + ): { convId: string; agentContinuationId: string } | undefined { + const row = this.db + .prepare( + 'SELECT conv_id, agent_continuation_id FROM conversations WHERE workspace = ? AND fingerprint = ?' + ) + .get(workspace, fingerprint) as + | { conv_id: string; agent_continuation_id: string | null } + | undefined + if (row) { + this.db + .prepare('UPDATE conversations SET last_used = ? WHERE workspace = ? AND fingerprint = ?') + .run(Date.now(), workspace, fingerprint) + } + return row + ? { convId: row.conv_id, agentContinuationId: row.agent_continuation_id || '' } + : undefined + } + + deleteConversationId(workspace: string, fingerprint: string): void { + this.db + .prepare('DELETE FROM conversations WHERE workspace = ? AND fingerprint = ?') + .run(workspace, fingerprint) + } + + /** + * Persist a conversationId and agentContinuationId, clean up entries older than ttlDays (default 7). + */ + setConversationId( + workspace: string, + fingerprint: string, + convId: string, + agentContinuationId: string, + ttlDays = 7 + ): void { + const now = Date.now() + const cutoff = now - ttlDays * 24 * 60 * 60 * 1000 + this.db.run('BEGIN TRANSACTION') + try { + this.db + .prepare( + `INSERT INTO conversations (workspace, fingerprint, conv_id, agent_continuation_id, last_used) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(workspace, fingerprint) DO UPDATE SET conv_id = excluded.conv_id, agent_continuation_id = excluded.agent_continuation_id, last_used = excluded.last_used` + ) + .run(workspace, fingerprint, convId, agentContinuationId, now) + this.db.prepare('DELETE FROM conversations WHERE last_used < ?').run(cutoff) + this.db.run('COMMIT') + } catch (e) { + this.db.run('ROLLBACK') + throw e + } + } } export function createDatabase(path?: string): KiroDatabase { diff --git a/src/plugin/streaming/sdk-stream-transformer.ts b/src/plugin/streaming/sdk-stream-transformer.ts index 6787c3e..02ba9d8 100644 --- a/src/plugin/streaming/sdk-stream-transformer.ts +++ b/src/plugin/streaming/sdk-stream-transformer.ts @@ -1,4 +1,6 @@ import { parseBracketToolCalls } from '../../infrastructure/transformers/tool-call-parser.js' +import { deduplicateToolCallsByContent } from '../../infrastructure/transformers/tool-transformer.js' +import * as logger from '../logger.js' import { getContextWindowSize } from '../models.js' import { estimateTokens } from '../response.js' import { convertToOpenAI } from './openai-converter.js' @@ -9,7 +11,8 @@ import { StreamState, THINKING_END_TAG, THINKING_START_TAG, ToolCallState } from export async function* transformSdkStream( sdkResponse: any, model: string, - conversationId: string + conversationId: string, + toolNameMapper?: (name: string) => string ): AsyncGenerator { const thinkingRequested = true @@ -29,6 +32,8 @@ export async function* transformSdkStream( let outputTokens = 0 let inputTokens = 0 let contextUsagePercentage: number | null = null + let realInputTokens: number | undefined + let realOutputTokens: number | undefined const toolCalls: ToolCallState[] = [] let currentToolCall: ToolCallState | null = null @@ -129,14 +134,14 @@ export async function* transformSdkStream( if (tc.name) totalContent += tc.name if (tc.input) totalContent += tc.input - if (tc.name && tc.toolUseId) { + if (tc.toolUseId) { if (currentToolCall && currentToolCall.toolUseId === tc.toolUseId) { currentToolCall.input += tc.input || '' - } else { + } else if (tc.name) { if (currentToolCall) toolCalls.push(currentToolCall) currentToolCall = { toolUseId: tc.toolUseId, - name: tc.name, + name: toolNameMapper ? toolNameMapper(tc.name) : tc.name, input: tc.input || '' } } @@ -149,11 +154,21 @@ export async function* transformSdkStream( if (event.metadataEvent.contextUsagePercentage) { contextUsagePercentage = event.metadataEvent.contextUsagePercentage } + if (event.metadataEvent.tokenUsage) { + const tu = event.metadataEvent.tokenUsage + if (typeof tu.inputTokens === 'number') realInputTokens = tu.inputTokens + if (typeof tu.outputTokens === 'number') realOutputTokens = tu.outputTokens + } } else if ((event as any).contextUsageEvent) { const cue = (event as any).contextUsageEvent if (cue.contextUsagePercentage) { contextUsagePercentage = cue.contextUsagePercentage } + } else if ((event as any).meteringEvent) { + const me = (event as any).meteringEvent + logger.debug( + `[CREDITS] usage=${me.usage} ${me.unit || 'credit'}${me.usage !== 1 ? 's' : ''}` + ) } } @@ -202,10 +217,12 @@ export async function* transformSdkStream( } } - if (toolCalls.length > 0) { + const dedupedToolCalls = deduplicateToolCallsByContent(toolCalls) + + if (dedupedToolCalls.length > 0) { const baseIndex = streamState.nextBlockIndex - for (let i = 0; i < toolCalls.length; i++) { - const tc = toolCalls[i] + for (let i = 0; i < dedupedToolCalls.length; i++) { + const tc = dedupedToolCalls[i] if (!tc) continue const blockIndex = baseIndex + i @@ -232,6 +249,9 @@ export async function* transformSdkStream( const parsed = JSON.parse(tc.input) inputJson = JSON.stringify(parsed) } catch (e) { + logger.debug( + `[TOOL_CALL] Invalid JSON for tool "${tc.name}" (id=${tc.toolUseId}): ${tc.input.slice(0, 500)}` + ) inputJson = tc.input } @@ -270,11 +290,15 @@ export async function* transformSdkStream( inputTokens = Math.max(0, totalTokens - outputTokens) } + // Real token counts from Kiro's metadata win over the context-% estimate. + if (realInputTokens !== undefined) inputTokens = realInputTokens + if (realOutputTokens !== undefined) outputTokens = realOutputTokens + { const _c = convertToOpenAI( { type: 'message_delta', - delta: { stop_reason: toolCalls.length > 0 ? 'tool_use' : 'end_turn' }, + delta: { stop_reason: dedupedToolCalls.length > 0 ? 'tool_use' : 'end_turn' }, usage: { input_tokens: inputTokens, output_tokens: outputTokens, @@ -293,6 +317,14 @@ export async function* transformSdkStream( if (_c !== null) yield _c } } catch (e) { + logger.debug( + `[STREAM] Error in transformSdkStream: ${e instanceof Error ? e.message : String(e)}` + ) + if (currentToolCall) { + logger.debug( + `[STREAM] Incomplete tool call: name=${currentToolCall.name} id=${currentToolCall.toolUseId} inputLen=${currentToolCall.input.length}` + ) + } throw e } } diff --git a/src/plugin/streaming/stream-transformer.ts b/src/plugin/streaming/stream-transformer.ts index 7036cb9..af9e839 100644 --- a/src/plugin/streaming/stream-transformer.ts +++ b/src/plugin/streaming/stream-transformer.ts @@ -43,7 +43,12 @@ export async function* transformKiroStream( try { while (true) { const { done, value } = await reader.read() - if (done) break + if (done) { + // Flush remaining multi-byte UTF-8 from TextDecoder + const tail = decoder.decode() + if (tail) rawBuffer += tail + break + } const chunk = decoder.decode(value, { stream: true }) rawBuffer += chunk diff --git a/src/plugin/sync/kiro-cli.ts b/src/plugin/sync/kiro-cli.ts index 6e46b01..3883871 100644 --- a/src/plugin/sync/kiro-cli.ts +++ b/src/plugin/sync/kiro-cli.ts @@ -122,8 +122,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) @@ -196,6 +221,10 @@ export async function syncFromKiroCli() { limitCount, lastSync: Date.now() }) + + if (authMethod === 'idc' && profileArn) { + await kiroDb.deleteStaleIdcDuplicates(id, resolvedEmail, profileArn) + } } } cliDb.close() diff --git a/src/plugin/types.ts b/src/plugin/types.ts index e8b05d3..06dfdca 100644 --- a/src/plugin/types.ts +++ b/src/plugin/types.ts @@ -82,6 +82,8 @@ export interface CodeWhispererMessage { export interface CodeWhispererRequest { conversationState: { + agentContinuationId?: string + agentTaskType?: string chatTriggerType: string conversationId: string history?: CodeWhispererMessage[] @@ -118,7 +120,9 @@ export interface SdkPreparedRequest { streaming: boolean effectiveModel: string conversationId: string + conversationKey: { workspace: string; fingerprint: string } region: string + toolNameMapper?: (name: string) => string } export type AccountSelectionStrategy = 'sticky' | 'round-robin' | 'lowest-usage' 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,