@rigkit/provider-freestyle 0.2.3 → 0.2.4

package/src/index.ts

@@ -3,30 +3,37 @@ import {
   type WorkflowProviderDefinition,
 } from "@rigkit/sdk";
 import type { BaseProviderPlugin } from "@rigkit/engine";
-import type { WorkflowProviderController } from "@rigkit/engine";
-import { Freestyle } from "freestyle";
 import * as z from "zod/v4-mini";
 import { freestyleIdentityId, freestyleToken, freestyleTokenId } from "./auth.ts";
+import {
+  createFreestyleAuthenticatedClient,
+  createFreestyleProxyFetch,
+  type FreestyleProviderAuthConfig,
+} from "./host-auth.ts";
 import {
   FREESTYLE_PROVIDER_ID,
   FREESTYLE_TERMINAL_PROVIDER_ID,
   createFreestyleTerminalController,
   createFreestyleWorkflowProvider,
-  isFreestyleVmSnapshotRef,
 } from "./provider.ts";
 import type { FreestyleRuntime, FreestyleTerminalRuntime } from "./provider.ts";
-import { createFreestyleStore } from "./store.ts";
 
 const freestyleProviderConfigSchema = z.object({
-  apiKey: z.string().check(z.minLength(1)),
-  image: z.string().check(z.minLength(1)),
-  cpu: z.optional(z.number()),
-  memory: z.optional(z.union([z.string(), z.number()])),
-  disk: z.optional(z.union([z.string(), z.number()])),
-  idleTimeoutSeconds: z.optional(z.nullable(z.number())),
+  auth: z.optional(z.object({
+    apiKey: z.optional(z.string().check(z.minLength(1))),
+    profile: z.optional(z.string().check(z.minLength(1))),
+    teamId: z.optional(z.string().check(z.minLength(1))),
+    apiUrl: z.optional(z.string().check(z.minLength(1))),
+    dashboardUrl: z.optional(z.string().check(z.minLength(1))),
+    stackApiUrl: z.optional(z.string().check(z.minLength(1))),
+    stackAppUrl: z.optional(z.string().check(z.minLength(1))),
+    stackProjectId: z.optional(z.string().check(z.minLength(1))),
+    stackPublishableClientKey: z.optional(z.string().check(z.minLength(1))),
+  })),
 });
 
 export type FreestyleProviderConfig = z.output<typeof freestyleProviderConfigSchema>;
+export type { FreestyleProviderAuthConfig };
 
 export type FreestyleProviderDefinition = WorkflowProviderDefinition<
   typeof FREESTYLE_PROVIDER_ID,
@@ -41,7 +48,7 @@ export type FreestyleTerminalProviderDefinition = WorkflowProviderDefinition<
 >;
 
 export function provider(
-  config: FreestyleProviderDefinition["config"],
+  config: FreestyleProviderDefinition["config"] = {},
 ): FreestyleProviderDefinition {
   return defineProvider(FREESTYLE_PROVIDER_ID, config, freestyleProviderPlugin);
 }
@@ -59,50 +66,18 @@ export const defineFreestyleProvider = provider;
 
 export const freestyleProviderPlugin: BaseProviderPlugin = {
   providerId: FREESTYLE_PROVIDER_ID,
-  createProvider({ provider, storage }) {
+  async createProvider({ provider, hostStorage, local }) {
     const config = parseFreestyleProviderConfig(provider.config);
-    const { apiKey, ...vm } = config;
-    let controller: Promise<WorkflowProviderController<FreestyleRuntime>> | undefined;
-
-    const load = async () => {
-      controller ??= create();
-      return await controller;
-    };
-
-    const create = async () => {
-      const store = createFreestyleStore(storage);
-      const savedIdentity = store.getIdentity();
-      if (savedIdentity) {
-        return createFreestyleWorkflowProvider({
-          apiKey,
-          identityId: savedIdentity.identityId,
-          token: savedIdentity.token,
-          vm,
-        });
-      }
-
-      const client = new Freestyle({ apiKey });
-      const { identity, identityId } = await client.identities.create();
-      const { token, tokenId } = await identity.tokens.create();
-      const createdIdentity = store.saveIdentity({
-        identityId: freestyleIdentityId(identityId),
-        tokenId: freestyleTokenId(tokenId),
-        token: freestyleToken(token),
-      });
-
-      return createFreestyleWorkflowProvider({
-        apiKey,
-        identityId: createdIdentity.identityId,
-        token: createdIdentity.token,
-        vm,
-      });
-    };
-
-    return {
-      providerId: FREESTYLE_PROVIDER_ID,
-      runtime: async (context) => await (await load()).runtime(context),
-      validateArtifact: (ref) => isFreestyleVmSnapshotRef(ref),
-    } satisfies WorkflowProviderController<FreestyleRuntime>;
+    const authenticated = await createFreestyleAuthenticatedClient({
+      auth: config.auth,
+      hostStorage,
+      local,
+    });
+    return createFreestyleWorkflowProvider({
+      client: authenticated.client,
+      identityId: authenticated.identityId,
+      token: authenticated.token,
+    });
   },
 };
 
@@ -113,6 +88,10 @@ export const freestyleTerminalPlugin: BaseProviderPlugin = {
   },
 };
 
+export {
+  createFreestyleAuthenticatedClient,
+  createFreestyleProxyFetch,
+} from "./host-auth.ts";
 export {
   freestyleIdentityId,
   freestyleToken,
@@ -124,24 +103,23 @@ export {
 export {
   FREESTYLE_PROVIDER_ID,
   FREESTYLE_TERMINAL_PROVIDER_ID,
-  createFreestyleProvider,
   createFreestyleTerminalController,
   createFreestyleWorkflowController,
   createFreestyleWorkflowProvider,
-  isFreestyleVmSnapshotRef,
 } from "./provider.ts";
 export { createFreestyleStore } from "./store.ts";
 export { createFreestyleTerminalSession } from "./terminal-session.ts";
 export { RIGKIT_PROVIDER_FREESTYLE_VERSION } from "./version.ts";
+export { Freestyle, VmBaseImage, VmSpec, VmWith, VmWithInstance } from "freestyle";
+export type { CreateVmOptions } from "freestyle";
 export type {
   FreestyleCmuxSshOptions,
   FreestyleCmuxSshOptionsInput,
   FreestyleRuntime,
+  FreestyleSdkVm,
+  FreestyleSshInput,
   FreestyleTerminalRuntime,
   FreestyleVscodeUrlOptions,
-  FreestyleVmConfig,
-  FreestyleVmRuntime,
-  FreestyleVmSnapshotRef,
 } from "./provider.ts";
 export type { FreestyleGitRelationship, FreestyleIdentity } from "./store.ts";
 

package/src/provider.test.ts

@@ -1,82 +1,61 @@
-import { describe, expect, test } from "bun:test";
-import type { ExecOptions, ExecOutputChunk, ExecResult } from "@rigkit/sdk";
-import type {
-  BaseDevMachineProvider,
-  ProviderRuntimeContext,
-  SshConnection,
-  VmHandle,
-  WorkflowEvent,
-} from "@rigkit/engine";
-import { createFreestyleWorkflowController, wrapCommand } from "./provider.ts";
-
-describe("Freestyle provider command wrapper", () => {
-  test("sets a root HOME fallback for exec commands", () => {
-    expect(wrapCommand("printf '%s\n' \"$HOME\"")).toContain("export HOME=${HOME:-/root}");
-  });
-
-  test("allows callers to override HOME explicitly", () => {
-    const wrapped = wrapCommand("pwd", {
-      env: { HOME: "/workspace/home" },
-    });
-
-    expect(wrapped).toContain("export HOME=${HOME:-/root}");
-    expect(wrapped).toContain("export HOME='\\''/workspace/home'\\''");
-  });
+import { afterEach, describe, expect, test } from "bun:test";
+import { Freestyle } from "freestyle";
+import type { ProviderInteractionSession, ProviderRuntimeContext } from "@rigkit/engine";
+import { freestyleIdentityId, freestyleToken } from "./auth.ts";
+import {
+  buildInteractiveSshCommand,
+  createFreestyleTerminalController,
+  createFreestyleWorkflowController,
+} from "./provider.ts";
+
+const previousFetch = globalThis.fetch;
+
+afterEach(() => {
+  globalThis.fetch = previousFetch;
+});
 
-  test("streams VM command output through run events without replaying buffered output", async () => {
-    const chunks: ExecOutputChunk[] = [];
-    const events: WorkflowEvent[] = [];
-    const provider = new StreamingProvider();
-    const controller = createFreestyleWorkflowController(provider);
-    const runtime = await controller.runtime(providerContext(events));
-    const vm = runtime.vms.fromId("vm-stream");
-
-    const result = await vm.exec("printf ready", {
-      name: "stream command",
-      onOutput: (chunk) => {
-        chunks.push(chunk);
-      },
+describe("Freestyle provider host adapters", () => {
+  test("creates SSH options and grants VM access internally", async () => {
+    const requests: string[] = [];
+    globalThis.fetch = (async (resource, init) => {
+      requests.push(`${init?.method ?? "GET"} ${String(resource)}`);
+      return Response.json({});
+    }) as typeof fetch;
+
+    const runtime = await createFreestyleWorkflowController({
+      client: new Freestyle({ apiKey: "test-key" }),
+      identityId: freestyleIdentityId("identity-stream"),
+      token: freestyleToken("token"),
+    }).runtime(providerContext());
+
+    await expect(runtime.createSSHOptions({ vmId: "vm-stream" })).resolves.toEqual({
+      kind: "ssh",
+      host: "vm-ssh.freestyle.sh",
+      username: "vm-stream+root",
+      auth: { type: "token", token: "token" },
+      command: "ssh vm-stream+root:token@vm-ssh.freestyle.sh",
     });
-
-    expect(result.stdout).toBe("ready\n");
-    expect(chunks).toEqual([{ stream: "stdout", data: "ready\n" }]);
-    expect(events).toEqual([
-      {
-        type: "command.started",
-        nodePath: "workflow.step",
-        commandName: "stream command",
-        command: "printf ready",
-      },
-      {
-        type: "command.output",
-        nodePath: "workflow.step",
-        commandName: "stream command",
-        stream: "stdout",
-        data: "ready\n",
-      },
-      {
-        type: "command.completed",
-        nodePath: "workflow.step",
-        commandName: "stream command",
-        exitCode: 0,
-      },
-    ]);
+    expect(requests).toContain("POST https://api.freestyle.sh/identity/v1/identities/identity-stream/permissions/vm/vm-stream");
   });
 
   test("creates cmux ssh options with Freestyle-owned ssh settings", async () => {
-    const provider = new StreamingProvider();
-    const controller = createFreestyleWorkflowController(provider);
-    const runtime = await controller.runtime(providerContext([]));
-    const vm = runtime.vms.fromId("vm-stream");
+    globalThis.fetch = (async () => Response.json({})) as unknown as typeof fetch;
 
-    const ssh = await runtime.cmux.createSshOptions(vm, {
+    const runtime = await createFreestyleWorkflowController({
+      client: new Freestyle({ apiKey: "test-key" }),
+      identityId: freestyleIdentityId("identity-stream"),
+      token: freestyleToken("token"),
+    }).runtime(providerContext());
+
+    const ssh = await runtime.cmux.createSshOptions({
+      vmId: "vm-stream",
       sshOptions: ["ServerAliveInterval=15"],
       skipDaemonBootstrap: true,
     });
 
     expect(ssh).toEqual({
       kind: "ssh",
-      destination: "root,token@localhost",
+      destination: "vm-stream+root,token@vm-ssh.freestyle.sh",
       skipDaemonBootstrap: true,
       sshOptions: [
         "StrictHostKeyChecking=no",
@@ -90,71 +69,151 @@ describe("Freestyle provider command wrapper", () => {
     });
   });
 
+  test("treats existing VM permissions as idempotent for cmux ssh options", async () => {
+    const calls: string[] = [];
+    const runtime = await createFreestyleWorkflowController({
+      client: {
+        identities: {
+          ref: () => ({
+            permissions: {
+              vms: {
+                grant: async () => {
+                  calls.push("grant");
+                  throw new Error("PERMISSION_ALREADY_EXISTS: Permission already exists");
+                },
+                update: async () => {
+                  calls.push("update");
+                },
+              },
+            },
+          }),
+        },
+      } as unknown as Freestyle,
+      identityId: freestyleIdentityId("identity-stream"),
+      token: freestyleToken("token"),
+    }).runtime(providerContext());
+
+    await expect(runtime.cmux.createSshOptions({ vmId: "vm-stream" })).resolves.toMatchObject({
+      destination: "vm-stream+root,token@vm-ssh.freestyle.sh",
+    });
+    expect(calls).toEqual(["grant", "update"]);
+  });
+
   test("creates VS Code URLs using the Freestyle ssh authority", async () => {
-    const provider = new StreamingProvider();
-    const controller = createFreestyleWorkflowController(provider);
-    const runtime = await controller.runtime(providerContext([]));
-    const vm = runtime.vms.fromId("vm-stream");
+    globalThis.fetch = (async () => Response.json({})) as unknown as typeof fetch;
+
+    const runtime = await createFreestyleWorkflowController({
+      client: new Freestyle({ apiKey: "test-key" }),
+      identityId: freestyleIdentityId("identity-stream"),
+      token: freestyleToken("token"),
+    }).runtime(providerContext());
 
-    const url = await runtime.vscode.createUrl(vm, { cwd: "/workspace/site" });
+    const url = await runtime.vscode.createUrl({ vmId: "vm-stream", cwd: "/workspace/site" });
 
     expect(url).toBe(
-      "vscode://vscode-remote/ssh-remote+root%3Atoken%40localhost/workspace/site?windowId=_blank",
+      "vscode://vscode-remote/ssh-remote+vm-stream%2Broot%3Atoken%40vm-ssh.freestyle.sh/workspace/site?windowId=_blank",
     );
   });
-});
 
-const sshConnection: SshConnection = {
-  kind: "ssh",
-  host: "localhost",
-  username: "root",
-  auth: { type: "token", token: "token" },
-  command: "ssh vm-stream",
-};
+  test("honors explicit SSH users", async () => {
+    globalThis.fetch = (async () => Response.json({})) as unknown as typeof fetch;
 
-class StreamingProvider implements BaseDevMachineProvider {
-  readonly providerId = "freestyle";
+    const runtime = await createFreestyleWorkflowController({
+      client: new Freestyle({ apiKey: "test-key" }),
+      identityId: freestyleIdentityId("identity-stream"),
+      token: freestyleToken("token"),
+    }).runtime(providerContext());
 
-  async createVm(): Promise<VmHandle> {
-    return { vmId: "vm-stream" };
-  }
+    await expect(runtime.createSSHOptions({ vmId: "vm-stream", user: "ubuntu" })).resolves.toMatchObject({
+      username: "vm-stream+ubuntu",
+      command: "ssh vm-stream+ubuntu:token@vm-ssh.freestyle.sh",
+    });
+  });
+
+  test("runs terminal commands as SSH remote commands instead of typed startup input", async () => {
+    let html = "";
+    const runtime = await createFreestyleTerminalController().runtime({
+      ...providerContext(),
+      interaction: {
+        present: async <Result>(session: ProviderInteractionSession<Result>) => {
+          const response = await fetch(session.url);
+          html = await response.text();
+          session.stop();
+          return { finished: true } as Result;
+        },
+      },
+    });
 
-  async createVmFromSnapshot(): Promise<VmHandle> {
-    return { vmId: "vm-stream" };
-  }
+    await expect(runtime.open("GitHub auth", {
+      ssh: {
+        kind: "ssh",
+        host: "vm-ssh.freestyle.sh",
+        username: "vm-stream+root",
+        auth: { type: "token", token: "token" },
+        command: "ssh vm-stream+root:token@vm-ssh.freestyle.sh",
+      },
+      command: "gh auth login --hostname github.com",
+    })).resolves.toEqual({ finished: true });
 
-  async exec(_vm: VmHandle, _command: string, options?: ExecOptions): Promise<ExecResult> {
-    await options?.onOutput?.({ stream: "stdout", data: "ready\n" });
-    return { stdout: "ready\n", stderr: "", exitCode: 0, ok: true };
-  }
+    expect(html).toContain("gh auth login --hostname github.com");
+    expect(html).toContain("const startupInput = null;");
+    expect(html).toContain("const canFinishWhileRunning = false;");
+  });
 
-  async readFile(): Promise<string> {
-    return "";
-  }
+  test("can keep an SSH terminal open after a successful remote command", () => {
+    const command = buildInteractiveSshCommand(
+      {
+        kind: "ssh",
+        host: "vm-ssh.freestyle.sh",
+        username: "vm-stream+root",
+        auth: { type: "token", token: "token" },
+        command: "ssh vm-stream+root:token@vm-ssh.freestyle.sh",
+      },
+      "gh auth status -h github.com",
+      { keepOpenAfterCommand: true },
+    );
 
-  async writeFile(): Promise<void> {}
+    expect(command).toContain("gh auth status -h github.com");
+    expect(command).toContain("status=$?");
+    expect(command).toContain('if [ "$status" -ne 0 ]; then exit "$status"; fi');
+    expect(command).toContain('exec "${SHELL:-/bin/bash}" -l');
+  });
 
-  async snapshot(): Promise<{ snapshotId: string; sourceVmId: string }> {
-    return { snapshotId: "snap-stream", sourceVmId: "vm-stream" };
-  }
+  test("allows finishing while a keep-open SSH command is running", async () => {
+    let html = "";
+    const runtime = await createFreestyleTerminalController().runtime({
+      ...providerContext(),
+      interaction: {
+        present: async <Result>(session: ProviderInteractionSession<Result>) => {
+          const response = await fetch(session.url);
+          html = await response.text();
+          session.stop();
+          return { finished: true } as Result;
+        },
+      },
+    });
 
-  async ssh(): Promise<SshConnection> {
-    return sshConnection;
-  }
+    await runtime.open("GitHub auth", {
+      ssh: {
+        kind: "ssh",
+        host: "vm-ssh.freestyle.sh",
+        username: "vm-stream+root",
+        auth: { type: "token", token: "token" },
+        command: "ssh vm-stream+root:token@vm-ssh.freestyle.sh",
+      },
+      command: "gh auth login --hostname github.com",
+      keepOpenAfterCommand: true,
+    });
 
-  async deleteVm(): Promise<void> {}
-}
+    expect(html).toContain("const canFinishWhileRunning = true;");
+  });
+});
 
-function providerContext(
-  events: WorkflowEvent[],
-  local: Partial<ProviderRuntimeContext["local"]> = {},
-): ProviderRuntimeContext {
+function providerContext(): ProviderRuntimeContext {
   return {
     workflow: "workflow",
     nodePath: "workflow.step",
-    emit: (event) => {
-      events.push(event);
-    },
+    emit: () => {},
     interaction: {
       present: async () => {
         throw new Error("unexpected interaction");
@@ -162,7 +221,6 @@ function providerContext(
     },
     local: {
       open: async () => {},
-      ...local,
     },
     metadata: () => {},
   };