Fix custom image terminal CI failures

Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/__tests__/environment-custom-image-routes.test.ts

@@ -170,6 +170,7 @@ function createSession(overrides: Record<string, unknown> = {}) {
     },
     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,
     });

server/src/__tests__/environment-custom-image-terminal-ws.test.ts

@@ -86,6 +86,7 @@ function createSession(overrides: Record<string, unknown> = {}) {
     provider: "daytona",
     status: "waiting_for_user",
     expiresAt: futureDate(),
+    metadata: { setupRpcCompanyId: "company-1" },
     ...overrides,
   };
 }

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<string, unknown> | null;
 }
 
 interface CustomImageTerminalService {
@@ -162,6 +162,15 @@ function requireFutureSetupExpiry(session: { expiresAt: Date | string | null },
   return expiresAt;
 }
 
+function readSetupSessionCompanyId(session: {
+  metadata?: Record<string, unknown> | 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) {

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<string, unknown> | 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<string, unknown> | null },
+  ): Promise<string> {
+    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",
     });