Merge main into model-registry
This commit is contained in:
Mario Zechner
@@ -5,6 +5,7 @@ import {
|
||||
type Model,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { estimateTokens } from "../../src/core/compaction/index.ts";
|
||||
import { createHarness, type Harness } from "./harness.ts";
|
||||
|
||||
type SessionWithCompactionInternals = {
|
||||
@@ -67,19 +68,20 @@ function useSummaryStreamFn(harness: Harness, summary: string): () => number {
|
||||
}
|
||||
|
||||
function seedCompactableSession(harness: Harness): void {
|
||||
harness.settingsManager.applyOverrides({ compaction: { keepRecentTokens: 1 } });
|
||||
const now = Date.now();
|
||||
harness.sessionManager.appendMessage({
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "message to compact" }],
|
||||
timestamp: now - 1000,
|
||||
});
|
||||
harness.sessionManager.appendMessage(
|
||||
createAssistant(harness, {
|
||||
stopReason: "stop",
|
||||
totalTokens: 100,
|
||||
timestamp: now - 500,
|
||||
}),
|
||||
);
|
||||
const assistant = createAssistant(harness, {
|
||||
stopReason: "stop",
|
||||
totalTokens: 100,
|
||||
timestamp: now - 500,
|
||||
});
|
||||
assistant.content = [{ type: "text", text: "assistant response to compact" }];
|
||||
harness.sessionManager.appendMessage(assistant);
|
||||
harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages;
|
||||
}
|
||||
|
||||
@@ -96,6 +98,7 @@ describe("AgentSession compaction characterization", () => {
|
||||
|
||||
it("manually compacts using an extension-provided summary", async () => {
|
||||
const harness = await createHarness({
|
||||
settings: { compaction: { keepRecentTokens: 1 } },
|
||||
extensionFactories: [
|
||||
(pi) => {
|
||||
pi.on("session_before_compact", async (event) => ({
|
||||
@@ -116,8 +119,10 @@ describe("AgentSession compaction characterization", () => {
|
||||
|
||||
const result = await harness.session.compact();
|
||||
const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction");
|
||||
const estimatedTokensAfter = harness.session.messages.reduce((sum, message) => sum + estimateTokens(message), 0);
|
||||
|
||||
expect(result.summary).toBe("summary from extension");
|
||||
expect(result.estimatedTokensAfter).toBe(estimatedTokensAfter);
|
||||
expect(compactionEntries).toHaveLength(1);
|
||||
expect(harness.session.messages[0]?.role).toBe("compactionSummary");
|
||||
});
|
||||
@@ -145,7 +150,7 @@ describe("AgentSession compaction characterization", () => {
|
||||
|
||||
const result = await harness.session.compact();
|
||||
|
||||
expect(result.summary).toBe("summary from custom stream");
|
||||
expect(result.summary).toContain("summary from custom stream");
|
||||
expect(getStreamCallCount()).toBe(1);
|
||||
});
|
||||
|
||||
@@ -159,12 +164,15 @@ describe("AgentSession compaction characterization", () => {
|
||||
await sessionInternals._runAutoCompaction("threshold", false);
|
||||
|
||||
const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction");
|
||||
const compactionEnd = harness.eventsOfType("compaction_end").at(-1);
|
||||
expect(compactionEntries).toHaveLength(1);
|
||||
expect(compactionEnd?.result?.estimatedTokensAfter).toBeGreaterThan(0);
|
||||
expect(getStreamCallCount()).toBe(1);
|
||||
});
|
||||
|
||||
it("cancels in-progress manual compaction when abortCompaction is called", async () => {
|
||||
const harness = await createHarness({
|
||||
settings: { compaction: { keepRecentTokens: 1 } },
|
||||
extensionFactories: [
|
||||
(pi) => {
|
||||
pi.on("session_before_compact", async (event) => {
|
||||
@@ -248,6 +256,37 @@ describe("AgentSession compaction characterization", () => {
|
||||
);
|
||||
});
|
||||
|
||||
it("compacts successful overflow responses without retrying", async () => {
|
||||
const harness = await createHarness({
|
||||
settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } },
|
||||
models: [{ id: "faux-1", contextWindow: 1, maxTokens: 100 }],
|
||||
extensionFactories: [
|
||||
(pi) => {
|
||||
pi.on("session_before_compact", async (event) => ({
|
||||
compaction: {
|
||||
summary: "successful overflow compacted",
|
||||
firstKeptEntryId: event.preparation.firstKeptEntryId,
|
||||
tokensBefore: event.preparation.tokensBefore,
|
||||
details: {},
|
||||
},
|
||||
}));
|
||||
},
|
||||
],
|
||||
});
|
||||
harnesses.push(harness);
|
||||
harness.setResponses([fauxAssistantMessage("completed answer")]);
|
||||
|
||||
await expect(harness.session.prompt("hello")).resolves.toBeUndefined();
|
||||
|
||||
const compactionEnd = harness.eventsOfType("compaction_end").at(-1);
|
||||
expect(compactionEnd).toMatchObject({
|
||||
reason: "overflow",
|
||||
aborted: false,
|
||||
willRetry: false,
|
||||
});
|
||||
expect(harness.faux.state.callCount).toBe(1);
|
||||
});
|
||||
|
||||
it("ignores stale pre-compaction assistant usage on pre-prompt checks", async () => {
|
||||
const harness = await createHarness();
|
||||
harnesses.push(harness);
|
||||
|
||||
+2
-1
@@ -21,6 +21,7 @@ type ShutdownThis = {
|
||||
unregisterSignalHandlers: () => void;
|
||||
runtimeHost: { dispose: () => Promise<void> };
|
||||
ui: { terminal: { drainInput: (ms: number) => Promise<void> } };
|
||||
themeController: { disableAutoSync: () => void };
|
||||
stop: () => void;
|
||||
sessionManager: SessionManager;
|
||||
};
|
||||
@@ -81,6 +82,7 @@ function createContext(order: string[], sessionManager = createSessionManager())
|
||||
}),
|
||||
},
|
||||
},
|
||||
themeController: { disableAutoSync: vi.fn() },
|
||||
stop: vi.fn(() => {
|
||||
order.push("stop");
|
||||
}),
|
||||
@@ -116,7 +118,6 @@ describe("InteractiveMode.shutdown ordering (#5080)", () => {
|
||||
|
||||
expect(order).toEqual(["dispose", "drainInput", "stop"]);
|
||||
expect(context.isShuttingDown).toBe(true);
|
||||
expect(context.unregisterSignalHandlers).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
test("interactive quit stops the TUI before emitting session_shutdown", async () => {
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { type BashOperations, createBashTool } from "../../../src/core/tools/bash.ts";
|
||||
|
||||
function getTextOutput(result: { content?: Array<{ type: string; text?: string }> }): string {
|
||||
return (
|
||||
result.content
|
||||
?.filter((block) => block.type === "text")
|
||||
.map((block) => block.text ?? "")
|
||||
.join("\n") ?? ""
|
||||
);
|
||||
}
|
||||
|
||||
describe("regression #5208: late bash output callbacks", () => {
|
||||
it("ignores output callbacks after bash operations resolve", async () => {
|
||||
const operations: BashOperations = {
|
||||
exec: async (_command, _cwd, { onData }) => {
|
||||
onData(Buffer.from("before\n", "utf-8"));
|
||||
setTimeout(() => onData(Buffer.from("late\n", "utf-8")), 0);
|
||||
return { exitCode: 0 };
|
||||
},
|
||||
};
|
||||
const bash = createBashTool(process.cwd(), { operations });
|
||||
|
||||
const result = await bash.execute("test-call-late-output", { command: "late-output" });
|
||||
await new Promise((resolve) => setTimeout(resolve, 20));
|
||||
|
||||
expect(getTextOutput(result).trim()).toBe("before");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,95 @@
|
||||
import { fauxAssistantMessage } from "@earendil-works/pi-ai";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import type { ExtensionFactory } from "../../../src/index.ts";
|
||||
import { createHarness, type Harness } from "../harness.ts";
|
||||
|
||||
type SessionWithCompactionInternals = {
|
||||
_runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise<boolean>;
|
||||
};
|
||||
|
||||
interface RecordedCompactionEvent {
|
||||
type: "session_before_compact" | "session_compact";
|
||||
reason: "manual" | "threshold" | "overflow";
|
||||
willRetry: boolean;
|
||||
}
|
||||
|
||||
function recordingExtension(recorded: RecordedCompactionEvent[]): ExtensionFactory {
|
||||
return (pi) => {
|
||||
pi.on("session_before_compact", async (event) => {
|
||||
recorded.push({ type: event.type, reason: event.reason, willRetry: event.willRetry });
|
||||
return {
|
||||
compaction: {
|
||||
summary: "summary from extension",
|
||||
firstKeptEntryId: event.preparation.firstKeptEntryId,
|
||||
tokensBefore: event.preparation.tokensBefore,
|
||||
details: {},
|
||||
},
|
||||
};
|
||||
});
|
||||
pi.on("session_compact", async (event) => {
|
||||
recorded.push({ type: event.type, reason: event.reason, willRetry: event.willRetry });
|
||||
});
|
||||
};
|
||||
}
|
||||
|
||||
async function createCompactionHarness(recorded: RecordedCompactionEvent[]): Promise<Harness> {
|
||||
const harness = await createHarness({
|
||||
settings: { compaction: { keepRecentTokens: 1 } },
|
||||
extensionFactories: [recordingExtension(recorded)],
|
||||
});
|
||||
harness.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two")]);
|
||||
await harness.session.prompt("first");
|
||||
await harness.session.prompt("second");
|
||||
return harness;
|
||||
}
|
||||
|
||||
describe("issue #5217 compaction reason on extension events", () => {
|
||||
const harnesses: Harness[] = [];
|
||||
|
||||
afterEach(() => {
|
||||
while (harnesses.length > 0) {
|
||||
harnesses.pop()?.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it("reports manual reason for compact()", async () => {
|
||||
const recorded: RecordedCompactionEvent[] = [];
|
||||
const harness = await createCompactionHarness(recorded);
|
||||
harnesses.push(harness);
|
||||
|
||||
await harness.session.compact();
|
||||
|
||||
expect(recorded).toEqual([
|
||||
{ type: "session_before_compact", reason: "manual", willRetry: false },
|
||||
{ type: "session_compact", reason: "manual", willRetry: false },
|
||||
]);
|
||||
});
|
||||
|
||||
it("reports threshold reason for auto-compaction", async () => {
|
||||
const recorded: RecordedCompactionEvent[] = [];
|
||||
const harness = await createCompactionHarness(recorded);
|
||||
harnesses.push(harness);
|
||||
const sessionInternals = harness.session as unknown as SessionWithCompactionInternals;
|
||||
|
||||
await sessionInternals._runAutoCompaction("threshold", false);
|
||||
|
||||
expect(recorded).toEqual([
|
||||
{ type: "session_before_compact", reason: "threshold", willRetry: false },
|
||||
{ type: "session_compact", reason: "threshold", willRetry: false },
|
||||
]);
|
||||
});
|
||||
|
||||
it("reports overflow reason and willRetry for overflow recovery", async () => {
|
||||
const recorded: RecordedCompactionEvent[] = [];
|
||||
const harness = await createCompactionHarness(recorded);
|
||||
harnesses.push(harness);
|
||||
const sessionInternals = harness.session as unknown as SessionWithCompactionInternals;
|
||||
|
||||
await sessionInternals._runAutoCompaction("overflow", true);
|
||||
|
||||
expect(recorded).toEqual([
|
||||
{ type: "session_before_compact", reason: "overflow", willRetry: true },
|
||||
{ type: "session_compact", reason: "overflow", willRetry: true },
|
||||
]);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,79 @@
|
||||
import type { ChildProcessByStdio } from "node:child_process";
|
||||
import type { Readable } from "node:stream";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import { spawnProcess, waitForChildProcess } from "../../../src/utils/child-process.ts";
|
||||
|
||||
/**
|
||||
* Regression test for https://github.com/earendil-works/pi/issues/5303
|
||||
*
|
||||
* waitForChildProcess armed a fixed 100ms timer on `exit` and destroyed the
|
||||
* stdio streams when it fired. When a short-lived detached descendant kept the
|
||||
* stdout pipe open, `close` never fired, so that timer was the only thing that
|
||||
* resolved the wait, and any output written more than 100ms after exit was
|
||||
* binned. In practice every git commit whose pre-commit hook runs lint-staged
|
||||
* came back truncated mid-listr2 output, read by the model as a hang.
|
||||
*
|
||||
* The fix re-arms the grace on each chunk, so an actively writing pipe keeps us
|
||||
* reading while a genuinely idle held-open handle still releases after the
|
||||
* grace elapses. Both behaviours are covered below.
|
||||
*/
|
||||
describe.skipIf(process.platform === "win32")("issue #5303 bash output truncation past exit", () => {
|
||||
let child: ChildProcessByStdio<null, Readable, Readable> | undefined;
|
||||
|
||||
afterEach(() => {
|
||||
if (child?.pid) {
|
||||
try {
|
||||
process.kill(-child.pid, "SIGKILL");
|
||||
} catch {
|
||||
// Already gone.
|
||||
}
|
||||
}
|
||||
child = undefined;
|
||||
});
|
||||
|
||||
it("captures output emitted after exit while a detached child holds stdout open", async () => {
|
||||
// The shell exits immediately, but a backgrounded subshell keeps the stdout
|
||||
// pipe open and emits ticks every 50ms, the last well past the 100ms grace.
|
||||
const command = 'printf "HEAD\\n"; ( for i in 1 2 3 4 5 6; do sleep 0.05; printf "TICK$i\\n"; done ) &';
|
||||
child = spawnProcess("/bin/sh", ["-c", command], {
|
||||
stdio: ["ignore", "pipe", "pipe"],
|
||||
detached: true,
|
||||
}) as ChildProcessByStdio<null, Readable, Readable>;
|
||||
|
||||
let output = "";
|
||||
child.stdout.on("data", (chunk: Buffer) => {
|
||||
output += chunk.toString();
|
||||
});
|
||||
|
||||
const exitCode = await waitForChildProcess(child);
|
||||
|
||||
expect(exitCode).toBe(0);
|
||||
expect(output).toContain("HEAD");
|
||||
expect(output).toContain("TICK6");
|
||||
});
|
||||
|
||||
it("resolves promptly when a detached child holds stdout open but stays quiet", async () => {
|
||||
// The shell exits, but a backgrounded sleeper inherits the stdout pipe and
|
||||
// keeps it open for a long time without writing. `close` never fires, so we
|
||||
// must still release via the idle grace rather than hang on the open handle.
|
||||
const command = 'printf "DONE\\n"; ( sleep 30 ) &';
|
||||
child = spawnProcess("/bin/sh", ["-c", command], {
|
||||
stdio: ["ignore", "pipe", "pipe"],
|
||||
detached: true,
|
||||
}) as ChildProcessByStdio<null, Readable, Readable>;
|
||||
|
||||
let output = "";
|
||||
child.stdout.on("data", (chunk: Buffer) => {
|
||||
output += chunk.toString();
|
||||
});
|
||||
|
||||
const start = Date.now();
|
||||
const exitCode = await waitForChildProcess(child);
|
||||
const elapsed = Date.now() - start;
|
||||
|
||||
expect(exitCode).toBe(0);
|
||||
expect(output).toContain("DONE");
|
||||
// Must not wait for the 30s sleeper; the idle grace releases us in well under a second.
|
||||
expect(elapsed).toBeLessThan(2000);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,89 @@
|
||||
import { existsSync, mkdtempSync, rmSync } from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { Agent } from "@earendil-works/pi-agent-core";
|
||||
import { fauxAssistantMessage, registerFauxProvider } from "@earendil-works/pi-ai";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import { AgentSession } from "../../../src/core/agent-session.ts";
|
||||
import { AuthStorage } from "../../../src/core/auth-storage.ts";
|
||||
import { convertToLlm } from "../../../src/core/messages.ts";
|
||||
import { ModelRegistry } from "../../../src/core/model-registry.ts";
|
||||
import { SessionManager } from "../../../src/core/session-manager.ts";
|
||||
import { SettingsManager } from "../../../src/core/settings-manager.ts";
|
||||
import { initTheme } from "../../../src/modes/interactive/theme/theme.ts";
|
||||
import { createTestResourceLoader } from "../../utilities.ts";
|
||||
|
||||
describe("regression #5596: missing configured theme export", () => {
|
||||
const cleanups: Array<() => void> = [];
|
||||
|
||||
afterEach(() => {
|
||||
while (cleanups.length > 0) {
|
||||
cleanups.pop()?.();
|
||||
}
|
||||
initTheme("dark");
|
||||
});
|
||||
|
||||
it("exports with the active fallback theme when the configured theme is missing", async () => {
|
||||
const tempDir = mkdtempSync(join(tmpdir(), "pi-5596-"));
|
||||
const faux = registerFauxProvider({
|
||||
models: [{ id: "faux-1", reasoning: false }],
|
||||
});
|
||||
faux.setResponses([fauxAssistantMessage("hello")]);
|
||||
|
||||
const model = faux.getModel();
|
||||
const authStorage = AuthStorage.inMemory();
|
||||
authStorage.setRuntimeApiKey(model.provider, "faux-key");
|
||||
const modelRegistry = ModelRegistry.inMemory(authStorage);
|
||||
modelRegistry.registerProvider(model.provider, {
|
||||
baseUrl: model.baseUrl,
|
||||
apiKey: "faux-key",
|
||||
api: faux.api,
|
||||
models: faux.models.map((registeredModel) => ({
|
||||
id: registeredModel.id,
|
||||
name: registeredModel.name,
|
||||
api: registeredModel.api,
|
||||
reasoning: registeredModel.reasoning,
|
||||
input: registeredModel.input,
|
||||
cost: registeredModel.cost,
|
||||
contextWindow: registeredModel.contextWindow,
|
||||
maxTokens: registeredModel.maxTokens,
|
||||
baseUrl: registeredModel.baseUrl,
|
||||
})),
|
||||
});
|
||||
|
||||
const settingsManager = SettingsManager.inMemory({ theme: "missing-theme" });
|
||||
const sessionManager = SessionManager.create(tempDir, join(tempDir, "sessions"));
|
||||
const agent = new Agent({
|
||||
getApiKey: () => "faux-key",
|
||||
initialState: {
|
||||
model,
|
||||
systemPrompt: "You are a test assistant.",
|
||||
tools: [],
|
||||
},
|
||||
convertToLlm,
|
||||
});
|
||||
const session = new AgentSession({
|
||||
agent,
|
||||
sessionManager,
|
||||
settingsManager,
|
||||
cwd: tempDir,
|
||||
modelRegistry,
|
||||
resourceLoader: createTestResourceLoader(),
|
||||
});
|
||||
cleanups.push(() => {
|
||||
session.dispose();
|
||||
faux.unregister();
|
||||
if (existsSync(tempDir)) {
|
||||
rmSync(tempDir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
await session.prompt("hi");
|
||||
initTheme(settingsManager.getTheme());
|
||||
|
||||
const outputPath = join(tempDir, "export.html");
|
||||
await expect(session.exportToHtml(outputPath)).resolves.toBe(outputPath);
|
||||
expect(existsSync(outputPath)).toBe(true);
|
||||
expect(settingsManager.getTheme()).toBe("missing-theme");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,91 @@
|
||||
import { readFileSync, writeFileSync } from "node:fs";
|
||||
import { join } from "node:path";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import { ENV_AGENT_DIR } from "../../../src/config.ts";
|
||||
import { AuthStorage } from "../../../src/core/auth-storage.ts";
|
||||
import { ModelRegistry } from "../../../src/core/model-registry.ts";
|
||||
import { runMigrations } from "../../../src/migrations.ts";
|
||||
import { createHarness } from "../harness.ts";
|
||||
|
||||
describe("regression #5661: uppercase models.json header values", () => {
|
||||
const cleanups: Array<() => void> = [];
|
||||
|
||||
afterEach(() => {
|
||||
while (cleanups.length > 0) {
|
||||
cleanups.pop()?.();
|
||||
}
|
||||
});
|
||||
|
||||
function withAgentDir(agentDir: string, fn: () => void): void {
|
||||
const previousAgentDir = process.env[ENV_AGENT_DIR];
|
||||
process.env[ENV_AGENT_DIR] = agentDir;
|
||||
try {
|
||||
fn();
|
||||
} finally {
|
||||
if (previousAgentDir === undefined) {
|
||||
delete process.env[ENV_AGENT_DIR];
|
||||
} else {
|
||||
process.env[ENV_AGENT_DIR] = previousAgentDir;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
it("keeps uppercase header strings as literals during startup migrations", async () => {
|
||||
const harness = await createHarness({ withConfiguredAuth: false });
|
||||
cleanups.push(harness.cleanup);
|
||||
|
||||
const envKeys = ["CUSTOM_API_KEY", "BEARER"];
|
||||
const savedEnv: Record<string, string | undefined> = {};
|
||||
for (const key of envKeys) {
|
||||
savedEnv[key] = process.env[key];
|
||||
process.env[key] = `env-${key}`;
|
||||
}
|
||||
cleanups.push(() => {
|
||||
for (const key of envKeys) {
|
||||
if (savedEnv[key] === undefined) {
|
||||
delete process.env[key];
|
||||
} else {
|
||||
process.env[key] = savedEnv[key];
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
const modelsPath = join(harness.tempDir, "models.json");
|
||||
writeFileSync(
|
||||
modelsPath,
|
||||
`${JSON.stringify(
|
||||
{
|
||||
providers: {
|
||||
"my-provider": {
|
||||
baseUrl: "https://example.com/v1",
|
||||
apiKey: "CUSTOM_API_KEY",
|
||||
api: "openai-completions",
|
||||
headers: { Authorization: "BEARER" },
|
||||
models: [{ id: "my-model" }],
|
||||
},
|
||||
},
|
||||
},
|
||||
null,
|
||||
2,
|
||||
)}\n`,
|
||||
"utf-8",
|
||||
);
|
||||
|
||||
withAgentDir(harness.tempDir, () => runMigrations(harness.tempDir));
|
||||
|
||||
const migrated = JSON.parse(readFileSync(modelsPath, "utf-8")) as {
|
||||
providers: Record<string, { apiKey?: string; headers?: Record<string, string> }>;
|
||||
};
|
||||
expect(migrated.providers["my-provider"]?.apiKey).toBe("CUSTOM_API_KEY");
|
||||
expect(migrated.providers["my-provider"]?.headers?.Authorization).toBe("BEARER");
|
||||
|
||||
const registry = ModelRegistry.create(AuthStorage.create(join(harness.tempDir, "auth.json")), modelsPath);
|
||||
const model = registry.find("my-provider", "my-model");
|
||||
expect(model).toBeDefined();
|
||||
expect(await registry.getApiKeyAndHeaders(model!)).toMatchObject({
|
||||
ok: true,
|
||||
apiKey: "CUSTOM_API_KEY",
|
||||
headers: { Authorization: "BEARER" },
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,94 @@
|
||||
import { afterEach, describe, expect, test, vi } from "vitest";
|
||||
import { InteractiveMode } from "../../../src/modes/interactive/interactive-mode.ts";
|
||||
|
||||
// Regression for https://github.com/earendil-works/pi/issues/5724
|
||||
//
|
||||
// `proper-lockfile` installs `signal-exit`, whose signal listener re-sends
|
||||
// SIGTERM/SIGHUP when it observes no other process listeners during the same
|
||||
// signal dispatch. InteractiveMode must therefore keep its signal handlers
|
||||
// registered until async terminal cleanup has completed.
|
||||
|
||||
type ShutdownThis = {
|
||||
isShuttingDown: boolean;
|
||||
unregisterSignalHandlers: () => void;
|
||||
runtimeHost: { dispose: () => Promise<void> };
|
||||
ui: { terminal: { drainInput: (ms: number) => Promise<void> } };
|
||||
themeController: { disableAutoSync: () => void };
|
||||
stop: () => void;
|
||||
};
|
||||
|
||||
type InteractiveModePrototypeWithShutdown = {
|
||||
shutdown(this: ShutdownThis, options?: { fromSignal?: boolean }): Promise<void>;
|
||||
};
|
||||
|
||||
const interactiveModePrototype = InteractiveMode.prototype as unknown;
|
||||
|
||||
class ProcessExitError extends Error {}
|
||||
|
||||
function deferred(): { promise: Promise<void>; resolve: () => void } {
|
||||
let resolve: (() => void) | undefined;
|
||||
const promise = new Promise<void>((res) => {
|
||||
resolve = res;
|
||||
});
|
||||
return {
|
||||
promise,
|
||||
resolve: () => resolve?.(),
|
||||
};
|
||||
}
|
||||
|
||||
async function callShutdown(context: ShutdownThis, options?: { fromSignal?: boolean }): Promise<void> {
|
||||
try {
|
||||
await (interactiveModePrototype as InteractiveModePrototypeWithShutdown).shutdown.call(context, options);
|
||||
} catch (error) {
|
||||
if (!(error instanceof ProcessExitError)) throw error;
|
||||
}
|
||||
}
|
||||
|
||||
describe("InteractiveMode SIGTERM shutdown with signal-exit (#5724)", () => {
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
test("keeps signal handlers registered while signal-triggered cleanup is pending", async () => {
|
||||
vi.spyOn(process, "exit").mockImplementation((() => {
|
||||
throw new ProcessExitError();
|
||||
}) as typeof process.exit);
|
||||
|
||||
const order: string[] = [];
|
||||
const dispose = deferred();
|
||||
const context: ShutdownThis = {
|
||||
isShuttingDown: false,
|
||||
unregisterSignalHandlers: vi.fn(() => {
|
||||
order.push("unregister");
|
||||
}),
|
||||
runtimeHost: {
|
||||
dispose: vi.fn(() => {
|
||||
order.push("dispose");
|
||||
return dispose.promise;
|
||||
}),
|
||||
},
|
||||
ui: {
|
||||
terminal: {
|
||||
drainInput: vi.fn(async () => {
|
||||
order.push("drainInput");
|
||||
}),
|
||||
},
|
||||
},
|
||||
themeController: { disableAutoSync: vi.fn() },
|
||||
stop: vi.fn(() => {
|
||||
order.push("stop");
|
||||
}),
|
||||
};
|
||||
|
||||
const shutdownPromise = callShutdown(context, { fromSignal: true });
|
||||
await Promise.resolve();
|
||||
|
||||
expect(order).toEqual(["dispose"]);
|
||||
expect(context.unregisterSignalHandlers).not.toHaveBeenCalled();
|
||||
|
||||
dispose.resolve();
|
||||
await shutdownPromise;
|
||||
|
||||
expect(order).toEqual(["dispose", "drainInput", "stop"]);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,113 @@
|
||||
import { afterEach, describe, expect, test, vi } from "vitest";
|
||||
import type { AgentSessionRuntime } from "../../../src/core/agent-session-runtime.ts";
|
||||
import { runRpcMode } from "../../../src/modes/rpc/rpc-mode.ts";
|
||||
import { createHarness, type Harness } from "../harness.ts";
|
||||
|
||||
// Regression for https://github.com/earendil-works/pi/issues/5868
|
||||
|
||||
const rpcIo = vi.hoisted(() => ({
|
||||
outputLines: [] as string[],
|
||||
lineHandler: undefined as ((line: string) => void) | undefined,
|
||||
}));
|
||||
|
||||
vi.mock("../../../src/core/output-guard.js", () => ({
|
||||
flushRawStdout: vi.fn(async () => {}),
|
||||
takeOverStdout: vi.fn(),
|
||||
waitForRawStdoutBackpressure: vi.fn(async () => {}),
|
||||
writeRawStdout: (line: string) => {
|
||||
rpcIo.outputLines.push(line);
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock("../../../src/modes/interactive/theme/theme.js", () => ({ theme: {} }));
|
||||
|
||||
vi.mock("../../../src/modes/rpc/jsonl.js", () => ({
|
||||
attachJsonlLineReader: vi.fn((_stream: NodeJS.ReadableStream, onLine: (line: string) => void) => {
|
||||
rpcIo.lineHandler = onLine;
|
||||
return () => {
|
||||
rpcIo.lineHandler = undefined;
|
||||
};
|
||||
}),
|
||||
serializeJsonLine: (value: unknown) => `${JSON.stringify(value)}\n`,
|
||||
}));
|
||||
|
||||
type NodeListener = Parameters<typeof process.on>[1];
|
||||
|
||||
type ListenerSnapshot = {
|
||||
stdinEnd: NodeListener[];
|
||||
signals: Map<NodeJS.Signals, NodeListener[]>;
|
||||
};
|
||||
|
||||
function takeListenerSnapshot(): ListenerSnapshot {
|
||||
const signals: NodeJS.Signals[] = process.platform === "win32" ? ["SIGTERM"] : ["SIGTERM", "SIGHUP"];
|
||||
return {
|
||||
stdinEnd: process.stdin.listeners("end") as NodeListener[],
|
||||
signals: new Map(signals.map((signal) => [signal, process.listeners(signal) as NodeListener[]])),
|
||||
};
|
||||
}
|
||||
|
||||
function restoreListeners(snapshot: ListenerSnapshot): void {
|
||||
for (const listener of process.stdin.listeners("end") as NodeListener[]) {
|
||||
if (!snapshot.stdinEnd.includes(listener)) {
|
||||
process.stdin.off("end", listener);
|
||||
}
|
||||
}
|
||||
|
||||
for (const [signal, previousListeners] of snapshot.signals) {
|
||||
for (const listener of process.listeners(signal) as NodeListener[]) {
|
||||
if (!previousListeners.includes(listener)) {
|
||||
process.off(signal, listener);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function parseOutputLines(): Array<Record<string, unknown>> {
|
||||
return rpcIo.outputLines
|
||||
.flatMap((line) => line.split("\n"))
|
||||
.filter((line) => line.trim().length > 0)
|
||||
.map((line) => JSON.parse(line) as Record<string, unknown>);
|
||||
}
|
||||
|
||||
function createRuntimeHost(harness: Harness): AgentSessionRuntime {
|
||||
return {
|
||||
session: harness.session,
|
||||
newSession: vi.fn(async () => ({ cancelled: true })),
|
||||
switchSession: vi.fn(async () => ({ cancelled: true })),
|
||||
fork: vi.fn(async () => ({ cancelled: true, selectedText: "" })),
|
||||
dispose: vi.fn(async () => {}),
|
||||
setRebindSession: vi.fn(),
|
||||
} as unknown as AgentSessionRuntime;
|
||||
}
|
||||
|
||||
describe("RPC unknown command responses (#5868)", () => {
|
||||
afterEach(() => {
|
||||
rpcIo.outputLines = [];
|
||||
rpcIo.lineHandler = undefined;
|
||||
});
|
||||
|
||||
test("preserves the request id on unknown command errors", async () => {
|
||||
const listenerSnapshot = takeListenerSnapshot();
|
||||
const harness = await createHarness();
|
||||
|
||||
try {
|
||||
void runRpcMode(createRuntimeHost(harness));
|
||||
await vi.waitFor(() => expect(rpcIo.lineHandler).toBeDefined());
|
||||
|
||||
rpcIo.lineHandler?.(JSON.stringify({ id: "test", type: "foobar" }));
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(parseOutputLines()).toContainEqual({
|
||||
id: "test",
|
||||
type: "response",
|
||||
command: "foobar",
|
||||
success: false,
|
||||
error: "Unknown command: foobar",
|
||||
});
|
||||
});
|
||||
} finally {
|
||||
harness.cleanup();
|
||||
restoreListeners(listenerSnapshot);
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,430 @@
|
||||
import { fauxAssistantMessage } from "@earendil-works/pi-ai";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import type { AgentSessionEvent } from "../../../src/core/agent-session.ts";
|
||||
import type { ExtensionUIContext } from "../../../src/core/extensions/index.ts";
|
||||
import { InteractiveMode } from "../../../src/modes/interactive/interactive-mode.ts";
|
||||
import { initTheme, type Theme, theme } from "../../../src/modes/interactive/theme/theme.ts";
|
||||
import { createHarness } from "../harness.ts";
|
||||
|
||||
function createUiContext(
|
||||
onNotify: (message: string, type: "info" | "warning" | "error" | undefined) => void,
|
||||
): ExtensionUIContext {
|
||||
return {
|
||||
select: async () => undefined,
|
||||
confirm: async () => false,
|
||||
input: async () => undefined,
|
||||
notify: onNotify,
|
||||
onTerminalInput: () => () => {},
|
||||
setStatus: () => {},
|
||||
setWorkingMessage: () => {},
|
||||
setWorkingVisible: () => {},
|
||||
setWorkingIndicator: () => {},
|
||||
setHiddenThinkingLabel: () => {},
|
||||
setWidget: () => {},
|
||||
setFooter: () => {},
|
||||
setHeader: () => {},
|
||||
setTitle: () => {},
|
||||
custom: async <T>() => undefined as T,
|
||||
pasteToEditor: () => {},
|
||||
setEditorText: () => {},
|
||||
getEditorText: () => "",
|
||||
editor: async () => undefined,
|
||||
addAutocompleteProvider: () => {},
|
||||
setEditorComponent: () => {},
|
||||
getEditorComponent: () => undefined,
|
||||
get theme() {
|
||||
return theme;
|
||||
},
|
||||
getAllThemes: () => [],
|
||||
getTheme: () => undefined,
|
||||
setTheme: (_theme: string | Theme) => ({ success: false, error: "Theme switching not available in tests" }),
|
||||
getToolsExpanded: () => false,
|
||||
setToolsExpanded: () => {},
|
||||
};
|
||||
}
|
||||
|
||||
type RebindContext = {
|
||||
unsubscribe?: () => void;
|
||||
applyRuntimeSettings: () => void;
|
||||
renderCurrentSessionState: () => void;
|
||||
bindCurrentSessionExtensions: () => Promise<void>;
|
||||
subscribeToAgent: () => void;
|
||||
updateAvailableProviderCount: () => Promise<void>;
|
||||
updateEditorBorderColor: () => void;
|
||||
updateTerminalTitle: () => void;
|
||||
};
|
||||
|
||||
type ReloadCommandContext = {
|
||||
hideThinkingBlock: boolean;
|
||||
session: {
|
||||
isStreaming: boolean;
|
||||
isCompacting: boolean;
|
||||
reload: (options?: { beforeSessionStart?: () => void | Promise<void> }) => Promise<void>;
|
||||
resourceLoader: { getThemes: () => { themes: [] } };
|
||||
extensionRunner: unknown;
|
||||
modelRegistry: { getError: () => string | undefined };
|
||||
};
|
||||
settingsManager: {
|
||||
getHttpIdleTimeoutMs: () => number;
|
||||
getHideThinkingBlock: () => boolean;
|
||||
getEditorPaddingX: () => number;
|
||||
getAutocompleteMaxVisible: () => number;
|
||||
getShowHardwareCursor: () => boolean;
|
||||
getClearOnShrink: () => boolean;
|
||||
};
|
||||
keybindings: { reload: () => void };
|
||||
customHeader?: unknown;
|
||||
builtInHeader?: unknown;
|
||||
editorContainer: { clear: () => void; addChild: (component: unknown) => void };
|
||||
ui: {
|
||||
setFocus: (component: unknown) => void;
|
||||
requestRender: (force?: boolean) => void;
|
||||
setShowHardwareCursor: (enabled: boolean) => void;
|
||||
setClearOnShrink: (enabled: boolean) => void;
|
||||
};
|
||||
editor: unknown;
|
||||
defaultEditor: { setPaddingX: (padding: number) => void; setAutocompleteMaxVisible: (maxVisible: number) => void };
|
||||
themeController: { applyFromSettings: () => Promise<void> };
|
||||
resetExtensionUI: () => void;
|
||||
rebuildChatFromMessages: () => void;
|
||||
setupAutocompleteProvider: () => void;
|
||||
setupExtensionShortcuts: (runner: unknown) => void;
|
||||
showLoadedResources: (options: unknown) => void;
|
||||
maybeSaveImplicitProjectTrustAfterReload: () => boolean;
|
||||
showStatus: (message: string) => void;
|
||||
showWarning: (message: string) => void;
|
||||
showError: (message: string) => void;
|
||||
};
|
||||
|
||||
type InteractiveModePrototype = {
|
||||
rebindCurrentSession(this: RebindContext, options?: { renderBeforeBind?: boolean }): Promise<void>;
|
||||
handleReloadCommand(this: ReloadCommandContext): Promise<void>;
|
||||
};
|
||||
|
||||
const interactiveModePrototype = InteractiveMode.prototype as unknown as InteractiveModePrototype;
|
||||
|
||||
type ReloadCommandContextOverrides = Omit<
|
||||
Partial<ReloadCommandContext>,
|
||||
"session" | "settingsManager" | "keybindings" | "editorContainer" | "ui" | "defaultEditor" | "themeController"
|
||||
> & {
|
||||
session?: Partial<ReloadCommandContext["session"]>;
|
||||
settingsManager?: Partial<ReloadCommandContext["settingsManager"]>;
|
||||
keybindings?: Partial<ReloadCommandContext["keybindings"]>;
|
||||
editorContainer?: Partial<ReloadCommandContext["editorContainer"]>;
|
||||
ui?: Partial<ReloadCommandContext["ui"]>;
|
||||
defaultEditor?: Partial<ReloadCommandContext["defaultEditor"]>;
|
||||
themeController?: Partial<ReloadCommandContext["themeController"]>;
|
||||
};
|
||||
|
||||
function createReloadCommandContext(overrides: ReloadCommandContextOverrides = {}): ReloadCommandContext {
|
||||
const editor = overrides.editor ?? {};
|
||||
return {
|
||||
hideThinkingBlock: overrides.hideThinkingBlock ?? false,
|
||||
session: {
|
||||
isStreaming: false,
|
||||
isCompacting: false,
|
||||
reload: async (options) => {
|
||||
await options?.beforeSessionStart?.();
|
||||
},
|
||||
resourceLoader: { getThemes: () => ({ themes: [] }) },
|
||||
extensionRunner: {},
|
||||
modelRegistry: { getError: () => undefined },
|
||||
...overrides.session,
|
||||
},
|
||||
settingsManager: {
|
||||
getHttpIdleTimeoutMs: () => 0,
|
||||
getHideThinkingBlock: () => false,
|
||||
getEditorPaddingX: () => 1,
|
||||
getAutocompleteMaxVisible: () => 10,
|
||||
getShowHardwareCursor: () => false,
|
||||
getClearOnShrink: () => false,
|
||||
...overrides.settingsManager,
|
||||
},
|
||||
keybindings: { reload: () => {}, ...overrides.keybindings },
|
||||
editorContainer: { clear: () => {}, addChild: () => {}, ...overrides.editorContainer },
|
||||
ui: {
|
||||
setFocus: () => {},
|
||||
requestRender: () => {},
|
||||
setShowHardwareCursor: () => {},
|
||||
setClearOnShrink: () => {},
|
||||
...overrides.ui,
|
||||
},
|
||||
editor,
|
||||
defaultEditor: { setPaddingX: () => {}, setAutocompleteMaxVisible: () => {}, ...overrides.defaultEditor },
|
||||
themeController: { applyFromSettings: async () => {}, ...overrides.themeController },
|
||||
customHeader: overrides.customHeader,
|
||||
builtInHeader: overrides.builtInHeader,
|
||||
resetExtensionUI: overrides.resetExtensionUI ?? (() => {}),
|
||||
rebuildChatFromMessages: overrides.rebuildChatFromMessages ?? (() => {}),
|
||||
setupAutocompleteProvider: overrides.setupAutocompleteProvider ?? (() => {}),
|
||||
setupExtensionShortcuts: overrides.setupExtensionShortcuts ?? (() => {}),
|
||||
showLoadedResources: overrides.showLoadedResources ?? (() => {}),
|
||||
maybeSaveImplicitProjectTrustAfterReload: overrides.maybeSaveImplicitProjectTrustAfterReload ?? (() => false),
|
||||
showStatus: overrides.showStatus ?? (() => {}),
|
||||
showWarning: overrides.showWarning ?? (() => {}),
|
||||
showError: overrides.showError ?? (() => {}),
|
||||
};
|
||||
}
|
||||
|
||||
type MessageEvent = Extract<AgentSessionEvent, { type: "message_start" | "message_end" }>;
|
||||
|
||||
function getMessageText(event: MessageEvent): string {
|
||||
const message = event.message;
|
||||
if (!("content" in message)) {
|
||||
return "";
|
||||
}
|
||||
const content = message.content;
|
||||
if (typeof content === "string") {
|
||||
return content;
|
||||
}
|
||||
return content
|
||||
.filter((part): part is { type: "text"; text: string } => part.type === "text")
|
||||
.map((part) => part.text)
|
||||
.join("");
|
||||
}
|
||||
|
||||
describe("regression #5943: session_start transient UI", () => {
|
||||
it("renders replacement session state before session_start handlers can notify", async () => {
|
||||
const events: string[] = [];
|
||||
const harness = await createHarness({
|
||||
extensionFactories: [
|
||||
(pi) => {
|
||||
pi.on("session_start", (_event, ctx) => {
|
||||
ctx.ui.notify("Hello Error", "error");
|
||||
});
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
try {
|
||||
const context: RebindContext = {
|
||||
applyRuntimeSettings: () => events.push("apply"),
|
||||
renderCurrentSessionState: () => events.push("render"),
|
||||
bindCurrentSessionExtensions: async () => {
|
||||
events.push("bind");
|
||||
await harness.session.bindExtensions({
|
||||
uiContext: createUiContext((message) => events.push(`notify:${message}`)),
|
||||
mode: "tui",
|
||||
});
|
||||
},
|
||||
subscribeToAgent: () => events.push("subscribe"),
|
||||
updateAvailableProviderCount: async () => {},
|
||||
updateEditorBorderColor: () => {},
|
||||
updateTerminalTitle: () => {},
|
||||
};
|
||||
|
||||
await interactiveModePrototype.rebindCurrentSession.call(context, { renderBeforeBind: true });
|
||||
|
||||
expect(events).toEqual(["apply", "render", "subscribe", "bind", "notify:Hello Error"]);
|
||||
} finally {
|
||||
harness.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it("subscribes before replacement session_start handlers send messages", async () => {
|
||||
const events: string[] = [];
|
||||
const harness = await createHarness({
|
||||
extensionFactories: [
|
||||
(pi) => {
|
||||
pi.on("session_start", () => {
|
||||
pi.sendMessage({
|
||||
customType: "session-start",
|
||||
content: "custom from start",
|
||||
display: true,
|
||||
});
|
||||
});
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
try {
|
||||
const context: RebindContext = {
|
||||
applyRuntimeSettings: () => {},
|
||||
renderCurrentSessionState: () => events.push("render"),
|
||||
bindCurrentSessionExtensions: async () => {
|
||||
events.push("bind");
|
||||
await harness.session.bindExtensions({
|
||||
uiContext: createUiContext(() => {}),
|
||||
mode: "tui",
|
||||
});
|
||||
},
|
||||
subscribeToAgent: () => {
|
||||
events.push("subscribe");
|
||||
harness.session.subscribe((event) => {
|
||||
if (event.type !== "message_start" && event.type !== "message_end") {
|
||||
return;
|
||||
}
|
||||
events.push(`${event.type}:${event.message.role}:${getMessageText(event)}`);
|
||||
});
|
||||
},
|
||||
updateAvailableProviderCount: async () => {},
|
||||
updateEditorBorderColor: () => {},
|
||||
updateTerminalTitle: () => {},
|
||||
};
|
||||
|
||||
await interactiveModePrototype.rebindCurrentSession.call(context, { renderBeforeBind: true });
|
||||
|
||||
expect(events).toEqual([
|
||||
"render",
|
||||
"subscribe",
|
||||
"bind",
|
||||
"message_start:custom:custom from start",
|
||||
"message_end:custom:custom from start",
|
||||
]);
|
||||
} finally {
|
||||
harness.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it("subscribes before replacement session_start handlers send user messages", async () => {
|
||||
const events: string[] = [];
|
||||
const harness = await createHarness({
|
||||
extensionFactories: [
|
||||
(pi) => {
|
||||
pi.on("session_start", () => {
|
||||
pi.sendUserMessage("user from start");
|
||||
});
|
||||
},
|
||||
],
|
||||
});
|
||||
harness.setResponses([fauxAssistantMessage("assistant from start")]);
|
||||
|
||||
try {
|
||||
const context: RebindContext = {
|
||||
applyRuntimeSettings: () => {},
|
||||
renderCurrentSessionState: () => events.push("render"),
|
||||
bindCurrentSessionExtensions: async () => {
|
||||
events.push("bind");
|
||||
await harness.session.bindExtensions({
|
||||
uiContext: createUiContext(() => {}),
|
||||
mode: "tui",
|
||||
});
|
||||
},
|
||||
subscribeToAgent: () => {
|
||||
events.push("subscribe");
|
||||
harness.session.subscribe((event) => {
|
||||
if (event.type !== "message_start" && event.type !== "message_end") {
|
||||
return;
|
||||
}
|
||||
events.push(`${event.type}:${event.message.role}:${getMessageText(event)}`);
|
||||
});
|
||||
},
|
||||
updateAvailableProviderCount: async () => {},
|
||||
updateEditorBorderColor: () => {},
|
||||
updateTerminalTitle: () => {},
|
||||
};
|
||||
|
||||
await interactiveModePrototype.rebindCurrentSession.call(context, { renderBeforeBind: true });
|
||||
await harness.session.agent.waitForIdle();
|
||||
|
||||
expect(events.slice(0, 3)).toEqual(["render", "subscribe", "bind"]);
|
||||
expect(events).toContain("message_start:user:user from start");
|
||||
expect(events).toContain("message_end:user:user from start");
|
||||
expect(events).toContain("message_end:assistant:assistant from start");
|
||||
} finally {
|
||||
harness.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it("runs the reload render hook before reload session_start handlers can notify", async () => {
|
||||
const events: string[] = [];
|
||||
const beforeSessionStart = vi.fn(() => {
|
||||
events.push("render");
|
||||
});
|
||||
const harness = await createHarness({
|
||||
extensionFactories: [
|
||||
(pi) => {
|
||||
pi.on("session_start", (event, ctx) => {
|
||||
events.push(`start:${event.reason}`);
|
||||
ctx.ui.notify(`notify:${event.reason}`, "error");
|
||||
});
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
try {
|
||||
await harness.session.bindExtensions({
|
||||
uiContext: createUiContext((message) => events.push(message)),
|
||||
mode: "tui",
|
||||
});
|
||||
expect(events).toEqual(["start:startup", "notify:startup"]);
|
||||
|
||||
events.length = 0;
|
||||
await harness.session.reload({ beforeSessionStart });
|
||||
|
||||
expect(beforeSessionStart).toHaveBeenCalledTimes(1);
|
||||
expect(events).toEqual(["render", "start:reload", "notify:reload"]);
|
||||
} finally {
|
||||
harness.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it("refreshes hideThinkingBlock before rebuilding chat during reload", async () => {
|
||||
initTheme("dark", false);
|
||||
const events: string[] = [];
|
||||
let context: ReloadCommandContext;
|
||||
context = createReloadCommandContext({
|
||||
settingsManager: { getHideThinkingBlock: () => true },
|
||||
session: {
|
||||
reload: async (options) => {
|
||||
events.push("reload");
|
||||
await options?.beforeSessionStart?.();
|
||||
events.push(`start:${context.hideThinkingBlock}`);
|
||||
},
|
||||
},
|
||||
rebuildChatFromMessages: () => {
|
||||
events.push(`rebuild:${context.hideThinkingBlock}`);
|
||||
},
|
||||
});
|
||||
|
||||
await interactiveModePrototype.handleReloadCommand.call(context);
|
||||
|
||||
expect(context.hideThinkingBlock).toBe(true);
|
||||
expect(events).toEqual(["reload", "rebuild:true", "start:true"]);
|
||||
});
|
||||
|
||||
it("keeps the reload blocker focused until async reload completes", async () => {
|
||||
initTheme("dark", false);
|
||||
const editor = {};
|
||||
let focused: unknown;
|
||||
let chatRestored = false;
|
||||
let markReloadWaiting!: () => void;
|
||||
let finishReload!: () => void;
|
||||
const reloadWaiting = new Promise<void>((resolve) => {
|
||||
markReloadWaiting = resolve;
|
||||
});
|
||||
const reloadFinished = new Promise<void>((resolve) => {
|
||||
finishReload = resolve;
|
||||
});
|
||||
|
||||
const context = createReloadCommandContext({
|
||||
editor,
|
||||
session: {
|
||||
reload: async (options) => {
|
||||
await options?.beforeSessionStart?.();
|
||||
markReloadWaiting();
|
||||
await reloadFinished;
|
||||
},
|
||||
},
|
||||
ui: {
|
||||
setFocus: (component) => {
|
||||
focused = component;
|
||||
},
|
||||
},
|
||||
rebuildChatFromMessages: () => {
|
||||
chatRestored = true;
|
||||
},
|
||||
});
|
||||
|
||||
const reloadPromise = interactiveModePrototype.handleReloadCommand.call(context);
|
||||
await reloadWaiting;
|
||||
|
||||
expect(chatRestored).toBe(true);
|
||||
expect(focused).not.toBe(editor);
|
||||
|
||||
finishReload();
|
||||
await reloadPromise;
|
||||
|
||||
expect(focused).toBe(editor);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,131 @@
|
||||
import { existsSync, mkdirSync, rmSync, writeFileSync } from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||
import { clearExtensionCache, loadExtensions, loadExtensionsCached } from "../../../src/core/extensions/loader.ts";
|
||||
import { DefaultResourceLoader } from "../../../src/core/resource-loader.ts";
|
||||
|
||||
interface TestState {
|
||||
moduleLoads?: number;
|
||||
factoryRuns?: number;
|
||||
}
|
||||
|
||||
function state(): TestState {
|
||||
const global = globalThis as typeof globalThis & { __extensionFactoryCacheTest?: TestState };
|
||||
if (!global.__extensionFactoryCacheTest) {
|
||||
global.__extensionFactoryCacheTest = {};
|
||||
}
|
||||
return global.__extensionFactoryCacheTest;
|
||||
}
|
||||
|
||||
function resetState(): void {
|
||||
delete (globalThis as typeof globalThis & { __extensionFactoryCacheTest?: TestState }).__extensionFactoryCacheTest;
|
||||
}
|
||||
|
||||
function writeCountingExtension(filePath: string): void {
|
||||
writeFileSync(
|
||||
filePath,
|
||||
`
|
||||
const state = (globalThis.__extensionFactoryCacheTest ??= {});
|
||||
state.moduleLoads = (state.moduleLoads ?? 0) + 1;
|
||||
|
||||
export default function () {
|
||||
state.factoryRuns = (state.factoryRuns ?? 0) + 1;
|
||||
}
|
||||
`,
|
||||
"utf-8",
|
||||
);
|
||||
}
|
||||
|
||||
describe("extension factory cache", () => {
|
||||
const roots: string[] = [];
|
||||
|
||||
function fixture(name: string) {
|
||||
const root = join(tmpdir(), `pi-extension-cache-${name}-${Date.now()}-${Math.random().toString(36).slice(2)}`);
|
||||
const cwd = join(root, "project");
|
||||
const agentDir = join(root, "agent");
|
||||
mkdirSync(cwd, { recursive: true });
|
||||
mkdirSync(agentDir, { recursive: true });
|
||||
roots.push(root);
|
||||
return { root, cwd, agentDir };
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
resetState();
|
||||
clearExtensionCache();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
while (roots.length > 0) {
|
||||
const root = roots.pop();
|
||||
if (root && existsSync(root)) {
|
||||
rmSync(root, { recursive: true, force: true });
|
||||
}
|
||||
}
|
||||
resetState();
|
||||
clearExtensionCache();
|
||||
});
|
||||
|
||||
it("caches extension modules for cached same-cwd loads but reruns factories", async () => {
|
||||
const { root, cwd } = fixture("same-cwd");
|
||||
const extensionPath = join(root, "counting.ts");
|
||||
writeCountingExtension(extensionPath);
|
||||
|
||||
const first = await loadExtensionsCached([extensionPath], cwd);
|
||||
const second = await loadExtensionsCached([extensionPath], cwd);
|
||||
|
||||
expect(state().moduleLoads).toBe(1);
|
||||
expect(state().factoryRuns).toBe(2);
|
||||
expect(first.extensions[0]).not.toBe(second.extensions[0]);
|
||||
expect(first.runtime).not.toBe(second.runtime);
|
||||
});
|
||||
|
||||
it("does not cache direct loadExtensions calls", async () => {
|
||||
const { root, cwd } = fixture("direct");
|
||||
const extensionPath = join(root, "counting.ts");
|
||||
writeCountingExtension(extensionPath);
|
||||
|
||||
await loadExtensions([extensionPath], cwd);
|
||||
await loadExtensions([extensionPath], cwd);
|
||||
|
||||
expect(state().moduleLoads).toBe(2);
|
||||
expect(state().factoryRuns).toBe(2);
|
||||
});
|
||||
|
||||
it("clears the cache on resource loader reload", async () => {
|
||||
const { cwd, agentDir } = fixture("reload");
|
||||
const extensionDir = join(agentDir, "extensions");
|
||||
mkdirSync(extensionDir, { recursive: true });
|
||||
writeCountingExtension(join(extensionDir, "counting.ts"));
|
||||
const loader = new DefaultResourceLoader({
|
||||
cwd,
|
||||
agentDir,
|
||||
noSkills: true,
|
||||
noPromptTemplates: true,
|
||||
noThemes: true,
|
||||
});
|
||||
|
||||
await loader.reload();
|
||||
await loader.reload();
|
||||
|
||||
expect(state().moduleLoads).toBe(2);
|
||||
expect(state().factoryRuns).toBe(2);
|
||||
});
|
||||
|
||||
it("keeps the cache scoped to one cwd", async () => {
|
||||
const { root } = fixture("cross-cwd");
|
||||
const firstCwd = join(root, "first");
|
||||
const secondCwd = join(root, "second");
|
||||
mkdirSync(firstCwd, { recursive: true });
|
||||
mkdirSync(secondCwd, { recursive: true });
|
||||
const extensionPath = join(root, "counting.ts");
|
||||
writeCountingExtension(extensionPath);
|
||||
|
||||
await loadExtensionsCached([extensionPath], firstCwd);
|
||||
await loadExtensionsCached([extensionPath], secondCwd);
|
||||
await loadExtensionsCached([extensionPath], secondCwd);
|
||||
|
||||
expect(state().moduleLoads).toBe(2);
|
||||
expect(state().factoryRuns).toBe(3);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user