| |
| import { EventEmitter } from "node:events"; |
| import { afterEach, describe, expect, it, vi } from "vitest"; |
| import { WebSocket } from "ws"; |
| import { |
| GATEWAY_CLIENT_IDS, |
| GATEWAY_CLIENT_MODES, |
| } from "../../../../packages/gateway-protocol/src/client-info.js"; |
| import { createDeferredCore } from "../../../shared/deferred.js"; |
| import { NodeRegistry } from "../../node-registry.js"; |
| import type { GatewayRequestOptions } from "../../server-methods/types.js"; |
| import type { GatewayWsClient } from "../ws-types.js"; |
| import { |
| createDispatchTestHarness, |
| createOperatorWsClient, |
| } from "./authenticated-request-dispatch.test-support.js"; |
|
|
| const handleGatewayRequest = vi.hoisted(() => vi.fn()); |
|
|
| vi.mock("./authenticated-request-dispatch.server-methods.runtime.js", () => ({ |
| handleGatewayRequest, |
| })); |
|
|
| const activeRegistries = new Set<NodeRegistry>(); |
|
|
| afterEach(() => { |
| for (const registry of activeRegistries) { |
| registry.unregister("paired-node-connection"); |
| } |
| activeRegistries.clear(); |
| handleGatewayRequest.mockReset(); |
| }); |
|
|
| function createPairedNode() { |
| const frames: string[] = []; |
| let frameArrived = createDeferredCore(); |
| const waitForFrameCount = async (count: number) => { |
| while (frames.length < count) { |
| await frameArrived.promise; |
| } |
| }; |
| const registry = new NodeRegistry(); |
| activeRegistries.add(registry); |
| registry.register( |
| { |
| connId: "paired-node-connection", |
| usesSharedGatewayAuth: false, |
| socket: { |
| readyState: WebSocket.OPEN, |
| send(frame: unknown) { |
| if (typeof frame === "string") { |
| frames.push(frame); |
| frameArrived.resolve(); |
| frameArrived = createDeferredCore(); |
| } |
| }, |
| }, |
| connect: { |
| minProtocol: 1, |
| maxProtocol: 1, |
| client: { |
| id: GATEWAY_CLIENT_IDS.NODE_HOST, |
| version: "1.0.0", |
| platform: "linux", |
| mode: "node", |
| }, |
| device: { |
| id: "paired-node", |
| publicKey: "public-key", |
| signature: "signature", |
| signedAt: 1, |
| nonce: "nonce", |
| }, |
| caps: ["local-inference"], |
| commands: ["ollama.chat"], |
| }, |
| } as GatewayWsClient, |
| { pairingIdentity: "paired-node-identity" }, |
| ); |
| return { registry, frames, waitForFrameCount }; |
| } |
|
|
| function createDispatcher( |
| socket: EventEmitter, |
| clientInfo: Pick<GatewayWsClient["connect"]["client"], "id" | "mode"> = { |
| id: GATEWAY_CLIENT_IDS.CLI, |
| mode: GATEWAY_CLIENT_MODES.CLI, |
| }, |
| ) { |
| const client = createOperatorWsClient({ |
| connId: "operator-connection", |
| socket, |
| clientInfo, |
| scopes: ["operator.write"], |
| }); |
| const harness = createDispatchTestHarness({ connId: client.connId }); |
| return { client, ...harness }; |
| } |
|
|
| describe("authenticated WebSocket request cancellation", () => { |
| it("forwards CLI socket closure to the actual first-party node cancel event", async () => { |
| const socket = new EventEmitter(); |
| const { registry, frames, waitForFrameCount } = createPairedNode(); |
| const { awaitResponseFrame, client, dispatcher } = createDispatcher(socket); |
| handleGatewayRequest.mockImplementation(async (options: GatewayRequestOptions) => { |
| const result = await registry.invoke({ |
| nodeId: "paired-node", |
| command: "ollama.chat", |
| timeoutMs: 10_000, |
| signal: options.signal, |
| }); |
| options.respond( |
| result.ok, |
| result.payload, |
| result.error |
| ? { |
| code: result.error.code ?? "UNAVAILABLE", |
| message: result.error.message ?? "node invocation failed", |
| } |
| : undefined, |
| ); |
| }); |
|
|
| const dispatch = dispatcher.dispatch( |
| { |
| type: "req", |
| id: "paired-node-cli-inference", |
| method: "node.invoke", |
| params: { |
| nodeId: "paired-node", |
| command: "ollama.chat", |
| idempotencyKey: "paired-node-cli-inference", |
| }, |
| }, |
| client, |
| ); |
| try { |
| await waitForFrameCount(1); |
| expect(socket.listenerCount("close")).toBe(1); |
| socket.emit("close", 1000, Buffer.alloc(0)); |
| await waitForFrameCount(2); |
| const request = JSON.parse(frames[0] ?? "{}") as { payload?: { id?: string } }; |
| expect(JSON.parse(frames[1] ?? "{}")).toMatchObject({ |
| event: "node.invoke.cancel", |
| payload: { invokeId: request.payload?.id, nodeId: "paired-node" }, |
| }); |
| await awaitResponseFrame("paired-node-cli-inference"); |
| } finally { |
| socket.emit("close", 1000, Buffer.alloc(0)); |
| await dispatch; |
| } |
| expect(socket.listenerCount("close")).toBe(0); |
| }); |
|
|
| it.each(["test.trace", "sessions.cleanup", "agents.delete"])( |
| "keeps authenticated %s work alive after its socket disconnects", |
| async (method) => { |
| const socket = new EventEmitter(); |
| const { awaitResponseFrame, client, dispatcher } = createDispatcher(socket); |
| const invoked = createDeferredCore(); |
| const completion = createDeferredCore(); |
| let completed = false; |
| handleGatewayRequest.mockImplementation(async (options: GatewayRequestOptions) => { |
| invoked.resolve(); |
| await completion.promise; |
| completed = true; |
| options.respond(true, { ok: true }); |
| }); |
|
|
| const dispatch = dispatcher.dispatch( |
| { type: "req", id: "ordinary-request", method, params: {} }, |
| client, |
| ); |
| try { |
| await invoked.promise; |
| expect(handleGatewayRequest).toHaveBeenCalledOnce(); |
|
|
| expect(handleGatewayRequest.mock.calls[0]?.[0]).not.toHaveProperty("signal"); |
| expect(socket.listenerCount("close")).toBe(0); |
| socket.emit("close", 1006, Buffer.alloc(0)); |
| expect(completed).toBe(false); |
| } finally { |
| completion.resolve(); |
| await awaitResponseFrame("ordinary-request"); |
| await dispatch; |
| } |
| expect(completed).toBe(true); |
| }, |
| ); |
|
|
| it("cancels a session companion ask when its authenticated socket closes", async () => { |
| const socket = new EventEmitter(); |
| const { client, dispatcher } = createDispatcher(socket, { |
| id: GATEWAY_CLIENT_IDS.CONTROL_UI, |
| mode: GATEWAY_CLIENT_MODES.UI, |
| }); |
| const invoked = createDeferredCore(); |
| const abortObserved = createDeferredCore(); |
| let observedSignal: AbortSignal | undefined; |
| handleGatewayRequest.mockImplementation(async (options: GatewayRequestOptions) => { |
| observedSignal = options.signal; |
| options.signal?.addEventListener("abort", () => abortObserved.resolve(), { once: true }); |
| invoked.resolve(); |
| await abortObserved.promise; |
| }); |
|
|
| const dispatch = dispatcher.dispatch( |
| { |
| type: "req", |
| id: "session-companion", |
| method: "sessions.companion.ask", |
| params: { sessionKey: "agent:main:main", question: "What changed?" }, |
| }, |
| client, |
| ); |
| try { |
| |
| await invoked.promise; |
| expect(socket.listenerCount("close")).toBe(1); |
| socket.emit("close", 1000, Buffer.alloc(0)); |
| await abortObserved.promise; |
| expect(observedSignal?.aborted).toBe(true); |
| } finally { |
| socket.emit("close", 1000, Buffer.alloc(0)); |
| await dispatch; |
| } |
| expect(socket.listenerCount("close")).toBe(0); |
| }); |
|
|
| it.each([ |
| { |
| label: "control UI", |
| id: GATEWAY_CLIENT_IDS.CONTROL_UI, |
| mode: GATEWAY_CLIENT_MODES.UI, |
| }, |
| { |
| label: "gateway SDK", |
| id: GATEWAY_CLIENT_IDS.GATEWAY_CLIENT, |
| mode: GATEWAY_CLIENT_MODES.BACKEND, |
| }, |
| { |
| label: "native macOS app", |
| id: GATEWAY_CLIENT_IDS.MACOS_APP, |
| mode: GATEWAY_CLIENT_MODES.UI, |
| }, |
| ])( |
| "does not cancel an authenticated $label node invocation on socket close", |
| async (identity) => { |
| const socket = new EventEmitter(); |
| const { registry, frames, waitForFrameCount } = createPairedNode(); |
| const { client, dispatcher } = createDispatcher(socket, identity); |
| handleGatewayRequest.mockImplementation(async (options: GatewayRequestOptions) => { |
| const result = await registry.invoke({ |
| nodeId: "paired-node", |
| command: "ollama.chat", |
| timeoutMs: 10_000, |
| signal: options.signal, |
| }); |
| options.respond(result.ok, result.payload); |
| }); |
|
|
| const dispatch = dispatcher.dispatch( |
| { |
| type: "req", |
| id: "paired-node-ordinary-inference", |
| method: "node.invoke", |
| params: { |
| nodeId: "paired-node", |
| command: "ollama.chat", |
| idempotencyKey: "paired-node-ordinary-inference", |
| }, |
| }, |
| client, |
| ); |
| try { |
| await waitForFrameCount(1); |
| expect(handleGatewayRequest.mock.calls[0]?.[0]).not.toHaveProperty("signal"); |
| expect(socket.listenerCount("close")).toBe(0); |
| socket.emit("close", 1000, Buffer.alloc(0)); |
| expect(frames).toHaveLength(1); |
| const request = JSON.parse(frames[0] ?? "{}") as { payload?: { id?: string } }; |
| expect( |
| registry.handleInvokeResult({ |
| id: request.payload?.id ?? "", |
| nodeId: "paired-node", |
| connId: "paired-node-connection", |
| ok: true, |
| }), |
| ).toBe(true); |
| } finally { |
| registry.unregister("paired-node-connection"); |
| await dispatch; |
| } |
| }, |
| ); |
| }); |
|
|
