@vellumai/cli 0.5.16 → 0.6.0

package/src/__tests__/teleport.test.ts

@@ -108,6 +108,37 @@ const platformImportBundleMock = mock(async () => ({
     },
   } as Record<string, unknown>,
 }));
+const platformRequestUploadUrlMock = mock(async () => ({
+  uploadUrl: "https://storage.googleapis.com/bucket/signed-upload-url",
+  bundleKey: "bundle-key-123",
+  expiresAt: new Date(Date.now() + 3600_000).toISOString(),
+}));
+const platformUploadToSignedUrlMock = mock(async () => {});
+const platformImportPreflightFromGcsMock = mock(async () => ({
+  statusCode: 200,
+  body: {
+    can_import: true,
+    summary: {
+      files_to_create: 2,
+      files_to_overwrite: 1,
+      files_unchanged: 0,
+      total_files: 3,
+    },
+  } as Record<string, unknown>,
+}));
+const platformImportBundleFromGcsMock = mock(async () => ({
+  statusCode: 200,
+  body: {
+    success: true,
+    summary: {
+      total_files: 3,
+      files_created: 2,
+      files_overwritten: 1,
+      files_skipped: 0,
+      backups_created: 1,
+    },
+  } as Record<string, unknown>,
+}));
 
 mock.module("../lib/platform-client.js", () => ({
   readPlatformToken: readPlatformTokenMock,
@@ -119,6 +150,10 @@ mock.module("../lib/platform-client.js", () => ({
   platformDownloadExport: platformDownloadExportMock,
   platformImportPreflight: platformImportPreflightMock,
   platformImportBundle: platformImportBundleMock,
+  platformRequestUploadUrl: platformRequestUploadUrlMock,
+  platformUploadToSignedUrl: platformUploadToSignedUrlMock,
+  platformImportPreflightFromGcs: platformImportPreflightFromGcsMock,
+  platformImportBundleFromGcs: platformImportBundleFromGcsMock,
 }));
 
 const hatchLocalMock = mock(async () => {});
@@ -256,6 +291,41 @@ beforeEach(() => {
       },
     },
   });
+  platformRequestUploadUrlMock.mockReset();
+  platformRequestUploadUrlMock.mockResolvedValue({
+    uploadUrl: "https://storage.googleapis.com/bucket/signed-upload-url",
+    bundleKey: "bundle-key-123",
+    expiresAt: new Date(Date.now() + 3600_000).toISOString(),
+  });
+  platformUploadToSignedUrlMock.mockReset();
+  platformUploadToSignedUrlMock.mockResolvedValue(undefined);
+  platformImportPreflightFromGcsMock.mockReset();
+  platformImportPreflightFromGcsMock.mockResolvedValue({
+    statusCode: 200,
+    body: {
+      can_import: true,
+      summary: {
+        files_to_create: 2,
+        files_to_overwrite: 1,
+        files_unchanged: 0,
+        total_files: 3,
+      },
+    },
+  });
+  platformImportBundleFromGcsMock.mockReset();
+  platformImportBundleFromGcsMock.mockResolvedValue({
+    statusCode: 200,
+    body: {
+      success: true,
+      summary: {
+        total_files: 3,
+        files_created: 2,
+        files_overwritten: 1,
+        files_skipped: 0,
+        backups_created: 1,
+      },
+    },
+  });
 
   hatchLocalMock.mockReset();
   hatchLocalMock.mockResolvedValue(undefined);
@@ -687,7 +757,10 @@ describe("resolveOrHatchTarget", () => {
     findAssistantByNameMock.mockReturnValue(null);
 
     const result = await resolveOrHatchTarget("platform", "nonexistent");
-    expect(hatchAssistantMock).toHaveBeenCalledWith("platform-token");
+    expect(hatchAssistantMock).toHaveBeenCalledWith(
+      "platform-token",
+      "org-123",
+    );
     expect(saveAssistantEntryMock).toHaveBeenCalledWith(
       expect.objectContaining({
         assistantId: "platform-new-id",
@@ -998,3 +1071,373 @@ describe("teleport full flow", () => {
     );
   });
 });
+
+// ---------------------------------------------------------------------------
+// Signed-URL upload tests
+// ---------------------------------------------------------------------------
+
+describe("signed-URL upload flow", () => {
+  test("happy path: signed URL upload succeeds → GCS-based import used", async () => {
+    setArgv("--from", "my-local", "--platform");
+
+    const localEntry = makeEntry("my-local", { cloud: "local" });
+
+    findAssistantByNameMock.mockImplementation((name: string) => {
+      if (name === "my-local") return localEntry;
+      return null;
+    });
+
+    const originalFetch = globalThis.fetch;
+    globalThis.fetch = createFetchMock() as unknown as typeof globalThis.fetch;
+
+    try {
+      await teleport();
+
+      // Signed-URL flow should be used
+      expect(platformRequestUploadUrlMock).toHaveBeenCalled();
+      expect(platformUploadToSignedUrlMock).toHaveBeenCalled();
+      expect(platformImportBundleFromGcsMock).toHaveBeenCalledWith(
+        "bundle-key-123",
+        "platform-token",
+        "org-123",
+        "https://platform.vellum.ai",
+      );
+      // Inline import should NOT be called
+      expect(platformImportBundleMock).not.toHaveBeenCalled();
+      expect(consoleLogSpy).toHaveBeenCalledWith(
+        expect.stringContaining("Teleport complete"),
+      );
+    } finally {
+      globalThis.fetch = originalFetch;
+    }
+  });
+
+  test("happy path dry-run: signed URL upload succeeds → GCS-based preflight used", async () => {
+    setArgv(
+      "--from",
+      "my-local",
+      "--platform",
+      "existing-platform",
+      "--dry-run",
+    );
+
+    const localEntry = makeEntry("my-local", { cloud: "local" });
+    const platformEntry = makeEntry("existing-platform", {
+      cloud: "vellum",
+      runtimeUrl: "https://platform.vellum.ai",
+    });
+
+    findAssistantByNameMock.mockImplementation((name: string) => {
+      if (name === "my-local") return localEntry;
+      if (name === "existing-platform") return platformEntry;
+      return null;
+    });
+
+    const originalFetch = globalThis.fetch;
+    globalThis.fetch = createFetchMock() as unknown as typeof globalThis.fetch;
+
+    try {
+      await teleport();
+
+      // Signed-URL flow should be used for preflight
+      expect(platformRequestUploadUrlMock).toHaveBeenCalled();
+      expect(platformUploadToSignedUrlMock).toHaveBeenCalled();
+      expect(platformImportPreflightFromGcsMock).toHaveBeenCalledWith(
+        "bundle-key-123",
+        "platform-token",
+        "org-123",
+        "https://platform.vellum.ai",
+      );
+      // Inline preflight should NOT be called
+      expect(platformImportPreflightMock).not.toHaveBeenCalled();
+    } finally {
+      globalThis.fetch = originalFetch;
+    }
+  });
+
+  test("fallback: platformRequestUploadUrl throws 503 → falls back to inline import", async () => {
+    setArgv("--from", "my-local", "--platform");
+
+    const localEntry = makeEntry("my-local", { cloud: "local" });
+
+    findAssistantByNameMock.mockImplementation((name: string) => {
+      if (name === "my-local") return localEntry;
+      return null;
+    });
+
+    // Simulate 503 — "not available" in the error message
+    platformRequestUploadUrlMock.mockRejectedValue(
+      new Error("Signed uploads are not available on this platform instance"),
+    );
+
+    const originalFetch = globalThis.fetch;
+    globalThis.fetch = createFetchMock() as unknown as typeof globalThis.fetch;
+
+    try {
+      await teleport();
+
+      // Should fall back to inline import
+      expect(platformRequestUploadUrlMock).toHaveBeenCalled();
+      expect(platformUploadToSignedUrlMock).not.toHaveBeenCalled();
+      expect(platformImportBundleFromGcsMock).not.toHaveBeenCalled();
+      expect(platformImportBundleMock).toHaveBeenCalled();
+      expect(consoleLogSpy).toHaveBeenCalledWith(
+        expect.stringContaining("Teleport complete"),
+      );
+    } finally {
+      globalThis.fetch = originalFetch;
+    }
+  });
+
+  test("fallback: platformRequestUploadUrl throws 404 → falls back to inline import", async () => {
+    setArgv("--from", "my-local", "--platform");
+
+    const localEntry = makeEntry("my-local", { cloud: "local" });
+
+    findAssistantByNameMock.mockImplementation((name: string) => {
+      if (name === "my-local") return localEntry;
+      return null;
+    });
+
+    // Simulate 404 — endpoint doesn't exist on older platform versions
+    platformRequestUploadUrlMock.mockRejectedValue(
+      new Error("Signed uploads are not available on this platform instance"),
+    );
+
+    const originalFetch = globalThis.fetch;
+    globalThis.fetch = createFetchMock() as unknown as typeof globalThis.fetch;
+
+    try {
+      await teleport();
+
+      // Should fall back to inline import
+      expect(platformRequestUploadUrlMock).toHaveBeenCalled();
+      expect(platformUploadToSignedUrlMock).not.toHaveBeenCalled();
+      expect(platformImportBundleFromGcsMock).not.toHaveBeenCalled();
+      expect(platformImportBundleMock).toHaveBeenCalled();
+      expect(consoleLogSpy).toHaveBeenCalledWith(
+        expect.stringContaining("Teleport complete"),
+      );
+    } finally {
+      globalThis.fetch = originalFetch;
+    }
+  });
+
+  test("upload error: platformUploadToSignedUrl throws → error propagates", async () => {
+    setArgv("--from", "my-local", "--platform");
+
+    const localEntry = makeEntry("my-local", { cloud: "local" });
+
+    findAssistantByNameMock.mockImplementation((name: string) => {
+      if (name === "my-local") return localEntry;
+      return null;
+    });
+
+    // Upload succeeds at getting URL but fails during PUT
+    platformUploadToSignedUrlMock.mockRejectedValue(
+      new Error("Upload to signed URL failed: 500 Internal Server Error"),
+    );
+
+    const originalFetch = globalThis.fetch;
+    globalThis.fetch = createFetchMock() as unknown as typeof globalThis.fetch;
+
+    try {
+      await expect(teleport()).rejects.toThrow(
+        "Upload to signed URL failed: 500 Internal Server Error",
+      );
+      // Should NOT fall back to inline import
+      expect(platformImportBundleMock).not.toHaveBeenCalled();
+      expect(platformImportBundleFromGcsMock).not.toHaveBeenCalled();
+    } finally {
+      globalThis.fetch = originalFetch;
+    }
+  });
+
+  test("413 from GCS import: error message includes 'too large'", async () => {
+    setArgv("--from", "my-local", "--platform");
+
+    const localEntry = makeEntry("my-local", { cloud: "local" });
+
+    findAssistantByNameMock.mockImplementation((name: string) => {
+      if (name === "my-local") return localEntry;
+      return null;
+    });
+
+    // GCS import returns 413
+    platformImportBundleFromGcsMock.mockRejectedValue(
+      new Error("Bundle too large to import"),
+    );
+
+    const originalFetch = globalThis.fetch;
+    globalThis.fetch = createFetchMock() as unknown as typeof globalThis.fetch;
+
+    try {
+      await expect(teleport()).rejects.toThrow("too large");
+    } finally {
+      globalThis.fetch = originalFetch;
+    }
+  });
+});
+
+// ---------------------------------------------------------------------------
+// Platform teleport org ID and reordered flow tests
+// ---------------------------------------------------------------------------
+
+describe("platform teleport org ID and reordered flow", () => {
+  test("hatchAssistant is called with the org ID from fetchOrganizationId", async () => {
+    setArgv("--from", "my-local", "--platform");
+
+    const localEntry = makeEntry("my-local", { cloud: "local" });
+
+    findAssistantByNameMock.mockImplementation((name: string) => {
+      if (name === "my-local") return localEntry;
+      return null;
+    });
+
+    fetchOrganizationIdMock.mockResolvedValue("test-org-456");
+
+    const originalFetch = globalThis.fetch;
+    globalThis.fetch = createFetchMock() as unknown as typeof globalThis.fetch;
+
+    try {
+      await teleport();
+
+      // fetchOrganizationId should have been called
+      expect(fetchOrganizationIdMock).toHaveBeenCalled();
+      // hatchAssistant must be called with (token, orgId)
+      expect(hatchAssistantMock).toHaveBeenCalledWith(
+        "platform-token",
+        "test-org-456",
+      );
+    } finally {
+      globalThis.fetch = originalFetch;
+    }
+  });
+
+  test("upload to GCS happens before hatchAssistant for platform targets", async () => {
+    setArgv("--from", "my-local", "--platform");
+
+    const localEntry = makeEntry("my-local", { cloud: "local" });
+
+    findAssistantByNameMock.mockImplementation((name: string) => {
+      if (name === "my-local") return localEntry;
+      return null;
+    });
+
+    const callOrder: string[] = [];
+
+    platformRequestUploadUrlMock.mockImplementation(async () => {
+      callOrder.push("platformRequestUploadUrl");
+      return {
+        uploadUrl: "https://storage.googleapis.com/bucket/signed-upload-url",
+        bundleKey: "bundle-key-123",
+        expiresAt: new Date(Date.now() + 3600_000).toISOString(),
+      };
+    });
+
+    platformUploadToSignedUrlMock.mockImplementation(async () => {
+      callOrder.push("platformUploadToSignedUrl");
+    });
+
+    hatchAssistantMock.mockImplementation(async () => {
+      callOrder.push("hatchAssistant");
+      return { id: "platform-new-id", name: "platform-new", status: "active" };
+    });
+
+    const originalFetch = globalThis.fetch;
+    globalThis.fetch = createFetchMock() as unknown as typeof globalThis.fetch;
+
+    try {
+      await teleport();
+
+      // Verify ordering: upload steps come before hatch
+      const uploadUrlIdx = callOrder.indexOf("platformRequestUploadUrl");
+      const uploadIdx = callOrder.indexOf("platformUploadToSignedUrl");
+      const hatchIdx = callOrder.indexOf("hatchAssistant");
+
+      expect(uploadUrlIdx).toBeGreaterThanOrEqual(0);
+      expect(uploadIdx).toBeGreaterThanOrEqual(0);
+      expect(hatchIdx).toBeGreaterThanOrEqual(0);
+      expect(uploadUrlIdx).toBeLessThan(hatchIdx);
+      expect(uploadIdx).toBeLessThan(hatchIdx);
+    } finally {
+      globalThis.fetch = originalFetch;
+    }
+  });
+
+  test("signed-URL fallback: when platformRequestUploadUrl throws 'not available', falls back to inline upload via importToAssistant", async () => {
+    setArgv("--from", "my-local", "--platform");
+
+    const localEntry = makeEntry("my-local", { cloud: "local" });
+
+    findAssistantByNameMock.mockImplementation((name: string) => {
+      if (name === "my-local") return localEntry;
+      return null;
+    });
+
+    // Simulate 503 — signed uploads not available
+    platformRequestUploadUrlMock.mockRejectedValue(
+      new Error("Signed uploads are not available on this platform instance"),
+    );
+
+    const originalFetch = globalThis.fetch;
+    globalThis.fetch = createFetchMock() as unknown as typeof globalThis.fetch;
+
+    try {
+      await teleport();
+
+      // Upload URL was attempted but failed
+      expect(platformRequestUploadUrlMock).toHaveBeenCalled();
+      // No signed URL upload should have happened
+      expect(platformUploadToSignedUrlMock).not.toHaveBeenCalled();
+      // Should NOT use GCS-based import
+      expect(platformImportBundleFromGcsMock).not.toHaveBeenCalled();
+      // Should fall back to inline import
+      expect(platformImportBundleMock).toHaveBeenCalled();
+      // Hatch should still succeed
+      expect(hatchAssistantMock).toHaveBeenCalled();
+      expect(consoleLogSpy).toHaveBeenCalledWith(
+        expect.stringContaining("Teleport complete"),
+      );
+    } finally {
+      globalThis.fetch = originalFetch;
+    }
+  });
+
+  test("bundleKey from pre-upload is forwarded to platformImportBundleFromGcs", async () => {
+    setArgv("--from", "my-local", "--platform");
+
+    const localEntry = makeEntry("my-local", { cloud: "local" });
+
+    findAssistantByNameMock.mockImplementation((name: string) => {
+      if (name === "my-local") return localEntry;
+      return null;
+    });
+
+    // Return a specific bundle key from the pre-upload step
+    platformRequestUploadUrlMock.mockResolvedValue({
+      uploadUrl: "https://storage.googleapis.com/bucket/signed-upload-url",
+      bundleKey: "pre-uploaded-key-789",
+      expiresAt: new Date(Date.now() + 3600_000).toISOString(),
+    });
+
+    const originalFetch = globalThis.fetch;
+    globalThis.fetch = createFetchMock() as unknown as typeof globalThis.fetch;
+
+    try {
+      await teleport();
+
+      // The bundle key from the pre-upload step should be forwarded to GCS import
+      expect(platformImportBundleFromGcsMock).toHaveBeenCalledWith(
+        "pre-uploaded-key-789",
+        "platform-token",
+        expect.any(String),
+        expect.any(String),
+      );
+      // Inline import should NOT be used since signed upload succeeded
+      expect(platformImportBundleMock).not.toHaveBeenCalled();
+    } finally {
+      globalThis.fetch = originalFetch;
+    }
+  });
+});

package/src/commands/hatch.ts

@@ -21,6 +21,7 @@ import { hatchGcp } from "../lib/gcp";
 import type { PollResult, WatchHatchingResult } from "../lib/gcp";
 import { hatchLocal } from "../lib/hatch-local";
 import {
+  fetchOrganizationId,
   getPlatformUrl,
   hatchAssistant,
   readPlatformToken,
@@ -593,7 +594,19 @@ async function hatchVellumPlatform(): Promise<void> {
   console.log("   Hatching assistant on Vellum platform...");
   console.log("");
 
-  const result = await hatchAssistant(token);
+  let orgId: string;
+  try {
+    orgId = await fetchOrganizationId(token);
+  } catch (err) {
+    const msg = err instanceof Error ? err.message : String(err);
+    if (msg.includes("401") || msg.includes("403")) {
+      console.error("Authentication failed. Run 'vellum login' to refresh.");
+      process.exit(1);
+    }
+    throw err;
+  }
+
+  const result = await hatchAssistant(token, orgId);
 
   const platformUrl = getPlatformUrl();
 

package/src/commands/rollback.ts

@@ -352,6 +352,11 @@ export async function rollback(): Promise<void> {
       `   Captured ${Object.keys(capturedEnv).length} env var(s) from ${res.assistantContainer}\n`,
     );
 
+    // Capture GUARDIAN_BOOTSTRAP_SECRET from the gateway container (it is only
+    // set on gateway, not assistant) so it persists across container restarts.
+    const gatewayEnv = await captureContainerEnv(res.gatewayContainer);
+    const bootstrapSecret = gatewayEnv["GUARDIAN_BOOTSTRAP_SECRET"];
+
     // Extract CES_SERVICE_TOKEN from captured env, or generate fresh one
     const cesServiceToken =
       capturedEnv["CES_SERVICE_TOKEN"] || randomBytes(32).toString("hex");
@@ -430,6 +435,7 @@ export async function rollback(): Promise<void> {
     await startContainers(
       {
         signingKey,
+        bootstrapSecret,
         cesServiceToken,
         extraAssistantEnv,
         gatewayPort,