@vellumai/cli 0.6.5 → 0.7.0

package/src/commands/teleport.ts

@@ -1,6 +1,7 @@
 import {
   findAssistantByName,
   loadAllAssistants,
+  getDaemonPidPath,
   removeAssistantEntry,
   saveAssistantEntry,
   setActiveAssistant,
@@ -17,15 +18,10 @@ import {
   hatchAssistant,
   checkExistingPlatformAssistant,
   platformInitiateExport,
-  platformPollExportStatus,
-  platformDownloadExport,
-  platformImportPreflight,
-  platformImportBundle,
-  platformRequestUploadUrl,
-  platformUploadToSignedUrl,
-  platformImportPreflightFromGcs,
+  platformPollJobStatus,
   platformImportBundleFromGcs,
-  platformPollImportStatus,
+  platformImportPreflightFromGcs,
+  platformRequestSignedUrl,
   ensureSelfHostedLocalRegistration,
   readGatewayCredential,
   reprovisionAssistantApiKey,
@@ -33,6 +29,13 @@ import {
   fetchCurrentUser,
   fetchOrganizationId,
 } from "../lib/platform-client.js";
+import {
+  localRuntimeExportToGcs,
+  localRuntimeImportFromGcs,
+  localRuntimePollJobStatus,
+  MigrationInProgressError,
+} from "../lib/local-runtime-client.js";
+import { pollJobUntilDone } from "../lib/job-polling.js";
 import {
   hatchDocker,
   retireDocker,
@@ -221,11 +224,17 @@ async function getAccessToken(
   runtimeUrl: string,
   assistantId: string,
   displayName: string,
+  options?: { forceRefresh?: boolean },
 ): Promise<string> {
-  const tokenData = loadGuardianToken(assistantId);
-
-  if (tokenData && new Date(tokenData.accessTokenExpiresAt) > new Date()) {
-    return tokenData.accessToken;
+  // When forceRefresh is set (e.g. after a runtime 401 on the cached token)
+  // we skip the cache and lease a brand-new token from the gateway, so a
+  // stale-but-unexpired token can't keep failing on every retry.
+  if (!options?.forceRefresh) {
+    const tokenData = loadGuardianToken(assistantId);
+
+    if (tokenData && new Date(tokenData.accessTokenExpiresAt) > new Date()) {
+      return tokenData.accessToken;
+    }
   }
 
   try {
@@ -244,336 +253,49 @@ async function getAccessToken(
   }
 }
 
-// ---------------------------------------------------------------------------
-// HTTP-based export/import helpers (shared by local and docker)
-// ---------------------------------------------------------------------------
-
-async function exportViaHttp(
-  entry: AssistantEntry,
-): Promise<Uint8Array<ArrayBuffer>> {
-  let accessToken = await getAccessToken(
-    entry.runtimeUrl,
-    entry.assistantId,
-    entry.assistantId,
-  );
-
-  let response: Response;
-  try {
-    response = await fetch(`${entry.runtimeUrl}/v1/migrations/export`, {
-      method: "POST",
-      headers: {
-        Authorization: `Bearer ${accessToken}`,
-        "Content-Type": "application/json",
-      },
-      body: JSON.stringify({ description: "teleport export" }),
-      signal: AbortSignal.timeout(120_000),
-    });
-
-    // Retry once with a fresh token on 401
-    if (response.status === 401) {
-      let refreshedToken: string | null = null;
-      try {
-        const freshToken = await leaseGuardianToken(
-          entry.runtimeUrl,
-          entry.assistantId,
-        );
-        refreshedToken = freshToken.accessToken;
-      } catch {
-        // If token refresh fails, fall through to the error handler below
-      }
-      if (refreshedToken) {
-        accessToken = refreshedToken;
-        response = await fetch(`${entry.runtimeUrl}/v1/migrations/export`, {
-          method: "POST",
-          headers: {
-            Authorization: `Bearer ${accessToken}`,
-            "Content-Type": "application/json",
-          },
-          body: JSON.stringify({ description: "teleport export" }),
-          signal: AbortSignal.timeout(120_000),
-        });
-      }
-    }
-  } catch (err) {
-    if (err instanceof Error && err.name === "TimeoutError") {
-      console.error("Error: Export request timed out after 2 minutes.");
-      process.exit(1);
-    }
-    const msg = err instanceof Error ? err.message : String(err);
-    if (msg.includes("ECONNREFUSED") || msg.includes("fetch failed")) {
-      console.error(
-        `Error: Could not connect to assistant '${entry.assistantId}'. Is it running?`,
-      );
-      console.error(`Try: vellum wake ${entry.assistantId}`);
-      process.exit(1);
-    }
-    throw err;
-  }
-
-  if (response.status === 401 || response.status === 403) {
-    console.error("Authentication failed.");
-    process.exit(1);
-  }
-
-  if (response.status === 404) {
-    console.error("Assistant not found or not running.");
-    process.exit(1);
-  }
-
-  if (
-    response.status === 502 ||
-    response.status === 503 ||
-    response.status === 504
-  ) {
-    console.error(
-      `Assistant is unreachable. Try 'vellum wake ${entry.assistantId}'.`,
-    );
-    process.exit(1);
-  }
-
-  if (!response.ok) {
-    const body = await response.text();
-    console.error(`Error: Export failed (${response.status}): ${body}`);
-    process.exit(1);
-  }
-
-  const arrayBuffer = await response.arrayBuffer();
-  return new Uint8Array(arrayBuffer);
+/**
+ * Detect a 401 Unauthorized raised by `localRuntimeExportToGcs` /
+ * `localRuntimeImportFromGcs`. Both throw Error with a message of the form
+ * `"Local runtime <op> failed (401): ..."` when the gateway rejects the
+ * cached guardian token.
+ */
+function isRuntime401(err: unknown): boolean {
+  const msg = err instanceof Error ? err.message : String(err);
+  return /Local runtime [^(]*failed \(401\)/.test(msg);
 }
 
-async function importViaHttp(
-  entry: AssistantEntry,
-  bundleData: Uint8Array<ArrayBuffer>,
-  dryRun: boolean,
-): Promise<void> {
-  let accessToken = await getAccessToken(
-    entry.runtimeUrl,
-    entry.assistantId,
-    entry.assistantId,
-  );
-
-  if (dryRun) {
-    console.log("Running preflight analysis...\n");
-
-    let response: Response;
-    try {
-      response = await fetch(
-        `${entry.runtimeUrl}/v1/migrations/import-preflight`,
-        {
-          method: "POST",
-          headers: {
-            Authorization: `Bearer ${accessToken}`,
-            "Content-Type": "application/octet-stream",
-          },
-          body: new Blob([bundleData]),
-          signal: AbortSignal.timeout(120_000),
-        },
-      );
-
-      // Retry once with a fresh token on 401
-      if (response.status === 401) {
-        let refreshedToken: string | null = null;
-        try {
-          const freshToken = await leaseGuardianToken(
-            entry.runtimeUrl,
-            entry.assistantId,
-          );
-          refreshedToken = freshToken.accessToken;
-        } catch {
-          // If token refresh fails, fall through to the error handler below
-        }
-        if (refreshedToken) {
-          accessToken = refreshedToken;
-          response = await fetch(
-            `${entry.runtimeUrl}/v1/migrations/import-preflight`,
-            {
-              method: "POST",
-              headers: {
-                Authorization: `Bearer ${accessToken}`,
-                "Content-Type": "application/octet-stream",
-              },
-              body: new Blob([bundleData]),
-              signal: AbortSignal.timeout(120_000),
-            },
-          );
-        }
-      }
-    } catch (err) {
-      if (err instanceof Error && err.name === "TimeoutError") {
-        console.error("Error: Preflight request timed out after 2 minutes.");
-        process.exit(1);
-      }
-      const msg = err instanceof Error ? err.message : String(err);
-      if (msg.includes("ECONNREFUSED") || msg.includes("fetch failed")) {
-        console.error(
-          `Error: Could not connect to assistant '${entry.assistantId}'. Is it running?`,
-        );
-        console.error(`Try: vellum wake ${entry.assistantId}`);
-        process.exit(1);
-      }
-      throw err;
-    }
-
-    handleLocalResponseErrors(response, entry.assistantId);
-
-    const result = (await response.json()) as PreflightResponse;
-    printPreflightSummary(result);
-    return;
-  }
-
-  // Actual import
-  console.log("Importing data...");
-
-  let response: Response;
+/**
+ * Run a runtime kickoff (`localRuntimeExportToGcs` / `localRuntimeImportFromGcs`)
+ * with a one-shot refresh-and-retry on 401. Matches the pre-rewrite
+ * `exportViaHttp`/`importViaHttp` behavior: if the cached guardian token is
+ * stale-but-unexpired and the runtime returns 401, we lease a fresh token
+ * and retry once. Any other error — or a repeated 401 on the refreshed token
+ * — propagates to the caller.
+ */
+async function callRuntimeWithAuthRetry<T>(
+  runtimeUrl: string,
+  assistantId: string,
+  fn: (token: string) => Promise<T>,
+): Promise<T> {
+  const firstToken = await getAccessToken(runtimeUrl, assistantId, assistantId);
   try {
-    response = await fetch(`${entry.runtimeUrl}/v1/migrations/import`, {
-      method: "POST",
-      headers: {
-        Authorization: `Bearer ${accessToken}`,
-        "Content-Type": "application/octet-stream",
-      },
-      body: new Blob([bundleData]),
-      signal: AbortSignal.timeout(300_000),
-    });
-
-    // Retry once with a fresh token on 401
-    if (response.status === 401) {
-      let refreshedToken: string | null = null;
-      try {
-        const freshToken = await leaseGuardianToken(
-          entry.runtimeUrl,
-          entry.assistantId,
-        );
-        refreshedToken = freshToken.accessToken;
-      } catch {
-        // If token refresh fails, fall through to the error handler below
-      }
-      if (refreshedToken) {
-        accessToken = refreshedToken;
-        response = await fetch(`${entry.runtimeUrl}/v1/migrations/import`, {
-          method: "POST",
-          headers: {
-            Authorization: `Bearer ${accessToken}`,
-            "Content-Type": "application/octet-stream",
-          },
-          body: new Blob([bundleData]),
-          signal: AbortSignal.timeout(300_000),
-        });
-      }
-    }
+    return await fn(firstToken);
   } catch (err) {
-    if (err instanceof Error && err.name === "TimeoutError") {
-      console.error("Error: Import request timed out after 5 minutes.");
-      process.exit(1);
-    }
-    const msg = err instanceof Error ? err.message : String(err);
-    if (msg.includes("ECONNREFUSED") || msg.includes("fetch failed")) {
-      console.error(
-        `Error: Could not connect to assistant '${entry.assistantId}'. Is it running?`,
-      );
-      console.error(`Try: vellum wake ${entry.assistantId}`);
-      process.exit(1);
-    }
-    throw err;
-  }
-
-  handleLocalResponseErrors(response, entry.assistantId);
-
-  const result = (await response.json()) as ImportResponse;
-  printImportSummary(result);
-}
-
-// ---------------------------------------------------------------------------
-// Export from source assistant
-// ---------------------------------------------------------------------------
-
-async function exportFromAssistant(
-  entry: AssistantEntry,
-  cloud: string,
-): Promise<Uint8Array<ArrayBuffer>> {
-  if (cloud === "vellum") {
-    // Platform source — use Django async export
-    const token = readPlatformToken();
-    if (!token) {
-      console.error("Not logged in. Run 'vellum login' first.");
-      process.exit(1);
+    if (!isRuntime401(err)) {
+      throw err;
     }
-
-    // Initiate export job
-    const { jobId } = await platformInitiateExport(
-      token,
-      "teleport export",
-      entry.runtimeUrl,
+    const refreshedToken = await getAccessToken(
+      runtimeUrl,
+      assistantId,
+      assistantId,
+      { forceRefresh: true },
     );
-
-    console.log(`Export started (job ${jobId})...`);
-
-    // Poll for completion
-    const POLL_INTERVAL_MS = 2_000;
-    const TIMEOUT_MS = 5 * 60 * 1_000;
-    const deadline = Date.now() + TIMEOUT_MS;
-    let downloadUrl: string | undefined;
-
-    while (Date.now() < deadline) {
-      let status: { status: string; downloadUrl?: string; error?: string };
-      try {
-        status = await platformPollExportStatus(jobId, token, entry.runtimeUrl);
-      } catch (err) {
-        const msg = err instanceof Error ? err.message : String(err);
-        if (msg.includes("not found")) {
-          throw err;
-        }
-        // Re-throw permanent 4xx errors (auth, forbidden, etc.)
-        // but retry transient 5xx errors
-        const statusMatch = msg.match(/status check failed: (\d+)/);
-        if (statusMatch) {
-          const statusCode = parseInt(statusMatch[1], 10);
-          if (statusCode >= 400 && statusCode < 500) {
-            throw err;
-          }
-        }
-        // Transient error — retry
-        console.warn(`Polling failed, retrying... (${msg})`);
-        await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
-        continue;
-      }
-
-      if (status.status === "complete") {
-        downloadUrl = status.downloadUrl;
-        break;
-      }
-
-      if (status.status === "failed") {
-        console.error(`Export failed: ${status.error ?? "unknown error"}`);
-        process.exit(1);
-      }
-
-      await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
-    }
-
-    if (!downloadUrl) {
-      console.error("Export timed out after 5 minutes.");
-      process.exit(1);
-    }
-
-    // Download the bundle
-    const response = await platformDownloadExport(downloadUrl);
-    const arrayBuffer = await response.arrayBuffer();
-    return new Uint8Array(arrayBuffer);
+    return await fn(refreshedToken);
   }
-
-  if (cloud === "local" || cloud === "docker") {
-    return exportViaHttp(entry);
-  }
-
-  console.error(
-    "Teleport only supports local, docker, and platform assistants as source.",
-  );
-  process.exit(1);
 }
 
 // ---------------------------------------------------------------------------
-// Import into target assistant
+// Summary response shapes (reused by the GCS job result payload)
 // ---------------------------------------------------------------------------
 
 interface PreflightFileEntry {
@@ -625,86 +347,198 @@ interface ImportResponse {
   };
 }
 
-async function importToAssistant(
+// ---------------------------------------------------------------------------
+// Export from source — unified GCS flow
+//
+// Every source (local, docker, platform) produces a `bundleKey` referring to
+// a bundle sitting in GCS. The CLI never holds the bundle bytes.
+// ---------------------------------------------------------------------------
+
+async function exportFromAssistant(
   entry: AssistantEntry,
   cloud: string,
-  bundleData: Uint8Array<ArrayBuffer>,
-  dryRun: boolean,
-  preUploadedBundleKey?: string | null,
-): Promise<void> {
-  if (cloud === "vellum") {
-    // Platform target
-    const token = readPlatformToken();
-    if (!token) {
-      console.error("Not logged in. Run 'vellum login' first.");
-      process.exit(1);
+  bundlePlatformUrl?: string,
+): Promise<{ bundleKey: string }> {
+  const platformToken = readPlatformToken();
+  if (!platformToken) {
+    console.error(
+      "Not logged in. Run 'vellum login' first (required for GCS-based teleport).",
+    );
+    process.exit(1);
+  }
+
+  if (cloud === "local" || cloud === "docker") {
+    // Request a signed upload URL from the platform instance that will
+    // eventually own the bundle (i.e. the one the importer will read from).
+    // Passing the target's runtime URL here keeps upload and download on
+    // the same platform — otherwise a non-default/stale platform URL would
+    // cause the import to look at an empty object.
+    const { url: uploadUrl, bundleKey } = await platformRequestSignedUrl(
+      { operation: "upload" },
+      platformToken,
+      bundlePlatformUrl,
+    );
+
+    // Wrap the kickoff in a one-shot refresh-and-retry helper so a stale-but-
+    // unexpired cached guardian token surfaces as 401 → re-lease → retry
+    // rather than a terminal failure. `accessToken` below is whichever token
+    // succeeded on the kickoff; we reuse it for polling so the runtime sees a
+    // consistent credential throughout the migration.
+    let jobId: string;
+    let accessToken: string;
+    try {
+      const result = await callRuntimeWithAuthRetry(
+        entry.runtimeUrl,
+        entry.assistantId,
+        async (token) => {
+          const r = await localRuntimeExportToGcs(entry.runtimeUrl, token, {
+            uploadUrl,
+            description: "teleport export",
+          });
+          return { jobId: r.jobId, token };
+        },
+      );
+      jobId = result.jobId;
+      accessToken = result.token;
+    } catch (err) {
+      if (err instanceof MigrationInProgressError) {
+        // Fail fast — the existing job is writing to a different GCS object
+        // (its caller's signed URL, not ours), so polling it would leave us
+        // pointing at an empty/unrelated bundle. Surface the existing job id
+        // so the user can decide whether to wait or investigate.
+        console.error(
+          `Error: Another teleport export is already in progress on '${entry.assistantId}' (job ${err.existingJobId}). Wait for it to finish or check its status, then re-run.`,
+        );
+        process.exit(1);
+      }
+      throw err;
     }
 
-    // Use pre-uploaded bundle key if provided (string), skip upload if null
-    // (signals signed URLs were already tried and unavailable), or try
-    // signed-URL upload if undefined (never attempted).
-    let bundleKey: string | undefined =
-      preUploadedBundleKey === null ? undefined : preUploadedBundleKey;
-    if (preUploadedBundleKey === undefined) {
-      try {
-        const { uploadUrl, bundleKey: key } = await platformRequestUploadUrl(
-          token,
+    console.log(`Export started (job ${jobId})...`);
+
+    const terminal = await pollJobUntilDone({
+      label: "local-runtime export",
+      poll: () =>
+        localRuntimePollJobStatus(entry.runtimeUrl, accessToken, jobId),
+      // Large exports can take longer than a guardian-token lease. If the
+      // runtime returns 401 mid-poll, re-lease a fresh token and rebind the
+      // closure variable so the next poll uses it.
+      refreshOn401: async () => {
+        accessToken = await getAccessToken(
           entry.runtimeUrl,
+          entry.assistantId,
+          entry.assistantId,
+          { forceRefresh: true },
         );
-        bundleKey = key;
-        console.log("Uploading bundle...");
-        await platformUploadToSignedUrl(uploadUrl, bundleData);
-      } catch (err) {
-        // If signed uploads unavailable (503), fall back to inline upload
-        const msg = err instanceof Error ? err.message : String(err);
-        if (msg.includes("not available")) {
-          bundleKey = undefined;
-        } else {
-          throw err;
+      },
+    });
+
+    if (terminal.status === "failed") {
+      console.error(`Export failed: ${terminal.error}`);
+      process.exit(1);
+    }
+
+    return { bundleKey };
+  }
+
+  if (cloud === "vellum") {
+    // Platform source — initiate a server-side export. The platform writes
+    // the bundle to its own `exports/<org>/<id>.vbundle` key; we discover
+    // that key via the unified job-status endpoint's `bundle_key` field.
+    const { jobId } = await platformInitiateExport(
+      platformToken,
+      "teleport export",
+      entry.runtimeUrl,
+    );
+
+    console.log(`Export started (job ${jobId})...`);
+
+    let exportPlatformToken = platformToken;
+    const terminal = await pollJobUntilDone({
+      label: "platform export",
+      poll: () =>
+        platformPollJobStatus(jobId, exportPlatformToken, entry.runtimeUrl),
+      // The platform token is normally static per-process, but re-reading the
+      // on-disk credential covers the case where the user ran `vellum login`
+      // in another terminal during a long migration. A persistent 401 after
+      // a re-read surfaces to the caller with a clear next step.
+      refreshOn401: async () => {
+        const refreshed = readPlatformToken();
+        if (!refreshed) {
+          throw new Error(
+            "Platform auth expired during export and no credential was found on disk. Run 'vellum login' and retry.",
+          );
         }
-      }
+        exportPlatformToken = refreshed;
+      },
+    });
+
+    if (terminal.status === "failed") {
+      console.error(`Export failed: ${terminal.error}`);
+      process.exit(1);
     }
 
+    if (!terminal.bundleKey) {
+      console.error(
+        "Export completed but the platform did not return a bundle_key. Is the platform up to date?",
+      );
+      process.exit(1);
+    }
+
+    return { bundleKey: terminal.bundleKey };
+  }
+
+  console.error(
+    "Teleport only supports local, docker, and platform assistants as source.",
+  );
+  process.exit(1);
+}
+
+// ---------------------------------------------------------------------------
+// Import into target — unified GCS flow
+// ---------------------------------------------------------------------------
+
+async function importToAssistant(
+  entry: AssistantEntry,
+  cloud: string,
+  bundleKey: string,
+  dryRun: boolean,
+  bundlePlatformUrl?: string,
+): Promise<void> {
+  const platformToken = readPlatformToken();
+  if (!platformToken) {
+    console.error(
+      "Not logged in. Run 'vellum login' first (required for GCS-based teleport).",
+    );
+    process.exit(1);
+  }
+
+  if (cloud === "vellum") {
+    // Platform target — the bundle is already in GCS; kick off preflight or
+    // async import via the unified job-status endpoint.
     if (dryRun) {
       console.log("Running preflight analysis...\n");
 
-      let preflightResult: {
-        statusCode: number;
-        body: Record<string, unknown>;
-      };
-      try {
-        preflightResult = bundleKey
-          ? await platformImportPreflightFromGcs(
-              bundleKey,
-              token,
-              entry.runtimeUrl,
-            )
-          : await platformImportPreflight(bundleData, token, entry.runtimeUrl);
-      } catch (err) {
-        if (err instanceof Error && err.name === "TimeoutError") {
-          console.error("Error: Preflight request timed out after 2 minutes.");
-          process.exit(1);
-        }
-        throw err;
-      }
+      const preflight = await platformImportPreflightFromGcs(
+        bundleKey,
+        platformToken,
+        entry.runtimeUrl,
+      );
 
-      if (
-        preflightResult.statusCode === 401 ||
-        preflightResult.statusCode === 403
-      ) {
+      if (preflight.statusCode === 401 || preflight.statusCode === 403) {
         console.error("Authentication failed. Run 'vellum login' to refresh.");
         process.exit(1);
       }
 
-      if (preflightResult.statusCode === 404) {
+      if (preflight.statusCode === 404) {
         console.error("Assistant not found or not running.");
         process.exit(1);
       }
 
       if (
-        preflightResult.statusCode === 502 ||
-        preflightResult.statusCode === 503 ||
-        preflightResult.statusCode === 504
+        preflight.statusCode === 502 ||
+        preflight.statusCode === 503 ||
+        preflight.statusCode === 504
       ) {
         console.error(
           `Assistant is unreachable. Try 'vellum wake ${entry.assistantId}'.`,
@@ -712,35 +546,53 @@ async function importToAssistant(
         process.exit(1);
       }
 
-      if (preflightResult.statusCode !== 200) {
+      if (preflight.statusCode !== 200) {
         console.error(
-          `Error: Preflight check failed (${preflightResult.statusCode}): ${JSON.stringify(preflightResult.body)}`,
+          `Error: Preflight check failed (${preflight.statusCode}): ${JSON.stringify(preflight.body)}`,
         );
         process.exit(1);
       }
 
-      const result = preflightResult.body as unknown as PreflightResponse;
+      const result = preflight.body as unknown as PreflightResponse;
       printPreflightSummary(result);
       return;
     }
 
-    // Actual import
     console.log("Importing data...");
 
-    let importResult: { statusCode: number; body: Record<string, unknown> };
-    try {
-      importResult = bundleKey
-        ? await platformImportBundleFromGcs(bundleKey, token, entry.runtimeUrl)
-        : await platformImportBundle(bundleData, token, entry.runtimeUrl);
-    } catch (err) {
-      if (err instanceof Error && err.name === "TimeoutError") {
-        console.error("Error: Import request timed out.");
-        process.exit(1);
-      }
-      throw err;
+    const importResult = await platformImportBundleFromGcs(
+      bundleKey,
+      platformToken,
+      entry.runtimeUrl,
+    );
+
+    if (importResult.statusCode === 401 || importResult.statusCode === 403) {
+      console.error("Authentication failed. Run 'vellum login' to refresh.");
+      process.exit(1);
+    }
+
+    if (importResult.statusCode === 404) {
+      console.error("Assistant not found or not running.");
+      process.exit(1);
     }
 
-    handleImportStatusErrors(importResult.statusCode, entry.assistantId);
+    if (
+      importResult.statusCode === 502 ||
+      importResult.statusCode === 503 ||
+      importResult.statusCode === 504
+    ) {
+      console.error(
+        `Assistant is unreachable. Try 'vellum wake ${entry.assistantId}'.`,
+      );
+      process.exit(1);
+    }
+
+    if (importResult.statusCode !== 202 && importResult.statusCode !== 200) {
+      console.error(`Error: Import failed (${importResult.statusCode})`);
+      process.exit(1);
+    }
+
+    let finalBody: Record<string, unknown> = importResult.body;
 
     if (importResult.statusCode === 202) {
       const jobId = (importResult.body as { job_id?: string }).job_id;
@@ -749,74 +601,107 @@ async function importToAssistant(
         process.exit(1);
       }
 
-      const POLL_INTERVAL_MS = 5_000;
-      const TIMEOUT_MS = 10 * 60 * 1_000; // 10 minutes (platform staleness is 930s)
-      const startTime = Date.now();
-      const deadline = startTime + TIMEOUT_MS;
-
-      while (Date.now() < deadline) {
-        await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
-
-        let status: {
-          status: string;
-          result?: Record<string, unknown>;
-          error?: string;
-        };
-        try {
-          status = await platformPollImportStatus(
-            jobId,
-            token,
-            entry.runtimeUrl,
-          );
-        } catch (err) {
-          const msg = err instanceof Error ? err.message : String(err);
-          if (msg.includes("not found")) {
-            throw err;
-          }
-          // Re-throw permanent 4xx errors (auth, forbidden, etc.)
-          // but retry transient 5xx errors
-          const statusMatch = msg.match(/status check failed: (\d+)/);
-          if (statusMatch) {
-            const statusCode = parseInt(statusMatch[1], 10);
-            if (statusCode >= 400 && statusCode < 500) {
-              throw err;
-            }
+      let importPlatformToken = platformToken;
+      const terminal = await pollJobUntilDone({
+        label: "platform import",
+        poll: () =>
+          platformPollJobStatus(jobId, importPlatformToken, entry.runtimeUrl),
+        refreshOn401: async () => {
+          const refreshed = readPlatformToken();
+          if (!refreshed) {
+            throw new Error(
+              "Platform auth expired during import and no credential was found on disk. Run 'vellum login' and retry.",
+            );
           }
-          // Transient error — retry
-          console.warn(`Polling failed, retrying... (${msg})`);
-          continue;
-        }
-
-        if (status.status === "complete") {
-          importResult = { statusCode: 200, body: status.result ?? {} };
-          break;
-        }
-
-        if (status.status === "failed") {
-          console.error(`Import failed: ${status.error ?? "unknown error"}`);
-          process.exit(1);
-        }
-
-        const elapsed = Math.round((Date.now() - startTime) / 1000);
-        process.stdout.write(`\rImporting... ${elapsed}s elapsed`);
-      }
-
-      // Clear the progress line
-      process.stdout.write("\r" + " ".repeat(40) + "\r");
+          importPlatformToken = refreshed;
+        },
+      });
 
-      if (importResult.statusCode === 202) {
-        console.error("Import timed out after 10 minutes.");
+      if (terminal.status === "failed") {
+        console.error(`Import failed: ${terminal.error}`);
         process.exit(1);
       }
+
+      finalBody = (terminal.result as Record<string, unknown>) ?? {};
     }
 
-    const result = importResult.body as unknown as ImportResponse;
+    const result = finalBody as unknown as ImportResponse;
     printImportSummary(result);
     return;
   }
 
   if (cloud === "local" || cloud === "docker") {
-    await importViaHttp(entry, bundleData, dryRun);
+    if (dryRun) {
+      // TODO(cli): support dry-run against local targets
+      console.error(
+        "Error: --dry-run is not yet supported for local or docker targets (no preflight-from-gcs endpoint on the runtime).",
+      );
+      process.exit(1);
+    }
+
+    // Ask the platform for a signed download URL and hand it to the local
+    // runtime. The runtime streams the bundle straight out of GCS — the CLI
+    // never touches the bytes. The URL must target the same platform the
+    // bundle was uploaded to; otherwise the object won't exist on this
+    // platform's GCS bucket.
+    const { url: bundleUrl } = await platformRequestSignedUrl(
+      { operation: "download", bundleKey },
+      platformToken,
+      bundlePlatformUrl,
+    );
+
+    console.log("Importing data...");
+
+    let jobId: string;
+    let accessToken: string;
+    try {
+      const result = await callRuntimeWithAuthRetry(
+        entry.runtimeUrl,
+        entry.assistantId,
+        async (token) => {
+          const r = await localRuntimeImportFromGcs(entry.runtimeUrl, token, {
+            bundleUrl,
+          });
+          return { jobId: r.jobId, token };
+        },
+      );
+      jobId = result.jobId;
+      accessToken = result.token;
+    } catch (err) {
+      if (err instanceof MigrationInProgressError) {
+        // Fail fast — the existing job is importing someone else's bundle
+        // (the original caller's), not ours. Polling it would report success
+        // on an import that wasn't the one we just kicked off.
+        console.error(
+          `Error: Another teleport import is already in progress on '${entry.assistantId}' (job ${err.existingJobId}). Wait for it to finish or check its status, then re-run.`,
+        );
+        process.exit(1);
+      }
+      throw err;
+    }
+
+    const terminal = await pollJobUntilDone({
+      label: "local-runtime import",
+      poll: () =>
+        localRuntimePollJobStatus(entry.runtimeUrl, accessToken, jobId),
+      refreshOn401: async () => {
+        accessToken = await getAccessToken(
+          entry.runtimeUrl,
+          entry.assistantId,
+          entry.assistantId,
+          { forceRefresh: true },
+        );
+      },
+    });
+
+    if (terminal.status === "failed") {
+      console.error(`Import failed: ${terminal.error}`);
+      process.exit(1);
+    }
+
+    const result = ((terminal.result as Record<string, unknown>) ??
+      {}) as unknown as ImportResponse;
+    printImportSummary(result);
     return;
   }
 
@@ -952,68 +837,6 @@ export async function resolveOrHatchTarget(
   process.exit(1);
 }
 
-// ---------------------------------------------------------------------------
-// Error handling helpers
-// ---------------------------------------------------------------------------
-
-function handleLocalResponseErrors(
-  response: Response,
-  assistantName: string,
-): void {
-  if (response.status === 401 || response.status === 403) {
-    console.error("Authentication failed.");
-    process.exit(1);
-  }
-
-  if (response.status === 404) {
-    console.error("Assistant not found or not running.");
-    process.exit(1);
-  }
-
-  if (
-    response.status === 502 ||
-    response.status === 503 ||
-    response.status === 504
-  ) {
-    console.error(
-      `Assistant is unreachable. Try 'vellum wake ${assistantName}'.`,
-    );
-    process.exit(1);
-  }
-
-  if (!response.ok) {
-    console.error(`Error: Request failed (${response.status})`);
-    process.exit(1);
-  }
-}
-
-function handleImportStatusErrors(
-  statusCode: number,
-  assistantName: string,
-): void {
-  if (statusCode === 401 || statusCode === 403) {
-    console.error("Authentication failed. Run 'vellum login' to refresh.");
-    process.exit(1);
-  }
-
-  if (statusCode === 404) {
-    console.error("Assistant not found or not running.");
-    process.exit(1);
-  }
-
-  if (statusCode === 502 || statusCode === 503 || statusCode === 504) {
-    console.error(
-      `Assistant is unreachable. Try 'vellum wake ${assistantName}'.`,
-    );
-    process.exit(1);
-  }
-
-  if (statusCode < 200 || statusCode >= 300) {
-    console.error(`Error: Import failed (${statusCode})`);
-    process.exit(1);
-  }
-}
-
 // ---------------------------------------------------------------------------
 // Summary printing — matches restore.ts format
 // ---------------------------------------------------------------------------
@@ -1270,6 +1093,19 @@ export async function teleport(): Promise<void> {
         process.exit(1);
       }
 
+      // Dry-run feasibility check — reject local/docker targets BEFORE any
+      // export work. The local runtime has no preflight-from-gcs endpoint yet,
+      // so we can't actually run a dry-run against it; burning a GCS upload
+      // just to fail afterwards would be wasteful.
+      // TODO(cli): support dry-run against local targets (needs a
+      // preflight-from-gcs endpoint on the runtime).
+      if (toCloud === "local" || toCloud === "docker") {
+        console.error(
+          "Error: --dry-run is not yet supported for local or docker targets (no preflight-from-gcs endpoint on the runtime).",
+        );
+        process.exit(1);
+      }
+
       // Version guard: block platform→non-platform when target is behind
       if (fromCloud === "vellum" && toCloud !== "vellum") {
         const [sourceVersion, targetVersion] = await Promise.all([
@@ -1292,41 +1128,49 @@ export async function teleport(): Promise<void> {
         }
       }
 
+      // Pin both upload and download to the same platform instance. For
+      // platform targets the bundle is owned by the target platform; for
+      // platform sources it's owned by the source platform. Only one of
+      // these branches applies at a time (same-env was rejected earlier).
+      const bundlePlatformUrl =
+        toCloud === "vellum"
+          ? existingTarget.runtimeUrl
+          : fromCloud === "vellum"
+            ? fromEntry.runtimeUrl
+            : undefined;
+
       console.log(`Exporting from ${from} (${fromCloud})...`);
-      const bundleData = await exportFromAssistant(fromEntry, fromCloud);
+      const { bundleKey } = await exportFromAssistant(
+        fromEntry,
+        fromCloud,
+        bundlePlatformUrl,
+      );
       console.log(`Importing to ${existingTarget.assistantId} (${toCloud})...`);
-      await importToAssistant(existingTarget, toCloud, bundleData, true);
+      await importToAssistant(
+        existingTarget,
+        toCloud,
+        bundleKey,
+        true,
+        bundlePlatformUrl,
+      );
     } else {
       // No existing target — just describe what would happen
       console.log("Dry run summary:");
       console.log(`  Would export data from: ${from} (${fromCloud})`);
-      if (targetEnv === "platform") {
-        // For platform targets, reflect the reordered flow
-        console.log(`  Would upload bundle via signed URL (if available)`);
-        console.log(
-          `  Would hatch a new ${targetEnv} assistant${targetName ? ` named '${targetName}'` : ""}`,
-        );
-        console.log(`  Would import data into the new assistant`);
-      } else {
-        console.log(
-          `  Would hatch a new ${targetEnv} assistant${targetName ? ` named '${targetName}'` : ""}`,
-        );
-        console.log(`  Would import data into the new assistant`);
-      }
+      console.log(`  Would upload bundle via signed URL`);
+      console.log(
+        `  Would hatch a new ${targetEnv} assistant${targetName ? ` named '${targetName}'` : ""}`,
+      );
+      console.log(`  Would import data into the new assistant`);
     }
 
     console.log(`Dry run complete — no changes were made.`);
     return;
   }
 
-  // Export from source
-  console.log(`Exporting from ${from} (${fromCloud})...`);
-  const bundleData = await exportFromAssistant(fromEntry, fromCloud);
-
   // Platform target: reordered flow — upload to GCS before hatching so that
-  // if upload fails, no empty assistant is left dangling on the platform.
+  // if export/upload fails, no empty assistant is left dangling on the platform.
   if (targetEnv === "platform") {
-    // Step B — Auth
     const token = readPlatformToken();
     if (!token) {
       console.error("Not logged in. Run 'vellum login' first.");
@@ -1334,7 +1178,7 @@ export async function teleport(): Promise<void> {
     }
 
     // If targeting an existing assistant, validate cloud match early — before
-    // uploading — so we don't waste a GCS upload on an invalid command.
+    // exporting — so we don't waste work on an invalid command.
     const existingTarget = targetName ? findAssistantByName(targetName) : null;
     if (existingTarget) {
       const existingCloud = resolveCloud(existingTarget);
@@ -1347,12 +1191,12 @@ export async function teleport(): Promise<void> {
       }
     }
 
-    // Use the existing target's runtimeUrl for all platform calls so upload
-    // and import hit the same instance.
+    // Use the existing target's runtimeUrl for all platform calls so the
+    // export, upload, and import all hit the same instance.
     const targetPlatformUrl = existingTarget?.runtimeUrl;
 
-    // Step B2 — Pre-check: block if the user already has a platform assistant.
-    // This runs BEFORE the expensive GCS upload so we don't waste bandwidth.
+    // Pre-check: block if the user already has a platform assistant. This
+    // runs BEFORE the expensive export so we don't waste the upload.
     if (!existingTarget) {
       const existing = await checkExistingPlatformAssistant(
         token,
@@ -1376,43 +1220,40 @@ export async function teleport(): Promise<void> {
       }
     }
 
-    // Step C — Upload to GCS
-    // bundleKey: string = uploaded successfully, null = tried but unavailable,
-    // undefined would mean "never tried" (not used here).
-    let bundleKey: string | null = null;
-    try {
-      const { uploadUrl, bundleKey: key } = await platformRequestUploadUrl(
-        token,
-        targetPlatformUrl,
-      );
-      bundleKey = key;
-      console.log("Uploading bundle to GCS...");
-      await platformUploadToSignedUrl(uploadUrl, bundleData);
-    } catch (err) {
-      // If signed uploads unavailable (503), fall back to inline upload later
-      const msg = err instanceof Error ? err.message : String(err);
-      if (msg.includes("not available")) {
-        bundleKey = null;
-      } else {
-        throw err;
-      }
-    }
+    // Export — for local/docker sources this uploads straight into GCS via
+    // the platform's signed URL; for platform sources this runs a server-side
+    // export and we read the resulting bundle_key.
+    // The signed upload URL must be requested from the same platform instance
+    // where the import will run. For existing targets that's the lockfile's
+    // runtimeUrl; for fresh hatches it's getPlatformUrl() (which is what
+    // resolveOrHatchTarget writes to the new entry).
+    console.log(`Exporting from ${from} (${fromCloud})...`);
+    const bundlePlatformUrl = targetPlatformUrl ?? getPlatformUrl();
+    const { bundleKey } = await exportFromAssistant(
+      fromEntry,
+      fromCloud,
+      bundlePlatformUrl,
+    );
 
-    // Step D — Hatch (upload succeeded or fallback to inline — safe to hatch)
+    // Hatch (export succeeded — safe to create the target)
     const toEntry = await resolveOrHatchTarget(targetEnv, targetName);
     const toCloud = resolveCloud(toEntry);
 
-    // Step E — Import from GCS (or inline fallback)
-    // Pass bundleKey (string) or null to signal "already tried, use inline".
+    // Import from GCS
     console.log(`Importing to ${toEntry.assistantId} (${toCloud})...`);
-    await importToAssistant(toEntry, toCloud, bundleData, false, bundleKey);
+    await importToAssistant(
+      toEntry,
+      toCloud,
+      bundleKey,
+      false,
+      bundlePlatformUrl,
+    );
 
-    // Success summary
     console.log(`Teleport complete: ${from} → ${toEntry.assistantId}`);
     return;
   }
 
-  // Non-platform targets (local/docker): existing flow unchanged
+  // Non-platform targets (local/docker)
   // For local<->docker transfers, stop (sleep) the source to free up ports
   // before hatching the target. We do NOT retire yet — if hatch or import
   // fails, the user can recover by running `vellum wake <source>`.
@@ -1448,16 +1289,34 @@ export async function teleport(): Promise<void> {
     }
   }
 
+  // Pin the bundle's platform instance so upload and download land on the
+  // same GCS bucket. For platform sources the bundle is owned by the source
+  // platform. For local/docker→local/docker the bundle lives on whatever
+  // platform getPlatformUrl() currently resolves to — we resolve it once
+  // here so a lockfile change mid-teleport can't split export and import.
+  const bundlePlatformUrl =
+    fromCloud === "vellum" ? fromEntry.runtimeUrl : getPlatformUrl();
+
+  // Export from source (bundle lives in GCS after this returns).
+  console.log(`Exporting from ${from} (${fromCloud})...`);
+  const { bundleKey } = await exportFromAssistant(
+    fromEntry,
+    fromCloud,
+    bundlePlatformUrl,
+  );
+
   if (sourceIsLocalOrDocker && targetIsLocalOrDocker && !keepSource) {
     console.log(`Stopping source assistant '${from}' to free ports...`);
     if (fromCloud === "docker") {
       const res = dockerResourceNames(fromEntry.assistantId);
       await sleepContainers(res);
     } else if (fromEntry.resources) {
-      const pidFile = fromEntry.resources.pidFile;
       const vellumDir = join(fromEntry.resources.instanceDir, ".vellum");
       const gatewayPidFile = join(vellumDir, "gateway.pid");
-      await stopProcessByPidFile(pidFile, "assistant");
+      await stopProcessByPidFile(
+        getDaemonPidPath(fromEntry.resources),
+        "assistant",
+      );
       await stopProcessByPidFile(gatewayPidFile, "gateway", undefined, 7000);
     }
     console.log(`Source assistant '${from}' stopped.`);
@@ -1513,9 +1372,15 @@ export async function teleport(): Promise<void> {
     }
   }
 
-  // Import to target
+  // Import to target (also GCS-driven)
   console.log(`Importing to ${toEntry.assistantId} (${toCloud})...`);
-  await importToAssistant(toEntry, toCloud, bundleData, false);
+  await importToAssistant(
+    toEntry,
+    toCloud,
+    bundleKey,
+    false,
+    bundlePlatformUrl,
+  );
 
   // After successful import, inject fresh platform credentials if the
   // user is logged in — replaces the source's stale vellum:* credentials
@@ -1540,6 +1405,5 @@ export async function teleport(): Promise<void> {
     }
   }
 
-  // Success summary
   console.log(`Teleport complete: ${from} → ${toEntry.assistantId}`);
 }