diff --git a/packages/ai/src/providers/anthropic.ts b/packages/ai/src/providers/anthropic.ts index 7d11a1289..629197328 100644 --- a/packages/ai/src/providers/anthropic.ts +++ b/packages/ai/src/providers/anthropic.ts @@ -1298,7 +1298,11 @@ export const streamAnthropic: StreamFunction<"anthropic-messages"> = ( } providerRetryAttempt++; const delayMs = PROVIDER_BASE_DELAY_MS * 2 ** (providerRetryAttempt - 1); - await scheduler.wait(delayMs, { signal: options?.signal }); + if (options?.providerRetryWait) { + await options.providerRetryWait(delayMs, options.signal); + } else { + await scheduler.wait(delayMs, { signal: options?.signal }); + } output.content.length = 0; output.responseId = undefined; output.errorMessage = strictFallbackErrorMessage; diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 096b221bb..51a295257 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -296,6 +296,10 @@ export interface StreamOptions { * Set to 0 to disable the inter-event idle watchdog for this request. */ streamIdleTimeoutMs?: number; + /** + * Optional retry delay hook for tests and transports that need custom scheduling. + */ + providerRetryWait?: (delayMs: number, signal?: AbortSignal) => Promise; /** * Optional `fetch` implementation override. Providers route every HTTP * request — direct calls, SDK clients, and retry helpers — through this diff --git a/packages/ai/src/usage.ts b/packages/ai/src/usage.ts index ab4c9e91d..2aa4531e1 100644 --- a/packages/ai/src/usage.ts +++ b/packages/ai/src/usage.ts @@ -156,6 +156,7 @@ export interface UsageFetchParams { export interface UsageFetchContext { fetch: typeof fetch; logger?: UsageLogger; + retryWait?: (delayMs: number, signal?: AbortSignal) => Promise; } /** Provider implementation for fetching usage information. */ diff --git a/packages/ai/src/usage/claude.ts b/packages/ai/src/usage/claude.ts index e4331c31e..8ae269d6b 100644 --- a/packages/ai/src/usage/claude.ts +++ b/packages/ai/src/usage/claude.ts @@ -149,11 +149,21 @@ function retryDelayMs(attempt: number, retryAfter: string | null): number { return Number.isFinite(dateDelay) ? Math.max(baseline, Math.max(0, dateDelay)) : baseline; } -async function waitBeforeRetry(attempt: number, retryAfter: string | null, signal?: AbortSignal): Promise { +async function waitBeforeRetry( + attempt: number, + retryAfter: string | null, + signal?: AbortSignal, + retryWait?: UsageFetchContext["retryWait"], +): Promise { if (signal?.aborted) return false; if (attempt >= MAX_ATTEMPTS - 1) return false; try { - await scheduler.wait(retryDelayMs(attempt, retryAfter), { signal }); + const delayMs = retryDelayMs(attempt, retryAfter); + if (retryWait) { + await retryWait(delayMs, signal); + } else { + await scheduler.wait(delayMs, { signal }); + } return !signal?.aborted; } catch (error) { if (isAbortError(error, signal)) return false; @@ -186,7 +196,8 @@ async function fetchUsagePayload( willRetry: retryable && attempt < MAX_ATTEMPTS - 1, }); if (!retryable) return null; - if (!(await waitBeforeRetry(attempt, response.headers.get("retry-after"), signal))) break; + const retryAfter = response.headers.get("retry-after"); + if (!(await waitBeforeRetry(attempt, retryAfter, signal, ctx.retryWait))) break; continue; } @@ -201,7 +212,7 @@ async function fetchUsagePayload( attempt, willRetry: attempt < MAX_ATTEMPTS - 1, }); - if (!(await waitBeforeRetry(attempt, null, signal))) break; + if (!(await waitBeforeRetry(attempt, null, signal, ctx.retryWait))) break; } catch (error) { if (isAbortError(error, signal)) return null; ctx.logger?.warn("Claude usage fetch error", { @@ -209,7 +220,7 @@ async function fetchUsagePayload( attempt, willRetry: attempt < MAX_ATTEMPTS - 1, }); - if (!(await waitBeforeRetry(attempt, null, signal))) break; + if (!(await waitBeforeRetry(attempt, null, signal, ctx.retryWait))) break; } } diff --git a/packages/ai/src/utils/oauth/github-copilot.ts b/packages/ai/src/utils/oauth/github-copilot.ts index 7430846ff..15c421816 100644 --- a/packages/ai/src/utils/oauth/github-copilot.ts +++ b/packages/ai/src/utils/oauth/github-copilot.ts @@ -165,10 +165,12 @@ async function pollForGitHubAccessToken( intervalSeconds: number, expiresIn: number, signal?: AbortSignal, + pollIntervalFloorMs = 1000, + pollIntervalScaleMs = 1000, ) { const urls = getUrls(domain); const deadline = Date.now() + expiresIn * 1000; - let intervalMs = Math.max(1000, Math.floor(intervalSeconds * 1000)); + let intervalMs = Math.max(pollIntervalFloorMs, Math.floor(intervalSeconds * pollIntervalScaleMs)); let intervalMultiplier = INITIAL_POLL_INTERVAL_MULTIPLIER; let slowDownResponses = 0; @@ -212,7 +214,9 @@ async function pollForGitHubAccessToken( if (error === "slow_down") { slowDownResponses += 1; intervalMs = - typeof interval === "number" && interval > 0 ? interval * 1000 : Math.max(1000, intervalMs + 5000); + typeof interval === "number" && interval > 0 + ? Math.max(pollIntervalFloorMs, interval * pollIntervalScaleMs) + : Math.max(pollIntervalFloorMs, intervalMs + 5 * pollIntervalScaleMs); intervalMultiplier = SLOW_DOWN_POLL_INTERVAL_MULTIPLIER; continue; } @@ -308,6 +312,8 @@ export async function loginGitHubCopilot(options: { onPrompt: (prompt: { message: string; placeholder?: string; allowEmpty?: boolean }) => Promise; onProgress?: (message: string) => void; signal?: AbortSignal; + pollIntervalFloorMs?: number; + pollIntervalScaleMs?: number; }): Promise { const input = await options.onPrompt({ message: "GitHub Enterprise URL/domain (blank for github.com)", @@ -337,6 +343,8 @@ export async function loginGitHubCopilot(options: { device.interval, device.expires_in, options.signal, + options.pollIntervalFloorMs, + options.pollIntervalScaleMs, ); // With opencode OAuth, the GitHub token is used directly for all API requests diff --git a/packages/ai/src/utils/retry.ts b/packages/ai/src/utils/retry.ts index 1a04263f6..732f54914 100644 --- a/packages/ai/src/utils/retry.ts +++ b/packages/ai/src/utils/retry.ts @@ -34,11 +34,12 @@ const COPILOT_MODEL_RETRY_BASE_DELAY_MS = 400; */ export async function callWithCopilotModelRetry( fn: () => Promise, - options: { provider: string; signal?: AbortSignal }, + options: { provider: string; signal?: AbortSignal; retryBaseDelayMs?: number }, ): Promise { if (options.provider !== "github-copilot") return fn(); let lastError: unknown; + const retryBaseDelayMs = options.retryBaseDelayMs ?? COPILOT_MODEL_RETRY_BASE_DELAY_MS; for (let attempt = 0; attempt < COPILOT_MODEL_RETRY_MAX_ATTEMPTS; attempt++) { try { return await fn(); @@ -46,7 +47,7 @@ export async function callWithCopilotModelRetry( lastError = error; if (!isCopilotTransientModelError(error) && !isRetryableError(error)) throw error; if (attempt === COPILOT_MODEL_RETRY_MAX_ATTEMPTS - 1) break; - await scheduler.wait(COPILOT_MODEL_RETRY_BASE_DELAY_MS * (attempt + 1), { signal: options.signal }); + await scheduler.wait(retryBaseDelayMs * (attempt + 1), { signal: options.signal }); } } throw lastError; diff --git a/packages/ai/test/anthropic-stream-envelope.test.ts b/packages/ai/test/anthropic-stream-envelope.test.ts index 19be1b7c8..f4f8acca9 100644 --- a/packages/ai/test/anthropic-stream-envelope.test.ts +++ b/packages/ai/test/anthropic-stream-envelope.test.ts @@ -1,6 +1,6 @@ import { afterEach, describe, expect, it, vi } from "bun:test"; +import { scheduler } from "node:timers/promises"; import { Messages } from "@anthropic-ai/sdk/resources/messages/messages"; -import * as z from "zod/v4"; import { streamAnthropic } from "../src/providers/anthropic"; import type { AssistantMessageEvent, Context, Model, ProviderSessionState } from "../src/types"; @@ -20,6 +20,17 @@ const model: Model<"anthropic-messages"> = { const context: Context = { messages: [{ role: "user", content: "Say hi", timestamp: Date.now() }], }; +const queryObjectSchema = { + type: "object", + properties: { query: { type: "string" } }, + required: ["query"], +}; + +const cityObjectSchema = { + type: "object", + properties: { city: { type: "string" } }, + required: ["city"], +}; type MockAnthropicEvent = Record; type MockAnthropicStream = AsyncIterable; @@ -111,8 +122,11 @@ function getStrictFlags(params: unknown): boolean[] { return tools.map(tool => tool.strict === true); } -function createTextSuccessEvents(text: string): MockAnthropicEvent[] { - return [ +function createTextSuccessEvents( + text: string, + options: { duplicateMessageStart?: boolean } = {}, +): MockAnthropicEvent[] { + const events: MockAnthropicEvent[] = [ { type: "message_start", message: { @@ -126,7 +140,6 @@ function createTextSuccessEvents(text: string): MockAnthropicEvent[] { }, }, { type: "content_block_start", index: 0, content_block: { type: "text", text: "" } }, - { type: "message_start", message: { id: "msg_duplicate", usage: { input_tokens: 99, output_tokens: 99 } } }, { type: "content_block_delta", index: 0, delta: { type: "text_delta", text } }, { type: "content_block_stop", index: 0 }, { @@ -141,6 +154,13 @@ function createTextSuccessEvents(text: string): MockAnthropicEvent[] { }, { type: "message_stop" }, ]; + if (options.duplicateMessageStart) { + events.splice(2, 0, { + type: "message_start", + message: { id: "msg_duplicate", usage: { input_tokens: 99, output_tokens: 99 } }, + }); + } + return events; } function createTextSuccessEventsWithPreamble(text: string, preambleEvents: MockAnthropicEvent[]): MockAnthropicEvent[] { @@ -190,7 +210,7 @@ afterEach(() => { describe("anthropic stream envelope handling", () => { it("ignores duplicate message_start envelopes without resetting streamed text", async () => { vi.spyOn(Messages.prototype, "create").mockImplementation( - () => createMockRequest(createTextSuccessEvents("hello")) as never, + () => createMockRequest(createTextSuccessEvents("hello", { duplicateMessageStart: true })) as never, ); const stream = streamAnthropic(model, context, { apiKey: "sk-ant-test" }); @@ -269,6 +289,7 @@ describe("anthropic stream envelope handling", () => { attempt === 1 ? createMalformedPreMessageStartEvents() : createTextSuccessEvents("recovered"), ) as never; }); + vi.spyOn(scheduler, "wait").mockResolvedValue(undefined); const stream = streamAnthropic(model, context, { apiKey: "sk-ant-test" }); const events: AssistantMessageEvent[] = []; @@ -294,7 +315,7 @@ describe("anthropic stream envelope handling", () => { name: "edit", description: "Edit a value", strict: true, - parameters: z.object({ query: z.string() }), + parameters: queryObjectSchema, }, ], }; @@ -350,7 +371,7 @@ describe("anthropic stream envelope handling", () => { name: "edit", description: "Edit a value", strict: true, - parameters: z.object({ query: z.string() }), + parameters: queryObjectSchema, }, ], }; @@ -464,7 +485,7 @@ describe("anthropic stream envelope handling", () => { sseFrame("content_block_start", successEvents[1]), sseRawFrame("content_block_delta", malformedTextDelta), sseFrame("content_block_stop", { type: "content_block_stop", index: 0 }), - sseFrame("message_delta", successEvents[5]), + sseFrame("message_delta", successEvents[4]), sseFrame("message_stop", { type: "message_stop" }), ]; vi.spyOn(Messages.prototype, "create").mockImplementation(() => createRawSseRequest(frames) as never); @@ -486,7 +507,7 @@ describe("anthropic stream envelope handling", () => { { name: "lookup_weather", description: "Lookup weather", - parameters: z.object({ city: z.string() }), + parameters: cityObjectSchema, }, ], }; diff --git a/packages/ai/test/anthropic-stream-timeout.test.ts b/packages/ai/test/anthropic-stream-timeout.test.ts index d07734a06..210ecba50 100644 --- a/packages/ai/test/anthropic-stream-timeout.test.ts +++ b/packages/ai/test/anthropic-stream-timeout.test.ts @@ -1,4 +1,4 @@ -import { afterEach, describe, expect, it } from "bun:test"; +import { afterEach, describe, expect, it, vi } from "bun:test"; import type Anthropic from "@anthropic-ai/sdk"; import { streamAnthropic } from "../src/providers/anthropic"; import type { Context, Model } from "../src/types"; @@ -147,13 +147,16 @@ describe("anthropic first-event timeout retries", () => { }) as never; }) as unknown as Anthropic["messages"]["create"]; const client = { messages: { create } } as Anthropic; + const providerRetryWait = vi.fn(async () => {}); const result = await streamAnthropic(model, context, { client, - streamFirstEventTimeoutMs: 20, + streamFirstEventTimeoutMs: 1, + providerRetryWait, }).result(); expect(attempt).toBe(2); + expect(providerRetryWait).toHaveBeenCalledWith(2000, undefined); expect(result.stopReason).toBe("stop"); expect(result.content).toEqual([{ type: "text", text: "retry recovered" }]); expect(result.responseId).toBe("msg_retry_success"); @@ -163,7 +166,7 @@ describe("anthropic first-event timeout retries", () => { const create = ((_body: unknown, requestOptions?: { signal?: AbortSignal }) => { return createAnthropicMockStream({ signal: requestOptions?.signal, - connectDelayMs: 30, + connectDelayMs: 2, events: createSuccessfulAnthropicEvents("delayed connect"), }) as never; }) as unknown as Anthropic["messages"]["create"]; @@ -171,7 +174,7 @@ describe("anthropic first-event timeout retries", () => { const result = await streamAnthropic(model, context, { client, - streamFirstEventTimeoutMs: 20, + streamFirstEventTimeoutMs: 1, }).result(); expect(result.stopReason).toBe("stop"); @@ -187,12 +190,12 @@ describe("anthropic first-event timeout retries", () => { const client = { messages: { create } } as Anthropic; const controller = new AbortController(); - setTimeout(() => controller.abort(), 5); + setTimeout(() => controller.abort(), 1); const result = await streamAnthropic(model, context, { client, signal: controller.signal, - streamFirstEventTimeoutMs: 50, + streamFirstEventTimeoutMs: 10, }).result(); expect(attempt).toBe(1); @@ -237,8 +240,8 @@ describe("anthropic first-event timeout retries", () => { const result = await streamAnthropic(model, context, { client, - streamFirstEventTimeoutMs: 1_000, - streamIdleTimeoutMs: 20, + streamFirstEventTimeoutMs: 10, + streamIdleTimeoutMs: 1, }).result(); expect(attempt).toBe(1); diff --git a/packages/ai/test/auth-gateway-anthropic-caching.test.ts b/packages/ai/test/auth-gateway-anthropic-caching.test.ts index 36921fa74..5ef393c13 100644 --- a/packages/ai/test/auth-gateway-anthropic-caching.test.ts +++ b/packages/ai/test/auth-gateway-anthropic-caching.test.ts @@ -13,9 +13,7 @@ * with the gateway live (`omp auth-gateway serve` or pm2). */ import { describe, expect, it } from "bun:test"; -import * as os from "node:os"; -import * as path from "node:path"; -import { isEnoent } from "@oh-my-pi/pi-utils"; +import { AUTH_GATEWAY_E2E_URL, checkAuthGatewayE2EAvailable } from "./helpers"; interface AnthropicUsage { input_tokens: number; @@ -32,30 +30,9 @@ interface AnthropicResponse { error?: { type: string; message: string }; } -const GATEWAY_URL = Bun.env.OMP_E2E_GATEWAY_URL ?? "http://127.0.0.1:4000"; -const TOKEN_PATH = path.join(os.homedir(), ".omp", "auth-gateway.token"); const MODEL = Bun.env.OMP_E2E_ANTHROPIC_MODEL ?? "claude-sonnet-4-5"; -async function checkGatewayAvailable(): Promise<{ ok: boolean; token?: string; reason?: string }> { - let token: string; - try { - token = (await Bun.file(TOKEN_PATH).text()).trim(); - } catch (err) { - if (isEnoent(err)) return { ok: false, reason: `no token at ${TOKEN_PATH}` }; - throw err; - } - if (!token) return { ok: false, reason: `empty token at ${TOKEN_PATH}` }; - try { - const res = await fetch(`${GATEWAY_URL}/healthz`, { signal: AbortSignal.timeout(2_000) }); - if (!res.ok) return { ok: false, reason: `healthz returned ${res.status}` }; - } catch (err) { - const msg = err instanceof Error ? err.message : String(err); - return { ok: false, reason: `healthz unreachable: ${msg}` }; - } - return { ok: true, token }; -} - -const gateway = await checkGatewayAvailable(); +const gateway = await checkAuthGatewayE2EAvailable(); // Build a system prompt that comfortably exceeds Anthropic's 1024-token cache // floor for Sonnet. Using a deterministic repeated paragraph so cache keys are @@ -84,7 +61,7 @@ interface MessageBlock { } async function callGateway(body: unknown, token: string): Promise { - const res = await fetch(`${GATEWAY_URL}/v1/messages`, { + const res = await fetch(`${AUTH_GATEWAY_E2E_URL}/v1/messages`, { method: "POST", headers: { "Content-Type": "application/json", @@ -134,7 +111,7 @@ describe.skipIf(!gateway.ok)("auth-gateway: anthropic prompt caching e2e", () => const turn1 = await callGateway( { model: MODEL, - max_tokens: 32, + max_tokens: 4, system, messages: turn1Messages, }, @@ -162,7 +139,7 @@ describe.skipIf(!gateway.ok)("auth-gateway: anthropic prompt caching e2e", () => const turn2 = await callGateway( { model: MODEL, - max_tokens: 32, + max_tokens: 4, system, messages: turn2Messages, }, diff --git a/packages/ai/test/auth-gateway-anthropic-to-codex-caching.test.ts b/packages/ai/test/auth-gateway-anthropic-to-codex-caching.test.ts index e96d13257..97ef4f1a9 100644 --- a/packages/ai/test/auth-gateway-anthropic-to-codex-caching.test.ts +++ b/packages/ai/test/auth-gateway-anthropic-to-codex-caching.test.ts @@ -25,9 +25,7 @@ * To run: `bun --cwd packages/ai test test/auth-gateway-anthropic-to-codex-caching.test.ts` */ import { describe, expect, it } from "bun:test"; -import * as os from "node:os"; -import * as path from "node:path"; -import { isEnoent } from "@oh-my-pi/pi-utils"; +import { AUTH_GATEWAY_E2E_URL, checkAuthGatewayE2EAvailable } from "./helpers"; interface AnthropicUsage { input_tokens: number; @@ -44,30 +42,9 @@ interface AnthropicResponse { error?: { type: string; message: string }; } -const GATEWAY_URL = Bun.env.OMP_E2E_GATEWAY_URL ?? "http://127.0.0.1:4000"; -const TOKEN_PATH = path.join(os.homedir(), ".omp", "auth-gateway.token"); const MODEL = Bun.env.OMP_E2E_CODEX_MODEL ?? "gpt-5.3-codex"; -async function checkGatewayAvailable(): Promise<{ ok: boolean; token?: string; reason?: string }> { - let token: string; - try { - token = (await Bun.file(TOKEN_PATH).text()).trim(); - } catch (err) { - if (isEnoent(err)) return { ok: false, reason: `no token at ${TOKEN_PATH}` }; - throw err; - } - if (!token) return { ok: false, reason: `empty token at ${TOKEN_PATH}` }; - try { - const res = await fetch(`${GATEWAY_URL}/healthz`, { signal: AbortSignal.timeout(2_000) }); - if (!res.ok) return { ok: false, reason: `healthz returned ${res.status}` }; - } catch (err) { - const msg = err instanceof Error ? err.message : String(err); - return { ok: false, reason: `healthz unreachable: ${msg}` }; - } - return { ok: true, token }; -} - -const gateway = await checkGatewayAvailable(); +const gateway = await checkAuthGatewayE2EAvailable(); // Long deterministic instructions, repeated to clear Codex's 1024-token // cache floor with headroom. @@ -94,7 +71,7 @@ interface MessageBlock { } async function callGateway(body: unknown, token: string): Promise { - const res = await fetch(`${GATEWAY_URL}/v1/messages`, { + const res = await fetch(`${AUTH_GATEWAY_E2E_URL}/v1/messages`, { method: "POST", headers: { "Content-Type": "application/json", @@ -142,7 +119,7 @@ describe.skipIf(!gateway.ok)("auth-gateway: anthropic-messages → openai-codex const turn1 = await callGateway( { model: MODEL, - max_tokens: 32, + max_tokens: 4, system, messages: turn1Messages, }, @@ -167,7 +144,7 @@ describe.skipIf(!gateway.ok)("auth-gateway: anthropic-messages → openai-codex const turn2 = await callGateway( { model: MODEL, - max_tokens: 32, + max_tokens: 4, system, messages: turn2Messages, }, diff --git a/packages/ai/test/auth-gateway-cross-protocol-caching.test.ts b/packages/ai/test/auth-gateway-cross-protocol-caching.test.ts index b2736140a..fc16388b4 100644 --- a/packages/ai/test/auth-gateway-cross-protocol-caching.test.ts +++ b/packages/ai/test/auth-gateway-cross-protocol-caching.test.ts @@ -27,9 +27,7 @@ * with the gateway live (`omp auth-gateway serve` or pm2). */ import { describe, expect, it } from "bun:test"; -import * as os from "node:os"; -import * as path from "node:path"; -import { isEnoent } from "@oh-my-pi/pi-utils"; +import { AUTH_GATEWAY_E2E_URL, checkAuthGatewayE2EAvailable } from "./helpers"; interface OpenAIResponsesUsage { input_tokens: number; @@ -49,30 +47,9 @@ interface OpenAIResponse { error?: { type?: string; message: string }; } -const GATEWAY_URL = Bun.env.OMP_E2E_GATEWAY_URL ?? "http://127.0.0.1:4000"; -const TOKEN_PATH = path.join(os.homedir(), ".omp", "auth-gateway.token"); const MODEL = Bun.env.OMP_E2E_ANTHROPIC_MODEL ?? "claude-sonnet-4-5"; -async function checkGatewayAvailable(): Promise<{ ok: boolean; token?: string; reason?: string }> { - let token: string; - try { - token = (await Bun.file(TOKEN_PATH).text()).trim(); - } catch (err) { - if (isEnoent(err)) return { ok: false, reason: `no token at ${TOKEN_PATH}` }; - throw err; - } - if (!token) return { ok: false, reason: `empty token at ${TOKEN_PATH}` }; - try { - const res = await fetch(`${GATEWAY_URL}/healthz`, { signal: AbortSignal.timeout(2_000) }); - if (!res.ok) return { ok: false, reason: `healthz returned ${res.status}` }; - } catch (err) { - const msg = err instanceof Error ? err.message : String(err); - return { ok: false, reason: `healthz unreachable: ${msg}` }; - } - return { ok: true, token }; -} - -const gateway = await checkGatewayAvailable(); +const gateway = await checkAuthGatewayE2EAvailable(); // Long deterministic instructions, repeated to clear Anthropic's 1024-token // cache floor for Sonnet. @@ -99,7 +76,7 @@ interface ResponseInputMessage { } async function callGateway(body: unknown, token: string): Promise { - const res = await fetch(`${GATEWAY_URL}/v1/responses`, { + const res = await fetch(`${AUTH_GATEWAY_E2E_URL}/v1/responses`, { method: "POST", headers: { "Content-Type": "application/json", @@ -147,7 +124,7 @@ describe.skipIf(!gateway.ok)("auth-gateway: openai-responses → anthropic cachi const turn1 = await callGateway( { model: MODEL, - max_output_tokens: 64, + max_output_tokens: 4, instructions: instructionsWithNonce, input: turn1Input, }, @@ -173,7 +150,7 @@ describe.skipIf(!gateway.ok)("auth-gateway: openai-responses → anthropic cachi const turn2 = await callGateway( { model: MODEL, - max_output_tokens: 64, + max_output_tokens: 4, instructions: instructionsWithNonce, input: turn2Input, }, diff --git a/packages/ai/test/auth-storage-credential-disabled-event.test.ts b/packages/ai/test/auth-storage-credential-disabled-event.test.ts index 01184eb53..7a67faf50 100644 --- a/packages/ai/test/auth-storage-credential-disabled-event.test.ts +++ b/packages/ai/test/auth-storage-credential-disabled-event.test.ts @@ -1,12 +1,10 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; -import * as fs from "node:fs/promises"; -import * as os from "node:os"; -import * as path from "node:path"; import { + type AuthCredential, type AuthCredentialStore, AuthStorage, type CredentialDisabledEvent, - SqliteAuthCredentialStore, + type StoredAuthCredential, } from "../src/auth-storage"; import * as oauthUtils from "../src/utils/oauth"; @@ -29,33 +27,105 @@ const failOAuthRefresh = (message = 'HTTP 400 invalid_grant {"error":"invalid_gr }); }; +class MemoryAuthCredentialStore implements AuthCredentialStore { + #rows: StoredAuthCredential[] = []; + #nextId = 1; + + close(): void {} + + listAuthCredentials(provider?: string): StoredAuthCredential[] { + return this.#rows.filter(row => row.disabledCause === null && (!provider || row.provider === provider)); + } + + updateAuthCredential(id: number, credential: AuthCredential): void { + const row = this.#rows.find(entry => entry.id === id); + if (row) row.credential = credential; + } + + deleteAuthCredential(id: number, disabledCause: string): void { + const row = this.#rows.find(entry => entry.id === id); + if (row) row.disabledCause = disabledCause; + } + + tryDisableAuthCredentialIfMatches(id: number, expectedData: string, disabledCause: string): boolean { + const row = this.#rows.find(entry => entry.id === id && entry.disabledCause === null); + if (!row || serializeTestCredential(row.credential) !== expectedData) return false; + row.disabledCause = disabledCause; + return true; + } + + replaceAuthCredentialsForProvider(provider: string, credentials: AuthCredential[]): StoredAuthCredential[] { + for (const row of this.#rows) { + if (row.provider === provider && row.disabledCause === null) { + row.disabledCause = "replaced by newer credential"; + } + } + const rows = credentials.map( + (credential): StoredAuthCredential => ({ + id: this.#nextId++, + provider, + credential, + disabledCause: null, + }), + ); + this.#rows.push(...rows); + return rows; + } + + upsertAuthCredentialForProvider(provider: string, credential: AuthCredential): StoredAuthCredential[] { + return this.replaceAuthCredentialsForProvider(provider, [credential]); + } + + deleteAuthCredentialsForProvider(provider: string, disabledCause: string): void { + for (const row of this.#rows) { + if (row.provider === provider && row.disabledCause === null) row.disabledCause = disabledCause; + } + } + + getCache(): string | null { + return null; + } + + setCache(): void {} + + cleanExpiredCache(): void {} +} + +function serializeTestCredential(credential: AuthCredential): string { + if (credential.type === "api_key") return JSON.stringify({ key: credential.key }); + if (credential.type === "oauth") { + const { type: _type, ...rest } = credential; + return JSON.stringify(rest); + } + return ""; +} + +function disableCredential(authStorage: AuthStorage, id: number, provider = "anthropic"): void { + expect(authStorage.disableCredentialById(id, "oauth refresh failed: invalid_grant")).toBe(true); + expect(authStorage.list()).not.toContain(provider); +} + describe("AuthStorage credential_disabled subscriptions", () => { - let tempDir = ""; const stores: AuthCredentialStore[] = []; - const openStorage = async (options?: ConstructorParameters[1]): Promise => { - const store = await SqliteAuthCredentialStore.open(path.join(tempDir, `agent-${stores.length}.db`)); + const openStorage = (options?: ConstructorParameters[1]): AuthStorage => { + const store = new MemoryAuthCredentialStore(); stores.push(store); return new AuthStorage(store, options); }; - beforeEach(async () => { - tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-credential-disabled-subs-")); + beforeEach(() => { for (const key of SUPPRESS_ANTHROPIC_ENV) { savedEnv[key] = process.env[key]; delete process.env[key]; } }); - afterEach(async () => { + afterEach(() => { vi.restoreAllMocks(); for (const store of stores.splice(0)) { store.close(); } - if (tempDir) { - await fs.rm(tempDir, { recursive: true, force: true }); - tempDir = ""; - } for (const key of SUPPRESS_ANTHROPIC_ENV) { if (savedEnv[key] === undefined) { delete process.env[key]; @@ -69,7 +139,7 @@ describe("AuthStorage credential_disabled subscriptions", () => { describe("constructor `onCredentialDisabled` option", () => { test("fires when an OAuth credential is disabled by a definitive refresh failure", async () => { const events: CredentialDisabledEvent[] = []; - const authStorage = await openStorage({ + const authStorage = openStorage({ onCredentialDisabled: event => { events.push(event); }, @@ -87,7 +157,7 @@ describe("AuthStorage credential_disabled subscriptions", () => { test("does not fire for transient (non-definitive) refresh failures", async () => { const events: CredentialDisabledEvent[] = []; - const authStorage = await openStorage({ + const authStorage = openStorage({ onCredentialDisabled: event => { events.push(event); }, @@ -100,21 +170,18 @@ describe("AuthStorage credential_disabled subscriptions", () => { }); test("swallows synchronous handler exceptions so the disable still completes", async () => { - const authStorage = await openStorage({ + const authStorage = openStorage({ onCredentialDisabled: () => { throw new Error("subscriber exploded"); }, }); await authStorage.set("anthropic", [expiredOAuth()]); - failOAuthRefresh("invalid_grant"); - - await expect(authStorage.getApiKey("anthropic", "session-handler-throws")).resolves.toBeUndefined(); - expect(authStorage.list()).not.toContain("anthropic"); + disableCredential(authStorage, 1); }); test("swallows async handler rejections so the disable path still completes", async () => { const settled = Promise.withResolvers(); - const authStorage = await openStorage({ + const authStorage = openStorage({ onCredentialDisabled: async () => { // Yield so the rejection lands on the microtask queue, not synchronously. await Promise.resolve(); @@ -123,7 +190,6 @@ describe("AuthStorage credential_disabled subscriptions", () => { }, }); await authStorage.set("anthropic", [expiredOAuth()]); - failOAuthRefresh("invalid_grant"); const unhandled: unknown[] = []; const onUnhandled = (reason: unknown): void => { @@ -131,10 +197,9 @@ describe("AuthStorage credential_disabled subscriptions", () => { }; process.on("unhandledRejection", onUnhandled); try { - await expect(authStorage.getApiKey("anthropic", "session-async-handler-throws")).resolves.toBeUndefined(); + disableCredential(authStorage, 1); await settled.promise; await Bun.sleep(0); - expect(authStorage.list()).not.toContain("anthropic"); expect(unhandled).toHaveLength(0); } finally { process.off("unhandledRejection", onUnhandled); @@ -146,7 +211,7 @@ describe("AuthStorage credential_disabled subscriptions", () => { test("registers an additional subscriber alongside the constructor handler — both fire", async () => { const constructorEvents: CredentialDisabledEvent[] = []; const runtimeEvents: CredentialDisabledEvent[] = []; - const authStorage = await openStorage({ + const authStorage = openStorage({ onCredentialDisabled: event => { constructorEvents.push(event); }, @@ -156,9 +221,7 @@ describe("AuthStorage credential_disabled subscriptions", () => { }); await authStorage.set("anthropic", [expiredOAuth()]); - failOAuthRefresh("invalid_grant"); - - await authStorage.getApiKey("anthropic", "session-both-fire"); + disableCredential(authStorage, 1); expect(constructorEvents).toHaveLength(1); expect(runtimeEvents).toHaveLength(1); expect(constructorEvents[0]?.provider).toBe("anthropic"); @@ -168,7 +231,7 @@ describe("AuthStorage credential_disabled subscriptions", () => { test("fans out every event to every subscriber", async () => { const aEvents: CredentialDisabledEvent[] = []; const bEvents: CredentialDisabledEvent[] = []; - const authStorage = await openStorage(); + const authStorage = openStorage(); authStorage.onCredentialDisabled(event => { aEvents.push(event); }); @@ -177,17 +240,15 @@ describe("AuthStorage credential_disabled subscriptions", () => { }); await authStorage.set("anthropic", [expiredOAuth()]); await authStorage.set("openai", [expiredOAuth()]); - failOAuthRefresh("invalid_grant"); - - await authStorage.getApiKey("anthropic", "session-fanout-anthropic"); - await authStorage.getApiKey("openai", "session-fanout-openai"); + disableCredential(authStorage, 1); + disableCredential(authStorage, 2, "openai"); expect(aEvents.map(event => event.provider)).toEqual(["anthropic", "openai"]); expect(bEvents.map(event => event.provider)).toEqual(["anthropic", "openai"]); }); test("unsubscribe removes only that listener; others continue to fire", async () => { - const authStorage = await openStorage(); + const authStorage = openStorage(); const aEvents: CredentialDisabledEvent[] = []; const bEvents: CredentialDisabledEvent[] = []; const unsubscribeA = authStorage.onCredentialDisabled(event => { @@ -199,21 +260,20 @@ describe("AuthStorage credential_disabled subscriptions", () => { await authStorage.set("anthropic", [expiredOAuth()]); await authStorage.set("openai", [expiredOAuth()]); - failOAuthRefresh("invalid_grant"); - await authStorage.getApiKey("anthropic", "session-pre-unsubscribe"); + disableCredential(authStorage, 1); expect(aEvents).toHaveLength(1); expect(bEvents).toHaveLength(1); unsubscribeA(); - await authStorage.getApiKey("openai", "session-post-unsubscribe"); + disableCredential(authStorage, 2, "openai"); expect(aEvents).toHaveLength(1); expect(bEvents).toHaveLength(2); }); test("unsubscribe is idempotent: a second call is a no-op and does not affect other listeners", async () => { - const authStorage = await openStorage(); + const authStorage = openStorage(); const aEvents: CredentialDisabledEvent[] = []; const bEvents: CredentialDisabledEvent[] = []; const unsubscribeA = authStorage.onCredentialDisabled(event => { @@ -227,15 +287,14 @@ describe("AuthStorage credential_disabled subscriptions", () => { unsubscribeA(); await authStorage.set("anthropic", [expiredOAuth()]); - failOAuthRefresh("invalid_grant"); - await authStorage.getApiKey("anthropic", "session-idempotent-unsub"); + disableCredential(authStorage, 1); expect(aEvents).toHaveLength(0); expect(bEvents).toHaveLength(1); }); test("a throwing subscriber does not block other subscribers from receiving the event", async () => { - const authStorage = await openStorage(); + const authStorage = openStorage(); const tailEvents: CredentialDisabledEvent[] = []; authStorage.onCredentialDisabled(() => { throw new Error("first subscriber exploded"); @@ -245,14 +304,13 @@ describe("AuthStorage credential_disabled subscriptions", () => { }); await authStorage.set("anthropic", [expiredOAuth()]); - failOAuthRefresh("invalid_grant"); - await expect(authStorage.getApiKey("anthropic", "session-throw-isolation")).resolves.toBeUndefined(); + disableCredential(authStorage, 1); expect(tailEvents).toHaveLength(1); }); test("an async-rejecting subscriber does not trip unhandledRejection and does not block others", async () => { - const authStorage = await openStorage(); + const authStorage = openStorage(); const tailEvents: CredentialDisabledEvent[] = []; const settled = Promise.withResolvers(); authStorage.onCredentialDisabled(async () => { @@ -265,7 +323,6 @@ describe("AuthStorage credential_disabled subscriptions", () => { }); await authStorage.set("anthropic", [expiredOAuth()]); - failOAuthRefresh("invalid_grant"); const unhandled: unknown[] = []; const onUnhandled = (reason: unknown): void => { @@ -273,7 +330,7 @@ describe("AuthStorage credential_disabled subscriptions", () => { }; process.on("unhandledRejection", onUnhandled); try { - await authStorage.getApiKey("anthropic", "session-async-throw-isolation"); + disableCredential(authStorage, 1); await settled.promise; await Bun.sleep(0); expect(tailEvents).toHaveLength(1); @@ -286,11 +343,10 @@ describe("AuthStorage credential_disabled subscriptions", () => { describe("buffer-and-replay for events fired with no subscribers", () => { test("replays buffered events to the first subscriber that triggers the empty→non-empty transition", async () => { - const authStorage = await openStorage(); + const authStorage = openStorage(); await authStorage.set("anthropic", [expiredOAuth()]); - failOAuthRefresh("invalid_grant"); - await authStorage.getApiKey("anthropic", "session-pre-subscribe"); + disableCredential(authStorage, 1); const replayed: CredentialDisabledEvent[] = []; authStorage.onCredentialDisabled(event => { @@ -305,11 +361,10 @@ describe("AuthStorage credential_disabled subscriptions", () => { }); test("drains once: a later subscriber attached after the first does not re-receive past events", async () => { - const authStorage = await openStorage(); + const authStorage = openStorage(); await authStorage.set("anthropic", [expiredOAuth()]); - failOAuthRefresh("invalid_grant"); - await authStorage.getApiKey("anthropic", "session-pre-first-listener"); + disableCredential(authStorage, 1); const firstEvents: CredentialDisabledEvent[] = []; authStorage.onCredentialDisabled(event => { @@ -328,21 +383,20 @@ describe("AuthStorage credential_disabled subscriptions", () => { }); test("after every subscriber unsubscribes, subsequent events buffer until the next subscribe", async () => { - const authStorage = await openStorage(); + const authStorage = openStorage(); const events: CredentialDisabledEvent[] = []; const unsubscribe = authStorage.onCredentialDisabled(event => { events.push(event); }); await authStorage.set("anthropic", [expiredOAuth()]); - failOAuthRefresh("invalid_grant"); - await authStorage.getApiKey("anthropic", "session-pre-unsubscribe"); + disableCredential(authStorage, 1); expect(events).toHaveLength(1); unsubscribe(); // No subscribers; the next disable goes to the buffer. await authStorage.set("openai", [expiredOAuth()]); - await authStorage.getApiKey("openai", "session-during-gap"); + disableCredential(authStorage, 2, "openai"); expect(events).toHaveLength(1); const replayed: CredentialDisabledEvent[] = []; diff --git a/packages/ai/test/auth-storage-usage-cache.test.ts b/packages/ai/test/auth-storage-usage-cache.test.ts index f967134bb..b57d28484 100644 --- a/packages/ai/test/auth-storage-usage-cache.test.ts +++ b/packages/ai/test/auth-storage-usage-cache.test.ts @@ -237,7 +237,9 @@ describe("AuthStorage usage cache: last-good failure fallback", () => { describe("AuthStorage usage cache: jitter", () => { it("writes per-credential cache TTLs with ±25% jitter so refreshes decorrelate", async () => { const store = makeStore([oauthRow(1, "a@example.com"), oauthRow(2, "b@example.com")]); - const storage = new AuthStorage(store); + const storage = new AuthStorage(store, { + usageProviderResolver: provider => (provider === "anthropic" ? claudeUsage.claudeUsageProvider : undefined), + }); await storage.reload(); try { const goldA = makeReport("a@example.com"); diff --git a/packages/ai/test/claude-usage-retry.test.ts b/packages/ai/test/claude-usage-retry.test.ts index 73a1cd8d4..3425f1c2b 100644 --- a/packages/ai/test/claude-usage-retry.test.ts +++ b/packages/ai/test/claude-usage-retry.test.ts @@ -1,5 +1,4 @@ import { afterEach, describe, expect, it, vi } from "bun:test"; -import { setTimeout as setTimeoutCb } from "node:timers"; import type { UsageFetchContext } from "../src/usage"; import { claudeUsageProvider } from "../src/usage/claude"; @@ -14,8 +13,8 @@ function jsonResponse(status: number, body: unknown, headers: Record { vi.restoreAllMocks(); }); + const instantRetryWait: UsageFetchContext["retryWait"] = async () => {}; + it("retries on 429 and succeeds on a later attempt", async () => { let attempt = 0; const fetchMock = (async () => { @@ -44,7 +45,7 @@ describe("claudeUsageProvider retry contract", () => { return jsonResponse(200, VALID_PAYLOAD); }) as unknown as typeof fetch; - const report = await claudeUsageProvider.fetchUsage(baseParams(), makeContext(fetchMock)); + const report = await claudeUsageProvider.fetchUsage(baseParams(), makeContext(fetchMock, instantRetryWait)); expect(report).not.toBeNull(); expect(attempt).toBe(3); expect(report?.limits[0]?.amount.used).toBe(42); @@ -58,7 +59,7 @@ describe("claudeUsageProvider retry contract", () => { return jsonResponse(200, VALID_PAYLOAD); }) as unknown as typeof fetch; - const report = await claudeUsageProvider.fetchUsage(baseParams(), makeContext(fetchMock)); + const report = await claudeUsageProvider.fetchUsage(baseParams(), makeContext(fetchMock, instantRetryWait)); expect(report).not.toBeNull(); expect(attempt).toBe(2); }); @@ -94,35 +95,28 @@ describe("claudeUsageProvider retry contract", () => { return jsonResponse(429, { error: "rate_limited" }); }) as unknown as typeof fetch; - // Provider's MAX_RETRIES is 3; provider sleeps BASE_RETRY_DELAY_MS * 2^attempt - // between attempts — total worst-case ~1.5s, well within our test budget. - const report = await claudeUsageProvider.fetchUsage(baseParams(), makeContext(fetchMock)); + const report = await claudeUsageProvider.fetchUsage(baseParams(), makeContext(fetchMock, instantRetryWait)); expect(report).toBeNull(); expect(attempt).toBe(3); }); it("honours Retry-After when retrying a 429", async () => { let attempt = 0; - const callTimes: number[] = []; + const retryWait = vi.fn(async () => {}); const fetchMock = (async () => { attempt += 1; - callTimes.push(Date.now()); if (attempt === 1) { - // Retry-After: 1 second. Provider must wait ~1s before re-attempting. + // Retry-After: 1 second. Provider must compute a 1s backoff before re-attempting. return jsonResponse(429, { error: "rate_limited" }, { "retry-after": "1" }); } return jsonResponse(200, VALID_PAYLOAD); }) as unknown as typeof fetch; - const t0 = Date.now(); - const report = await claudeUsageProvider.fetchUsage(baseParams(), makeContext(fetchMock)); - const elapsed = Date.now() - t0; + const report = await claudeUsageProvider.fetchUsage(baseParams(), makeContext(fetchMock, retryWait)); expect(report).not.toBeNull(); expect(attempt).toBe(2); - // Allow generous slop (Bun scheduling jitter) but ensure we actually waited - // closer to the Retry-After than to the default 500ms backoff. - expect(elapsed).toBeGreaterThanOrEqual(800); - expect(callTimes[1] - callTimes[0]).toBeGreaterThanOrEqual(800); + expect(retryWait).toHaveBeenCalledTimes(1); + expect(retryWait.mock.calls[0]?.[0]).toBe(1000); }); it("aborts the retry sleep when the signal fires mid-backoff", async () => { @@ -140,16 +134,26 @@ describe("claudeUsageProvider retry contract", () => { }) as unknown as typeof fetch; const controller = new AbortController(); - setTimeoutCb(() => controller.abort(), 150); + const retryWait = vi.fn(async (delayMs: number, signal?: AbortSignal) => { + expect(delayMs).toBe(60_000); + if (signal?.aborted) throw new Error("AbortError"); + const { promise, reject } = Promise.withResolvers(); + const onAbort = () => reject(new Error("AbortError")); + signal?.addEventListener("abort", onAbort, { once: true }); + queueMicrotask(() => controller.abort()); + try { + await promise; + } finally { + signal?.removeEventListener("abort", onAbort); + } + }); - const t0 = Date.now(); const report = await claudeUsageProvider.fetchUsage( { ...baseParams(), signal: controller.signal }, - makeContext(fetchMock), + makeContext(fetchMock, retryWait), ); - const elapsed = Date.now() - t0; expect(report).toBeNull(); - expect(elapsed).toBeLessThan(3_000); + expect(retryWait).toHaveBeenCalledTimes(1); expect(attempt).toBe(1); }); @@ -168,7 +172,7 @@ describe("claudeUsageProvider retry contract", () => { return jsonResponse(429, { error: "rate_limited" }); }) as unknown as typeof fetch; - const report = await claudeUsageProvider.fetchUsage(baseParams(), makeContext(fetchMock)); + const report = await claudeUsageProvider.fetchUsage(baseParams(), makeContext(fetchMock, instantRetryWait)); // The 200 set lastPayload but had no usage data; 429s mean no further // successes. lastPayload survives but has no usage data → no limits. // Specifically: report is null (since lastPayload has nothing to expose). diff --git a/packages/ai/test/copilot-retry.test.ts b/packages/ai/test/copilot-retry.test.ts index 28881b14d..42ebd6319 100644 --- a/packages/ai/test/copilot-retry.test.ts +++ b/packages/ai/test/copilot-retry.test.ts @@ -83,7 +83,7 @@ describe("callWithCopilotModelRetry", () => { calls += 1; throw err; }, - { provider: "github-copilot" }, + { provider: "github-copilot", retryBaseDelayMs: 0 }, ), ).rejects.toBe(err); expect(calls).toBe(3); @@ -99,7 +99,7 @@ describe("callWithCopilotModelRetry", () => { } return "ok" as const; }, - { provider: "github-copilot" }, + { provider: "github-copilot", retryBaseDelayMs: 0 }, ); expect(result).toBe("ok"); expect(calls).toBe(2); @@ -130,7 +130,7 @@ describe("callWithCopilotModelRetry", () => { calls += 1; throw copilotError({ status: 400, code: "model_not_supported", message: "transient" }); }, - { provider: "github-copilot", signal: controller.signal }, + { provider: "github-copilot", signal: controller.signal, retryBaseDelayMs: 0 }, ), ).rejects.toBeDefined(); // fn runs once; scheduler.wait rejects before a second attempt. diff --git a/packages/ai/test/github-copilot-login.test.ts b/packages/ai/test/github-copilot-login.test.ts index 2e24dfea4..b4ff72e75 100644 --- a/packages/ai/test/github-copilot-login.test.ts +++ b/packages/ai/test/github-copilot-login.test.ts @@ -2,6 +2,7 @@ import { afterEach, describe, expect, it, vi } from "bun:test"; import { loginGitHubCopilot } from "../src/utils/oauth/github-copilot"; const originalFetch = global.fetch; +const FAST_POLL_OPTIONS = { pollIntervalFloorMs: 0, pollIntervalScaleMs: 1 } as const; afterEach(() => { global.fetch = originalFetch; @@ -59,6 +60,7 @@ describe("loginGitHubCopilot", () => { const onAuth = vi.fn(); const credentials = await loginGitHubCopilot({ + ...FAST_POLL_OPTIONS, onAuth, onPrompt: mockOnPrompt(""), }); @@ -94,6 +96,7 @@ describe("loginGitHubCopilot", () => { global.fetch = fetchMock as unknown as typeof fetch; const credentials = await loginGitHubCopilot({ + ...FAST_POLL_OPTIONS, onAuth: vi.fn(), onPrompt: mockOnPrompt("ghe.example.com"), }); @@ -125,6 +128,7 @@ describe("loginGitHubCopilot", () => { global.fetch = fetchMock as unknown as typeof fetch; const credentials = await loginGitHubCopilot({ + ...FAST_POLL_OPTIONS, onAuth: vi.fn(), onPrompt: mockOnPrompt(" "), }); @@ -191,6 +195,7 @@ describe("loginGitHubCopilot", () => { global.fetch = fetchMock as unknown as typeof fetch; const credentials = await loginGitHubCopilot({ + ...FAST_POLL_OPTIONS, onAuth: vi.fn(), onPrompt: mockOnPrompt(""), }); @@ -247,6 +252,7 @@ describe("loginGitHubCopilot", () => { await expect( loginGitHubCopilot({ + ...FAST_POLL_OPTIONS, onAuth: vi.fn(), onPrompt: mockOnPrompt(""), }), @@ -276,6 +282,7 @@ describe("loginGitHubCopilot", () => { global.fetch = fetchMock as unknown as typeof fetch; const credentials = await loginGitHubCopilot({ + ...FAST_POLL_OPTIONS, onAuth: vi.fn(), onPrompt: mockOnPrompt(""), }); diff --git a/packages/ai/test/helpers/index.ts b/packages/ai/test/helpers/index.ts index 7f6b1158d..700f673b2 100644 --- a/packages/ai/test/helpers/index.ts +++ b/packages/ai/test/helpers/index.ts @@ -1,5 +1,8 @@ +import * as os from "node:os"; +import * as path from "node:path"; import { enrichModelThinking } from "@oh-my-pi/pi-ai/model-thinking"; import type { Model } from "@oh-my-pi/pi-ai/types"; +import { isEnoent } from "@oh-my-pi/pi-utils"; export async function withEnv( overrides: Record, @@ -65,3 +68,44 @@ export function createCodexModel(id: string): Model<"openai-codex-responses"> { maxTokens: 128000, }); } + +export interface AuthGatewayE2EStatus { + ok: boolean; + token?: string; + reason?: string; +} + +export const AUTH_GATEWAY_E2E_URL = Bun.env.OMP_E2E_GATEWAY_URL ?? "http://127.0.0.1:4000"; + +const AUTH_GATEWAY_TOKEN_PATH = path.join(os.homedir(), ".omp", "auth-gateway.token"); +const AUTH_GATEWAY_HEALTH_TIMEOUT_MS = 500; + +let authGatewayE2EStatus: Promise | undefined; + +export function checkAuthGatewayE2EAvailable(): Promise { + authGatewayE2EStatus ??= readAuthGatewayE2EStatus(); + return authGatewayE2EStatus; +} + +async function readAuthGatewayE2EStatus(): Promise { + if (!Bun.env.E2E) return { ok: false, reason: "E2E env not set" }; + let token: string; + try { + token = (await Bun.file(AUTH_GATEWAY_TOKEN_PATH).text()).trim(); + } catch (err) { + if (isEnoent(err)) return { ok: false, reason: `no token at ${AUTH_GATEWAY_TOKEN_PATH}` }; + throw err; + } + if (!token) return { ok: false, reason: `empty token at ${AUTH_GATEWAY_TOKEN_PATH}` }; + + try { + const res = await fetch(`${AUTH_GATEWAY_E2E_URL}/healthz`, { + signal: AbortSignal.timeout(AUTH_GATEWAY_HEALTH_TIMEOUT_MS), + }); + if (!res.ok) return { ok: false, reason: `healthz returned ${res.status}` }; + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + return { ok: false, reason: `healthz unreachable: ${msg}` }; + } + return { ok: true, token }; +} diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index 6b2fff79d..2aa63bd2b 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/controllers/input-controller.ts @@ -2,7 +2,7 @@ import * as fs from "node:fs/promises"; import { type AgentMessage, ThinkingLevel } from "@oh-my-pi/pi-agent-core"; import type { AutocompleteProvider, SlashCommand } from "@oh-my-pi/pi-tui"; import { $env, sanitizeText } from "@oh-my-pi/pi-utils"; -import { settings } from "../../config/settings"; +import { isSettingsInitialized, settings } from "../../config/settings"; import { expandEmoticons } from "../../modes/emoji-autocomplete"; import { createPromptActionAutocompleteProvider } from "../../modes/prompt-action-autocomplete"; import { theme } from "../../modes/theme/theme"; @@ -187,7 +187,7 @@ export class InputController { setupEditorSubmitHandler(): void { this.ctx.editor.onSubmit = async (text: string) => { text = text.trim(); - if (settings.get("emojiAutocomplete")) text = expandEmoticons(text); + if ((!isSettingsInitialized() || settings.get("emojiAutocomplete")) && text) text = expandEmoticons(text); // Empty submit while streaming with queued messages: flush queues immediately if (!text && this.ctx.session.isStreaming && this.ctx.session.queuedMessageCount > 0) { diff --git a/packages/coding-agent/src/modes/prompt-action-autocomplete.ts b/packages/coding-agent/src/modes/prompt-action-autocomplete.ts index 6a379aa28..2dc6211e9 100644 --- a/packages/coding-agent/src/modes/prompt-action-autocomplete.ts +++ b/packages/coding-agent/src/modes/prompt-action-autocomplete.ts @@ -6,7 +6,7 @@ import { type SlashCommand, } from "@oh-my-pi/pi-tui"; import { formatKeyHints, type KeybindingsManager } from "../config/keybindings"; -import { settings } from "../config/settings"; +import { isSettingsInitialized, settings } from "../config/settings"; import { applyEmojiCompletion, getEmojiSuggestions, isEmojiPrefix, tryEmojiInlineReplace } from "./emoji-autocomplete"; interface PromptActionDefinition { @@ -128,7 +128,7 @@ export class PromptActionAutocompleteProvider implements AutocompleteProvider { } } - if (settings.get("emojiAutocomplete")) { + if (!isSettingsInitialized() || settings.get("emojiAutocomplete")) { const emojiSuggestions = getEmojiSuggestions(textBeforeCursor); if (emojiSuggestions) return emojiSuggestions; } @@ -183,7 +183,7 @@ export class PromptActionAutocompleteProvider implements AutocompleteProvider { return this.#baseProvider.trySyncSlashCompletion?.(textBeforeCursor) ?? null; } trySyncInlineReplace(textBeforeCursor: string): { replaceLen: number; insert: string } | null { - if (!settings.get("emojiAutocomplete")) return null; + if (isSettingsInitialized() && !settings.get("emojiAutocomplete")) return null; return tryEmojiInlineReplace(textBeforeCursor); } } diff --git a/packages/coding-agent/test/acp-stdout-hygiene.test.ts b/packages/coding-agent/test/acp-stdout-hygiene.test.ts index 97d37532e..c9e1cc74e 100644 --- a/packages/coding-agent/test/acp-stdout-hygiene.test.ts +++ b/packages/coding-agent/test/acp-stdout-hygiene.test.ts @@ -10,21 +10,73 @@ import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; +type AcpProc = Bun.Subprocess<"pipe", "pipe", "pipe">; + const repoRoot = path.resolve(import.meta.dir, "..", "..", ".."); const cliEntry = path.join(repoRoot, "packages", "coding-agent", "src", "cli.ts"); const cleanupRoots: string[] = []; -let activeProc: ReturnType | undefined; +let activeProc: AcpProc | undefined; + +/** + * Tear the child down hard. SIGTERM first so the process gets a chance to + * unwind, but force-kill quickly if it hasn't reaped — `omp acp` blocks on + * stdin reads and won't notice SIGTERM until we close the pipes. We bound + * the entire shutdown to ~2s so a stuck child never trips Bun's 5s hook + * timeout (which is what produced the "afterEach hook timed out" flakes). + */ +async function teardown(proc: AcpProc): Promise { + // Close stdin so any blocking read in the child wakes up. + try { + proc.stdin.end(); + } catch { + // already closed + } + // Best-effort detach from stdout/stderr so the child's pipe writes don't + // block on a full buffer once we stop draining. + for (const stream of [proc.stdout, proc.stderr] as Array | undefined>) { + if (!stream) continue; + try { + await stream.cancel(); + } catch { + // reader may already be detached + } + } + + try { + proc.kill("SIGTERM"); + } catch { + // already exited + } + + // Race the natural exit against a short grace, then escalate to SIGKILL + // and race again against a hard cap. `await proc.exited` after SIGKILL + // always returns promptly on Darwin/Linux. + const graceMs = 200; + const hardCapMs = 1500; + const exited = proc.exited; + const raced = await Promise.race([ + exited.then(() => "exited" as const), + Bun.sleep(graceMs).then(() => "grace" as const), + ]); + if (raced === "exited") return; + try { + proc.kill("SIGKILL"); + } catch { + // already exited between the SIGTERM and SIGKILL + } + await Promise.race([exited, Bun.sleep(hardCapMs)]); +} afterEach(async () => { if (activeProc) { - try { - activeProc.kill(); - await activeProc.exited; - } catch { - // ignore - } + const proc = activeProc; activeProc = undefined; + try { + await teardown(proc); + } catch { + // teardown is already best-effort; never let cleanup fail a test + } } for (const root of cleanupRoots.splice(0)) { await fs.promises.rm(root, { recursive: true, force: true }); @@ -53,13 +105,16 @@ describe("ACP stdout hygiene", () => { it("emits a JSON-RPC initialize response as the first bytes on stdout", async () => { const root = await fs.promises.mkdtemp(path.join(os.tmpdir(), "omp-acp-stdout-")); cleanupRoots.push(root); - const home = path.join(root, "home"); const xdg = path.join(root, "xdg"); const agentDir = path.join(root, "agent"); - await fs.promises.mkdir(home, { recursive: true }); await fs.promises.mkdir(xdg, { recursive: true }); await fs.promises.mkdir(agentDir, { recursive: true }); + // NOTE: we intentionally do NOT override HOME. Bun keys its transpile + // cache at `$HOME/.bun/install/cache`; pointing HOME at a fresh tmp + // dir forces a full re-transpile of the CLI's module graph on every + // run (~12s cold vs ~0.4s warm). XDG_* and PI_CODING_AGENT_DIR + // already isolate PI's on-disk state for this smoke test. const proc = Bun.spawn(["bun", cliEntry, "acp"], { cwd: repoRoot, stdin: "pipe", @@ -67,15 +122,38 @@ describe("ACP stdout hygiene", () => { stderr: "pipe", env: { ...process.env, - HOME: home, XDG_DATA_HOME: xdg, XDG_CONFIG_HOME: xdg, PI_CODING_AGENT_DIR: agentDir, PI_NO_TITLE: "1", + NO_COLOR: "1", }, }); activeProc = proc; + // Buffer stderr in the background so we can assert no JSON-RPC frame + // leaks onto it. The pump exits the moment stderr closes, which + // happens during teardown — we never wait on it from the test body. + const stderrChunks: Uint8Array[] = []; + const stderrPump = (async () => { + const reader = proc.stderr.getReader(); + try { + while (true) { + const { value, done } = await reader.read(); + if (done) break; + if (value) stderrChunks.push(value); + } + } catch { + // reader cancelled by teardown — expected + } finally { + try { + reader.releaseLock(); + } catch { + // already released + } + } + })(); + const initRequest = { jsonrpc: "2.0", id: 1, @@ -85,27 +163,7 @@ describe("ACP stdout hygiene", () => { proc.stdin.write(new TextEncoder().encode(`${JSON.stringify(initRequest)}\n`)); proc.stdin.flush(); - // Capture stderr in parallel so we can verify it does not carry any - // JSON-RPC frame. ACP owns stdout; banners, progress text, or stray - // protocol bytes on stderr indicate a misroute. - const stderrChunks: Uint8Array[] = []; - const stderrPump = (async () => { - const reader = (proc.stderr as ReadableStream).getReader(); - try { - while (true) { - const { value, done } = await reader.read(); - if (done) break; - if (value) stderrChunks.push(value); - // Stop once the first stdout frame arrives so the pump terminates - // alongside the test rather than waiting for process exit. - if (stderrChunks.length > 32) break; - } - } finally { - reader.releaseLock(); - } - })(); - - const firstLine = await readFirstFrame(proc.stdout as ReadableStream); + const firstLine = await readFirstFrame(proc.stdout); expect(firstLine.length).toBeGreaterThan(0); expect(firstLine[0]).toBe("{"); @@ -126,18 +184,19 @@ describe("ACP stdout hygiene", () => { ]), ); - // Terminate the process so the stderr pump promise resolves. Race with a - // short timeout in case stderr is empty (common path). - try { - proc.kill(); - } catch { - // process may already be exiting - } - await Promise.race([stderrPump, new Promise(resolve => setTimeout(resolve, 500))]); - const stderrText = new TextDecoder().decode(new Uint8Array(stderrChunks.flatMap(chunk => Array.from(chunk)))); - // Guard against JSON-RPC frames sneaking onto stderr. We allow normal - // stderr output (warnings, telemetry, etc.) but reject anything that - // parses as a JSON-RPC envelope on the wrong channel. + // First frame is good. Tear the child down now so the test body's + // wall time is bounded by "boot + first frame", not by waiting for + // stderr or a delayed shutdown. teardown() closes stdin/stdout/stderr + // and escalates SIGTERM→SIGKILL, which both stops the child and + // resolves stderrPump. + await teardown(proc); + activeProc = undefined; + await stderrPump; + + const stderrText = Buffer.concat(stderrChunks).toString("utf8"); + // Guard against JSON-RPC frames sneaking onto stderr. Normal stderr + // output (warnings, telemetry, etc.) is allowed, but anything that + // parses as a JSON-RPC envelope on the wrong channel is a misroute. for (const line of stderrText.split("\n")) { const trimmed = line.trim(); if (!trimmed.startsWith("{")) continue; diff --git a/packages/coding-agent/test/agent-session-concurrent.test.ts b/packages/coding-agent/test/agent-session-concurrent.test.ts index 1da8dec3e..71f21ef13 100644 --- a/packages/coding-agent/test/agent-session-concurrent.test.ts +++ b/packages/coding-agent/test/agent-session-concurrent.test.ts @@ -64,14 +64,15 @@ describe("AgentSession concurrent prompt guard", () => { const stream = new AssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); - const checkAbort = () => { - if (abortSignal?.aborted) { - stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") }); - } else { - setTimeout(checkAbort, 5); - } - }; - checkAbort(); + if (abortSignal) { + abortSignal.addEventListener( + "abort", + () => { + stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") }); + }, + { once: true }, + ); + } }); return stream; }, @@ -110,11 +111,7 @@ describe("AgentSession concurrent prompt guard", () => { // Start first prompt (don't await, it will block until abort) const firstPrompt = session.prompt("First message"); - // Wait a tick for isStreaming to be set - await Bun.sleep(10); - - // Verify we're streaming - expect(session.isStreaming).toBe(true); + await waitFor(() => session.isStreaming); // Second prompt should reject await expect(session.prompt("Second message")).rejects.toBeInstanceOf(AgentBusyError); @@ -129,7 +126,7 @@ describe("AgentSession concurrent prompt guard", () => { // Start first prompt const firstPrompt = session.prompt("First message"); - await Bun.sleep(10); + await waitFor(() => session.isStreaming); // steer should work while streaming expect(() => session.steer("Steering message")).not.toThrow(); @@ -145,7 +142,7 @@ describe("AgentSession concurrent prompt guard", () => { // Start first prompt const firstPrompt = session.prompt("First message"); - await Bun.sleep(10); + await waitFor(() => session.isStreaming); // followUp should work while streaming expect(() => session.followUp("Follow-up message")).not.toThrow(); @@ -293,6 +290,15 @@ describe("AgentSession TTSR resume gate", () => { } }); + async function waitFor(predicate: () => boolean, timeoutMs = 500): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (predicate()) return; + await Bun.sleep(10); + } + + throw new Error("Timed out waiting for condition"); + } const testRule: Rule = { name: "no-unwrap", path: "/tmp/no-unwrap.md", @@ -322,18 +328,16 @@ describe("AgentSession TTSR resume gate", () => { } function pushContinuationStream(stream: AssistantMessageEventStream, onComplete: () => void): void { - setTimeout(() => { + queueMicrotask(() => { const partial = makeMsg(""); stream.push({ type: "start", partial }); - setTimeout(() => { - onComplete(); - stream.push({ - type: "done", - reason: "stop", - message: makeMsg('Fixed: let val = result.expect("msg")'), - }); - }, 80); - }, 10); + onComplete(); + stream.push({ + type: "done", + reason: "stop", + message: makeMsg('Fixed: let val = result.expect("msg")'), + }); + }); } function pushAbortableTtsrStream(stream: AssistantMessageEventStream, signal: AbortSignal | undefined): void { @@ -346,19 +350,19 @@ describe("AgentSession TTSR resume gate", () => { delta: "let val = result.unwrap(", partial: makeMsg("let val = result.unwrap("), }); - // TTSR abort should fire synchronously; poll for it - const checkAbort = () => { - if (signal?.aborted) { - stream.push({ - type: "error", - reason: "aborted", - error: makeMsg("let val = result.unwrap(", "aborted"), - }); - } else { - setTimeout(checkAbort, 2); - } - }; - checkAbort(); + if (signal) { + signal.addEventListener( + "abort", + () => { + stream.push({ + type: "error", + reason: "aborted", + error: makeMsg("let val = result.unwrap(", "aborted"), + }); + }, + { once: true }, + ); + } }); } @@ -525,18 +529,19 @@ describe("AgentSession TTSR resume gate", () => { delta: "result.unwrap(", partial: makeMsg("result.unwrap("), }); - const checkAbort = () => { - if (signal?.aborted) { - stream.push({ - type: "error", - reason: "aborted", - error: makeMsg("result.unwrap(", "aborted"), - }); - } else { - setTimeout(checkAbort, 2); - } - }; - checkAbort(); + if (signal) { + signal.addEventListener( + "abort", + () => { + stream.push({ + type: "error", + reason: "aborted", + error: makeMsg("result.unwrap(", "aborted"), + }); + }, + { once: true }, + ); + } }); return stream; @@ -560,9 +565,7 @@ describe("AgentSession TTSR resume gate", () => { // Start prompt (will trigger TTSR and create resume gate) const promptPromise = session.prompt("Write some Rust code"); - - // Wait for TTSR abort to be pending - await Bun.sleep(20); + await waitFor(() => session.isStreaming); // Abort session — prompt() should unblock await session.abort(); @@ -592,7 +595,6 @@ describe("AgentSession TTSR resume gate", () => { description: "A mock edit tool", parameters: z.object({}), execute: async () => { - await Bun.sleep(100); toolExecutionFinished = true; return { content: [{ type: "text" as const, text: "edit applied" }] }; }, @@ -638,19 +640,19 @@ describe("AgentSession TTSR resume gate", () => { pushAbortableTtsrStream(stream, signal); } else if (streamCallCount === 2) { // Continuation: return assistant message with a tool call - setTimeout(() => { + queueMicrotask(() => { const msg = makeToolCallMsg(); stream.push({ type: "start", partial: msg }); stream.push({ type: "done", reason: "toolUse", message: msg }); - }, 10); + }); } else { // After tool execution: return final response - setTimeout(() => { + queueMicrotask(() => { allTurnsCompleted = true; const msg = makeMsg('Fixed: let val = result.expect("msg")'); stream.push({ type: "start", partial: msg }); stream.push({ type: "done", reason: "stop", message: msg }); - }, 10); + }); } return stream; @@ -756,11 +758,11 @@ describe("AgentSession TTSR resume gate", () => { }); } else { // Continuation after tool result; finish cleanly. - setTimeout(() => { + queueMicrotask(() => { const done = makeMsg("ok"); stream.push({ type: "start", partial: done }); stream.push({ type: "done", reason: "stop", message: done }); - }, 10); + }); } return stream; }, @@ -892,11 +894,11 @@ describe("AgentSession TTSR resume gate", () => { stream.push({ type: "done", reason: "toolUse", message: partial }); }); } else { - setTimeout(() => { + queueMicrotask(() => { const done = makeMsg("ok"); stream.push({ type: "start", partial: done }); stream.push({ type: "done", reason: "stop", message: done }); - }, 10); + }); } return stream; }, @@ -997,12 +999,12 @@ describe("AgentSession TTSR resume gate", () => { stream.push({ type: "error", reason: "error", error: message }); }); } else { - setTimeout(() => { + queueMicrotask(() => { continuationCompleted = true; const message = makeSuccessMessage(); stream.push({ type: "start", partial: message }); stream.push({ type: "done", reason: "stop", message }); - }, 80); + }); } return stream; }, diff --git a/packages/coding-agent/test/agent-session-python-cleanup.test.ts b/packages/coding-agent/test/agent-session-python-cleanup.test.ts index b2a2d30dd..126d8051b 100644 --- a/packages/coding-agent/test/agent-session-python-cleanup.test.ts +++ b/packages/coding-agent/test/agent-session-python-cleanup.test.ts @@ -1,4 +1,4 @@ -import { afterEach, describe, expect, it, vi } from "bun:test"; +import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; @@ -8,7 +8,7 @@ import * as pythonExecutor from "@oh-my-pi/pi-coding-agent/eval/py/executor"; import type { PythonKernel as PythonKernelInstance } from "@oh-my-pi/pi-coding-agent/eval/py/kernel"; import * as pythonKernel from "@oh-my-pi/pi-coding-agent/eval/py/kernel"; import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; -import { createAgentSession, type ExtensionFactory } from "@oh-my-pi/pi-coding-agent/sdk"; +import { createAgentSession, type ExtensionFactory, type WorkspaceTree } from "@oh-my-pi/pi-coding-agent/sdk"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import { Snowflake } from "@oh-my-pi/pi-utils"; @@ -75,6 +75,23 @@ const createTempProject = () => { return { tempDir, cwd }; }; +const emptyWorkspaceTree = (cwd: string): WorkspaceTree => ({ + rootPath: cwd, + rendered: ".", + truncated: false, + totalLines: 1, + agentsMdFiles: [], +}); + +const mockPositiveSleepsImmediate = () => { + const realSleep = Bun.sleep.bind(Bun); + return vi.spyOn(Bun, "sleep").mockImplementation((duration?: number | Date) => { + if (typeof duration === "number" && duration > 0) { + return Promise.resolve(); + } + return realSleep(duration ?? 0); + }); +}; const createSession = async ( tempDir: string, cwd: string, @@ -92,6 +109,7 @@ const createSession = async ( skills: [], contextFiles: [], promptTemplates: [], + workspaceTree: emptyWorkspaceTree(cwd), slashCommands: [], enableMCP: false, enableLsp: false, @@ -117,8 +135,20 @@ const createMockKernel = () => { describe("AgentSession python cleanup", () => { const tempDirs: string[] = []; + let originalNullPrompt: string | undefined; + + beforeEach(() => { + originalNullPrompt = Bun.env.NULL_PROMPT; + Bun.env.NULL_PROMPT = "true"; + }); afterEach(async () => { + if (originalNullPrompt === undefined) { + delete Bun.env.NULL_PROMPT; + } else { + Bun.env.NULL_PROMPT = originalNullPrompt; + } + originalNullPrompt = undefined; vi.restoreAllMocks(); await pythonExecutor.disposeAllKernelSessions(); for (const tempDir of tempDirs.splice(0)) { @@ -162,6 +192,7 @@ describe("AgentSession python cleanup", () => { enableMCP: false, enableLsp: false, toolNames: ["eval"], + workspaceTree: emptyWorkspaceTree(cwd), }), ).rejects.toThrow("Extension init failed"); @@ -227,6 +258,7 @@ describe("AgentSession python cleanup", () => { enableMCP: false, enableLsp: false, toolNames: ["eval"], + workspaceTree: emptyWorkspaceTree(cwd), agentRegistry: throwingRegistry, }), ).rejects.toThrow("Agent registry failed"); @@ -374,19 +406,17 @@ describe("AgentSession python cleanup", () => { toolExecutionSettled = true; }); await blockedExecuteStarted.promise; + const sleepSpy = mockPositiveSleepsImmediate(); let disposed = false; const disposeSession = session.dispose().then(() => { disposed = true; }); - await Bun.sleep(0); - - expect(disposed).toBe(false); - expect(toolExecutionSettled).toBe(false); - expect(executeSpy).toHaveBeenCalledTimes(1); const [toolResult] = await Promise.all([toolExecution, disposeSession]); + expect(sleepSpy).toHaveBeenCalledWith(3000); + expect(disposed).toBe(true); expect(toolExecutionSettled).toBe(true); expect(executeSpy).toHaveBeenCalledTimes(1); @@ -408,6 +438,8 @@ describe("AgentSession python cleanup", () => { kernel.abortBlockedExecution = false; vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const sleepSpy = vi.spyOn(Bun, "sleep").mockResolvedValue(undefined); + const startSpy = vi .spyOn(pythonKernel.PythonKernel, "start") .mockResolvedValue(kernel as unknown as PythonKernelInstance); @@ -428,6 +460,7 @@ describe("AgentSession python cleanup", () => { firstDisposed = true; }); await disposeFirst; + expect(sleepSpy).toHaveBeenCalledWith(3000); expect(firstDisposed).toBe(true); expect(firstExecutionSettled).toBe(false); @@ -666,9 +699,10 @@ describe("AgentSession python cleanup", () => { const firstExecution = session.executePython("print('first')"); await blockedExecutionStarted.promise; const secondExecution = session.executePython("print('second')"); - await Bun.sleep(0); + const sleepSpy = mockPositiveSleepsImmediate(); await session.dispose(); + expect(sleepSpy).toHaveBeenCalledWith(3000); const [firstResult, secondResult] = await Promise.all([firstExecution, secondExecution]); expect(firstResult.cancelled).toBe(true); diff --git a/packages/coding-agent/test/agent-session-retry-fallback.test.ts b/packages/coding-agent/test/agent-session-retry-fallback.test.ts index d7d337874..2df771f09 100644 --- a/packages/coding-agent/test/agent-session-retry-fallback.test.ts +++ b/packages/coding-agent/test/agent-session-retry-fallback.test.ts @@ -1,4 +1,4 @@ -import { afterEach, beforeEach, describe, expect, it } from "bun:test"; +import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import * as path from "node:path"; import { Agent } from "@oh-my-pi/pi-agent-core"; import { type AssistantMessage, Effort, getBundledModel, type Model, writeModelCache } from "@oh-my-pi/pi-ai"; @@ -83,6 +83,7 @@ describe("AgentSession retry fallback", () => { } authStorage.close(); tempDir.removeSync(); + vi.restoreAllMocks(); }); it("advances through a role-keyed fallback chain across retries", async () => { @@ -569,6 +570,8 @@ describe("AgentSession retry fallback", () => { settings, modelRegistry, }); + let now = Date.now(); + vi.spyOn(Date, "now").mockImplementation(() => now); await session.prompt("First prompt triggers fallback"); await session.waitForIdle(); @@ -589,7 +592,7 @@ describe("AgentSession retry fallback", () => { expect(session.model?.provider).toBe(fallbackModel.provider); expect(session.model?.id).toBe(fallbackModel.id); - await Bun.sleep(240); + now += 240; await session.prompt("Third prompt should lazily revert to primary"); await session.waitForIdle(); expect(requestedModels).toEqual([ @@ -629,6 +632,8 @@ describe("AgentSession retry fallback", () => { modelRegistry, thinkingLevel: Effort.High, }); + let now = Date.now(); + vi.spyOn(Date, "now").mockImplementation(() => now); await session.prompt("First prompt triggers bare-selector fallback"); await session.waitForIdle(); @@ -641,7 +646,7 @@ describe("AgentSession retry fallback", () => { expect(session.thinkingLevel).toBeUndefined(); session.setThinkingLevel(Effort.Low); - await Bun.sleep(240); + now += 240; await session.prompt("Second prompt should restore model but preserve user thinking change"); await session.waitForIdle(); expect(requestedModels).toEqual([ diff --git a/packages/coding-agent/test/bash-acp-terminal.test.ts b/packages/coding-agent/test/bash-acp-terminal.test.ts index d97115e70..0323dec17 100644 --- a/packages/coding-agent/test/bash-acp-terminal.test.ts +++ b/packages/coding-agent/test/bash-acp-terminal.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it, spyOn } from "bun:test"; +import { afterEach, describe, expect, it, mock, spyOn } from "bun:test"; import type { ClientBridge, ClientBridgeTerminalHandle } from "../src/session/client-bridge"; import type { ToolSession } from "../src/tools"; import { BashTool } from "../src/tools/bash"; @@ -29,6 +29,10 @@ function makeSession(bridge: ClientBridge): ToolSession { } as unknown as ToolSession; } +afterEach(() => { + mock.restore(); +}); + describe("BashTool ACP terminal routing", () => { it("routes through bridge, emits terminalId update, and releases the handle", async () => { const stubText = "hello from terminal\n"; @@ -140,6 +144,8 @@ describe("BashTool ACP terminal routing", () => { const killSpy = spyOn(handle, "kill"); const releaseSpy = spyOn(handle, "release"); + spyOn(Bun, "sleep").mockImplementation(async () => {}); + const tool = new BashTool(makeSession(bridge)); await expect(tool.execute("call-timeout", { command: "sleep 60", timeout: 1 })).rejects.toThrow( diff --git a/packages/coding-agent/test/bash-executor.test.ts b/packages/coding-agent/test/bash-executor.test.ts index 730331bd2..f728efac4 100644 --- a/packages/coding-agent/test/bash-executor.test.ts +++ b/packages/coding-agent/test/bash-executor.test.ts @@ -10,6 +10,9 @@ import * as shellSnapshot from "@oh-my-pi/pi-coding-agent/utils/shell-snapshot"; // Matches the schema default for `tools.artifactHeadBytes` (20 KB) used by // OutputSink when bash-executor pulls settings via resolveOutputSinkHeadBytes. const ARTIFACT_HEAD_BYTES_DEFAULT = 20 * 1024; +const BACKGROUND_COMPLETION_RACE_MS = 750; +const KILL_MARKER_DELAY_SECONDS = "0.4"; +const KILL_MARKER_ASSERTION_WAIT_MS = 900; function makeTempDir(): string { return fs.mkdtempSync(path.join(os.tmpdir(), "omp-bash-exec-")); @@ -95,13 +98,18 @@ describe("executeBash", () => { if (process.platform === "win32") { return; } - const start = Date.now(); - const result = await executeBash("{ sleep 5; } & echo fg", { + const runPromise = executeBash("{ sleep 2; } & echo fg", { cwd: tempDir, timeout: 5000, }); - expect(result.output).toContain("fg"); - expect(Date.now() - start).toBeLessThan(3000); + const timed = await Promise.race([ + runPromise.then(result => ({ type: "result" as const, result })), + Bun.sleep(BACKGROUND_COMPLETION_RACE_MS).then(() => ({ type: "timeout" as const })), + ]); + expect(timed.type).toBe("result"); + if (timed.type === "result") { + expect(timed.result.output).toContain("fg"); + } }); it("returns a real PID for background external commands", async () => { @@ -369,13 +377,13 @@ describe("executeBash", () => { it("completes even when background job keeps stdout pipe open", async () => { if (process.platform === "win32") return; - const runPromise = executeBash("{ sleep 3; echo late; } & echo immediate", { + const runPromise = executeBash("{ sleep 2; echo late; } & echo immediate", { cwd: tempDir, timeout: 5000, }); const timed = await Promise.race([ runPromise.then(result => ({ type: "result" as const, result })), - Bun.sleep(1500).then(() => ({ type: "timeout" as const })), + Bun.sleep(BACKGROUND_COMPLETION_RACE_MS).then(() => ({ type: "timeout" as const })), ]); expect(timed.type).toBe("result"); @@ -389,17 +397,18 @@ describe("executeBash", () => { if (process.platform === "win32") return; const marker = path.join(tempDir, "marker.txt"); + const markerEscaped = marker.replace(/'/g, "'\\''"); - // Command creates marker after 2s, but we timeout after 100ms - const result = await executeBash(`sleep 2 && echo done > ${marker}`, { + // Command creates marker after a short delay, but we timeout before then. + const result = await executeBash(`sleep ${KILL_MARKER_DELAY_SECONDS} && echo done > '${markerEscaped}'`, { cwd: tempDir, timeout: 100, }); expect(result.cancelled).toBe(true); - // Wait longer than the command would have taken - await Bun.sleep(3000); + // Wait longer than the command would have needed to create the marker. + await Bun.sleep(KILL_MARKER_ASSERTION_WAIT_MS); // If process was killed (not orphaned), marker should NOT exist expect(fs.existsSync(marker)).toBe(false); @@ -411,14 +420,17 @@ describe("executeBash", () => { const marker = path.join(tempDir, "marker-bg.txt"); const markerEscaped = marker.replace(/'/g, "'\\''"); - const result = await executeBash(`{ sleep 2; echo done > '${markerEscaped}'; } & sleep 10`, { - cwd: tempDir, - timeout: 100, - }); + const result = await executeBash( + `{ sleep ${KILL_MARKER_DELAY_SECONDS}; echo done > '${markerEscaped}'; } & sleep 10`, + { + cwd: tempDir, + timeout: 100, + }, + ); expect(result.cancelled).toBe(true); - await Bun.sleep(3000); + await Bun.sleep(KILL_MARKER_ASSERTION_WAIT_MS); expect(fs.existsSync(marker)).toBe(false); }); @@ -429,19 +441,23 @@ describe("executeBash", () => { const markerEscaped = marker.replace(/'/g, "'\\''"); const controller = new AbortController(); - const promise = executeBash(`{ sleep 2; echo done > '${markerEscaped}'; } & sleep 10`, { - cwd: tempDir, - timeout: 10000, - signal: controller.signal, - }); + const promise = executeBash( + `{ sleep ${KILL_MARKER_DELAY_SECONDS}; echo done > '${markerEscaped}'; } & sleep 10`, + { + cwd: tempDir, + timeout: 10000, + signal: controller.signal, + }, + ); await Bun.sleep(100); controller.abort(); const result = await promise; expect(result.cancelled).toBe(true); + expect(result.output).toContain("Command cancelled"); - await Bun.sleep(3000); + await Bun.sleep(KILL_MARKER_ASSERTION_WAIT_MS); expect(fs.existsSync(marker)).toBe(false); }); @@ -449,24 +465,26 @@ describe("executeBash", () => { if (process.platform === "win32") return; const marker = path.join(tempDir, "marker.txt"); + const markerEscaped = marker.replace(/'/g, "'\\''"); const controller = new AbortController(); - // Command creates marker after 2s - const promise = executeBash(`sleep 2 && echo done > ${marker}`, { + // Command creates marker after a short delay. + const promise = executeBash(`sleep ${KILL_MARKER_DELAY_SECONDS} && echo done > '${markerEscaped}'`, { cwd: tempDir, timeout: 10000, signal: controller.signal, }); - // Abort after 100ms + // Abort before the command can create the marker. await Bun.sleep(100); controller.abort(); const result = await promise; expect(result.cancelled).toBe(true); + expect(result.output).toContain("Command cancelled"); - // Wait longer than the command would have taken - await Bun.sleep(3000); + // Wait longer than the command would have needed to create the marker. + await Bun.sleep(KILL_MARKER_ASSERTION_WAIT_MS); // If process was killed (not orphaned), marker should NOT exist expect(fs.existsSync(marker)).toBe(false); diff --git a/packages/coding-agent/test/core/js-executor.test.ts b/packages/coding-agent/test/core/js-executor.test.ts index d424437e6..3d33d369d 100644 --- a/packages/coding-agent/test/core/js-executor.test.ts +++ b/packages/coding-agent/test/core/js-executor.test.ts @@ -1,4 +1,4 @@ -import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; +import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "bun:test"; import * as path from "node:path"; import type { AgentTool, AgentToolResult } from "@oh-my-pi/pi-agent-core"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; @@ -42,7 +42,7 @@ describe("executeJs", () => { let sessionFile: string; let sessionId: string; - beforeEach(async () => { + beforeAll(async () => { tempDir = TempDir.createSync("@js-executor-"); sessionFile = path.join(tempDir.path(), "session.jsonl"); sessionId = `session:${sessionFile}:cwd:${tempDir.path()}`; @@ -61,10 +61,13 @@ describe("executeJs", () => { await Bun.write(path.join(tempDir.path(), "config.yaml"), "name: demo\nenabled: true\n"); }); - afterEach(async () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + afterAll(async () => { await disposeAllVmContexts(); tempDir.removeSync(); - vi.restoreAllMocks(); }); it("persists bindings across calls and reset clears them", async () => { @@ -417,19 +420,6 @@ describe("executeJs", () => { expect(result.displayOutputs).toEqual([{ type: "json", data: { answer: 42, nested: { ok: true } } }]); }); - it("cancels execution when the timeout expires", async () => { - const result = await executeJs("await new Promise(() => {})", { - sessionId, - session, - sessionFile, - timeoutMs: 20, - }); - - expect(result.cancelled).toBe(true); - expect(result.exitCode).toBeUndefined(); - expect(result.output).toContain("Command timed out"); - }); - it('rewrites static `import { x } from "pkg"` to dynamic import', async () => { const result = await executeJs('import { join } from "node:path";\nreturn join("a", "b");', { sessionId, @@ -460,4 +450,17 @@ describe("executeJs", () => { // No JSON display because structuredClone fails on the embedded function. expect(result.displayOutputs.filter(o => o.type === "json")).toHaveLength(0); }); + + it("cancels execution when the timeout expires", async () => { + const result = await executeJs("await new Promise(() => {})", { + sessionId, + session, + sessionFile, + timeoutMs: 20, + }); + + expect(result.cancelled).toBe(true); + expect(result.exitCode).toBeUndefined(); + expect(result.output).toContain("Command timed out"); + }); }); diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index fe236e18a..8f9f76a50 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -644,25 +644,25 @@ describe("ExtensionRunner", () => { runner.onError(err => { errors.push(err); }); - testSetExtensionHandlerTimeoutMs(50); + testSetExtensionHandlerTimeoutMs(10); const startedAt = performance.now(); await runner.emit({ type: "session_start" }); const elapsedMs = performance.now() - startedAt; - expect(elapsedMs).toBeGreaterThanOrEqual(40); - expect(elapsedMs).toBeLessThan(250); + expect(elapsedMs).toBeGreaterThanOrEqual(8); + expect(elapsedMs).toBeLessThan(150); expect(fs.readFileSync(markerPath, "utf8")).toBe("fast\n"); expect(warnSpy).toHaveBeenCalledWith("Extension handler timed out", { extensionPath: hangExtensionPath, event: "session_start", - timeoutMs: 50, + timeoutMs: 10, }); expect(errors).toEqual([ { extensionPath: hangExtensionPath, event: "session_start", - error: "handler timed out after 50ms", + error: "handler timed out after 10ms", }, ]); @@ -936,7 +936,7 @@ describe("ExtensionRunner", () => { ); // Drain microtasks so the fire-and-forget emit() calls inside initialize() complete. - await new Promise(resolve => setTimeout(resolve, 50)); + for (let i = 0; i < 5; i++) await Promise.resolve(); const events = fs .readFileSync(eventsPath, "utf8") diff --git a/packages/coding-agent/test/history-storage-search.test.ts b/packages/coding-agent/test/history-storage-search.test.ts index b6e05c838..ad78fb419 100644 --- a/packages/coding-agent/test/history-storage-search.test.ts +++ b/packages/coding-agent/test/history-storage-search.test.ts @@ -1,4 +1,4 @@ -import { afterEach, beforeEach, describe, expect, it } from "bun:test"; +import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; @@ -13,17 +13,19 @@ async function freshStorage(): Promise { } async function seed(storage: HistoryStorage, prompts: string[]): Promise { - for (const prompt of prompts) { - await storage.add(prompt, "/tmp/test"); - } + const writes = prompts.map(prompt => storage.add(prompt, "/tmp/test")); + vi.advanceTimersByTime(100); + await Promise.all(writes); } beforeEach(() => { HistoryStorage.resetInstance(); + vi.useFakeTimers(); }); afterEach(async () => { HistoryStorage.resetInstance(); + vi.useRealTimers(); if (tempDir) { await fs.rm(tempDir, { recursive: true, force: true }); tempDir = ""; diff --git a/packages/coding-agent/test/sdk-credential-disabled-bridge.test.ts b/packages/coding-agent/test/sdk-credential-disabled-bridge.test.ts index 020711ef8..fe0bd9e0b 100644 --- a/packages/coding-agent/test/sdk-credential-disabled-bridge.test.ts +++ b/packages/coding-agent/test/sdk-credential-disabled-bridge.test.ts @@ -118,13 +118,17 @@ describe("createAgentSession credential_disabled subscription", () => { if (events.length > waiters.length) { return Promise.resolve(events[waiters.length] as CredentialDisabledEvent); } - return new Promise(resolve => { - waiters.push({ resolve }); - }); + const { promise, resolve } = Promise.withResolvers(); + waiters.push({ resolve }); + return promise; }; return { factory, events, next }; }; + const drainCredentialDisabledDispatch = async (): Promise => { + for (let i = 0; i < 5; i++) await Promise.resolve(); + }; + afterEach(() => { vi.restoreAllMocks(); for (const dir of tempDirs.splice(0)) { @@ -190,8 +194,8 @@ describe("createAgentSession credential_disabled subscription", () => { // Post-dispose: only the embedder fires; the extension's listener was unsubscribed. await authStorage.set("openai", [expiredOAuth()]); await authStorage.getApiKey("openai", "post-dispose"); - // Allow any (non-existent) async listener microtasks a chance to run before asserting absence. - await Bun.sleep(20); + // Drain async dispatch turns before asserting absence. + await drainCredentialDisabledDispatch(); expect(embedderEvents).toEqual([ { provider: "anthropic", disabledCause: expect.stringContaining("invalid_grant") }, @@ -242,7 +246,7 @@ describe("createAgentSession credential_disabled subscription", () => { const wait2 = Promise.all([ext2.next(), ext3.next()]); await authStorage.getApiKey("openai", "concurrent-2"); await wait2; - await Bun.sleep(20); + await drainCredentialDisabledDispatch(); expect(embedderEvents.map(e => e.provider)).toEqual(["anthropic", "openai"]); expect(ext1.events.map(e => e.provider)).toEqual(["anthropic"]); expect(ext2.events.map(e => e.provider)).toEqual(["anthropic", "openai"]); @@ -255,7 +259,7 @@ describe("createAgentSession credential_disabled subscription", () => { const wait3 = ext3.next(); await authStorage.getApiKey("google", "concurrent-3"); await wait3; - await Bun.sleep(20); + await drainCredentialDisabledDispatch(); expect(embedderEvents.map(e => e.provider)).toEqual(["anthropic", "openai", "google"]); expect(ext1.events.map(e => e.provider)).toEqual(["anthropic"]); expect(ext2.events.map(e => e.provider)).toEqual(["anthropic", "openai"]); @@ -266,7 +270,7 @@ describe("createAgentSession credential_disabled subscription", () => { await authStorage.set("anthropic", [expiredOAuth()]); await authStorage.getApiKey("anthropic", "concurrent-final"); - await Bun.sleep(20); + await drainCredentialDisabledDispatch(); expect(embedderEvents.map(e => e.provider)).toEqual(["anthropic", "openai", "google", "anthropic"]); expect(ext1.events).toHaveLength(1); expect(ext2.events).toHaveLength(2); @@ -291,7 +295,7 @@ describe("createAgentSession credential_disabled subscription", () => { await authStorage.set("anthropic", [expiredOAuth()]); failOAuthRefresh(); await authStorage.getApiKey("anthropic", "pre-init"); - await Bun.sleep(20); + await drainCredentialDisabledDispatch(); expect(ext.events).toHaveLength(0); // Initializing flushes the buffer through `emit()` with the now-populated @@ -330,7 +334,7 @@ describe("createAgentSession credential_disabled subscription", () => { await authStorage.set("anthropic", [expiredOAuth()]); failOAuthRefresh(); await authStorage.getApiKey("anthropic", "startup-with-embedder"); - await Bun.sleep(20); + await drainCredentialDisabledDispatch(); // Embedder fires immediately (sync push from AuthStorage's fan-out loop). The // extension still hasn't received it because the runner is uninitialized. @@ -380,7 +384,7 @@ describe("createAgentSession credential_disabled subscription", () => { failOAuthRefresh(); await authStorage.set("anthropic", [expiredOAuth()]); await authStorage.getApiKey("anthropic", "post-failure"); - await Bun.sleep(20); + await drainCredentialDisabledDispatch(); expect(embedderEvents).toEqual([ { provider: "anthropic", disabledCause: expect.stringContaining("invalid_grant") }, diff --git a/packages/coding-agent/test/sdk-mcp-discovery.test.ts b/packages/coding-agent/test/sdk-mcp-discovery.test.ts index 395a9b449..2d5d754c5 100644 --- a/packages/coding-agent/test/sdk-mcp-discovery.test.ts +++ b/packages/coding-agent/test/sdk-mcp-discovery.test.ts @@ -3,7 +3,8 @@ import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; import { ThinkingLevel } from "@oh-my-pi/pi-agent-core"; -import { Effort, getBundledModel, type Model } from "@oh-my-pi/pi-ai"; +import { AuthStorage, Effort, getBundledModel, type Model } from "@oh-my-pi/pi-ai"; +import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import type { CustomTool } from "@oh-my-pi/pi-coding-agent/extensibility/custom-tools/types"; import { createAgentSession } from "@oh-my-pi/pi-coding-agent/sdk"; @@ -41,15 +42,22 @@ function createReasoningModel(): Model<"openai-responses"> { }; } +const oldSessionMtime = new Date("2000-01-01T00:00:00.000Z"); + describe("createAgentSession MCP discovery prompt gating", () => { let tempDir: string; + let authStorage: AuthStorage; + let modelRegistry: ModelRegistry; - beforeEach(() => { + beforeEach(async () => { tempDir = path.join(os.tmpdir(), `pi-sdk-mcp-discovery-${Snowflake.next()}`); fs.mkdirSync(tempDir, { recursive: true }); + authStorage = await AuthStorage.create(path.join(tempDir, "auth.db")); + modelRegistry = new ModelRegistry(authStorage); }); afterEach(() => { + authStorage.close(); if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true, force: true }); } @@ -59,6 +67,7 @@ describe("createAgentSession MCP discovery prompt gating", () => { const { session } = await createAgentSession({ cwd: tempDir, agentDir: tempDir, + modelRegistry, sessionManager: SessionManager.inMemory(), settings: Settings.isolated({ "mcp.discoveryMode": true }), model: getBundledModel("openai", "gpt-4o-mini"), @@ -83,6 +92,7 @@ describe("createAgentSession MCP discovery prompt gating", () => { const { session } = await createAgentSession({ cwd: tempDir, agentDir: tempDir, + modelRegistry, sessionManager: SessionManager.inMemory(), settings: Settings.isolated({ "tools.discoveryMode": "all" }), model: getBundledModel("openai", "gpt-4o-mini"), @@ -106,6 +116,7 @@ describe("createAgentSession MCP discovery prompt gating", () => { const { session } = await createAgentSession({ cwd: tempDir, agentDir: tempDir, + modelRegistry, sessionManager: SessionManager.inMemory(), settings: Settings.isolated({ "mcp.discoveryMode": true }), model: getBundledModel("openai", "gpt-4o-mini"), @@ -139,6 +150,7 @@ describe("createAgentSession MCP discovery prompt gating", () => { const { session } = await createAgentSession({ cwd: tempDir, agentDir: tempDir, + modelRegistry, sessionManager: SessionManager.inMemory(), settings: Settings.isolated({ "mcp.discoveryMode": true, @@ -173,6 +185,7 @@ describe("createAgentSession MCP discovery prompt gating", () => { const { session } = await createAgentSession({ cwd: tempDir, agentDir: tempDir, + modelRegistry, sessionManager: SessionManager.inMemory(), settings: Settings.isolated({ "mcp.discoveryMode": true }), model: getBundledModel("openai", "gpt-4o-mini"), @@ -196,6 +209,7 @@ describe("createAgentSession MCP discovery prompt gating", () => { const { session } = await createAgentSession({ cwd: tempDir, agentDir: tempDir, + modelRegistry, sessionManager: SessionManager.inMemory(), settings: Settings.isolated({ "tools.discoveryMode": "all" }), model: getBundledModel("openai", "gpt-4o-mini"), @@ -223,6 +237,7 @@ describe("createAgentSession MCP discovery prompt gating", () => { const { session: firstSession } = await createAgentSession({ cwd: tempDir, agentDir: tempDir, + modelRegistry, sessionManager: firstManager, settings: Settings.isolated({ "mcp.discoveryMode": true, @@ -251,14 +266,15 @@ describe("createAgentSession MCP discovery prompt gating", () => { const sessionFile = firstSession.sessionFile; expect(sessionFile).toBeDefined(); await firstSession.sessionManager.rewriteEntries(); + fs.utimesSync(sessionFile!, oldSessionMtime, oldSessionMtime); const persistedBeforeResume = fs.readFileSync(sessionFile!, "utf8"); const persistedMtimeBeforeResume = fs.statSync(sessionFile!).mtimeMs; - await Bun.sleep(20); await firstSession.dispose(); const resumedManager = await SessionManager.open(sessionFile!, tempDir); const { session: resumedSession } = await createAgentSession({ cwd: tempDir, agentDir: tempDir, + modelRegistry, sessionManager: resumedManager, settings: Settings.isolated({ "mcp.discoveryMode": true, @@ -304,13 +320,14 @@ describe("createAgentSession MCP discovery prompt gating", () => { const sessionFile = sessionManager.getSessionFile(); expect(sessionFile).toBeDefined(); await sessionManager.rewriteEntries(); + fs.utimesSync(sessionFile!, oldSessionMtime, oldSessionMtime); const persistedBeforeResume = fs.readFileSync(sessionFile!, "utf8"); const persistedMtimeBeforeResume = fs.statSync(sessionFile!).mtimeMs; - await Bun.sleep(20); const resumedManager = await SessionManager.open(sessionFile!, tempDir); const { session } = await createAgentSession({ cwd: tempDir, agentDir: tempDir, + modelRegistry, sessionManager: resumedManager, settings: Settings.isolated({ "mcp.discoveryMode": true, @@ -352,6 +369,7 @@ describe("createAgentSession MCP discovery prompt gating", () => { const { session: firstSession } = await createAgentSession({ cwd: tempDir, agentDir: tempDir, + modelRegistry, sessionManager: firstManager, settings: Settings.isolated({ "mcp.discoveryMode": true }), model: getBundledModel("openai", "gpt-4o-mini"), @@ -379,6 +397,7 @@ describe("createAgentSession MCP discovery prompt gating", () => { const { session: resumedSession } = await createAgentSession({ cwd: tempDir, agentDir: tempDir, + modelRegistry, sessionManager: resumedManager, settings: Settings.isolated({ "mcp.discoveryMode": true }), model: getBundledModel("openai", "gpt-4o-mini"), diff --git a/packages/coding-agent/test/tools.test.ts b/packages/coding-agent/test/tools.test.ts index 33957fdc6..8b18dbb8f 100644 --- a/packages/coding-agent/test/tools.test.ts +++ b/packages/coding-agent/test/tools.test.ts @@ -1085,7 +1085,7 @@ function b() { const updates: string[] = []; const result = await bashTool.execute( "test-call-8-stream", - { command: "for i in 1 2 3; do echo $i; sleep 0.2; done" }, + { command: "for i in 1 2 3; do echo $i; sleep 0.05; done" }, undefined, update => { const text = update.content?.find(c => c.type === "text")?.text ?? ""; @@ -1155,7 +1155,7 @@ function b() { expect(getTextOutput(result)).toContain("short"); expect(result.details?.timeoutSeconds).toBe(300); expect(result.details?.async).toBeUndefined(); - await Bun.sleep(150); + await asyncJobManager.drainDeliveries({ timeoutMs: 1 }); expect(deliveries).toEqual([]); await asyncJobManager.dispose(); }); @@ -1174,7 +1174,7 @@ function b() { testDir, Settings.isolated({ "bash.autoBackground.enabled": true, - "bash.autoBackground.thresholdMs": 50, + "bash.autoBackground.thresholdMs": 10, }), { getSessionId: () => "test-session", @@ -1184,7 +1184,7 @@ function b() { ); const result = await autoBackgroundBashTool.execute("test-call-9-auto-running", { - command: "printf 'start\\n'; sleep 0.2; printf 'done\\n'", + command: "printf 'start\\n'; sleep 0.05; printf 'done\\n'", }); expect(result.details?.async?.state).toBe("running"); @@ -1199,7 +1199,7 @@ function b() { const runningJob = asyncJobManager.getJob(jobId); expect(runningJob?.status).toBe("running"); await runningJob?.promise; - await Bun.sleep(50); + await asyncJobManager.drainDeliveries({ timeoutMs: 1 }); expect(deliveries).toHaveLength(1); expect(deliveries[0]?.jobId).toBe(jobId); expect(deliveries[0]?.text).toContain("done"); @@ -1244,7 +1244,7 @@ function b() { const runningJob = asyncJobManager.getJob(jobId); expect(runningJob?.status).toBe("running"); await runningJob?.promise; - await Bun.sleep(50); + await asyncJobManager.drainDeliveries({ timeoutMs: 1 }); expect(deliveries).toHaveLength(1); expect(deliveries[0]?.jobId).toBe(jobId); expect(deliveries[0]?.text).toContain("Command timed out after 1 seconds"); @@ -1269,9 +1269,16 @@ function b() { it("should abort and recover for subsequent commands", async () => { const controller = new AbortController(); - const promise = bashTool.execute("test-call-10-abort", { command: "sleep 5" }, controller.signal); - await Bun.sleep(200); - controller.abort("test abort"); + const promise = bashTool.execute( + "test-call-10-abort", + { command: "printf 'started\\n'; sleep 5" }, + controller.signal, + update => { + if (update.content?.some(content => content.type === "text" && content.text.includes("started"))) { + controller.abort("test abort"); + } + }, + ); await expect(promise).rejects.toThrow(/abort|cancel|timed out/i); const result = await bashTool.execute("test-call-10-after-abort", { command: "echo ok" }); diff --git a/packages/coding-agent/test/tools/gh.test.ts b/packages/coding-agent/test/tools/gh.test.ts index ec923698d..fc5f6d658 100644 --- a/packages/coding-agent/test/tools/gh.test.ts +++ b/packages/coding-agent/test/tools/gh.test.ts @@ -725,9 +725,6 @@ describe("github tool", () => { it("treats git.remote.add as a no-op when the remote already exists with the same URL", async () => { const fixture = await createPrFixture(); try { - // Fixture already created `forksrc -> forkBare`. A second add with the - // same URL must succeed silently — this is the cross-process / leftover- - // state path that used to fail with `error: remote forksrc already exists`. await git.remote.add(fixture.repoRoot, "forksrc", fixture.forkBare); expect(runGit(fixture.repoRoot, ["remote", "get-url", "forksrc"])).toBe(fixture.forkBare); } finally { @@ -751,17 +748,18 @@ describe("github tool", () => { it("serializes concurrent git mutations through withRepoLock so callers don't race git's internal locks", async () => { const fixture = await createPrFixture(); try { - // Without serialization, ~20 concurrent `git config` invocations against - // the same `.git/config` produce "could not lock config file" failures - // (the lock is O_EXCL with no waiter). Wrapping each write in - // `withRepoLock` makes the queue per-repo so all 20 succeed. - const writes = Array.from({ length: 20 }, (_, idx) => + // Without serialization, concurrent `git config` invocations against the + // same `.git/config` produce "could not lock config file" failures (the + // lock is O_EXCL with no waiter). Wrapping each write in `withRepoLock` + // makes the queue per-repo so all writes succeed. + const writeCount = 8; + const writes = Array.from({ length: writeCount }, (_, idx) => git.withRepoLock(fixture.repoRoot, () => git.config.set(fixture.repoRoot, `branch.race-test.key${idx}`, `value-${idx}`), ), ); await Promise.all(writes); - for (let idx = 0; idx < 20; idx += 1) { + for (let idx = 0; idx < writeCount; idx += 1) { expect(runGit(fixture.repoRoot, ["config", "--get", `branch.race-test.key${idx}`])).toBe(`value-${idx}`); } } finally { diff --git a/packages/natives/test/native.test.ts b/packages/natives/test/native.test.ts index 29380df91..bdb1cb14c 100644 --- a/packages/natives/test/native.test.ts +++ b/packages/natives/test/native.test.ts @@ -522,14 +522,14 @@ describe("pi-natives", () => { await fs.rm(markerPath, { force: true }); const result = await executeShell({ - command: `{ sleep 2; echo done > '${markerEscaped}'; } & sleep 10`, + command: `{ sleep 0.15; echo done > '${markerEscaped}'; } & sleep 10`, cwd: testDir, - timeoutMs: 100, + timeoutMs: 50, }); expect(result.timedOut).toBe(true); - await Bun.sleep(3000); + await Bun.sleep(500); expect(await Bun.file(markerPath).exists()).toBe(false); }); }); diff --git a/packages/tui/test/overlay-scroll.test.ts b/packages/tui/test/overlay-scroll.test.ts index 21f160670..42cafbc67 100644 --- a/packages/tui/test/overlay-scroll.test.ts +++ b/packages/tui/test/overlay-scroll.test.ts @@ -81,6 +81,12 @@ function longestBlankRun(lines: string[]): number { return longest; } +async function flushRender(term: VirtualTerminal): Promise { + await new Promise(resolve => process.nextTick(resolve)); + await Bun.sleep(17); + await term.flush(); +} + describe("TUI overlays", () => { it("does not scroll the terminal when an overlay is shown with a large historical working area", async () => { const term = new VirtualTerminal(80, 24); @@ -89,16 +95,14 @@ describe("TUI overlays", () => { tui.addChild(new LineComponent("base-", 5)); tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); // Simulate a large historical working area (max lines ever rendered) without actually // rendering that many lines in the current view. (tui as unknown as { maxLinesRendered: number }).maxLinesRendered = 1500; tui.showOverlay(new LineComponent("overlay-", 3), { anchor: "center" }); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); // The scroll buffer should stay small; we should not have printed hundreds/thousands of blank lines. expect(term.getScrollBuffer().length).toBeLessThan(200); @@ -107,19 +111,17 @@ describe("TUI overlays", () => { it("clears preexisting terminal scrollback on startup full redraw", async () => { const term = new VirtualTerminal(40, 4); term.write("shell-0\r\nshell-1\r\nshell-2\r\nshell-3\r\nshell-4\r\n"); - await term.waitForRender(); + await flushRender(term); const tui = new TUI(term); const component = new MutableContentComponent(["ui-0", "ui-1", "ui-2", "ui-3", "ui-4", "ui-5"]); tui.addChild(component); tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); term.resize(39, 4); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const scrollback = term.getScrollBuffer().join("\n"); expect(scrollback.includes("shell-0")).toBeFalsy(); @@ -134,15 +136,13 @@ describe("TUI overlays", () => { tui.addChild(component); tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const before = term.getScrollBuffer().join("\n"); expect(before.includes("row-0")).toBeTruthy(); tui.requestRender(true); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const after = term.getScrollBuffer().join("\n"); expect(after.includes("row-0")).toBeTruthy(); @@ -152,19 +152,17 @@ describe("TUI overlays", () => { it("fully redraws on height increase to avoid stale viewport rows", async () => { const term = new VirtualTerminal(40, 4); term.write("shell-0\r\nshell-1\r\nshell-2\r\nshell-3\r\nshell-4\r\n"); - await term.waitForRender(); + await flushRender(term); const tui = new TUI(term); const component = new MutableContentComponent(["ui-0", "ui-1", "ui-2", "ui-3"]); tui.addChild(component); tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); term.resize(40, 8); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const viewport = term.getViewport().join("\n"); expect(viewport.includes("shell-")).toBeFalsy(); @@ -178,14 +176,12 @@ describe("TUI overlays", () => { tui.addChild(component); try { tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const before = term.getScrollBuffer().length; for (let i = 0; i < 8; i++) { term.resize(i % 2 === 0 ? 59 : 60, i % 2 === 0 ? 9 : 8); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); } const after = term.getScrollBuffer().length; @@ -202,12 +198,10 @@ describe("TUI overlays", () => { tui.addChild(component); try { tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); component.setLines(Array.from({ length: 140 }, (_v, i) => `row-${i}`)); term.resize(59, 9); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const viewport = term.getViewport(); expect(viewport.at(-1)?.includes("row-139")).toBeTruthy(); } finally { @@ -218,16 +212,14 @@ describe("TUI overlays", () => { it("keeps scrollback on viewport-only resize redraw", async () => { const term = new VirtualTerminal(40, 4); term.write("shell-0\r\nshell-1\r\nshell-2\r\nshell-3\r\n"); - await term.waitForRender(); + await flushRender(term); const tui = new TUI(term); tui.addChild(new MutableContentComponent(["ui-0", "ui-1", "ui-2", "ui-3", "ui-4"])); try { tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); term.resize(39, 4); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const scrollback = term.getScrollBuffer().join("\n"); expect(scrollback.includes("shell-0")).toBeFalsy(); } finally { @@ -242,21 +234,19 @@ describe("TUI overlays", () => { tui.addChild(component); try { tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); - for (let count = 5; count <= 45; count++) { + for (let count = 5; count <= 29; count++) { component.setLines(buildRows(count)); term.resize(40, count % 2 === 0 ? 4 : 5); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); } const scrollbackLines = term.getScrollBuffer().map(line => line.trim()); expect(scrollbackLines).toContain("row-0"); - expect(scrollbackLines).toContain("row-20"); + expect(scrollbackLines).toContain("row-12"); const viewport = term.getViewport().map(line => line.trim()); - expect(viewport.at(-1)).toBe("row-44"); + expect(viewport.at(-1)).toBe("row-28"); } finally { tui.stop(); } @@ -265,30 +255,27 @@ describe("TUI overlays", () => { it("stays anchored across shrink-grow cycles while overflowing viewport", async () => { const term = new VirtualTerminal(30, 6); const tui = new TUI(term); - const component = new MutableContentComponent(Array.from({ length: 120 }, (_v, i) => `row-${i}`)); + const component = new MutableContentComponent(Array.from({ length: 64 }, (_v, i) => `row-${i}`)); tui.addChild(component); try { tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); - for (let cycle = 0; cycle < 5; cycle++) { - component.setLines(Array.from({ length: 120 - cycle * 8 }, (_v, i) => `row-${i}`)); + for (let cycle = 0; cycle < 3; cycle++) { + component.setLines(Array.from({ length: 64 - cycle * 8 }, (_v, i) => `row-${i}`)); tui.requestRender(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); - component.setLines(Array.from({ length: 120 - cycle * 8 + 4 }, (_v, i) => `row-${i}`)); + component.setLines(Array.from({ length: 64 - cycle * 8 + 4 }, (_v, i) => `row-${i}`)); tui.requestRender(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); } const viewport = term.getViewport().map(line => line.trim()); expect(viewport.every(line => /^row-\d+$/.test(line))).toBeTruthy(); const viewportRows = viewport.map(line => Number.parseInt(line.slice(4), 10)); - expect(viewportRows.at(-1)).toBe(91); - expect(viewportRows[0]).toBeGreaterThanOrEqual(80); + expect(viewportRows.at(-1)).toBe(51); + expect(viewportRows[0]).toBeGreaterThanOrEqual(40); } finally { tui.stop(); } @@ -301,15 +288,13 @@ describe("TUI overlays", () => { tui.addChild(component); try { tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const before = term.getScrollBuffer().length; for (let col = 0; col <= 10; col++) { component.setCursorCol(col); tui.requestRender(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); } const viewport = term.getViewport(); @@ -323,26 +308,24 @@ describe("TUI overlays", () => { it("limits scrollback growth during resize oscillation with overflowing content", async () => { const term = new VirtualTerminal(60, 10); const tui = new TUI(term); - const component = new MutableContentComponent(buildRows(320)); + const component = new MutableContentComponent(buildRows(160)); tui.addChild(component); try { tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const before = term.getScrollBuffer().length; - for (let i = 0; i < 80; i++) { - component.setLines(buildRows(280 + (i % 6) * 15)); + for (let i = 0; i < 18; i++) { + component.setLines(buildRows(140 + (i % 6) * 8)); term.resize(i % 2 === 0 ? 59 : 60, i % 3 === 0 ? 11 : 10); tui.requestRender(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const viewportRows = viewportRowNumbers(term); expect(viewportRows.length).toBeGreaterThan(0); } const scrollback = term.getScrollBuffer(); - expect(scrollback.length - before).toBeLessThan(700); + expect(scrollback.length - before).toBeLessThan(220); expect(longestBlankRun(scrollback)).toBeLessThan(30); } finally { tui.stop(); @@ -352,34 +335,30 @@ describe("TUI overlays", () => { it("limits scrollback while toggling overlays over overflowing content", async () => { const term = new VirtualTerminal(60, 10); const tui = new TUI(term); - const component = new MutableContentComponent(buildRows(300)); + const component = new MutableContentComponent(buildRows(150)); tui.addChild(component); try { tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const before = term.getScrollBuffer().length; - for (let i = 0; i < 50; i++) { + for (let i = 0; i < 12; i++) { const handle = tui.showOverlay(new LineComponent(`overlay-${i}-`, 3), { anchor: "center" }); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); handle.hide(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); - if (i % 5 === 0) { - component.setLines(buildRows(280 + (i % 4) * 10)); + if (i % 4 === 0) { + component.setLines(buildRows(140 + (i % 4) * 10)); tui.requestRender(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); } expect(viewportRowNumbers(term).length).toBeGreaterThan(0); } const scrollback = term.getScrollBuffer(); - expect(scrollback.length - before).toBeLessThan(1200); + expect(scrollback.length - before).toBeLessThan(320); expect(longestBlankRun(scrollback)).toBeLessThan(50); } finally { tui.stop(); @@ -389,23 +368,21 @@ describe("TUI overlays", () => { it("keeps scrollback bounded under rapid micro-resize oscillation", async () => { const term = new VirtualTerminal(80, 12); const tui = new TUI(term); - const component = new MutableContentComponent(buildRows(360)); + const component = new MutableContentComponent(buildRows(180)); tui.addChild(component); try { tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const before = term.getScrollBuffer().length; - for (let i = 0; i < 120; i++) { + for (let i = 0; i < 24; i++) { term.resize(i % 2 === 0 ? 79 : 80, i % 3 === 0 ? 11 : 12); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); expect(viewportRowNumbers(term).length).toBeGreaterThan(0); } const scrollback = term.getScrollBuffer(); - expect(scrollback.length - before).toBeLessThan(1300); + expect(scrollback.length - before).toBeLessThan(320); expect(longestBlankRun(scrollback)).toBeLessThan(60); } finally { tui.stop(); @@ -415,17 +392,15 @@ describe("TUI overlays", () => { it("avoids scrollback growth on repeated no-op renders with overflowing content", async () => { const term = new VirtualTerminal(70, 10); const tui = new TUI(term); - tui.addChild(new MutableContentComponent(buildRows(260))); + tui.addChild(new MutableContentComponent(buildRows(130))); try { tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const before = term.getScrollBuffer().length; - for (let i = 0; i < 80; i++) { + for (let i = 0; i < 16; i++) { tui.requestRender(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); } const scrollback = term.getScrollBuffer(); @@ -437,25 +412,23 @@ describe("TUI overlays", () => { it("stays stable with direct row-delta movement", async () => { const term = new VirtualTerminal(50, 10); const tui = new TUI(term); - const component = new MutableContentComponent(buildRows(260)); + const component = new MutableContentComponent(buildRows(150)); tui.addChild(component); try { tui.start(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); const before = term.getScrollBuffer().length; - for (let i = 0; i < 60; i++) { - component.setLines(buildRows(220 + (i % 8) * 12)); + for (let i = 0; i < 18; i++) { + component.setLines(buildRows(120 + (i % 8) * 6)); term.resize(i % 2 === 0 ? 50 : 49, i % 3 === 0 ? 11 : 10); tui.requestRender(); - await Bun.sleep(0); - await term.waitForRender(); + await flushRender(term); expect(viewportRowNumbers(term).length).toBeGreaterThan(0); } const scrollback = term.getScrollBuffer(); - expect(scrollback.length - before).toBeLessThan(900); + expect(scrollback.length - before).toBeLessThan(260); expect(longestBlankRun(scrollback)).toBeLessThan(40); } finally { tui.stop(); diff --git a/packages/tui/test/render-regressions.test.ts b/packages/tui/test/render-regressions.test.ts index 34656101d..3aebebd9a 100644 --- a/packages/tui/test/render-regressions.test.ts +++ b/packages/tui/test/render-regressions.test.ts @@ -1,4 +1,4 @@ -import { afterEach, describe, expect, it, vi } from "bun:test"; +import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import { type Component, TUI } from "@oh-my-pi/pi-tui"; import { VirtualTerminal } from "./virtual-terminal"; @@ -25,7 +25,9 @@ function rows(prefix: string, count: number): string[] { } async function settle(term: VirtualTerminal): Promise { - await term.waitForRender(); + await new Promise(resolve => process.nextTick(resolve)); + await Bun.sleep(1); + await term.flush(); } function visible(term: VirtualTerminal): string[] { @@ -41,6 +43,21 @@ function countMatches(lines: string[], pattern: RegExp): number { } describe("TUI terminal-state regressions", () => { + let monotonicNow = 0; + // Keep TUI's 16ms render throttle deterministic without sleeping a real frame per render. + + beforeEach(() => { + monotonicNow = 0; + vi.spyOn(performance, "now").mockImplementation(() => { + monotonicNow += 20; + return monotonicNow; + }); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + describe("cursor + differential stability", () => { it("keeps stable output across repeated no-op renders", async () => { const term = new VirtualTerminal(40, 10);