diff --git a/packages/plugins/sdk/README.md b/packages/plugins/sdk/README.md index 80509728aa7..07bf77b16c5 100644 --- a/packages/plugins/sdk/README.md +++ b/packages/plugins/sdk/README.md @@ -1576,6 +1576,8 @@ Release and destroy must retire only that login allocation and preserve the computer. Omitted purpose retains the existing execution behavior. Core supplies `expiresAt` to the PTY open and an `inputId` for each input write. +Input, stop, and close retain the opening company context so provider callbacks +can use company-scoped host services during cleanup. A provider that retries delivery must retain the input identity and bytes, and reject reuse of the identity with different input. A failed input call ends the Core login; Core does not silently replay an unconfirmed write. diff --git a/packages/plugins/sdk/src/protocol.ts b/packages/plugins/sdk/src/protocol.ts index ee463fd84ba..ff3a43d51a8 100644 --- a/packages/plugins/sdk/src/protocol.ts +++ b/packages/plugins/sdk/src/protocol.ts @@ -1124,6 +1124,8 @@ export interface PluginLoginPtyOpenResult { /** The input request. It carries the worker session identifier and the raw input bytes. */ export interface PluginLoginPtyInputParams { + /** Host-bound company context for provider callbacks during this lifecycle call. */ + companyId?: string; /** Stable identity of this input write. Retries retain the identity and bytes. */ inputId?: string; /** The worker session identifier that the open reply returned. */ @@ -1148,12 +1150,16 @@ export interface PluginLoginPtyReadCredentialResult { /** The stop request. It carries the worker session identifier. */ export interface PluginLoginPtyStopParams { + /** Host-bound company context for provider callbacks during this lifecycle call. */ + companyId?: string; /** The worker session identifier that the open reply returned. */ workerSessionId: string; } /** The close request. The host route identifier is the authoritative key. */ export interface PluginLoginPtyCloseParams { + /** Host-bound company context for provider callbacks during this lifecycle call. */ + companyId?: string; /** * The host-owned opaque route identifier. This is the authoritative close key, * so the host closes the terminal even when no worker session identifier diff --git a/server/src/__tests__/fixtures/plugin-worker-login-pty.cjs b/server/src/__tests__/fixtures/plugin-worker-login-pty.cjs index 80ae3f750e8..87383b2aefb 100644 --- a/server/src/__tests__/fixtures/plugin-worker-login-pty.cjs +++ b/server/src/__tests__/fixtures/plugin-worker-login-pty.cjs @@ -146,6 +146,16 @@ function scriptedOutputLines(directive, hostRouteId, workerSessionId) { // The registered terminals, keyed by the host route id. Each entry records the // bound worker session id and the close directive. const routes = new Map(); +const pendingHost = new Map(); +let nextHostCall = 0; +function probeCompanyScope(message, entry) { + const id = `scope-${++nextHostCall}`; + return new Promise((resolve, reject) => { + pendingHost.set(id, { resolve, reject }); + send({ jsonrpc: "2.0", id, method: "companies.get", params: { companyId: entry.companyId, stage: message.method }, + ...(message.paperclipInvocation ? { paperclipInvocationId: message.paperclipInvocation.id } : {}) }); + }); +} function parseDirective(raw) { try { @@ -158,12 +168,24 @@ function parseDirective(raw) { const rl = readline.createInterface({ input: process.stdin, crlfDelay: Infinity }); -rl.on("line", (line) => { +rl.on("line", async (line) => { if (!line.trim()) return; const message = JSON.parse(line); + if (pendingHost.has(message.id)) { + const pending = pendingHost.get(message.id); pendingHost.delete(message.id); + if (message.error) pending.reject(new Error("Host company scope denied")); else pending.resolve(message.result); + return; + } const method = message && typeof message.method === "string" ? message.method : null; const params = message.params ?? {}; + if (["loginPtyInput", "loginPtyStop", "loginPtyClose"].includes(method)) { + const entry = method === "loginPtyClose" ? routes.get(params.hostRouteId) : [...routes.values()].find(value => value.workerSessionId === params.workerSessionId); + if (entry?.probeLifecycle) { + try { await probeCompanyScope(message, entry); } + catch { send({ jsonrpc: "2.0", id: message.id, error: { code: -32000, message: "Host company scope denied" } }); return; } + } + } if (method === "initialize") { send({ jsonrpc: "2.0", @@ -188,6 +210,7 @@ rl.on("line", (line) => { const closeMode = directive.closeMode ?? "ack"; routes.set(params.hostRouteId, { workerSessionId, + companyId: params.companyId, probeLifecycle: directive.probeLifecycle === true, closeMode, rejectInput: directive.rejectInput === true, pendingOutput: directive.emitOnInput === true diff --git a/server/src/__tests__/plugin-worker-manager.test.ts b/server/src/__tests__/plugin-worker-manager.test.ts index d5438e69955..55424e9f9b3 100644 --- a/server/src/__tests__/plugin-worker-manager.test.ts +++ b/server/src/__tests__/plugin-worker-manager.test.ts @@ -1336,6 +1336,27 @@ describe("plugin worker manager setup-token pty route gate", () => { } }); + it("retains the opening company for scoped host calls during input, stop, and close", async () => { + const stages: string[] = []; + const handle = makeLoginPtyHandle({ hostHandlers: { "companies.get": async ( + params: { companyId: string; stage: string }, context?: { invocationScope?: { companyId?: string } }, + ) => { + if (context?.invocationScope?.companyId !== params.companyId || params.companyId !== "company-1") throw new Error("Host company scope denied"); + stages.push(params.stage); + return { id: params.companyId }; + } } }); + try { + await handle.start(); + const session = await handle.openLoginPtySession(ptyOpenInput({ probeLifecycle: true })); + session.write("browser-code"); + await vi.waitFor(() => expect(stages).toEqual(["loginPtyInput"])); + session.kill(); + await vi.waitFor(() => expect(stages).toEqual(["loginPtyInput", "loginPtyStop"])); + await session.close(); + expect(stages).toEqual(["loginPtyInput", "loginPtyStop", "loginPtyClose"]); + } finally { await handle.stop().catch(() => undefined); } + }); + it("fails the login when input delivery is unconfirmed", async () => { const handle = makeLoginPtyHandle(); try { diff --git a/server/src/services/plugin-worker-manager.ts b/server/src/services/plugin-worker-manager.ts index 0cc50a5024c..459f676148e 100644 --- a/server/src/services/plugin-worker-manager.ts +++ b/server/src/services/plugin-worker-manager.ts @@ -1565,6 +1565,7 @@ export function createPluginWorkerHandle( type LoginPtyPreBindRecord = LoginPtyPreBindOutputRecord | LoginPtyPreBindExitRecord; interface LoginPtyRoute { + companyId: string; hostRouteId: string; state: LoginPtyRouteState; workerSessionId: string | null; @@ -1627,11 +1628,11 @@ export function createPluginWorkerHandle( // that carries the exact host route identifier. An absent, malformed, // mismatched, or timed-out acknowledgement returns false, so the caller fails // closed. - async function closeLoginPtyTerminal(hostRouteId: string): Promise { + async function closeLoginPtyTerminal(hostRouteId: string, companyId: string): Promise { try { const ack = await callInternal( "loginPtyClose", - { hostRouteId }, + { hostRouteId, companyId }, loginPtyCloseTimeoutMs, ); return isRecord(ack) && readNonEmptyString(ack.hostRouteId) === hostRouteId; @@ -1664,7 +1665,7 @@ export function createPluginWorkerHandle( // A terminalized route reports a null exit code, which the runner treats as a // failure. settleRouteWait(route, { exitCode: null }); - const confirmed = await closeLoginPtyTerminal(route.hostRouteId); + const confirmed = await closeLoginPtyTerminal(route.hostRouteId, route.companyId); if (!confirmed) { // The worker did not acknowledge the close, so the host cannot prove the // terminal is gone. Fail closed: retire the worker before any reuse. This @@ -1921,6 +1922,7 @@ export function createPluginWorkerHandle( settleWait = resolve; }); const route: LoginPtyRoute = { + companyId: input.companyId, hostRouteId, state: "reserved", workerSessionId: null, @@ -2003,7 +2005,7 @@ export function createPluginWorkerHandle( if (route.state !== "open" || !sid) return; void callInternal( "loginPtyInput", - { workerSessionId: sid, data, inputId: randomUUID() }, + { companyId: route.companyId, workerSessionId: sid, data, inputId: randomUUID() }, loginPtyOpenTimeoutMs, ).catch(() => terminalizeLoginPtyRoute(route)); }, @@ -2015,7 +2017,7 @@ export function createPluginWorkerHandle( if (!sid) return; void callInternal( "loginPtyStop", - { workerSessionId: sid }, + { companyId: route.companyId, workerSessionId: sid }, loginPtyOpenTimeoutMs, ).catch(() => {}); },