refactor: share chat abort test helpers

This commit is contained in:
Peter Steinberger 2026-03-13 16:55:23 +00:00
parent 4a00cefe63
commit 8de94abfbc
3 changed files with 132 additions and 130 deletions

View File

@ -1,68 +1,24 @@
import { describe, expect, it, vi } from "vitest";
import { describe, expect, it } from "vitest";
import {
createActiveRun,
createChatAbortContext,
invokeChatAbortHandler,
} from "./chat.abort.test-helpers.js";
import { chatHandlers } from "./chat.js";
function createActiveRun(sessionKey: string, owner?: { connId?: string; deviceId?: string }) {
const now = Date.now();
return {
controller: new AbortController(),
sessionId: `${sessionKey}-session`,
sessionKey,
startedAtMs: now,
expiresAtMs: now + 30_000,
ownerConnId: owner?.connId,
ownerDeviceId: owner?.deviceId,
};
}
function createContext(overrides: Record<string, unknown> = {}) {
return {
chatAbortControllers: new Map(),
chatRunBuffers: new Map(),
chatDeltaSentAt: new Map(),
chatAbortedRuns: new Map<string, number>(),
removeChatRun: vi
.fn()
.mockImplementation((run: string) => ({ sessionKey: "main", clientRunId: run })),
agentRunSeq: new Map<string, number>(),
broadcast: vi.fn(),
nodeSendToSession: vi.fn(),
logGateway: { warn: vi.fn() },
...overrides,
};
}
async function invokeChatAbort(params: {
context: ReturnType<typeof createContext>;
request: { sessionKey: string; runId?: string };
client?: {
connId?: string;
connect?: {
device?: { id?: string };
scopes?: string[];
};
} | null;
}) {
const respond = vi.fn();
await chatHandlers["chat.abort"]({
params: params.request,
respond: respond as never,
context: params.context as never,
req: {} as never,
client: (params.client ?? null) as never,
isWebchatConnect: () => false,
});
return respond;
}
describe("chat.abort authorization", () => {
it("rejects explicit run aborts from other clients", async () => {
const context = createContext({
const context = createChatAbortContext({
chatAbortControllers: new Map([
["run-1", createActiveRun("main", { connId: "conn-owner", deviceId: "dev-owner" })],
[
"run-1",
createActiveRun("main", { owner: { connId: "conn-owner", deviceId: "dev-owner" } }),
],
]),
});
const respond = await invokeChatAbort({
const respond = await invokeChatAbortHandler({
handler: chatHandlers["chat.abort"],
context,
request: { sessionKey: "main", runId: "run-1" },
client: {
@ -79,13 +35,14 @@ describe("chat.abort authorization", () => {
});
it("allows the same paired device to abort after reconnecting", async () => {
const context = createContext({
const context = createChatAbortContext({
chatAbortControllers: new Map([
["run-1", createActiveRun("main", { connId: "conn-old", deviceId: "dev-1" })],
["run-1", createActiveRun("main", { owner: { connId: "conn-old", deviceId: "dev-1" } })],
]),
});
const respond = await invokeChatAbort({
const respond = await invokeChatAbortHandler({
handler: chatHandlers["chat.abort"],
context,
request: { sessionKey: "main", runId: "run-1" },
client: {
@ -101,14 +58,15 @@ describe("chat.abort authorization", () => {
});
it("only aborts session-scoped runs owned by the requester", async () => {
const context = createContext({
const context = createChatAbortContext({
chatAbortControllers: new Map([
["run-mine", createActiveRun("main", { deviceId: "dev-1" })],
["run-other", createActiveRun("main", { deviceId: "dev-2" })],
["run-mine", createActiveRun("main", { owner: { deviceId: "dev-1" } })],
["run-other", createActiveRun("main", { owner: { deviceId: "dev-2" } })],
]),
});
const respond = await invokeChatAbort({
const respond = await invokeChatAbortHandler({
handler: chatHandlers["chat.abort"],
context,
request: { sessionKey: "main" },
client: {
@ -125,13 +83,17 @@ describe("chat.abort authorization", () => {
});
it("allows operator.admin clients to bypass owner checks", async () => {
const context = createContext({
const context = createChatAbortContext({
chatAbortControllers: new Map([
["run-1", createActiveRun("main", { connId: "conn-owner", deviceId: "dev-owner" })],
[
"run-1",
createActiveRun("main", { owner: { connId: "conn-owner", deviceId: "dev-owner" } }),
],
]),
});
const respond = await invokeChatAbort({
const respond = await invokeChatAbortHandler({
handler: chatHandlers["chat.abort"],
context,
request: { sessionKey: "main", runId: "run-1" },
client: {

View File

@ -3,6 +3,11 @@ import os from "node:os";
import path from "node:path";
import { CURRENT_SESSION_VERSION } from "@mariozechner/pi-coding-agent";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
createActiveRun,
createChatAbortContext,
invokeChatAbortHandler,
} from "./chat.abort.test-helpers.js";
type TranscriptLine = {
message?: Record<string, unknown>;
@ -31,17 +36,6 @@ vi.mock("../session-utils.js", async (importOriginal) => {
const { chatHandlers } = await import("./chat.js");
function createActiveRun(sessionKey: string, sessionId: string) {
const now = Date.now();
return {
controller: new AbortController(),
sessionId,
sessionKey,
startedAtMs: now,
expiresAtMs: now + 30_000,
};
}
async function writeTranscriptHeader(transcriptPath: string, sessionId: string) {
const header = {
type: "session",
@ -81,49 +75,6 @@ async function createTranscriptFixture(prefix: string) {
return { transcriptPath, sessionId };
}
function createChatAbortContext(overrides: Record<string, unknown> = {}): {
chatAbortControllers: Map<string, ReturnType<typeof createActiveRun>>;
chatRunBuffers: Map<string, string>;
chatDeltaSentAt: Map<string, number>;
chatAbortedRuns: Map<string, number>;
removeChatRun: ReturnType<typeof vi.fn>;
agentRunSeq: Map<string, number>;
broadcast: ReturnType<typeof vi.fn>;
nodeSendToSession: ReturnType<typeof vi.fn>;
logGateway: { warn: ReturnType<typeof vi.fn> };
dedupe?: { get: ReturnType<typeof vi.fn> };
} {
return {
chatAbortControllers: new Map(),
chatRunBuffers: new Map(),
chatDeltaSentAt: new Map(),
chatAbortedRuns: new Map<string, number>(),
removeChatRun: vi
.fn()
.mockImplementation((run: string) => ({ sessionKey: "main", clientRunId: run })),
agentRunSeq: new Map<string, number>(),
broadcast: vi.fn(),
nodeSendToSession: vi.fn(),
logGateway: { warn: vi.fn() },
...overrides,
};
}
async function invokeChatAbort(
context: ReturnType<typeof createChatAbortContext>,
params: { sessionKey: string; runId?: string },
respond: ReturnType<typeof vi.fn>,
) {
await chatHandlers["chat.abort"]({
params,
respond: respond as never,
context: context as never,
req: {} as never,
client: null,
isWebchatConnect: () => false,
});
}
afterEach(() => {
vi.restoreAllMocks();
});
@ -134,7 +85,7 @@ describe("chat abort transcript persistence", () => {
const runId = "idem-abort-run-1";
const respond = vi.fn();
const context = createChatAbortContext({
chatAbortControllers: new Map([[runId, createActiveRun("main", sessionId)]]),
chatAbortControllers: new Map([[runId, createActiveRun("main", { sessionId })]]),
chatRunBuffers: new Map([[runId, "Partial from run abort"]]),
chatDeltaSentAt: new Map([[runId, Date.now()]]),
removeChatRun: vi
@ -149,17 +100,27 @@ describe("chat abort transcript persistence", () => {
logGateway: { warn: vi.fn() },
});
await invokeChatAbort(context, { sessionKey: "main", runId }, respond);
await invokeChatAbortHandler({
handler: chatHandlers["chat.abort"],
context,
request: { sessionKey: "main", runId },
respond,
});
const [ok1, payload1] = respond.mock.calls.at(-1) ?? [];
expect(ok1).toBe(true);
expect(payload1).toMatchObject({ aborted: true, runIds: [runId] });
context.chatAbortControllers.set(runId, createActiveRun("main", sessionId));
context.chatAbortControllers.set(runId, createActiveRun("main", { sessionId }));
context.chatRunBuffers.set(runId, "Partial from run abort");
context.chatDeltaSentAt.set(runId, Date.now());
await invokeChatAbort(context, { sessionKey: "main", runId }, respond);
await invokeChatAbortHandler({
handler: chatHandlers["chat.abort"],
context,
request: { sessionKey: "main", runId },
respond,
});
const lines = await readTranscriptLines(transcriptPath);
const persisted = lines
@ -188,8 +149,8 @@ describe("chat abort transcript persistence", () => {
const respond = vi.fn();
const context = createChatAbortContext({
chatAbortControllers: new Map([
["run-a", createActiveRun("main", sessionId)],
["run-b", createActiveRun("main", sessionId)],
["run-a", createActiveRun("main", { sessionId })],
["run-b", createActiveRun("main", { sessionId })],
]),
chatRunBuffers: new Map([
["run-a", "Session abort partial"],
@ -201,7 +162,12 @@ describe("chat abort transcript persistence", () => {
]),
});
await invokeChatAbort(context, { sessionKey: "main" }, respond);
await invokeChatAbortHandler({
handler: chatHandlers["chat.abort"],
context,
request: { sessionKey: "main" },
respond,
});
const [ok, payload] = respond.mock.calls.at(-1) ?? [];
expect(ok).toBe(true);
@ -280,12 +246,17 @@ describe("chat abort transcript persistence", () => {
const runId = "idem-abort-run-blank";
const respond = vi.fn();
const context = createChatAbortContext({
chatAbortControllers: new Map([[runId, createActiveRun("main", sessionId)]]),
chatAbortControllers: new Map([[runId, createActiveRun("main", { sessionId })]]),
chatRunBuffers: new Map([[runId, " \n\t "]]),
chatDeltaSentAt: new Map([[runId, Date.now()]]),
});
await invokeChatAbort(context, { sessionKey: "main", runId }, respond);
await invokeChatAbortHandler({
handler: chatHandlers["chat.abort"],
context,
request: { sessionKey: "main", runId },
respond,
});
const [ok, payload] = respond.mock.calls.at(-1) ?? [];
expect(ok).toBe(true);

View File

@ -0,0 +1,69 @@
import { vi } from "vitest";
export function createActiveRun(
sessionKey: string,
params: {
sessionId?: string;
owner?: { connId?: string; deviceId?: string };
} = {},
) {
const now = Date.now();
return {
controller: new AbortController(),
sessionId: params.sessionId ?? `${sessionKey}-session`,
sessionKey,
startedAtMs: now,
expiresAtMs: now + 30_000,
ownerConnId: params.owner?.connId,
ownerDeviceId: params.owner?.deviceId,
};
}
export function createChatAbortContext(overrides: Record<string, unknown> = {}) {
return {
chatAbortControllers: new Map(),
chatRunBuffers: new Map(),
chatDeltaSentAt: new Map(),
chatAbortedRuns: new Map<string, number>(),
removeChatRun: vi
.fn()
.mockImplementation((run: string) => ({ sessionKey: "main", clientRunId: run })),
agentRunSeq: new Map<string, number>(),
broadcast: vi.fn(),
nodeSendToSession: vi.fn(),
logGateway: { warn: vi.fn() },
...overrides,
};
}
export async function invokeChatAbortHandler(params: {
handler: (args: {
params: { sessionKey: string; runId?: string };
respond: never;
context: never;
req: never;
client: never;
isWebchatConnect: () => boolean;
}) => Promise<void>;
context: ReturnType<typeof createChatAbortContext>;
request: { sessionKey: string; runId?: string };
client?: {
connId?: string;
connect?: {
device?: { id?: string };
scopes?: string[];
};
} | null;
respond?: ReturnType<typeof vi.fn>;
}) {
const respond = params.respond ?? vi.fn();
await params.handler({
params: params.request,
respond: respond as never,
context: params.context as never,
req: {} as never,
client: (params.client ?? null) as never,
isWebchatConnect: () => false,
});
return respond;
}