diff --git a/server/src/__tests__/fixtures/plugin-worker-duplex-channel.cjs b/server/src/__tests__/fixtures/plugin-worker-duplex-channel.cjs index 2cd1967d95..7c116b1dbd 100644 --- a/server/src/__tests__/fixtures/plugin-worker-duplex-channel.cjs +++ b/server/src/__tests__/fixtures/plugin-worker-duplex-channel.cjs @@ -18,12 +18,38 @@ // - `closeMode`: "ack" | "bad-ack" | "no-ack" (default "ack"). It controls the // close reply, so a test proves the host retires the worker on an unconfirmed // close. +// - `batchWithOpenReply`: when true, the fixture writes the open reply and the +// scripted data and exit in one stdout write. The host then reads the open +// reply and the notifications in one batch, so a test proves the host holds +// and replays a frame that arrives before the route binds. const readline = require("node:readline"); function send(message) { process.stdout.write(`${JSON.stringify(message)}\n`); } +// Serialize the scripted data and exit frames as newline-delimited lines. The +// batch mode writes these together with the open reply in one stdout write. +function scriptedFrameLines(directive, workerSessionId) { + const data = Array.isArray(directive.data) ? directive.data : []; + let lines = ""; + for (const entry of data) { + lines += `${JSON.stringify({ + jsonrpc: "2.0", + method: "duplexChannel.data", + params: { workerSessionId: entry.sid ?? workerSessionId, chunk: entry.chunk }, + })}\n`; + } + if (typeof directive.exitCode === "number") { + lines += `${JSON.stringify({ + jsonrpc: "2.0", + method: "duplexChannel.exit", + params: { workerSessionId, exitCode: directive.exitCode }, + })}\n`; + } + return lines; +} + // The registered channels, keyed by the host route id. Each entry records the // bound worker session id and the close directive. const routes = new Map(); @@ -84,8 +110,21 @@ rl.on("line", (line) => { return; } - const reply = () => - send({ jsonrpc: "2.0", id: message.id, result: { workerSessionId } }); + const openReplyLine = `${JSON.stringify({ + jsonrpc: "2.0", + id: message.id, + result: { workerSessionId }, + })}\n`; + + if (directive.batchWithOpenReply === true) { + // Write the open reply and the scripted frames in one stdout write. The + // host reads them in one batch, so a data or exit frame arrives before the + // route binds. The host must hold and replay the frame after the bind. + process.stdout.write(openReplyLine + scriptedFrameLines(directive, workerSessionId)); + return; + } + + const reply = () => process.stdout.write(openReplyLine); reply(); if (mode === "duplicate-open-reply") { // Send a second open reply for the same request id. The host drops it. @@ -95,24 +134,7 @@ rl.on("line", (line) => { // Emit the scripted data and the exit after the open reply, so the host // binds the route first. setImmediate(() => { - const data = Array.isArray(directive.data) ? directive.data : []; - for (const entry of data) { - send({ - jsonrpc: "2.0", - method: "duplexChannel.data", - params: { - workerSessionId: entry.sid ?? workerSessionId, - chunk: entry.chunk, - }, - }); - } - if (typeof directive.exitCode === "number") { - send({ - jsonrpc: "2.0", - method: "duplexChannel.exit", - params: { workerSessionId, exitCode: directive.exitCode }, - }); - } + process.stdout.write(scriptedFrameLines(directive, workerSessionId)); }); return; } diff --git a/server/src/__tests__/issue-dependency-wakeups-routes.test.ts b/server/src/__tests__/issue-dependency-wakeups-routes.test.ts index a5d22dc1da..a840a81a4a 100644 --- a/server/src/__tests__/issue-dependency-wakeups-routes.test.ts +++ b/server/src/__tests__/issue-dependency-wakeups-routes.test.ts @@ -2,6 +2,17 @@ import express from "express"; import request from "supertest"; import { beforeEach, describe, expect, it, vi } from "vitest"; +// The first test in this suite imports the large `routes/issues.ts` module +// through `vi.importActual` inside `createApp`. `vi.resetModules()` in +// `beforeEach` forces a fresh import each test, so the first test pays the +// one-time transform and execution cost of that module. Locally the first +// test takes about 3.7s while the later tests take about 0.13s each. Under +// the loaded serial shard (maxWorkers=1) this cold-start can cross vitest's +// default 5000ms test timeout and produce a flaky "Test timed out in 5000ms" +// failure. Give the suite generous headroom, far above the observed cold-start +// yet still below the 30s hook timeout. +vi.setConfig({ testTimeout: 15000 }); + const mockWakeup = vi.hoisted(() => vi.fn(async () => undefined)); const mockFindExistingIssueBlockersResolvedWakeForReadyState = vi.hoisted(() => vi.fn(async () => null)); const mockIssueService = vi.hoisted(() => ({ diff --git a/server/src/__tests__/plugin-worker-manager-duplex.test.ts b/server/src/__tests__/plugin-worker-manager-duplex.test.ts index 7a02b92b50..b6a5c86a37 100644 --- a/server/src/__tests__/plugin-worker-manager-duplex.test.ts +++ b/server/src/__tests__/plugin-worker-manager-duplex.test.ts @@ -460,6 +460,58 @@ describe("plugin worker manager duplex channel route", () => { } }); + // ------------------------------------------------------------------------- + // The open reply and a frame arrive in one read batch. + // ------------------------------------------------------------------------- + + it("holds and replays a data frame that arrives in the open-reply read batch", async () => { + const handle = makeDuplexHandle(); + try { + await handle.start(); + const session = await handle.openDuplexChannel( + duplexOpenInput({ + // The worker writes the open reply and the data and exit frames in one + // stdout write. The host reads them in one batch, so the data and exit + // frames arrive before the route binds. The host must hold the frames + // and replay them after the bind, not drop them. + batchWithOpenReply: true, + workerSessionId: "ws-A", + data: [{ chunk: "batched-one" }, { chunk: "batched-two" }], + exitCode: 0, + }), + ); + const chunks: string[] = []; + session.onData((chunk) => chunks.push(chunk)); + await expect(session.wait()).resolves.toEqual({ exitCode: 0 }); + expect(chunks).toEqual(["batched-one", "batched-two"]); + await session.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("ends the route when a batched frame passes the per-chunk limit before the bind", async () => { + const handle = makeDuplexHandle({ + duplexChannelLimits: { maxChunkChars: 4 }, + }); + try { + await handle.start(); + const session = await handle.openDuplexChannel( + duplexOpenInput({ + // The worker batches the data frame with the open reply, so the frame + // arrives before the route binds. The replay after the bind applies the + // per-chunk limit, so the one large chunk ends the route. + batchWithOpenReply: true, + workerSessionId: "ws-A", + data: [{ chunk: "this-one-chunk-is-too-large" }], + }), + ); + await expect(session.wait()).resolves.toEqual({ exitCode: null }); + } finally { + await handle.stop().catch(() => undefined); + } + }); + // ------------------------------------------------------------------------- // Authoritative closure and worker retirement. // ------------------------------------------------------------------------- diff --git a/server/src/__tests__/tool-gateway.test.ts b/server/src/__tests__/tool-gateway.test.ts index 8ce669a407..e923d275ff 100644 --- a/server/src/__tests__/tool-gateway.test.ts +++ b/server/src/__tests__/tool-gateway.test.ts @@ -889,7 +889,13 @@ describeEmbeddedPostgres("tool gateway acceptance", () => { effect: "include", toolName: gatewayToolName, }); + // Pin the clock so every request in this test shares one rate-limit + // window. The window boundary aligns to wall-clock time, so a real clock + // can advance past the boundary between two paired requests and reset the + // counter. That reset makes the second request return 200 instead of 429. + const fixedNow = Date.now(); const gateway = createTestToolGatewayService(db, { + now: () => fixedNow, mcpGatewayProtocolLimits: { gatewayRequests: { max: 1, windowMs: 60_000 }, tokenRequests: { max: 1, windowMs: 60_000 }, diff --git a/server/src/__tests__/workspace-runtime.test.ts b/server/src/__tests__/workspace-runtime.test.ts index 4b59464fd2..882ec1114a 100644 --- a/server/src/__tests__/workspace-runtime.test.ts +++ b/server/src/__tests__/workspace-runtime.test.ts @@ -6455,9 +6455,16 @@ describeEmbeddedPostgres("workspace runtime service control persistence", () => return result[0]!; })); - expect(first.port).toBe(basePort); - expect(second.port).not.toBe(first.port); - expect(second.port).toBeGreaterThan(basePort); + // The two lanes claim ports concurrently in-process. The allocator guarantees + // distinct ports, but not which lane wins the base port. It depends on the + // scheduling order. So assert order-independent invariants, not array index. + const lowerPort = Math.min(first.port!, second.port!); + const upperPort = Math.max(first.port!, second.port!); + const maxAllocatablePort = basePort + WORKSPACE_RUNTIME_PORT_ALLOCATION_ATTEMPTS - 1; + expect(first.port).not.toBe(second.port); + expect(lowerPort).toBe(basePort); + expect(upperPort).toBeGreaterThan(basePort); + expect(upperPort).toBeLessThanOrEqual(maxAllocatablePort); expect(first.url).toBe(`http://127.0.0.1:${first.port}`); expect(second.url).toBe(`http://127.0.0.1:${second.port}`); await expect(fetch(first.url!)).resolves.toMatchObject({ ok: true }); @@ -6913,6 +6920,10 @@ describeEmbeddedPostgres("workspace runtime startup reconciliation", () => { await db.delete(projectWorkspaces); await db.delete(projects); await db.delete(heartbeatRuns); + // The runtime service control path writes activity_log rows through + // logActivity. Those rows reference the company. Delete them before the + // company to respect the activity_log.company_id foreign key. + await db.delete(activityLog); await db.delete(agents); await db.delete(companies); }); diff --git a/server/src/services/plugin-worker-manager.ts b/server/src/services/plugin-worker-manager.ts index 3412290cb2..dfbe2a444a 100644 --- a/server/src/services/plugin-worker-manager.ts +++ b/server/src/services/plugin-worker-manager.ts @@ -1483,6 +1483,13 @@ export function createPluginWorkerHandle( listener: ((chunk: string) => void) | null; buffered: string[]; bufferedChars: number; + // Raw data and exit notifications that arrive before the route binds. The + // host reads the worker stdout line by line. The open reply and a data or + // exit notification can arrive in one read batch, so the host dispatches the + // notification before the deferred open-reply continuation flips the state to + // `open`. The host holds these frames here and replays them in order right + // after it binds the route, so a batched frame is never lost. + preOpen: JsonRpcNotification[]; pendingRequests: number; protocolErrors: number; totalDataBytes: number; @@ -1531,6 +1538,7 @@ export function createPluginWorkerHandle( route.listener = null; route.buffered = []; route.bufferedChars = 0; + route.preOpen = []; clearDuplexChannelLifetimeTimer(route); // A terminalized route reports a null exit code, which the caller treats as a // failure. @@ -1558,6 +1566,40 @@ export function createPluginWorkerHandle( } } + // Hold one data or exit notification that arrives before the route binds. The + // host replays the held frames in order after it binds the route. Bound the + // hold by the pre-bind frame count, so a worker that floods frames before it + // replies to the open cannot make the host hold an unbounded number of frames. + // Count one protocol error for each frame past the bound. + function bufferPreOpenDuplexChannelNotification( + route: DuplexChannelRoute, + notification: JsonRpcNotification, + ): void { + if (route.preOpen.length >= maxDuplexChannelPreBindFrames) { + recordDuplexChannelProtocolError(route); + return; + } + route.preOpen.push(notification); + } + + // Replay the held pre-open frames in order right after the route binds. The + // route is `open` now, so each frame passes through the normal per-frame bounds + // and the session-identifier match. A frame that ends the route terminalizes + // it, and every later frame in the replay is a no-op, because the routing + // functions drop a frame when the route is not `open`. + function drainPreOpenDuplexChannelNotifications(route: DuplexChannelRoute): void { + if (route.preOpen.length === 0) return; + const pending = route.preOpen; + route.preOpen = []; + for (const notification of pending) { + if (notification.method === DUPLEX_CHANNEL_DATA_NOTIFICATION) { + routeDuplexChannelData(notification); + } else if (notification.method === DUPLEX_CHANNEL_EXIT_NOTIFICATION) { + routeDuplexChannelExit(notification); + } + } + } + // Deliver one duplex channel chunk to the bound listener in isolation. A // listener that throws must not escape the worker stdout notification handler // or the buffered replay, so a throw here breaks neither the notification @@ -1586,7 +1628,13 @@ export function createPluginWorkerHandle( // attached yet. Never log the raw bytes. function routeDuplexChannelData(notification: JsonRpcNotification): void { const route = duplexChannelRoute; - if (!route || route.state !== "open") return; + if (!route || route.terminalized) return; + if (route.state === "reserved" || route.state === "opening") { + // The route did not bind yet. Hold the frame and replay it after the bind. + bufferPreOpenDuplexChannelNotification(route, notification); + return; + } + if (route.state !== "open") return; const params = isRecord(notification.params) ? notification.params : {}; const workerSessionId = readNonEmptyString(params.workerSessionId); const chunk = params.chunk; @@ -1638,7 +1686,13 @@ export function createPluginWorkerHandle( // session identifier. function routeDuplexChannelExit(notification: JsonRpcNotification): void { const route = duplexChannelRoute; - if (!route || route.state !== "open") return; + if (!route || route.terminalized) return; + if (route.state === "reserved" || route.state === "opening") { + // The route did not bind yet. Hold the frame and replay it after the bind. + bufferPreOpenDuplexChannelNotification(route, notification); + return; + } + if (route.state !== "open") return; const params = isRecord(notification.params) ? notification.params : {}; const workerSessionId = readNonEmptyString(params.workerSessionId); if (!workerSessionId || workerSessionId !== route.workerSessionId) return; @@ -1658,6 +1712,7 @@ export function createPluginWorkerHandle( route.listener = null; route.buffered = []; route.bufferedChars = 0; + route.preOpen = []; clearDuplexChannelLifetimeTimer(route); settleRouteWait(route, { exitCode: null }); } @@ -1686,6 +1741,7 @@ export function createPluginWorkerHandle( listener: null, buffered: [], bufferedChars: 0, + preOpen: [], pendingRequests: 0, protocolErrors: 0, totalDataBytes: 0, @@ -1728,13 +1784,22 @@ export function createPluginWorkerHandle( route.workerSessionId = workerSessionId; route.state = "open"; + // Replay any data or exit frame that arrived in the open-reply read batch, + // before the route bound. The route is `open` now, so each replayed frame + // passes through the normal per-frame bounds and the session match. + drainPreOpenDuplexChannelNotifications(route); + // Start the route lifetime timer now the route is open. The route ends when // the timer expires. Every terminal path and the worker-exit path clears the // timer. Unreference the timer so it never blocks the host process shutdown. - route.lifetimeTimer = setTimeout(() => { - void terminalizeDuplexChannelRoute(route); - }, maxDuplexChannelDurationMs); - route.lifetimeTimer.unref?.(); + // A replayed frame can end the route during the drain above, so start the + // timer only while the route is still open. + if (route.state === "open") { + route.lifetimeTimer = setTimeout(() => { + void terminalizeDuplexChannelRoute(route); + }, maxDuplexChannelDurationMs); + route.lifetimeTimer.unref?.(); + } // Send one host→worker request under the pending-request bound. End the route // when too many requests are in-flight, so a worker that never replies cannot