@zocomputer/agent-sdk 0.5.0

package/platform/agent-sandbox/ssh-connection.ts

@@ -0,0 +1,165 @@
+/** Scoped SSH credentials the control plane mints for one sandbox. */
+export interface SshSandboxAccess {
+  readonly sandboxId: string;
+  readonly sshHost: string;
+  /** Username = the scoped access token (Daytona's SSH gateway model). */
+  readonly sshUser: string;
+  /** ISO timestamp when the token expires; we re-mint before then. */
+  readonly expiresAt: string;
+}
+
+// The SSH connection lifecycle as an explicit, testable state machine, lifted
+// out of the session surface so its transitions live in one auditable place
+// (the runtime's recurring bugs all came from ad-hoc interleavings here).
+//
+// States: idle (no connection) → connecting (a connect in flight) → connected
+// (live) → back to idle on a socket drop, or disposed (terminal). The scoped
+// SSH token (`heldAccess`) is retained across a socket drop so a pure reconnect
+// reuses it instead of re-minting; it's dropped only when expired or on dispose.
+//
+// `connect` is injected so the manager is unit-tested without a live SSH socket;
+// the runtime passes a real ssh2-backed connector.
+
+/**
+ * Minimal live-connection handle the manager drives. Generic over the carried
+ * client `C` (the runtime carries an ssh2 `Client`, which `run` execs against;
+ * tests carry a stub) — the manager only ever uses `end`/`onClose`.
+ */
+export interface ManagedConn<C = unknown> {
+  /** The underlying client (ssh2 `Client` in the runtime); used by callers, not the manager. */
+  readonly client: C;
+  /** Close the connection (idempotent). */
+  end(): void;
+  /** Register a one-shot callback for when the connection drops. */
+  onClose(cb: () => void): void;
+}
+
+/** Opens a connection for the given scoped access, or rejects. */
+export type Connector<C> = (access: SshSandboxAccess) => Promise<ManagedConn<C>>;
+
+/** Is the access token expired (or unparsable → treat as expired, fail-safe)? */
+function isExpired(access: SshSandboxAccess, skewMs: number, now: number): boolean {
+  const expiry = Date.parse(access.expiresAt);
+  if (Number.isNaN(expiry)) return true;
+  return expiry - skewMs <= now;
+}
+
+export interface SshConnectionManagerOptions<C> {
+  /** Mint fresh scoped access (the runtime's control-plane call). */
+  readonly acquireAccess: () => Promise<SshSandboxAccess>;
+  /** Open a connection for an access (injectable; real impl wraps ssh2). */
+  readonly connect: Connector<C>;
+  /** Re-mint this many ms before the token's actual expiry. */
+  readonly expirySkewMs: number;
+  /** Clock, injectable for tests. Defaults to `Date.now`. */
+  readonly now?: () => number;
+}
+
+/**
+ * Owns one sandbox's SSH connection across its lifetime. `ensure()` returns a
+ * live connection, (re)connecting/(re)minting as the rules below require;
+ * `dispose()` shuts it down terminally.
+ */
+export class SshConnectionManager<C> {
+  private conn: ManagedConn<C> | null = null;
+  private access: SshSandboxAccess | null = null;
+  private connecting: Promise<ManagedConn<C>> | null = null;
+  private disposed = false;
+  private sandboxId: string | null = null;
+  private readonly opts: SshConnectionManagerOptions<C>;
+  private readonly now: () => number;
+
+  constructor(opts: SshConnectionManagerOptions<C>) {
+    this.opts = opts;
+    this.now = opts.now ?? Date.now;
+  }
+
+  /** The provider sandbox id once a connection has been established, else null. */
+  currentSandboxId(): string | null {
+    return this.sandboxId;
+  }
+
+  private expired(access: SshSandboxAccess): boolean {
+    return isExpired(access, this.opts.expirySkewMs, this.now());
+  }
+
+  /**
+   * Get a live connection. Reuses a live, non-expired one; coalesces concurrent
+   * callers onto one in-flight connect; reconnects with the held (still-valid)
+   * token on a pure socket drop, and only re-mints when the token is expired/
+   * absent. A connect that fails on a reused token (gateway may have revoked it
+   * on a sandbox restart) re-mints once and retries.
+   */
+  async ensure(): Promise<ManagedConn<C>> {
+    if (this.disposed) throw new Error("zo sandbox: session is disposed");
+    if (this.conn !== null && this.access !== null && !this.expired(this.access)) {
+      return this.conn;
+    }
+    if (this.connecting !== null) return await this.connecting;
+
+    const held = this.access;
+    this.conn?.end();
+    this.conn = null;
+    this.connecting = this.openConnection(held);
+    try {
+      return await this.connecting;
+    } finally {
+      this.connecting = null;
+    }
+  }
+
+  private async openConnection(held: SshSandboxAccess | null): Promise<ManagedConn<C>> {
+    // dispose() can land at any await below; bail at each step so a disposed
+    // session never spends a control-plane mint or opens a doomed connection.
+    const abortIfDisposed = (): void => {
+      if (this.disposed) throw new Error("zo sandbox: session disposed during connect");
+    };
+
+    // Reuse the held token if still valid (socket just dropped); else mint fresh.
+    let used: SshSandboxAccess;
+    if (held !== null && !this.expired(held)) {
+      used = held;
+    } else {
+      abortIfDisposed(); // don't mint for a disposed session
+      used = await this.opts.acquireAccess();
+    }
+    abortIfDisposed(); // don't open a connection for a disposed session
+    let c: ManagedConn<C>;
+    try {
+      c = await this.opts.connect(used);
+    } catch (error) {
+      // A reused token may have been revoked server-side (sandbox restart) —
+      // re-mint once and retry. A fresh-token failure is a real error.
+      if (used === held) {
+        abortIfDisposed();
+        used = await this.opts.acquireAccess();
+        abortIfDisposed();
+        c = await this.opts.connect(used);
+      } else {
+        throw error;
+      }
+    }
+    // dispose() during the connect → this connection is orphaned; close it.
+    if (this.disposed) {
+      c.end();
+      throw new Error("zo sandbox: session disposed during connect");
+    }
+    // On a socket drop, fall back to idle — but KEEP the token so the next
+    // ensure() reconnects with it rather than re-minting needlessly.
+    c.onClose(() => {
+      if (this.conn === c) this.conn = null;
+    });
+    this.conn = c;
+    this.access = used;
+    this.sandboxId = used.sandboxId;
+    return c;
+  }
+
+  /** Terminal shutdown: close the connection and refuse further ensure() calls. */
+  dispose(): void {
+    this.disposed = true;
+    this.conn?.end();
+    this.conn = null;
+    this.access = null;
+  }
+}

package/platform/agent-sandbox/ssh-exec.ts

@@ -0,0 +1,98 @@
+import type { ClientChannel } from "ssh2";
+
+// The one place ssh2 exec-channel completion is reconciled, so run/spawn/sftp
+// can't drift on it. Two concerns live here, both correct-by-construction:
+//   1. exit reconciliation — prefer the `exit` event's (code, signal) over what
+//      `close` carried; a signal → SIGNAL_EXIT_CODE, never a clean 0.
+//   2. abort precedence — if the abort signal is (or becomes) aborted, the
+//      command rejects with the abort reason, EVEN IF the channel also exited in
+//      the same tick. Abort structurally wins because the single settle() point
+//      checks `aborted` before it ever resolves success.
+
+/** Exit code we report when a command is killed by a signal (no numeric code). */
+export const SIGNAL_EXIT_CODE = 137; // 128 + SIGKILL(9), the conventional shell value
+
+/** The reconciled outcome of an exec channel that exited (not aborted/errored). */
+export interface ExecExit {
+  /** Reconciled exit code; a signal termination maps to SIGNAL_EXIT_CODE. */
+  readonly exitCode: number;
+  /** The terminating signal name, if the channel reported one (else null). */
+  readonly signal: string | null;
+}
+
+/** The abort reason as an Error (falls back to a generic one). */
+function abortError(signal: AbortSignal): Error {
+  return signal.reason instanceof Error
+    ? signal.reason
+    : new Error(typeof signal.reason === "string" ? signal.reason : "aborted");
+}
+
+/**
+ * Await an exec channel's completion, with abort taking precedence over a
+ * same-tick exit.
+ *
+ * - Resolves `ExecExit` when the channel exits/closes (and not aborted).
+ * - Rejects with the abort reason if `abortSignal` is or becomes aborted — and
+ *   closes the channel so the remote process is torn down. Wins a same-tick race
+ *   with exit because `settle()` re-checks `aborted` before resolving success.
+ * - Rejects with the channel error on a stream `error`.
+ *
+ * Resolves only on exit/close/error/abort; attach `data`/`stderr` separately.
+ */
+export function awaitCommand(
+  stream: ClientChannel,
+  abortSignal?: AbortSignal,
+): Promise<ExecExit> {
+  return new Promise<ExecExit>((resolve, reject) => {
+    let code: number | null = null;
+    let signal: string | null = null;
+    let sawExit = false;
+    let settled = false;
+
+    const onAbort = (): void => {
+      // close() tears down the remote process; settle() turns this into a reject.
+      try {
+        stream.close();
+      } catch {
+        // already closed
+      }
+      settle();
+    };
+
+    // The single settle point. Abort always wins: if the signal is aborted we
+    // reject with its reason regardless of whether the channel also exited.
+    const settle = (channelError?: Error): void => {
+      if (settled) return;
+      settled = true;
+      abortSignal?.removeEventListener("abort", onAbort);
+      if (abortSignal?.aborted) { reject(abortError(abortSignal)); return; }
+      if (channelError !== undefined) { reject(channelError); return; }
+      const sig = signal;
+      const reconciled = sig != null ? SIGNAL_EXIT_CODE : (code ?? 0);
+      resolve({ exitCode: reconciled, signal: sig });
+    };
+
+    // Already aborted before we wired up — close + reject without hanging.
+    if (abortSignal?.aborted) {
+      onAbort();
+      return;
+    }
+    abortSignal?.addEventListener("abort", onAbort, { once: true });
+
+    stream
+      .on("exit", (c: number | null, s?: string | null) => {
+        sawExit = true;
+        code = c;
+        signal = s ?? null;
+      })
+      .on("close", (closeCode?: number | null, closeSignal?: string | null) => {
+        // Prefer the authoritative `exit` values; fall back to what close carried.
+        if (!sawExit) {
+          code = closeCode ?? null;
+          signal = closeSignal ?? null;
+        }
+        settle();
+      })
+      .on("error", (e: Error) => settle(e));
+  });
+}