diff --git a/packages/db/src/environment-custom-images-schema.test.ts b/packages/db/src/environment-custom-images-schema.test.ts index a3761a9e32..040b58df12 100644 --- a/packages/db/src/environment-custom-images-schema.test.ts +++ b/packages/db/src/environment-custom-images-schema.test.ts @@ -14,33 +14,33 @@ function indexWhere(table: Parameters[0], indexName: stri } describe("environment customImage schema", () => { - it("scopes template lookup and active uniqueness by company and environment", () => { + it("scopes template lookup and active uniqueness by environment", () => { expect(indexColumns( environmentCustomImageTemplates, - "environment_custom_image_templates_company_environment_status_idx", - )).toEqual(["company_id", "environment_id", "status"]); + "environment_custom_image_templates_environment_status_idx", + )).toEqual(["environment_id", "status"]); expect(indexColumns( environmentCustomImageTemplates, - "environment_custom_image_templates_company_environment_active_uq", - )).toEqual(["company_id", "environment_id"]); + "environment_custom_image_templates_environment_active_uq", + )).toEqual(["environment_id"]); expect(indexWhere( environmentCustomImageTemplates, - "environment_custom_image_templates_company_environment_active_uq", + "environment_custom_image_templates_environment_active_uq", )).toBeDefined(); }); - it("scopes setup-session lookup and active uniqueness by company and environment", () => { + it("scopes setup-session lookup and active uniqueness by environment", () => { expect(indexColumns( environmentCustomImageSetupSessions, - "environment_custom_image_setup_sessions_company_environment_status_idx", - )).toEqual(["company_id", "environment_id", "status"]); + "environment_custom_image_setup_sessions_environment_status_idx", + )).toEqual(["environment_id", "status"]); expect(indexColumns( environmentCustomImageSetupSessions, - "environment_custom_image_setup_sessions_company_environment_active_uq", - )).toEqual(["company_id", "environment_id"]); + "environment_custom_image_setup_sessions_environment_active_uq", + )).toEqual(["environment_id"]); expect(indexWhere( environmentCustomImageSetupSessions, - "environment_custom_image_setup_sessions_company_environment_active_uq", + "environment_custom_image_setup_sessions_environment_active_uq", )).toBeDefined(); }); }); diff --git a/packages/db/src/migrations/0127_environment_custom_images_instance_scoped.sql b/packages/db/src/migrations/0127_environment_custom_images_instance_scoped.sql new file mode 100644 index 0000000000..060606af8d --- /dev/null +++ b/packages/db/src/migrations/0127_environment_custom_images_instance_scoped.sql @@ -0,0 +1,106 @@ +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 + FROM "environment_custom_image_templates" + WHERE "status" = 'active' + GROUP BY "environment_id" + HAVING count(*) > 1 + ) THEN + RAISE EXCEPTION 'Cannot migrate environment custom image templates to environment scope while multiple active templates exist for the same environment. Revoke or supersede the extra active templates before retrying.'; + END IF; +END $$; +--> statement-breakpoint +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 + FROM "environment_custom_image_setup_sessions" + WHERE "status" IN ('starting', 'waiting_for_user', 'capturing') + GROUP BY "environment_id" + HAVING count(*) > 1 + ) THEN + RAISE EXCEPTION 'Cannot migrate environment custom image setup sessions to environment scope while multiple active sessions exist for the same environment. Finish or cancel the extra sessions before retrying.'; + END IF; +END $$; +--> statement-breakpoint +UPDATE "environment_custom_image_templates" +SET + "metadata" = jsonb_set( + COALESCE("metadata", '{}'::jsonb), + '{setupRpcCompanyId}', + to_jsonb("company_id"::text), + true + ), + "updated_at" = now() +WHERE "company_id" IS NOT NULL + AND COALESCE("metadata" ->> 'setupRpcCompanyId', '') = ''; +--> statement-breakpoint +UPDATE "environment_custom_image_setup_sessions" +SET + "metadata" = jsonb_set( + COALESCE("metadata", '{}'::jsonb), + '{setupRpcCompanyId}', + to_jsonb("company_id"::text), + true + ), + "updated_at" = now() +WHERE "company_id" IS NOT NULL + AND COALESCE("metadata" ->> 'setupRpcCompanyId', '') = ''; +--> statement-breakpoint +DROP INDEX IF EXISTS "environment_custom_image_templates_company_environment_status_idx"; +--> statement-breakpoint +DROP INDEX IF EXISTS "environment_custom_image_templates_company_provider_status_idx"; +--> statement-breakpoint +DROP INDEX IF EXISTS "environment_custom_image_templates_company_environment_active_uq"; +--> statement-breakpoint +DROP INDEX IF EXISTS "environment_custom_image_templates_company_last_used_idx"; +--> statement-breakpoint +DROP INDEX IF EXISTS "environment_custom_image_setup_sessions_company_environment_status_idx"; +--> statement-breakpoint +DROP INDEX IF EXISTS "environment_custom_image_setup_sessions_company_environment_active_uq"; +--> statement-breakpoint +DROP INDEX IF EXISTS "environment_custom_image_setup_sessions_company_template_idx"; +--> statement-breakpoint +DROP INDEX IF EXISTS "environment_custom_image_setup_sessions_company_promoted_template_idx"; +--> statement-breakpoint +DROP INDEX IF EXISTS "environment_custom_image_setup_sessions_company_expires_idx"; +--> statement-breakpoint +ALTER TABLE "environment_custom_image_templates" + DROP CONSTRAINT IF EXISTS "environment_custom_image_templates_company_id_companies_id_fk"; +--> statement-breakpoint +ALTER TABLE "environment_custom_image_setup_sessions" + DROP CONSTRAINT IF EXISTS "environment_custom_image_setup_sessions_company_id_companies_id_fk"; +--> statement-breakpoint +ALTER TABLE "environment_custom_image_templates" DROP COLUMN IF EXISTS "company_id"; +--> statement-breakpoint +ALTER TABLE "environment_custom_image_setup_sessions" DROP COLUMN IF EXISTS "company_id"; +--> statement-breakpoint +CREATE INDEX "environment_custom_image_templates_environment_status_idx" + ON "environment_custom_image_templates" USING btree ("environment_id", "status"); +--> statement-breakpoint +CREATE INDEX "environment_custom_image_templates_environment_provider_status_idx" + ON "environment_custom_image_templates" USING btree ("environment_id", "provider", "status"); +--> statement-breakpoint +CREATE UNIQUE INDEX "environment_custom_image_templates_environment_active_uq" + ON "environment_custom_image_templates" USING btree ("environment_id") + WHERE "status" = 'active'; +--> statement-breakpoint +CREATE INDEX "environment_custom_image_templates_last_used_idx" + ON "environment_custom_image_templates" USING btree ("last_used_at"); +--> statement-breakpoint +CREATE INDEX "environment_custom_image_setup_sessions_environment_status_idx" + ON "environment_custom_image_setup_sessions" USING btree ("environment_id", "status"); +--> statement-breakpoint +CREATE UNIQUE INDEX "environment_custom_image_setup_sessions_environment_active_uq" + ON "environment_custom_image_setup_sessions" USING btree ("environment_id") + WHERE "status" IN ('starting', 'waiting_for_user', 'capturing'); +--> statement-breakpoint +CREATE INDEX "environment_custom_image_setup_sessions_template_idx" + ON "environment_custom_image_setup_sessions" USING btree ("template_id"); +--> statement-breakpoint +CREATE INDEX "environment_custom_image_setup_sessions_promoted_template_idx" + ON "environment_custom_image_setup_sessions" USING btree ("promoted_template_id"); +--> statement-breakpoint +CREATE INDEX "environment_custom_image_setup_sessions_expires_idx" + ON "environment_custom_image_setup_sessions" USING btree ("expires_at"); diff --git a/packages/db/src/migrations/meta/_journal.json b/packages/db/src/migrations/meta/_journal.json index fa5b9832c9..9cac3b4213 100644 --- a/packages/db/src/migrations/meta/_journal.json +++ b/packages/db/src/migrations/meta/_journal.json @@ -890,6 +890,13 @@ "when": 1782526400000, "tag": "0126_issue_comment_derived_attribution", "breakpoints": true + }, + { + "idx": 127, + "version": "7", + "when": 1782526500000, + "tag": "0127_environment_custom_images_instance_scoped", + "breakpoints": true } ] } diff --git a/packages/db/src/schema/environment_custom_image_setup_sessions.ts b/packages/db/src/schema/environment_custom_image_setup_sessions.ts index e88af26b89..d7a30bf9cb 100644 --- a/packages/db/src/schema/environment_custom_image_setup_sessions.ts +++ b/packages/db/src/schema/environment_custom_image_setup_sessions.ts @@ -5,7 +5,6 @@ import type { EnvironmentCustomImageSetupSessionStatus, } from "@paperclipai/shared"; import { agents } from "./agents.js"; -import { companies } from "./companies.js"; import { environmentLeases } from "./environment_leases.js"; import { environmentCustomImageTemplates } from "./environment_custom_image_templates.js"; import { environments } from "./environments.js"; @@ -14,7 +13,6 @@ export const environmentCustomImageSetupSessions = pgTable( "environment_custom_image_setup_sessions", { id: uuid("id").primaryKey().defaultRandom(), - companyId: uuid("company_id").notNull().references(() => companies.id, { onDelete: "cascade" }), environmentId: uuid("environment_id").notNull().references(() => environments.id, { onDelete: "cascade" }), templateId: uuid("template_id").references(() => environmentCustomImageTemplates.id, { onDelete: "set null" }), promotedTemplateId: uuid("promoted_template_id") @@ -36,26 +34,17 @@ export const environmentCustomImageSetupSessions = pgTable( updatedAt: timestamp("updated_at", { withTimezone: true }).notNull().defaultNow(), }, (table) => ({ - companyEnvironmentStatusIdx: index("environment_custom_image_setup_sessions_company_environment_status_idx").on( - table.companyId, + environmentStatusIdx: index("environment_custom_image_setup_sessions_environment_status_idx").on( table.environmentId, table.status, ), - companyEnvironmentActiveUq: uniqueIndex("environment_custom_image_setup_sessions_company_environment_active_uq") - .on(table.companyId, table.environmentId) + environmentActiveUq: uniqueIndex("environment_custom_image_setup_sessions_environment_active_uq") + .on(table.environmentId) .where(sql`${table.status} IN ('starting', 'waiting_for_user', 'capturing')`), - companyTemplateIdx: index("environment_custom_image_setup_sessions_company_template_idx").on( - table.companyId, - table.templateId, - ), - companyPromotedTemplateIdx: index("environment_custom_image_setup_sessions_company_promoted_template_idx").on( - table.companyId, - table.promotedTemplateId, - ), - companyExpiresIdx: index("environment_custom_image_setup_sessions_company_expires_idx").on( - table.companyId, - table.expiresAt, - ), + templateIdx: index("environment_custom_image_setup_sessions_template_idx").on(table.templateId), + promotedTemplateIdx: index("environment_custom_image_setup_sessions_promoted_template_idx") + .on(table.promotedTemplateId), + expiresIdx: index("environment_custom_image_setup_sessions_expires_idx").on(table.expiresAt), providerLeaseIdx: index("environment_custom_image_setup_sessions_provider_lease_idx").on( table.provider, table.providerLeaseId, diff --git a/packages/db/src/schema/environment_custom_image_templates.ts b/packages/db/src/schema/environment_custom_image_templates.ts index aadd8080ba..589906cfa7 100644 --- a/packages/db/src/schema/environment_custom_image_templates.ts +++ b/packages/db/src/schema/environment_custom_image_templates.ts @@ -11,14 +11,12 @@ import { } from "drizzle-orm/pg-core"; import type { EnvironmentCustomImageTemplateStatus } from "@paperclipai/shared"; import { agents } from "./agents.js"; -import { companies } from "./companies.js"; import { environments } from "./environments.js"; export const environmentCustomImageTemplates = pgTable( "environment_custom_image_templates", { id: uuid("id").primaryKey().defaultRandom(), - companyId: uuid("company_id").notNull().references(() => companies.id, { onDelete: "cascade" }), environmentId: uuid("environment_id").notNull().references(() => environments.id, { onDelete: "cascade" }), provider: text("provider").notNull(), templateKind: text("template_kind").notNull().default("unknown"), @@ -37,23 +35,19 @@ export const environmentCustomImageTemplates = pgTable( updatedAt: timestamp("updated_at", { withTimezone: true }).notNull().defaultNow(), }, (table) => ({ - companyEnvironmentStatusIdx: index("environment_custom_image_templates_company_environment_status_idx").on( - table.companyId, + environmentStatusIdx: index("environment_custom_image_templates_environment_status_idx").on( table.environmentId, table.status, ), - companyProviderStatusIdx: index("environment_custom_image_templates_company_provider_status_idx").on( - table.companyId, + environmentProviderStatusIdx: index("environment_custom_image_templates_environment_provider_status_idx").on( + table.environmentId, table.provider, table.status, ), - companyEnvironmentActiveUq: uniqueIndex("environment_custom_image_templates_company_environment_active_uq") - .on(table.companyId, table.environmentId) + environmentActiveUq: uniqueIndex("environment_custom_image_templates_environment_active_uq") + .on(table.environmentId) .where(sql`${table.status} = 'active'`), supersededByIdx: index("environment_custom_image_templates_superseded_by_idx").on(table.supersededByTemplateId), - companyLastUsedIdx: index("environment_custom_image_templates_company_last_used_idx").on( - table.companyId, - table.lastUsedAt, - ), + lastUsedIdx: index("environment_custom_image_templates_last_used_idx").on(table.lastUsedAt), }), ); diff --git a/packages/shared/src/environment-custom-images.test.ts b/packages/shared/src/environment-custom-images.test.ts index add76996c9..e102d95e5c 100644 --- a/packages/shared/src/environment-custom-images.test.ts +++ b/packages/shared/src/environment-custom-images.test.ts @@ -11,16 +11,14 @@ import { startEnvironmentCustomImageSetupSessionSchema, } from "./validators/environment-custom-images.js"; -const companyId = "11111111-1111-4111-8111-111111111111"; const environmentId = "22222222-2222-4222-8222-222222222222"; const templateId = "33333333-3333-4333-8333-333333333333"; const sessionId = "44444444-4444-4444-8444-444444444444"; describe("environment customImage validators", () => { - it("requires company and environment scope on templates and setup sessions", () => { + it("requires environment scope on templates and setup sessions", () => { const template = environmentCustomImageTemplateSchema.parse({ id: templateId, - companyId, environmentId, provider: "daytona", templateKind: "snapshot", @@ -38,12 +36,10 @@ describe("environment customImage validators", () => { updatedAt: "2026-06-25T12:00:00.000Z", }); - expect(template.companyId).toBe(companyId); expect(template.environmentId).toBe(environmentId); const session = environmentCustomImageSetupSessionSchema.parse({ id: sessionId, - companyId, environmentId, templateId, promotedTemplateId: null, @@ -69,7 +65,6 @@ describe("environment customImage validators", () => { updatedAt: "2026-06-25T12:00:00.000Z", }); - expect(session.companyId).toBe(companyId); expect(session.environmentId).toBe(environmentId); }); diff --git a/packages/shared/src/validators/environment-custom-images.ts b/packages/shared/src/validators/environment-custom-images.ts index 523f1d2155..bd1b76a6c1 100644 --- a/packages/shared/src/validators/environment-custom-images.ts +++ b/packages/shared/src/validators/environment-custom-images.ts @@ -32,7 +32,6 @@ export type EnvironmentCustomImageSetupConnectionSummary = export const environmentCustomImageTemplateSchema = z.object({ id: z.string().uuid(), - companyId: z.string().uuid(), environmentId: z.string().uuid(), provider: providerKeySchema, templateKind: environmentCustomImageTemplateKindSchema, @@ -54,7 +53,6 @@ export type EnvironmentCustomImageTemplate = export const environmentCustomImageSetupSessionSchema = z.object({ id: z.string().uuid(), - companyId: z.string().uuid(), environmentId: z.string().uuid(), templateId: z.string().uuid().nullable(), promotedTemplateId: z.string().uuid().nullable(), diff --git a/server/src/__tests__/environment-custom-image-routes.test.ts b/server/src/__tests__/environment-custom-image-routes.test.ts index 3d90806ea9..6d3268c78f 100644 --- a/server/src/__tests__/environment-custom-image-routes.test.ts +++ b/server/src/__tests__/environment-custom-image-routes.test.ts @@ -120,7 +120,6 @@ function createEnvironment(overrides: Record = {}) { function createTemplate(overrides: Record = {}) { return { id: "template-1", - companyId: "company-1", environmentId: "env-1", provider: "daytona", templateKind: "snapshot", @@ -143,7 +142,6 @@ function createTemplate(overrides: Record = {}) { function createSession(overrides: Record = {}) { return { id: "session-1", - companyId: "company-1", environmentId: "env-1", templateId: "template-1", promotedTemplateId: null, @@ -282,13 +280,12 @@ describe("environment customImage setup routes", () => { it("starts a setup session, returns the live payload, and logs redacted details", async () => { const res = await request(createApp(boardActor())) - .post("/api/environments/env-1/custom-image-setup-sessions?companyId=company-1") + .post("/api/environments/env-1/custom-image-setup-sessions") .send({ ttlSeconds: 3600 }); 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, @@ -320,7 +317,7 @@ describe("environment customImage setup routes", () => { it("denies agent API key actors before customImage state or payloads are read", async () => { const app = createApp(agentActor()); const start = await request(app) - .post("/api/environments/env-1/custom-image-setup-sessions?companyId=company-1") + .post("/api/environments/env-1/custom-image-setup-sessions") .send({}); const status = await request(app) .get("/api/environment-custom-image-setup-sessions/session-1"); @@ -346,18 +343,17 @@ describe("environment customImage setup routes", () => { expect(mockEnvironmentCustomImageService.refreshSetupSession).not.toHaveBeenCalled(); }); - it("denies single-company fallback when the board actor is not a member", async () => { + it("does not require a company fallback to start setup", async () => { mockInstanceSettingsService.listCompanyIds.mockResolvedValue(["company-2"]); - const res = await request(createApp(boardActor({ - companyIds: ["company-1"], - isInstanceAdmin: false, - }))) + const res = await request(createApp(boardActor())) .post("/api/environments/env-1/custom-image-setup-sessions") .send({}); - expect(res.status).toBe(403); - expect(mockEnvironmentCustomImageService.startSetupSession).not.toHaveBeenCalled(); + expect(res.status).toBe(201); + expect(mockEnvironmentCustomImageService.startSetupSession).toHaveBeenCalledWith(expect.objectContaining({ + environmentId: "env-1", + })); }); it("finishes and promotes a template while logging redacted template details", async () => { @@ -395,22 +391,20 @@ describe("environment customImage setup routes", () => { expect(loggedActivityJson()).not.toContain("lease-secret"); }); - it("rolls back and disables active templates through company-scoped routes", async () => { + it("rolls back and disables active templates through instance-scoped routes", async () => { const app = createApp(boardActor()); const rollback = await request(app) - .post("/api/environments/env-1/custom-image-template/rollback?companyId=company-1") + .post("/api/environments/env-1/custom-image-template/rollback") .send({}); const disable = await request(app) - .delete("/api/environments/env-1/custom-image-template?companyId=company-1&deleteProviderTemplate=true"); + .delete("/api/environments/env-1/custom-image-template?deleteProviderTemplate=true"); 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-images-service.test.ts b/server/src/__tests__/environment-custom-images-service.test.ts index c3cf2c10a0..883e4c7f32 100644 --- a/server/src/__tests__/environment-custom-images-service.test.ts +++ b/server/src/__tests__/environment-custom-images-service.test.ts @@ -157,12 +157,10 @@ describeEmbeddedPostgres("environmentCustomImageService", () => { async function seed() { const companyId = randomUUID(); - const otherCompanyId = randomUUID(); const environmentId = randomUUID(); - await db.insert(companies).values([ + await db.insert(companies).values( { id: companyId, name: "Acme", issuePrefix: `A${companyId.slice(0, 4)}` }, - { id: otherCompanyId, name: "Other", issuePrefix: `B${otherCompanyId.slice(0, 4)}` }, - ]); + ); await db.insert(environments).values({ id: environmentId, name: `Fake ${environmentId.slice(0, 8)}`, @@ -184,16 +182,15 @@ describeEmbeddedPostgres("environmentCustomImageService", () => { manifestJson: pluginManifest(), status: "ready", }); - return { companyId, otherCompanyId, environmentId }; + return { companyId, environmentId }; } it("starts, refreshes, finishes, refreshes again, and rolls back setup sessions", async () => { - const { companyId, environmentId } = await seed(); + const { environmentId } = await seed(); const workerManager = createWorkerManager(); const service = environmentCustomImageService(db, { pluginWorkerManager: workerManager }); const started = await service.startSetupSession({ - companyId, environmentId, actor: { userId: "user-1" }, ttlSeconds: 600, @@ -221,25 +218,69 @@ describeEmbeddedPostgres("environmentCustomImageService", () => { }); const refresh = await service.startSetupSession({ - companyId, environmentId, actor: { userId: "user-1" }, }); const replacement = await service.finishSetupSession({ sessionId: refresh.session.id }); expect(replacement.template.id).not.toBe(promoted.template.id); - const rollback = await service.rollbackTemplate({ companyId, environmentId }); + const rollback = await service.rollbackTemplate({ environmentId }); expect(rollback.activeTemplate.id).toBe(promoted.template.id); expect(rollback.supersededTemplate.id).toBe(replacement.template.id); }); - it("revokes the active template before deleting the provider template", async () => { + it("reuses the setup provider company context across lifecycle calls", async () => { const { companyId, environmentId } = await seed(); const workerManager = createWorkerManager(); const service = environmentCustomImageService(db, { pluginWorkerManager: workerManager }); const started = await service.startSetupSession({ - companyId, + environmentId, + actor: { userId: "user-1" }, + secretContextCompanyId: companyId, + }); + expect(started.session.metadata).toMatchObject({ + setupRpcCompanyId: companyId, + }); + + await service.refreshSetupSession({ + sessionId: started.session.id, + includeConnectionPayload: true, + }); + const promoted = await service.finishSetupSession({ sessionId: started.session.id }); + + const lifecycleCalls = workerManager.call.mock.calls + .filter(([, method]) => [ + "environmentStartInteractiveSetup", + "environmentGetInteractiveSetup", + "environmentCaptureTemplate", + "environmentCancelInteractiveSetup", + ].includes(method)) + .map(([, method, params]) => ({ + method, + companyId: (params as Record).companyId, + })); + + expect(lifecycleCalls).toEqual([ + { method: "environmentStartInteractiveSetup", companyId }, + { method: "environmentGetInteractiveSetup", companyId }, + { method: "environmentCaptureTemplate", companyId }, + { method: "environmentCancelInteractiveSetup", companyId }, + ]); + expect(promoted.session.metadata).toMatchObject({ + setupRpcCompanyId: companyId, + }); + expect(promoted.template.metadata).toMatchObject({ + setupRpcCompanyId: companyId, + }); + }); + + it("revokes the active template before deleting the provider template", async () => { + const { environmentId } = await seed(); + const workerManager = createWorkerManager(); + const service = environmentCustomImageService(db, { pluginWorkerManager: workerManager }); + + const started = await service.startSetupSession({ environmentId, actor: { userId: "user-1" }, }); @@ -258,7 +299,6 @@ describeEmbeddedPostgres("environmentCustomImageService", () => { }); const disabled = await service.disableTemplate({ - companyId, environmentId, deleteProviderTemplate: true, }); @@ -276,13 +316,53 @@ describeEmbeddedPostgres("environmentCustomImageService", () => { ); }); - it("cancels and times out setup sessions without changing the active template", async () => { - const { companyId, environmentId } = await seed(); + it("does not send provider template refs to a different current provider", async () => { + const { environmentId } = await seed(); + const workerManager = createWorkerManager(); + const service = environmentCustomImageService(db, { pluginWorkerManager: workerManager }); + + const started = await service.startSetupSession({ + environmentId, + actor: { userId: "user-1" }, + }); + const promoted = await service.finishSetupSession({ sessionId: started.session.id }); + const deleteCallsBefore = workerManager.call.mock.calls + .filter(([, method]) => method === "environmentDeleteTemplate") + .length; + + await db.update(environments) + .set({ + config: { + provider: "other-plugin", + image: "other:base", + reuseLease: false, + }, + }) + .where(eq(environments.id, environmentId)); + + await expect(service.disableTemplate({ + environmentId, + deleteProviderTemplate: true, + })).rejects.toThrow("Environment customImage provider changed"); + + const deleteCallsAfter = workerManager.call.mock.calls + .filter(([, method]) => method === "environmentDeleteTemplate") + .length; + const [templateRow] = await db + .select({ status: environmentCustomImageTemplates.status }) + .from(environmentCustomImageTemplates) + .where(eq(environmentCustomImageTemplates.id, promoted.template.id)); + + expect(deleteCallsAfter).toBe(deleteCallsBefore); + expect(templateRow?.status).toBe("active"); + }); + + it("cancels and times out setup sessions without changing the active template", async () => { + const { environmentId } = await seed(); const workerManager = createWorkerManager(); const service = environmentCustomImageService(db, { pluginWorkerManager: workerManager }); const started = await service.startSetupSession({ - companyId, environmentId, actor: { userId: "user-1" }, }); @@ -291,10 +371,9 @@ describeEmbeddedPostgres("environmentCustomImageService", () => { reason: "user_cancelled", }); expect(cancelled.status).toBe("cancelled"); - expect(await service.getActiveTemplate({ companyId, environmentId })).toBeNull(); + expect(await service.getActiveTemplate({ environmentId })).toBeNull(); const expired = await service.startSetupSession({ - companyId, environmentId, actor: { userId: "user-1" }, ttlSeconds: 60, @@ -308,21 +387,32 @@ describeEmbeddedPostgres("environmentCustomImageService", () => { expect(timedOut?.status).toBe("timed_out"); }); - it("rejects templates from another company or environment", async () => { - const { companyId, otherCompanyId, environmentId } = await seed(); + it("rejects templates from another environment", async () => { + const { environmentId } = await seed(); + const otherEnvironmentId = randomUUID(); + await db.insert(environments).values({ + id: otherEnvironmentId, + name: `Other ${otherEnvironmentId.slice(0, 8)}`, + driver: "sandbox", + status: "active", + config: { + provider: "fake-plugin", + image: "fake:base", + reuseLease: false, + }, + envVars: {}, + }); const workerManager = createWorkerManager(); const service = environmentCustomImageService(db, { pluginWorkerManager: workerManager }); const [otherTemplate] = await db.insert(environmentCustomImageTemplates).values({ - companyId: otherCompanyId, - environmentId, + environmentId: otherEnvironmentId, provider: "fake-plugin", templateKind: "snapshot", - templateRef: "snapshot-other-company", + templateRef: "snapshot-other-environment", status: "active", }).returning(); await expect(service.startSetupSession({ - companyId, environmentId, templateId: otherTemplate!.id, actor: { userId: "user-1" }, @@ -342,7 +432,6 @@ describeEmbeddedPostgres("environmentCustomImageService", () => { expect(fallback.config).toMatchObject({ image: "fake:base" }); await db.insert(environmentCustomImageTemplates).values({ - companyId, environmentId, provider: "fake-plugin", templateKind: "snapshot", @@ -380,7 +469,6 @@ describeEmbeddedPostgres("environmentCustomImageService", () => { const environment = await db.select().from(environments).where(eq(environments.id, environmentId)).then((rows) => rows[0]!); await db.insert(environmentCustomImageTemplates).values({ - companyId, environmentId, provider: "fake-plugin", templateKind: "snapshot", @@ -414,7 +502,6 @@ describeEmbeddedPostgres("environmentCustomImageService", () => { const service = environmentCustomImageService(db, { pluginWorkerManager: workerManager }); const started = await service.startSetupSession({ - companyId, environmentId, actor: { userId: "user-1" }, }); diff --git a/server/src/__tests__/environment-routes.test.ts b/server/src/__tests__/environment-routes.test.ts index a45b298012..d830c8cdaf 100644 --- a/server/src/__tests__/environment-routes.test.ts +++ b/server/src/__tests__/environment-routes.test.ts @@ -1346,6 +1346,7 @@ describe("environment routes", () => { expect(mockProbeEnvironment).toHaveBeenCalledWith(expect.anything(), environment, { companyId: null, pluginWorkerManager: undefined, + applyCustomImageTemplate: false, }); expect(mockLogActivity).toHaveBeenCalledWith( expect.anything(), @@ -1444,6 +1445,7 @@ describe("environment routes", () => { expect(mockProbeEnvironment).toHaveBeenCalledWith(expect.anything(), environment, { companyId: null, pluginWorkerManager: undefined, + applyCustomImageTemplate: true, }); expect(mockLogActivity).toHaveBeenCalledWith( expect.anything(), @@ -1460,6 +1462,48 @@ describe("environment routes", () => { ); }); + it("probes saved sandbox environments with the active custom image template without company context", async () => { + const environment = { + ...createEnvironment(), + id: "env-sandbox", + name: "Daytona Sandbox", + driver: "sandbox" as const, + config: { + provider: "daytona", + image: "ubuntu:24.04", + reuseLease: true, + }, + }; + mockEnvironmentService.getById.mockResolvedValue(environment); + mockProbeEnvironment.mockResolvedValue({ + ok: true, + driver: "sandbox", + summary: "Connected to Daytona sandbox.", + details: { + provider: "daytona", + snapshot: "captured-template", + }, + }); + const app = createApp({ + type: "board", + userId: "user-1", + source: "local_implicit", + runId: "run-1", + }); + + const res = await request(app) + .post(`/api/environments/${environment.id}/probe`) + .send({}); + + expect(res.status).toBe(200); + expect(res.body.driver).toBe("sandbox"); + expect(mockProbeEnvironment).toHaveBeenCalledWith(expect.anything(), environment, { + companyId: null, + pluginWorkerManager: undefined, + applyCustomImageTemplate: true, + }); + }); + it("probes unsaved provider config without persisting secrets", async () => { mockProbeEnvironment.mockResolvedValue({ ok: true, diff --git a/server/src/routes/environments.ts b/server/src/routes/environments.ts index b1f1887910..2ac76b9393 100644 --- a/server/src/routes/environments.ts +++ b/server/src/routes/environments.ts @@ -71,17 +71,6 @@ export function environmentRoutes( assertBoardOrgAccess(req); } - function assertCustomImageCompanyAccess(req: Request, companyId: string) { - if (req.actor.type !== "board") { - throw forbidden("Board access required"); - } - if (req.actor.source === "local_implicit" || req.actor.isInstanceAdmin) return; - const allowedCompanies = req.actor.companyIds ?? []; - if (!allowedCompanies.includes(companyId)) { - throw forbidden("User does not have access to this company"); - } - } - function canReadFullInstanceEnvironment(req: Request) { return req.actor.type === "board" && (req.actor.source === "local_implicit" || req.actor.isInstanceAdmin); @@ -139,47 +128,6 @@ export function environmentRoutes( ); } - async function logEnvironmentCustomImageActivity(input: { - actor: ReturnType; - companyId: string; - action: string; - entityId: string; - details: Record; - }) { - await logActivity(db, { - companyId: input.companyId, - actorType: input.actor.actorType, - actorId: input.actor.actorId, - agentId: input.actor.agentId, - runId: input.actor.runId, - action: input.action, - entityType: "environment", - entityId: input.entityId, - details: input.details, - }); - } - - async function resolveCustomImageCompanyId(req: Request): Promise { - const queryCompanyId = - typeof req.query.companyId === "string" && req.query.companyId.trim().length > 0 - ? req.query.companyId.trim() - : null; - if (queryCompanyId) { - assertCustomImageCompanyAccess(req, queryCompanyId); - return queryCompanyId; - } - if (req.actor.type === "board" && req.actor.companyIds?.length === 1) { - return req.actor.companyIds[0]!; - } - const companyIds = await instanceSettings.listCompanyIds(); - if (companyIds.length === 1 && companyIds[0]) { - const companyId = companyIds[0]; - assertCustomImageCompanyAccess(req, companyId); - return companyId; - } - throw unprocessable("companyId query parameter is required for environment customImage setup."); - } - async function resolveEnvironmentSecretContextCompanyId( req: Request, environmentId: string, @@ -339,9 +287,7 @@ export function environmentRoutes( router.get("/environments/:environmentId/custom-image-template", async (req, res) => { assertCanAccessInstanceEnvironments(req); - const companyId = await resolveCustomImageCompanyId(req); const overview = await customImages.getOverview({ - companyId, environmentId: req.params.environmentId as string, }); res.json(overview); @@ -352,11 +298,10 @@ export function environmentRoutes( validate(startEnvironmentCustomImageSetupSessionSchema), async (req, res) => { assertCanAccessInstanceEnvironments(req); - const companyId = await resolveCustomImageCompanyId(req); + const environmentId = req.params.environmentId as string; const actor = getActorInfo(req); const result = await customImages.startSetupSession({ - companyId, - environmentId: req.params.environmentId as string, + environmentId, templateId: req.body.templateId ?? null, ttlSeconds: req.body.ttlSeconds ?? null, actor: { @@ -364,9 +309,8 @@ export function environmentRoutes( agentId: actor.agentId, }, }); - await logEnvironmentCustomImageActivity({ + await logInstanceEnvironmentActivity({ actor, - companyId, action: "environment.custom_image_setup.started", entityId: result.session.environmentId, details: setupSessionActivityDetails(result.session), @@ -382,7 +326,6 @@ export function environmentRoutes( res.status(404).json({ error: "Environment customImage setup session not found" }); return; } - assertCustomImageCompanyAccess(req, session.companyId); const result = await customImages.refreshSetupSession({ sessionId: session.id, includeConnectionPayload: true, @@ -400,15 +343,13 @@ export function environmentRoutes( res.status(404).json({ error: "Environment customImage setup session not found" }); return; } - assertCustomImageCompanyAccess(req, session.companyId); const actor = getActorInfo(req); const result = await customImages.finishSetupSession({ sessionId: session.id, metadata: req.body.metadata, }); - await logEnvironmentCustomImageActivity({ + await logInstanceEnvironmentActivity({ actor, - companyId: session.companyId, action: "environment.custom_image_setup.finished", entityId: result.session.environmentId, details: { @@ -430,15 +371,13 @@ export function environmentRoutes( res.status(404).json({ error: "Environment customImage setup session not found" }); return; } - assertCustomImageCompanyAccess(req, session.companyId); const actor = getActorInfo(req); const cancelled = await customImages.cancelSetupSession({ sessionId: session.id, reason: req.body.reason ?? null, }); - await logEnvironmentCustomImageActivity({ + await logInstanceEnvironmentActivity({ actor, - companyId: session.companyId, action: "environment.custom_image_setup.cancelled", entityId: cancelled.environmentId, details: setupSessionActivityDetails(cancelled), @@ -449,15 +388,12 @@ export function environmentRoutes( router.post("/environments/:environmentId/custom-image-template/rollback", async (req, res) => { assertCanAccessInstanceEnvironments(req); - const companyId = await resolveCustomImageCompanyId(req); const actor = getActorInfo(req); const result = await customImages.rollbackTemplate({ - companyId, environmentId: req.params.environmentId as string, }); - await logEnvironmentCustomImageActivity({ + await logInstanceEnvironmentActivity({ actor, - companyId, action: "environment.custom_image_template.rolled_back", entityId: req.params.environmentId as string, details: { @@ -470,16 +406,13 @@ export function environmentRoutes( router.delete("/environments/:environmentId/custom-image-template", async (req, res) => { assertCanAccessInstanceEnvironments(req); - 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", }); - await logEnvironmentCustomImageActivity({ + await logInstanceEnvironmentActivity({ actor, - companyId, action: "environment.custom_image_template.disabled", entityId: req.params.environmentId as string, details: templateActivityDetails(template), @@ -715,6 +648,7 @@ export function environmentRoutes( const probe = await probeEnvironment(db, environment, { companyId: companyIdForSecrets, pluginWorkerManager: options.pluginWorkerManager, + applyCustomImageTemplate: environment.driver === "sandbox", }); await logInstanceEnvironmentActivity({ actor, diff --git a/server/src/services/environment-config.ts b/server/src/services/environment-config.ts index 2dfbb36228..05878ed276 100644 --- a/server/src/services/environment-config.ts +++ b/server/src/services/environment-config.ts @@ -556,7 +556,7 @@ export async function normalizeEnvironmentConfigForPersistence(input: { export async function resolveEnvironmentDriverConfigForRuntime( db: Db, - companyId: string, + companyId: string | null, environment: Pick & Partial>, context?: { issueId?: string | null; @@ -576,6 +576,9 @@ export async function resolveEnvironmentDriverConfigForRuntime( } if (parsed.driver === "ssh" && parsed.config.privateKeySecretRef) { + if (!companyId) { + throw unprocessable("Runtime secret resolution requires a companyId context"); + } return { driver: "ssh", config: { @@ -600,22 +603,31 @@ export async function resolveEnvironmentDriverConfigForRuntime( if (parsed.driver === "sandbox" && parsed.config.provider !== "fake") { const schema = await getSandboxProviderConfigSchema(db, parsed.config.provider); - const runtimeConfig = await resolveConfigSecretRefsForRuntime({ - db, - companyId, - config: parsed.config as Record, - schema, - context: { - consumerId: environmentId!, - issueId: context?.issueId ?? null, - heartbeatRunId: context?.heartbeatRunId ?? null, - }, - }) as SandboxEnvironmentConfig; + let runtimeConfig = parsed.config; + if (companyId) { + runtimeConfig = await resolveConfigSecretRefsForRuntime({ + db, + companyId, + config: parsed.config as Record, + schema, + context: { + consumerId: environmentId!, + issueId: context?.issueId ?? null, + heartbeatRunId: context?.heartbeatRunId ?? null, + }, + }) as SandboxEnvironmentConfig; + } else { + for (const path of collectSecretRefPaths(schema)) { + const current = readConfigValueAtPath(parsed.config as Record, path); + if (typeof current === "string" && isUuidSecretRef(current.trim())) { + throw unprocessable("Runtime secret resolution requires a companyId context"); + } + } + } return { driver: "sandbox", config: environmentId && (context?.issueId || context?.heartbeatRunId || context?.applyCustomImageTemplate) ? await resolveActiveEnvironmentCustomImageTemplateForRuntime(db, { - companyId, environmentId, baseConfig: parsed.config, runtimeConfig, diff --git a/server/src/services/environment-custom-image-runtime.ts b/server/src/services/environment-custom-image-runtime.ts index 0f938c1244..9c56321cd4 100644 --- a/server/src/services/environment-custom-image-runtime.ts +++ b/server/src/services/environment-custom-image-runtime.ts @@ -114,7 +114,6 @@ export function applyCustomImageTemplateToSandboxConfig( export function environmentCustomImageTemplateFromRow(row: TemplateRow): EnvironmentCustomImageTemplate { return { id: row.id, - companyId: row.companyId, environmentId: row.environmentId, provider: row.provider, templateKind: readEnvironmentCustomImageTemplateKind(row.templateKind), @@ -136,7 +135,6 @@ export function environmentCustomImageTemplateFromRow(row: TemplateRow): Environ export async function resolveActiveEnvironmentCustomImageTemplateForRuntime( db: Db, input: { - companyId: string; environmentId: string; baseConfig: SandboxEnvironmentConfig; runtimeConfig: SandboxEnvironmentConfig; @@ -147,7 +145,6 @@ export async function resolveActiveEnvironmentCustomImageTemplateForRuntime( .select() .from(environmentCustomImageTemplates) .where(and( - eq(environmentCustomImageTemplates.companyId, input.companyId), eq(environmentCustomImageTemplates.environmentId, input.environmentId), eq(environmentCustomImageTemplates.provider, input.baseConfig.provider), eq(environmentCustomImageTemplates.status, "active"), diff --git a/server/src/services/environment-custom-images.ts b/server/src/services/environment-custom-images.ts index a044d14e7f..ae3b2f0c6a 100644 --- a/server/src/services/environment-custom-images.ts +++ b/server/src/services/environment-custom-images.ts @@ -31,6 +31,7 @@ import { resolveSandboxProviderSecretRefPaths, stripSandboxProviderEnvelope, } from "./environment-config.js"; +import { secretService } from "./secrets.js"; import { resolvePluginExecuteRpcTimeoutMs, resolvePluginSandboxProviderDriverByKey, @@ -49,6 +50,8 @@ import type { PluginWorkerManager } from "./plugin-worker-manager.js"; const ACTIVE_SETUP_STATUSES = ["starting", "waiting_for_user", "capturing"] as const; const DEFAULT_SETUP_TTL_SECONDS = 60 * 60; const DEFAULT_CONNECTION_EXPIRES_IN_MINUTES = 15; +const SETUP_RPC_COMPANY_ID_METADATA_KEY = "setupRpcCompanyId"; +const SOURCE_ENVIRONMENT_CONFIG_FINGERPRINT_METADATA_KEY = "sourceEnvironmentConfigFingerprint"; type SetupSessionRow = typeof environmentCustomImageSetupSessions.$inferSelect; @@ -72,7 +75,6 @@ export interface EnvironmentCustomImageSetupCleanupResult { function toSession(row: SetupSessionRow): EnvironmentCustomImageSetupSession { return { id: row.id, - companyId: row.companyId, environmentId: row.environmentId, templateId: row.templateId ?? null, promotedTemplateId: row.promotedTemplateId ?? null, @@ -130,6 +132,42 @@ function normalizeProviderMetadata(metadata: Record | null | un return redactEnvironmentCustomImageValue(metadata); } +function metadataRecord(metadata: Record | null | undefined): Record { + return metadata && typeof metadata === "object" && !Array.isArray(metadata) ? metadata : {}; +} + +function normalizeSetupRpcCompanyId(value: unknown): string | null { + return typeof value === "string" && value.trim().length > 0 ? value.trim() : null; +} + +function readSetupRpcCompanyId(metadata: Record | null | undefined): string | null { + return normalizeSetupRpcCompanyId(metadataRecord(metadata)[SETUP_RPC_COMPANY_ID_METADATA_KEY]); +} + +function persistedSetupMetadata(metadata: Record | null | undefined): Record { + const record = metadataRecord(metadata); + const result: Record = {}; + const setupRpcCompanyId = normalizeSetupRpcCompanyId(record[SETUP_RPC_COMPANY_ID_METADATA_KEY]); + if (setupRpcCompanyId) { + result[SETUP_RPC_COMPANY_ID_METADATA_KEY] = setupRpcCompanyId; + } + const fingerprint = readString(record[SOURCE_ENVIRONMENT_CONFIG_FINGERPRINT_METADATA_KEY]); + if (fingerprint) { + result[SOURCE_ENVIRONMENT_CONFIG_FINGERPRINT_METADATA_KEY] = fingerprint; + } + return result; +} + +function mergeSetupSessionMetadata( + existing: Record | null | undefined, + providerMetadata: Record | null | undefined, +): Record | null { + const provider = normalizeProviderMetadata(providerMetadata) ?? {}; + const persisted = persistedSetupMetadata(existing); + const merged = { ...provider, ...persisted }; + return Object.keys(merged).length > 0 ? merged : null; +} + function normalizePersistedStatus( status: string, fallback: EnvironmentCustomImageSetupSessionStatus = "failed", @@ -179,10 +217,9 @@ function sourceTemplateFromConfig( async function resolveActiveTemplate( db: Db, - input: { companyId: string; environmentId: string; provider?: string | null }, + input: { environmentId: string; provider?: string | null }, ): Promise { const conditions = [ - eq(environmentCustomImageTemplates.companyId, input.companyId), eq(environmentCustomImageTemplates.environmentId, input.environmentId), eq(environmentCustomImageTemplates.status, "active"), ]; @@ -203,6 +240,7 @@ export function environmentCustomImageService( options: { pluginWorkerManager?: PluginWorkerManager } = {}, ) { const environments = environmentService(db); + const secrets = secretService(db); async function getTemplateById(id: string): Promise { const row = await db @@ -223,14 +261,12 @@ export function environmentCustomImageService( } async function getActiveSetupSession(input: { - companyId: string; environmentId: string; }): Promise { const row = await db .select() .from(environmentCustomImageSetupSessions) .where(and( - eq(environmentCustomImageSetupSessions.companyId, input.companyId), eq(environmentCustomImageSetupSessions.environmentId, input.environmentId), inArray(environmentCustomImageSetupSessions.status, [...ACTIVE_SETUP_STATUSES]), )) @@ -240,14 +276,12 @@ export function environmentCustomImageService( } async function getLatestSetupSession(input: { - companyId: string; environmentId: string; }): Promise { const row = await db .select() .from(environmentCustomImageSetupSessions) .where(and( - eq(environmentCustomImageSetupSessions.companyId, input.companyId), eq(environmentCustomImageSetupSessions.environmentId, input.environmentId), )) .orderBy(desc(environmentCustomImageSetupSessions.createdAt)) @@ -261,8 +295,25 @@ export function environmentCustomImageService( return environment; } + async function resolveSecretContextCompanyId( + environmentId: string, + explicitCompanyId?: string | null, + ): Promise { + if (explicitCompanyId) return explicitCompanyId; + const bindingCompanyIds = await secrets.listBindingCompanyIdsForTarget({ + targetType: "environment", + targetId: environmentId, + }); + if (bindingCompanyIds.length > 1) { + throw conflict("Environment secret bindings span multiple companies and require explicit companyId context."); + } + return bindingCompanyIds[0] ?? null; + } + async function resolveSetupProvider(input: { - companyId: string; + secretContextCompanyId?: string | null; + storedRpcCompanyId?: string | null; + storedProvider?: string | null; environment: Environment; requireCapture?: boolean; requireDelete?: boolean; @@ -270,9 +321,16 @@ export function environmentCustomImageService( if (!options.pluginWorkerManager) { throw unprocessable("Environment customImage setup requires a running plugin worker manager."); } + const storedRpcCompanyId = normalizeSetupRpcCompanyId(input.storedRpcCompanyId); + const secretContextCompanyId = storedRpcCompanyId + ? (storedRpcCompanyId === "instance" ? null : storedRpcCompanyId) + : await resolveSecretContextCompanyId( + input.environment.id, + input.secretContextCompanyId, + ); const parsed = await resolveEnvironmentDriverConfigForRuntime( db, - input.companyId, + secretContextCompanyId, input.environment, { issueId: null, heartbeatRunId: null }, ); @@ -280,6 +338,12 @@ export function environmentCustomImageService( throw unprocessable("Environment customImage setup is only supported for sandbox environments."); } const provider = parsed.config.provider; + const storedProvider = readString(input.storedProvider); + if (storedProvider && provider !== storedProvider) { + throw conflict( + `Environment customImage provider changed from "${storedProvider}" to "${provider}". Switch the environment back before continuing this customImage lifecycle operation.`, + ); + } const resolved = await resolvePluginSandboxProviderDriverByKey({ db, driverKey: provider, @@ -300,6 +364,7 @@ export function environmentCustomImageService( } return { provider, + rpcCompanyId: storedRpcCompanyId ?? secretContextCompanyId ?? "instance", pluginId: resolved.plugin.id, driver: resolved.driver, runtimeConfig: parsed.config, @@ -308,17 +373,20 @@ export function environmentCustomImageService( } async function callProviderStart(input: { - companyId: string; environment: Environment; sessionId: string; expiresAt: Date; sourceTemplateRef: string | null; sourceTemplateKind: EnvironmentCustomImageTemplateKind | null; + secretContextCompanyId?: string | null; }): Promise { - const provider = await resolveSetupProvider({ companyId: input.companyId, environment: input.environment }); + const provider = await resolveSetupProvider({ + secretContextCompanyId: input.secretContextCompanyId, + environment: input.environment, + }); return await options.pluginWorkerManager!.call(provider.pluginId, "environmentStartInteractiveSetup", { driverKey: provider.provider, - companyId: input.companyId, + companyId: provider.rpcCompanyId, environmentId: input.environment.id, issueId: null, config: provider.driverConfig, @@ -338,10 +406,14 @@ export function environmentCustomImageService( includeConnectionPayload: boolean; }): Promise { const environment = await requireEnvironment(input.session.environmentId); - const provider = await resolveSetupProvider({ companyId: input.session.companyId, environment }); + const provider = await resolveSetupProvider({ + environment, + storedProvider: input.session.provider, + storedRpcCompanyId: readSetupRpcCompanyId(input.session.metadata), + }); return await options.pluginWorkerManager!.call(provider.pluginId, "environmentGetInteractiveSetup", { driverKey: provider.provider, - companyId: input.session.companyId, + companyId: provider.rpcCompanyId, environmentId: environment.id, issueId: null, config: provider.driverConfig, @@ -361,13 +433,14 @@ export function environmentCustomImageService( }): Promise { const environment = await requireEnvironment(input.session.environmentId); const provider = await resolveSetupProvider({ - companyId: input.session.companyId, environment, + storedProvider: input.session.provider, + storedRpcCompanyId: readSetupRpcCompanyId(input.session.metadata), requireCapture: true, }); return await options.pluginWorkerManager!.call(provider.pluginId, "environmentCaptureTemplate", { driverKey: provider.provider, - companyId: input.session.companyId, + companyId: provider.rpcCompanyId, environmentId: environment.id, issueId: null, config: provider.driverConfig, @@ -388,10 +461,14 @@ export function environmentCustomImageService( reason: string | null; }): Promise { const environment = await requireEnvironment(input.session.environmentId); - const provider = await resolveSetupProvider({ companyId: input.session.companyId, environment }); + const provider = await resolveSetupProvider({ + environment, + storedProvider: input.session.provider, + storedRpcCompanyId: readSetupRpcCompanyId(input.session.metadata), + }); return await options.pluginWorkerManager!.call(provider.pluginId, "environmentCancelInteractiveSetup", { driverKey: provider.provider, - companyId: input.session.companyId, + companyId: provider.rpcCompanyId, environmentId: environment.id, issueId: null, config: provider.driverConfig, @@ -404,26 +481,37 @@ export function environmentCustomImageService( })); } + async function resolveTemplateDeleteProvider(template: EnvironmentCustomImageTemplate) { + if (!template.templateRef) { + throw unprocessable("Cannot delete an environment customImage template without a provider template ref."); + } + const environment = await requireEnvironment(template.environmentId); + return await resolveSetupProvider({ + environment, + storedProvider: template.provider, + storedRpcCompanyId: readSetupRpcCompanyId(template.metadata), + requireDelete: true, + }); + } + async function callProviderDeleteTemplate(input: { template: EnvironmentCustomImageTemplate; + provider: Awaited>; reason: string | null; }) { - if (!input.template.templateRef) { + const templateRef = input.template.templateRef; + if (!templateRef) { throw unprocessable("Cannot delete an environment customImage template without a provider template ref."); } const environment = await requireEnvironment(input.template.environmentId); - const provider = await resolveSetupProvider({ - companyId: input.template.companyId, - environment, - requireDelete: true, - }); + const provider = input.provider; return await options.pluginWorkerManager!.call(provider.pluginId, "environmentDeleteTemplate", { driverKey: provider.provider, - companyId: input.template.companyId, + companyId: provider.rpcCompanyId, environmentId: environment.id, issueId: null, config: provider.driverConfig, - templateRef: input.template.templateRef, + templateRef, templateKind: input.template.templateKind as PluginEnvironmentTemplateRefKind, metadata: input.template.metadata ?? undefined, reason: input.reason, @@ -434,7 +522,7 @@ export function environmentCustomImageService( } async function updateSessionFromProvider( - sessionId: string, + session: Pick, providerSession: PluginEnvironmentInteractiveSetupSession, fallbackStatus: EnvironmentCustomImageSetupSessionStatus = "failed", ): Promise { @@ -447,11 +535,11 @@ export function environmentCustomImageService( status, connectionSummary: normalizeConnectionSummary(providerSession.connectionSummary), expiresAt: providerSession.expiresAt ? new Date(providerSession.expiresAt) : undefined, - metadata: normalizeProviderMetadata(providerSession.metadata), + metadata: mergeSetupSessionMetadata(session.metadata, providerSession.metadata), failureReason: status === "failed" ? "Provider setup session failed or is missing." : null, updatedAt: now, }) - .where(eq(environmentCustomImageSetupSessions.id, sessionId)) + .where(eq(environmentCustomImageSetupSessions.id, session.id)) .returning() .then((rows) => rows[0] ?? null); if (!row) throw notFound("Environment customImage setup session not found"); @@ -509,7 +597,6 @@ export function environmentCustomImageService( return { getOverview: async (input: { - companyId: string; environmentId: string; }): Promise => { await requireEnvironment(input.environmentId); @@ -522,7 +609,6 @@ export function environmentCustomImageService( }, getActiveTemplate: async (input: { - companyId: string; environmentId: string; provider?: string | null; }): Promise => resolveActiveTemplate(db, input), @@ -530,15 +616,18 @@ export function environmentCustomImageService( getSessionById, startSetupSession: async (input: { - companyId: string; environmentId: string; templateId?: string | null; ttlSeconds?: number | null; actor: { userId?: string | null; agentId?: string | null }; + secretContextCompanyId?: string | null; now?: Date; }): Promise => { const environment = await requireEnvironment(input.environmentId); - const provider = await resolveSetupProvider({ companyId: input.companyId, environment }); + const provider = await resolveSetupProvider({ + secretContextCompanyId: input.secretContextCompanyId, + environment, + }); const activeSession = await getActiveSetupSession(input); if (activeSession) { throw conflict("An environment customImage setup session is already active for this environment."); @@ -546,7 +635,6 @@ export function environmentCustomImageService( const selectedTemplate = input.templateId ? await getTemplateById(input.templateId) : await resolveActiveTemplate(db, { - companyId: input.companyId, environmentId: input.environmentId, provider: provider.provider, }); @@ -556,13 +644,12 @@ export function environmentCustomImageService( if ( selectedTemplate && ( - selectedTemplate.companyId !== input.companyId || selectedTemplate.environmentId !== input.environmentId || selectedTemplate.provider !== provider.provider || selectedTemplate.status !== "active" ) ) { - throw unprocessable("Setup template must be the active template for this company and environment."); + throw unprocessable("Setup template must be the active template for this environment."); } const source = selectedTemplate ? { @@ -579,9 +666,12 @@ export function environmentCustomImageService( const expiresAt = addSeconds(now, ttlSeconds); const sessionId = randomUUID(); const fingerprint = fingerprintEnvironmentSandboxProviderConfig(provider.runtimeConfig); + const setupMetadata = { + [SOURCE_ENVIRONMENT_CONFIG_FINGERPRINT_METADATA_KEY]: fingerprint, + [SETUP_RPC_COMPANY_ID_METADATA_KEY]: provider.rpcCompanyId, + }; await db.insert(environmentCustomImageSetupSessions).values({ id: sessionId, - companyId: input.companyId, environmentId: input.environmentId, templateId: selectedTemplate?.id ?? null, provider: provider.provider, @@ -590,23 +680,21 @@ export function environmentCustomImageService( startedByAgentId: input.actor.agentId ?? null, baseTemplateRef: source.sourceTemplateRef, expiresAt, - metadata: { - sourceEnvironmentConfigFingerprint: fingerprint, - }, + metadata: setupMetadata, createdAt: now, updatedAt: now, }); try { const providerSession = await callProviderStart({ - companyId: input.companyId, environment, sessionId, expiresAt, sourceTemplateRef: source.sourceTemplateRef, sourceTemplateKind: source.sourceTemplateKind, + secretContextCompanyId: input.secretContextCompanyId, }); - const session = await updateSessionFromProvider(sessionId, providerSession); + const session = await updateSessionFromProvider({ id: sessionId, metadata: setupMetadata }, providerSession); return { session, connectionPayload: providerSession.connectionPayload ?? null, @@ -644,7 +732,7 @@ export function environmentCustomImageService( includeConnectionPayload: input.includeConnectionPayload ?? false, }); return { - session: await updateSessionFromProvider(session.id, providerSession), + session: await updateSessionFromProvider(session, providerSession), connectionPayload: providerSession.connectionPayload ?? null, }; }, @@ -666,7 +754,6 @@ export function environmentCustomImageService( } await markSessionStatus({ sessionId: session.id, status: "capturing" }); const currentActive = await resolveActiveTemplate(db, { - companyId: session.companyId, environmentId: session.environmentId, provider: session.provider, }); @@ -680,8 +767,9 @@ export function environmentCustomImageService( }) : null; const provider = await resolveSetupProvider({ - companyId: session.companyId, environment, + storedProvider: session.provider, + storedRpcCompanyId: readSetupRpcCompanyId(session.metadata), requireCapture: true, }); const runtimeConfigBinding = templateConfigBindingFromDriver({ @@ -699,7 +787,6 @@ export function environmentCustomImageService( updatedAt: now, }) .where(and( - eq(environmentCustomImageTemplates.companyId, session.companyId), eq(environmentCustomImageTemplates.environmentId, session.environmentId), eq(environmentCustomImageTemplates.provider, session.provider), eq(environmentCustomImageTemplates.status, "active"), @@ -709,7 +796,6 @@ export function environmentCustomImageService( .insert(environmentCustomImageTemplates) .values({ id: templateId, - companyId: session.companyId, environmentId: session.environmentId, provider: session.provider, templateKind: readTemplateKind(captured.templateKind), @@ -723,6 +809,7 @@ export function environmentCustomImageService( metadata: normalizeProviderMetadata({ ...(captured.metadata ?? {}), ...(input.metadata ? { userMetadata: input.metadata } : {}), + ...persistedSetupMetadata(session.metadata), [ENVIRONMENT_CUSTOM_IMAGE_RUNTIME_CONFIG_BINDING_METADATA_KEY]: runtimeConfigBinding, }), createdAt: now, @@ -790,7 +877,6 @@ export function environmentCustomImageService( }, rollbackTemplate: async (input: { - companyId: string; environmentId: string; now?: Date; }): Promise<{ @@ -804,7 +890,6 @@ export function environmentCustomImageService( .select() .from(environmentCustomImageTemplates) .where(and( - eq(environmentCustomImageTemplates.companyId, input.companyId), eq(environmentCustomImageTemplates.environmentId, input.environmentId), eq(environmentCustomImageTemplates.provider, active.provider), eq(environmentCustomImageTemplates.status, "superseded"), @@ -843,7 +928,6 @@ export function environmentCustomImageService( }, disableTemplate: async (input: { - companyId: string; environmentId: string; deleteProviderTemplate?: boolean; now?: Date; @@ -851,6 +935,9 @@ export function environmentCustomImageService( await requireEnvironment(input.environmentId); const active = await resolveActiveTemplate(db, input); if (!active) throw notFound("Active environment customImage template not found"); + const deleteProvider = input.deleteProviderTemplate + ? await resolveTemplateDeleteProvider(active) + : null; const now = input.now ?? new Date(); const row = await db .update(environmentCustomImageTemplates) @@ -860,8 +947,8 @@ export function environmentCustomImageService( .then((rows) => rows[0] ?? null); if (!row) throw notFound("Active environment customImage template not found"); const template = environmentCustomImageTemplateFromRow(row); - if (input.deleteProviderTemplate) { - await callProviderDeleteTemplate({ template, reason: "disabled" }); + if (deleteProvider) { + await callProviderDeleteTemplate({ template, provider: deleteProvider, reason: "disabled" }); } return template; }, diff --git a/server/src/services/environment-probe.ts b/server/src/services/environment-probe.ts index 8c0e39b31a..eaac5b22e8 100644 --- a/server/src/services/environment-probe.ts +++ b/server/src/services/environment-probe.ts @@ -18,12 +18,15 @@ export async function probeEnvironment( companyId?: string | null; pluginWorkerManager?: PluginWorkerManager; resolvedConfig?: ParsedEnvironmentConfig; + applyCustomImageTemplate?: boolean; } = {}, ): Promise { const resolvedCompanyId = options.companyId ?? null; const parsed = options.resolvedConfig ?? ( - resolvedCompanyId - ? await resolveEnvironmentDriverConfigForRuntime(db, resolvedCompanyId, environment) + resolvedCompanyId || options.applyCustomImageTemplate === true + ? await resolveEnvironmentDriverConfigForRuntime(db, resolvedCompanyId, environment, { + applyCustomImageTemplate: options.applyCustomImageTemplate === true, + }) : parseEnvironmentDriverConfig(environment) ); diff --git a/ui/src/api/environments.ts b/ui/src/api/environments.ts index 0d1a99b80f..7cfe039cb9 100644 --- a/ui/src/api/environments.ts +++ b/ui/src/api/environments.ts @@ -67,17 +67,14 @@ export const environmentsApi = { config?: Record; metadata?: Record | null; }) => api.post(`/companies/${companyId}/environments/probe-config`, body), - customImageTemplate: (environmentId: string, companyId: string) => - api.get( - `/environments/${environmentId}/custom-image-template?companyId=${encodeURIComponent(companyId)}`, - ), + customImageTemplate: (environmentId: string) => + api.get(`/environments/${environmentId}/custom-image-template`), startCustomImageSetupSession: ( environmentId: string, - companyId: string, body: StartEnvironmentCustomImageSetupSession = {}, ) => api.post( - `/environments/${environmentId}/custom-image-setup-sessions?companyId=${encodeURIComponent(companyId)}`, + `/environments/${environmentId}/custom-image-setup-sessions`, body, ), customImageSetupSession: (sessionId: string) => @@ -100,17 +97,16 @@ export const environmentsApi = { `/environment-custom-image-setup-sessions/${sessionId}/cancel`, body, ), - rollbackCustomImageTemplate: (environmentId: string, companyId: string) => + rollbackCustomImageTemplate: (environmentId: string) => api.post( - `/environments/${environmentId}/custom-image-template/rollback?companyId=${encodeURIComponent(companyId)}`, + `/environments/${environmentId}/custom-image-template/rollback`, {}, ), disableCustomImageTemplate: ( environmentId: string, - companyId: string, options: { deleteProviderTemplate?: boolean } = {}, ) => api.delete( - `/environments/${environmentId}/custom-image-template?companyId=${encodeURIComponent(companyId)}&deleteProviderTemplate=${options.deleteProviderTemplate === true ? "true" : "false"}`, + `/environments/${environmentId}/custom-image-template?deleteProviderTemplate=${options.deleteProviderTemplate === true ? "true" : "false"}`, ), }; diff --git a/ui/src/lib/queryKeys.ts b/ui/src/lib/queryKeys.ts index d4ca9c06df..2fbc9a9224 100644 --- a/ui/src/lib/queryKeys.ts +++ b/ui/src/lib/queryKeys.ts @@ -161,8 +161,8 @@ export const queryKeys = { environments: { list: (companyId: string) => ["environments", companyId] as const, capabilities: (companyId: string) => ["environment-capabilities", companyId] as const, - customImageTemplate: (companyId: string, environmentId: string) => - ["environments", companyId, environmentId, "custom-image-template"] as const, + customImageTemplate: (environmentId: string) => + ["environments", environmentId, "custom-image-template"] as const, customImageSetupSession: (sessionId: string) => ["environment-custom-image-setup-sessions", sessionId] as const, }, diff --git a/ui/src/pages/CompanyEnvironments.test.tsx b/ui/src/pages/CompanyEnvironments.test.tsx index 7305d6d21d..5585ab2e2f 100644 --- a/ui/src/pages/CompanyEnvironments.test.tsx +++ b/ui/src/pages/CompanyEnvironments.test.tsx @@ -121,7 +121,6 @@ function getOpenDialog(): HTMLElement | null { function createSession(overrides: Record = {}) { return { id: "session-1", - companyId: "company-1", environmentId: "env-1", templateId: null, promotedTemplateId: null, @@ -152,7 +151,6 @@ function createSession(overrides: Record = {}) { function createTemplate(overrides: Record = {}) { return { id: "template-1", - companyId: "company-1", environmentId: "env-1", provider: "daytona", templateKind: "snapshot", @@ -457,7 +455,7 @@ describe("CompanyEnvironments — test provider button", () => { await waitForAssertion(() => { expect(getOpenDialog()?.textContent).toContain("Configure image"); }); - expect(mockEnvironmentsApi.customImageTemplate).toHaveBeenCalledExactlyOnceWith("env-1", "company-1"); + expect(mockEnvironmentsApi.customImageTemplate).toHaveBeenCalledExactlyOnceWith("env-1"); await act(async () => click(findButton(document.body, "Cancel"))); await waitForAssertion(() => expect(getOpenDialog()).toBeNull()); @@ -656,7 +654,6 @@ describe("CompanyEnvironments — test provider button", () => { expect(mockEnvironmentsApi.startCustomImageSetupSession).toHaveBeenCalledWith( "env-1", - "company-1", { templateId: "template-active" }, ); }); diff --git a/ui/src/pages/CompanyEnvironments.tsx b/ui/src/pages/CompanyEnvironments.tsx index 2f475fddb3..6fd86efef4 100644 --- a/ui/src/pages/CompanyEnvironments.tsx +++ b/ui/src/pages/CompanyEnvironments.tsx @@ -245,12 +245,10 @@ function sessionStatusCopy(status: EnvironmentCustomImageSetupSession["status"]) } function EnvironmentImageTemplatePanel({ - companyId, environment, providerCapability, providerDisplayName, }: { - companyId: string; environment: Environment; providerCapability: EnvironmentProviderCapability | null | undefined; providerDisplayName: string; @@ -258,11 +256,11 @@ function EnvironmentImageTemplatePanel({ const { pushToast } = useToast(); const queryClient = useQueryClient(); const state = capabilityState(providerCapability); - const overviewKey = queryKeys.environments.customImageTemplate(companyId, environment.id); + const overviewKey = queryKeys.environments.customImageTemplate(environment.id); const overviewQuery = useQuery({ queryKey: overviewKey, - queryFn: () => environmentsApi.customImageTemplate(environment.id, companyId), + queryFn: () => environmentsApi.customImageTemplate(environment.id), enabled: state.kind === "supported", retry: false, }); @@ -290,7 +288,7 @@ function EnvironmentImageTemplatePanel({ const startSetupMutation = useMutation({ mutationFn: (input: { templateId?: string | null } = {}) => - environmentsApi.startCustomImageSetupSession(environment.id, companyId, { + environmentsApi.startCustomImageSetupSession(environment.id, { templateId: input.templateId ?? null, }), onSuccess: (result) => { @@ -368,7 +366,7 @@ function EnvironmentImageTemplatePanel({ }); const rollbackTemplateMutation = useMutation({ - mutationFn: () => environmentsApi.rollbackCustomImageTemplate(environment.id, companyId), + mutationFn: () => environmentsApi.rollbackCustomImageTemplate(environment.id), onSuccess: (result) => { queryClient.setQueryData(overviewKey, (current: typeof overviewQuery.data) => ({ activeTemplate: result.activeTemplate, @@ -392,7 +390,7 @@ function EnvironmentImageTemplatePanel({ }); const disableTemplateMutation = useMutation({ - mutationFn: () => environmentsApi.disableCustomImageTemplate(environment.id, companyId), + mutationFn: () => environmentsApi.disableCustomImageTemplate(environment.id), onSuccess: (template) => { queryClient.setQueryData(overviewKey, (current: typeof overviewQuery.data) => ({ activeTemplate: null, @@ -1246,7 +1244,6 @@ export function CompanyEnvironments() { running machine as a reusable image for future runs.