recreate pint client

Author: mohamedveron

src/PintClient/index.ts

@@ -0,0 +1,130 @@
+import { Port } from "../pitcher-protocol/messages/port";
+import { Emitter, EmitterSubscription, Event } from "../utils/event";
+import { SandboxSession } from "../types";
+import { Disposable } from "../utils/disposable";
+import { Client, createClient, createConfig } from "../api-clients/pint/client";
+import {
+  IAgentClient,
+  IAgentClientPorts,
+  IAgentClientShells,
+  IAgentClientState,
+  IAgentClientFS,
+  IAgentClientSetup,
+  IAgentClientTasks,
+  IAgentClientSystem,
+  PickRawFsResult,
+} from "../agent-client-interface";
+import {
+  listPorts,
+  PortInfo,
+  streamPortsList,
+} from "../api-clients/pint";
+
+
+
+function parseStreamEvent<T>(evt: unknown): T {
+  if (typeof evt !== "string") {
+    return evt as T;
+  }
+
+  const evtWithoutDataPrefix = evt.substring(5);
+
+  return JSON.parse(evtWithoutDataPrefix);
+}
+
+class PintPortsClient implements IAgentClientPorts {
+  private onPortsUpdatedEmitter = new EmitterSubscription<Port[]>((fire) => {
+    const abortController = new AbortController();
+
+    streamPortsList({
+      signal: abortController.signal,
+      headers: {
+        headers: { Accept: "text/event-stream" },
+      },
+    }).then(async ({ stream }) => {
+      for await (const evt of stream) {
+        const data = parseStreamEvent<PortInfo[]>(evt);
+
+        fire(
+          data.map((pintPort) => ({
+            port: pintPort.port,
+            url: pintPort.address,
+          }))
+        );
+      }
+    });
+
+    return Disposable.create(() => {
+      abortController.abort();
+    });
+  });
+  onPortsUpdated = this.onPortsUpdatedEmitter.event;
+
+  constructor(private apiClient: Client, private sandboxId: string) {}
+
+  async getPorts(): Promise<Port[]> {
+    const ports = await listPorts({
+      client: this.apiClient,
+    });
+
+    return (
+      ports.data?.ports.map((port) => ({
+        port: port.port,
+        url: `https://${this.sandboxId}-${port.port}.csb.app`,
+      })) ?? []
+    );
+  }
+}
+
+
+export class PintClient implements IAgentClient {
+  static async create(session: SandboxSession) {
+    return new PintClient(session);
+  }
+
+  readonly type = "pint" as const;
+
+  // Since there is no websocket connection or internal hibernation, the state
+  // will always be CONNECTED. No state change events will be triggered
+  readonly state = "CONNECTED";
+  private onStateChangeEmitter = new Emitter<IAgentClientState>();
+  onStateChange = this.onStateChangeEmitter.event;
+
+  sandboxId: string;
+  workspacePath: string;
+  isUpToDate: boolean;
+
+  ports: IAgentClientPorts;
+  shells: IAgentClientShells;
+  fs: IAgentClientFS;
+  setup: IAgentClientSetup;
+  tasks: IAgentClientTasks;
+  system: IAgentClientSystem;
+
+  constructor(session: SandboxSession) {
+    this.sandboxId = session.sandboxId;
+    this.workspacePath = session.workspacePath;
+    this.isUpToDate = true;
+
+    const apiClient = createClient(
+      createConfig({
+        baseUrl: session.pitcherURL,
+        headers: {
+          Authorization: `Bearer ${session.pitcherToken}`,
+        },
+      })
+    );
+
+    this.ports = new PintPortsClient(apiClient, this.sandboxId);
+    this.shells = {} as IAgentClientShells; // Not implemented for Pint
+    this.fs = {} as IAgentClientFS; // Not implemented for Pint  
+    this.tasks = {} as IAgentClientTasks; // Not implemented for Pint
+    this.setup = {} as IAgentClientSetup; // Not implemented for Pint
+    this.system = {} as IAgentClientSystem; // Not implemented for Pint
+  }
+
+  ping(): void {}
+  async reconnect(): Promise<void> {}
+  async disconnect(): Promise<void> {}
+  dispose(): void {}
+}

src/Sandbox.ts

@@ -151,13 +151,13 @@ export class Sandbox {
           return `export ${key}='${safe}'`;
         })
         .join("\n");
-      commands.push(
-        [
-          `cat << 'EOF' > "$HOME/.private/.env"`,
-          envStrings,
-          `EOF`,
-        ].join("\n")
-      );
+      const cmd = [
+        `mkdir -p "$HOME/.private"`,
+        `cat << 'EOF' > "$HOME/.private/.env"`,
+        envStrings,
+        `EOF`,
+      ].join("\n");
+      await client.commands.run(cmd);
     }
 
     if (customSession.git) {
@@ -188,8 +188,12 @@ export class Sandbox {
     pitcherManagerResponse: PitcherManagerResponse,
     customSession?: SessionCreateOptions
   ): Promise<SandboxSession> {
+    // HACK: we currently do not get a flag for pint, but this is a check we can use for now
+    const isPint = false;
+
     if (!customSession || !customSession.id) {
       return {
+        isPint: isPint,
         sandboxId: this.id,
         bootupType: this.bootupType,
         hostToken: customSession?.hostToken,
@@ -214,6 +218,7 @@ export class Sandbox {
     });
 
     return {
+      isPint: isPint,
       sandboxId: this.id,
       sessionId: customSession?.id,
       hostToken: customSession?.hostToken,

src/types.ts

@@ -1,6 +1,6 @@
 import { VMTier } from "./VMTier";
 import { HostToken } from "./HostTokens";
-import { Config } from "@hey-api/client-fetch";
+import { Config } from "./api-clients/client/client";
 import { Tracer } from "@opentelemetry/api";
 
 export interface PitcherManagerResponse {
@@ -235,6 +235,7 @@ export type SandboxOpts = {
 };
 
 export type SandboxSession = PitcherManagerResponse & {
+  isPint: boolean;
   sandboxId: string;
   sessionId?: string;
   hostToken?: HostToken;

src/utils/event.ts

@@ -130,3 +130,53 @@ export class AsyncEmitter<T> implements IDisposable {
     this.registeredListeners = new Set();
   }
 }
+
+/**
+ * EmitterSubscription provides an abstraction that manages a subscription lifecycle
+ * tied to the number of listeners on an emitter. The subscription is created when
+ * the first listener is added and disposed when the last listener is removed.
+ */
+export class EmitterSubscription<T> implements IDisposable {
+  private emitter = new Emitter<T>();
+  private subscription: IDisposable | undefined;
+  private listenerCount = 0;
+
+  constructor(
+    private createSubscription: (fire: (value: T) => void) => IDisposable
+  ) {}
+
+  get event(): Event<T> {
+    return (listener: (e: T) => void) => {
+      // Add listener to emitter
+      const listenerDisposable = this.emitter.event(listener);
+
+      // Create subscription if this is the first listener
+      if (this.listenerCount === 0) {
+        this.subscription = this.createSubscription((value) =>
+          this.emitter.fire(value)
+        );
+      }
+
+      this.listenerCount++;
+
+      // Return disposable that removes listener and cleans up subscription if needed
+      return Disposable.create(() => {
+        listenerDisposable.dispose();
+        this.listenerCount--;
+
+        // Dispose subscription when last listener is removed
+        if (this.listenerCount === 0 && this.subscription) {
+          this.subscription.dispose();
+          this.subscription = undefined;
+        }
+      });
+    };
+  }
+
+  dispose(): void {
+    this.subscription?.dispose();
+    this.subscription = undefined;
+    this.emitter.dispose();
+    this.listenerCount = 0;
+  }
+}