diff --git a/src/gateway/server-methods/chat.abort-authorization.test.ts b/src/gateway/server-methods/chat.abort-authorization.test.ts index 6fbf0478df3..607e80b58ff 100644 --- a/src/gateway/server-methods/chat.abort-authorization.test.ts +++ b/src/gateway/server-methods/chat.abort-authorization.test.ts @@ -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 = {}) { - return { - chatAbortControllers: new Map(), - chatRunBuffers: new Map(), - chatDeltaSentAt: new Map(), - chatAbortedRuns: new Map(), - removeChatRun: vi - .fn() - .mockImplementation((run: string) => ({ sessionKey: "main", clientRunId: run })), - agentRunSeq: new Map(), - broadcast: vi.fn(), - nodeSendToSession: vi.fn(), - logGateway: { warn: vi.fn() }, - ...overrides, - }; -} - -async function invokeChatAbort(params: { - context: ReturnType; - 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: { diff --git a/src/gateway/server-methods/chat.abort-persistence.test.ts b/src/gateway/server-methods/chat.abort-persistence.test.ts index b7add3740eb..31a00a3f186 100644 --- a/src/gateway/server-methods/chat.abort-persistence.test.ts +++ b/src/gateway/server-methods/chat.abort-persistence.test.ts @@ -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; @@ -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 = {}): { - chatAbortControllers: Map>; - chatRunBuffers: Map; - chatDeltaSentAt: Map; - chatAbortedRuns: Map; - removeChatRun: ReturnType; - agentRunSeq: Map; - broadcast: ReturnType; - nodeSendToSession: ReturnType; - logGateway: { warn: ReturnType }; - dedupe?: { get: ReturnType }; -} { - return { - chatAbortControllers: new Map(), - chatRunBuffers: new Map(), - chatDeltaSentAt: new Map(), - chatAbortedRuns: new Map(), - removeChatRun: vi - .fn() - .mockImplementation((run: string) => ({ sessionKey: "main", clientRunId: run })), - agentRunSeq: new Map(), - broadcast: vi.fn(), - nodeSendToSession: vi.fn(), - logGateway: { warn: vi.fn() }, - ...overrides, - }; -} - -async function invokeChatAbort( - context: ReturnType, - params: { sessionKey: string; runId?: string }, - respond: ReturnType, -) { - 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); diff --git a/src/gateway/server-methods/chat.abort.test-helpers.ts b/src/gateway/server-methods/chat.abort.test-helpers.ts new file mode 100644 index 00000000000..fe5cd324ccb --- /dev/null +++ b/src/gateway/server-methods/chat.abort.test-helpers.ts @@ -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 = {}) { + return { + chatAbortControllers: new Map(), + chatRunBuffers: new Map(), + chatDeltaSentAt: new Map(), + chatAbortedRuns: new Map(), + removeChatRun: vi + .fn() + .mockImplementation((run: string) => ({ sessionKey: "main", clientRunId: run })), + agentRunSeq: new Map(), + 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; + context: ReturnType; + request: { sessionKey: string; runId?: string }; + client?: { + connId?: string; + connect?: { + device?: { id?: string }; + scopes?: string[]; + }; + } | null; + respond?: ReturnType; +}) { + 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; +}