From f5b2564e33e4fcac79ee2d6bb4a73d98380d1705 Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Thu, 2 Jul 2026 20:40:22 -0700 Subject: [PATCH] Fix custom image terminal CI failures Co-Authored-By: Paperclip --- .../environment-custom-image-routes.test.ts | 5 +-- ...vironment-custom-image-terminal-ws.test.ts | 1 + .../environment-custom-image-terminal-ws.ts | 20 +++++++-- server/src/routes/environments.ts | 44 +++++++++++++------ 4 files changed, 51 insertions(+), 19 deletions(-) diff --git a/server/src/__tests__/environment-custom-image-routes.test.ts b/server/src/__tests__/environment-custom-image-routes.test.ts index 35aebec21e6..ef416a948bf 100644 --- a/server/src/__tests__/environment-custom-image-routes.test.ts +++ b/server/src/__tests__/environment-custom-image-routes.test.ts @@ -170,6 +170,7 @@ function createSession(overrides: Record = {}) { }, connectionSecretRef: null, metadata: { + setupRpcCompanyId: "company-1", safeLabel: "setup", connectUrl: "https://203.0.113.10/setup", }, @@ -298,7 +299,6 @@ describe("environment customImage setup routes", () => { expect(res.status).toBe(201); expect(res.body.connectionPayload.command).toContain("203.0.113.10"); expect(mockEnvironmentCustomImageService.startSetupSession).toHaveBeenCalledWith({ - companyId: "company-1", environmentId: "env-1", templateId: null, ttlSeconds: 3600, @@ -306,6 +306,7 @@ describe("environment customImage setup routes", () => { userId: "user-1", agentId: null, }, + secretContextCompanyId: "company-1", }); const activity = loggedActivityJson(); expect(activity).not.toContain("203.0.113.10"); @@ -624,11 +625,9 @@ describe("environment customImage setup routes", () => { expect(rollback.status).toBe(200); expect(disable.status).toBe(200); expect(mockEnvironmentCustomImageService.rollbackTemplate).toHaveBeenCalledWith({ - companyId: "company-1", environmentId: "env-1", }); expect(mockEnvironmentCustomImageService.disableTemplate).toHaveBeenCalledWith({ - companyId: "company-1", environmentId: "env-1", deleteProviderTemplate: true, }); diff --git a/server/src/__tests__/environment-custom-image-terminal-ws.test.ts b/server/src/__tests__/environment-custom-image-terminal-ws.test.ts index 03aa9329988..d8e3400e684 100644 --- a/server/src/__tests__/environment-custom-image-terminal-ws.test.ts +++ b/server/src/__tests__/environment-custom-image-terminal-ws.test.ts @@ -86,6 +86,7 @@ function createSession(overrides: Record = {}) { provider: "daytona", status: "waiting_for_user", expiresAt: futureDate(), + metadata: { setupRpcCompanyId: "company-1" }, ...overrides, }; } diff --git a/server/src/realtime/environment-custom-image-terminal-ws.ts b/server/src/realtime/environment-custom-image-terminal-ws.ts index 403fc6e3f17..31bd599152e 100644 --- a/server/src/realtime/environment-custom-image-terminal-ws.ts +++ b/server/src/realtime/environment-custom-image-terminal-ws.ts @@ -43,11 +43,11 @@ interface TerminalWsServer { interface SetupSessionSnapshot { id: string; - companyId: string; environmentId: string; provider: string; status: string; expiresAt: Date | string | null; + metadata?: Record | null; } interface CustomImageTerminalService { @@ -162,6 +162,15 @@ function requireFutureSetupExpiry(session: { expiresAt: Date | string | null }, return expiresAt; } +function readSetupSessionCompanyId(session: { + metadata?: Record | null; +}): string | null { + const value = session.metadata?.setupRpcCompanyId; + if (typeof value !== "string") return null; + const companyId = value.trim(); + return companyId && companyId !== "instance" ? companyId : null; +} + function statusLineForError(err: unknown) { const status = typeof err === "object" && err !== null && "status" in err ? Number((err as { status?: unknown }).status) @@ -274,8 +283,9 @@ async function validateTerminalUpgrade(input: { input.sessionStore.delete(terminalSession.id); throw unprocessable("Invalid terminal setup session."); } + const storedSetupCompanyId = readSetupSessionCompanyId(storedSetupSession); if ( - storedSetupSession.companyId !== terminalSession.companyId + storedSetupCompanyId !== terminalSession.companyId || storedSetupSession.environmentId !== terminalSession.environmentId || storedSetupSession.provider !== terminalSession.provider ) { @@ -289,7 +299,7 @@ async function validateTerminalUpgrade(input: { }); if ( refreshed.session.id !== terminalSession.setupSessionId - || refreshed.session.companyId !== terminalSession.companyId + || readSetupSessionCompanyId(refreshed.session) !== terminalSession.companyId || refreshed.session.environmentId !== terminalSession.environmentId || refreshed.session.provider !== terminalSession.provider ) { @@ -565,6 +575,10 @@ export function setupEnvironmentCustomImageTerminalWebSocketServer( connectionRegistry.closeAll("server_shutdown"); }); + if (typeof server.on !== "function") { + return wss; + } + server.on("close", () => { connectionRegistry.closeAll("server_shutdown"); for (const client of wss.clients) { diff --git a/server/src/routes/environments.ts b/server/src/routes/environments.ts index d4c590bbc7b..269609d3495 100644 --- a/server/src/routes/environments.ts +++ b/server/src/routes/environments.ts @@ -187,6 +187,27 @@ export function environmentRoutes( throw unprocessable("companyId query parameter is required for environment customImage setup."); } + function readSetupSessionCompanyId(session: { + metadata?: Record | null; + }): string | null { + const value = session.metadata?.setupRpcCompanyId; + if (typeof value !== "string") return null; + const companyId = value.trim(); + return companyId && companyId !== "instance" ? companyId : null; + } + + async function resolveCustomImageSessionCompanyId( + req: Request, + session: { metadata?: Record | null }, + ): Promise { + const metadataCompanyId = readSetupSessionCompanyId(session); + if (metadataCompanyId) { + assertCustomImageCompanyAccess(req, metadataCompanyId); + return metadataCompanyId; + } + return await resolveCustomImageCompanyId(req); + } + async function resolveEnvironmentSecretContextCompanyId( req: Request, environmentId: string, @@ -369,9 +390,8 @@ export function environmentRoutes( router.get("/environments/:environmentId/custom-image-template", async (req, res) => { assertCanAccessInstanceEnvironments(req); - const companyId = await resolveCustomImageCompanyId(req); + await resolveCustomImageCompanyId(req); const overview = await customImages.getOverview({ - companyId, environmentId: req.params.environmentId as string, }); res.json(overview); @@ -385,7 +405,6 @@ export function environmentRoutes( const companyId = await resolveCustomImageCompanyId(req); const actor = getActorInfo(req); const result = await customImages.startSetupSession({ - companyId, environmentId: req.params.environmentId as string, templateId: req.body.templateId ?? null, ttlSeconds: req.body.ttlSeconds ?? null, @@ -393,6 +412,7 @@ export function environmentRoutes( userId: actor.actorType === "user" ? actor.actorId : null, agentId: actor.agentId, }, + secretContextCompanyId: companyId, }); await logEnvironmentCustomImageActivity({ actor, @@ -412,7 +432,7 @@ export function environmentRoutes( res.status(404).json({ error: "Environment customImage setup session not found" }); return; } - assertCustomImageCompanyAccess(req, session.companyId); + await resolveCustomImageSessionCompanyId(req, session); const result = await customImages.refreshSetupSession({ sessionId: session.id, includeConnectionPayload: true, @@ -430,7 +450,7 @@ export function environmentRoutes( res.status(404).json({ error: "Environment customImage setup session not found" }); return; } - assertCustomImageCompanyAccess(req, session.companyId); + const companyId = await resolveCustomImageSessionCompanyId(req, session); const refreshed = await customImages.refreshSetupSession({ sessionId: session.id, @@ -448,7 +468,7 @@ export function environmentRoutes( const minted = environmentCustomImageTerminalSessionStore.create({ setupSessionId: refreshed.session.id, - companyId: refreshed.session.companyId, + companyId, environmentId: refreshed.session.environmentId, provider: refreshed.session.provider, ssh: payloadValidation.ssh, @@ -459,7 +479,7 @@ export function environmentRoutes( const actor = getActorInfo(req); await logEnvironmentCustomImageActivity({ actor, - companyId: refreshed.session.companyId, + companyId, action: "environment.custom_image_terminal_session_token.created", entityId: refreshed.session.environmentId, details: { @@ -495,7 +515,7 @@ export function environmentRoutes( res.status(404).json({ error: "Environment customImage setup session not found" }); return; } - assertCustomImageCompanyAccess(req, session.companyId); + const companyId = await resolveCustomImageSessionCompanyId(req, session); const actor = getActorInfo(req); const result = await customImages.finishSetupSession({ sessionId: session.id, @@ -505,7 +525,7 @@ export function environmentRoutes( environmentCustomImageTerminalConnectionRegistry.closeBySetupSessionId(session.id, "setup_finished"); await logEnvironmentCustomImageActivity({ actor, - companyId: session.companyId, + companyId, action: "environment.custom_image_setup.finished", entityId: result.session.environmentId, details: { @@ -527,7 +547,7 @@ export function environmentRoutes( res.status(404).json({ error: "Environment customImage setup session not found" }); return; } - assertCustomImageCompanyAccess(req, session.companyId); + const companyId = await resolveCustomImageSessionCompanyId(req, session); const actor = getActorInfo(req); const cancelled = await customImages.cancelSetupSession({ sessionId: session.id, @@ -537,7 +557,7 @@ export function environmentRoutes( environmentCustomImageTerminalConnectionRegistry.closeBySetupSessionId(session.id, "setup_cancelled"); await logEnvironmentCustomImageActivity({ actor, - companyId: session.companyId, + companyId, action: "environment.custom_image_setup.cancelled", entityId: cancelled.environmentId, details: setupSessionActivityDetails(cancelled), @@ -551,7 +571,6 @@ export function environmentRoutes( const companyId = await resolveCustomImageCompanyId(req); const actor = getActorInfo(req); const result = await customImages.rollbackTemplate({ - companyId, environmentId: req.params.environmentId as string, }); await logEnvironmentCustomImageActivity({ @@ -572,7 +591,6 @@ export function environmentRoutes( const companyId = await resolveCustomImageCompanyId(req); const actor = getActorInfo(req); const template = await customImages.disableTemplate({ - companyId, environmentId: req.params.environmentId as string, deleteProviderTemplate: req.query.deleteProviderTemplate === "true", });