v0.2.0: Project sync into sandboxes, background processes, and opencode auth login

.opencode/plugin/tensorlake/core/client.ts

@@ -1,16 +1,10 @@
-import { SandboxClient } from 'tensorlake'
-import type { Sandbox } from 'tensorlake'
+import { RemoteAPIError, Sandbox, SandboxConnectionError, SandboxNotFoundError } from 'tensorlake'
+import type { FileSystemMount } from 'tensorlake'
 import { execFileSync } from 'child_process'
 import { logger } from './logger.js'
 
 const MANAGEMENT_API = process.env.TENSORLAKE_API_URL ?? 'https://api.tensorlake.ai'
 
-function getSandboxProxyUrl(sandboxId: string): string {
-  const custom = process.env.TENSORLAKE_SANDBOX_PROXY_URL
-  if (custom) return custom
-  return `https://${sandboxId}.sandbox.tensorlake.ai`
-}
-
 export type SandboxInfo = {
   sandbox_id: string
   status: string
@@ -33,69 +27,168 @@ export type DirectoryEntry = {
   size?: number
 }
 
-export class TensorLakeClient {
-  private readonly sdk: SandboxClient
+export type ProcessStatusInfo = {
+  pid: number
+  status: string
+  exitCode?: number
+  signal?: number
+  command: string
+}
+
+export class TensorlakeClient {
+  // Connected handles keyed by sandboxId, so repeated operations reuse the
+  // resolved proxy routing instead of re-resolving on every call.
+  private readonly handles = new Map<string, Sandbox>()
+
+  // Credentials are resolved lazily on every use so a key added via
+  // `opencode auth login` after startup is picked up without a restart.
+  constructor(private readonly resolveKey: () => string | undefined) {}
+
+  hasApiKey(): boolean {
+    return (this.resolveKey() ?? '').length > 0
+  }
+
+  getApiKey(): string {
+    return this.resolveKey() ?? ''
+  }
 
-  constructor(private readonly apiKey: string) {
-    this.sdk = SandboxClient.forCloud({
-      apiKey,
+  private clientOptions() {
+    return {
+      apiKey: this.getApiKey(),
       apiUrl: MANAGEMENT_API,
-      organizationId: process.env.TENSORLAKE_ORGANIZATION_ID,
-      projectId: process.env.TENSORLAKE_PROJECT_ID,
+      ...(process.env.TENSORLAKE_ORGANIZATION_ID
+        ? { organizationId: process.env.TENSORLAKE_ORGANIZATION_ID }
+        : {}),
+      ...(process.env.TENSORLAKE_PROJECT_ID ? { projectId: process.env.TENSORLAKE_PROJECT_ID } : {}),
+    }
+  }
+
+  private async connectSandbox(sandboxId: string): Promise<Sandbox> {
+    const cached = this.handles.get(sandboxId)
+    if (cached) return cached
+    const proxyUrl = process.env.TENSORLAKE_SANDBOX_PROXY_URL
+    const sandbox = await Sandbox.connect({
+      sandboxId,
+      ...(proxyUrl ? { proxyUrl } : {}),
+      ...this.clientOptions(),
     })
+    this.handles.set(sandboxId, sandbox)
+    return sandbox
   }
 
-  hasApiKey(): boolean {
-    return this.apiKey.length > 0
+  // Errors suggesting the cached handle's proxy routing is stale — e.g. the
+  // sandbox was suspended and resumed outside this process, moving it to a
+  // different host. Only resume() on this handle refreshes routing; info()
+  // never does, so the only self-heal is dropping the handle and reconnecting.
+  private isStaleHandleError(err: unknown): boolean {
+    return (
+      err instanceof SandboxConnectionError ||
+      err instanceof SandboxNotFoundError ||
+      (err instanceof RemoteAPIError && [404, 502, 503].includes(err.statusCode))
+    )
+  }
+
+  // Runs a proxy-routed operation. On a stale-handle error the handle is
+  // dropped so the next call reconnects with fresh routing; idempotent ops
+  // (retry: true) additionally reconnect and retry once themselves. Command
+  // execution uses retry: false to avoid any chance of running a command twice.
+  private async withSandbox<T>(
+    sandboxId: string,
+    op: (sandbox: Sandbox) => Promise<T>,
+    opts: { retry: boolean } = { retry: true },
+  ): Promise<T> {
+    const sandbox = await this.connectSandbox(sandboxId)
+    try {
+      return await op(sandbox)
+    } catch (err: unknown) {
+      if (!this.isStaleHandleError(err)) throw err
+      this.dropHandle(sandboxId)
+      if (!opts.retry) throw err
+      logger.warn(`Sandbox ${sandboxId} call failed (${(err as Error)?.message ?? err}); reconnecting and retrying once`)
+      return op(await this.connectSandbox(sandboxId))
+    }
   }
 
-  async createSandbox(opts: { image?: string; name?: string; timeoutSecs?: number } = {}): Promise<CreateSandboxResponse> {
+  private dropHandle(sandboxId: string): void {
+    const handle = this.handles.get(sandboxId)
+    if (!handle) return
+    this.handles.delete(sandboxId)
+    try {
+      handle.close()
+    } catch {
+      // closing a stale handle is best-effort
+    }
+  }
+
+  async createSandbox(opts: { image?: string; name?: string; timeoutSecs?: number; fileSystems?: FileSystemMount[] } = {}): Promise<CreateSandboxResponse> {
     const cpus = parseFloat(process.env.TENSORLAKE_CPUS ?? '2')
     const memoryMb = parseInt(process.env.TENSORLAKE_MEMORY_MB ?? '4096', 10)
     const ephemeralDiskMb = parseInt(process.env.TENSORLAKE_DISK_MB ?? '10240', 10)
     logger.info(`Creating sandbox name=${opts.name ?? '(ephemeral)'} image=${opts.image ?? '(default)'} cpus=${cpus} memoryMb=${memoryMb} diskMb=${ephemeralDiskMb}`)
-    const sandbox = await this.sdk.createAndConnect({
+    const sandbox = await Sandbox.create({
       ...(opts.image ? { image: opts.image } : {}),
       cpus,
       memoryMb,
       diskMb: ephemeralDiskMb,
       ...(opts.name ? { name: opts.name } : {}),
       ...(opts.timeoutSecs ? { timeoutSecs: opts.timeoutSecs } : {}),
+      ...(opts.fileSystems?.length ? { fileSystems: opts.fileSystems } : {}),
+      ...this.clientOptions(),
     })
+    this.handles.set(sandbox.sandboxId, sandbox)
     return { sandbox_id: sandbox.sandboxId, status: 'running' }
   }
 
+  async listSandboxFileSystems(sandboxId: string): Promise<FileSystemMount[]> {
+    const sandbox = await this.connectSandbox(sandboxId)
+    const info = await sandbox.info()
+    return info.fileSystems ?? []
+  }
+
+  async attachFileSystem(sandboxId: string, fileSystemId: string, mountPath: string): Promise<void> {
+    const sandbox = await this.connectSandbox(sandboxId)
+    await sandbox.attachFileSystem(fileSystemId, mountPath)
+  }
+
   async getSandbox(sandboxId: string): Promise<SandboxInfo> {
-    const info = await this.sdk.get(sandboxId)
+    const sandbox = await this.connectSandbox(sandboxId)
+    const info = await sandbox.info()
     return { sandbox_id: info.sandboxId, status: info.status as unknown as string }
   }
 
   async deleteSandbox(sandboxId: string): Promise<void> {
     try {
-      await this.sdk.delete(sandboxId)
+      const sandbox = await this.connectSandbox(sandboxId)
+      await sandbox.terminate()
     } catch (err: unknown) {
-      if (String((err as Error)?.message ?? err).includes('404')) return
+      // Already deleted — treat as success.
+      if (err instanceof SandboxNotFoundError) return
+      if (err instanceof RemoteAPIError && err.statusCode === 404) return
       throw err
+    } finally {
+      this.dropHandle(sandboxId)
     }
   }
 
   async suspendSandbox(sandboxId: string): Promise<void> {
-    await this.sdk.suspend(sandboxId)
+    const sandbox = await this.connectSandbox(sandboxId)
+    await sandbox.suspend()
   }
 
   suspendSandboxSync(sandboxId: string): void {
+    const apiKey = this.getApiKey()
     execFileSync('curl', [
       '-s', '-X', 'POST',
       `${MANAGEMENT_API}/sandboxes/${sandboxId}/suspend`,
-      '-H', `Authorization: Bearer ${this.apiKey}`,
+      '-H', `Authorization: Bearer ${apiKey}`,
     ], { timeout: 10_000 })
     const deadline = Date.now() + 30_000
     while (Date.now() < deadline) {
       try {
         const out = execFileSync('curl', [
           '-s',
           `${MANAGEMENT_API}/sandboxes/${sandboxId}`,
-          '-H', `Authorization: Bearer ${this.apiKey}`,
+          '-H', `Authorization: Bearer ${apiKey}`,
         ], { timeout: 10_000 }).toString()
         const status = JSON.parse(out)?.status
         if (status === 'suspended' || status === 'terminated') return
@@ -107,7 +200,9 @@ export class TensorLakeClient {
   }
 
   async resumeSandbox(sandboxId: string): Promise<void> {
-    await this.sdk.resume(sandboxId)
+    // resume() waits for running and refreshes the handle's proxy routing.
+    const sandbox = await this.connectSandbox(sandboxId)
+    await sandbox.resume()
   }
 
   async waitForSuspended(sandboxId: string, timeoutMs = 30_000): Promise<void> {
@@ -130,43 +225,75 @@ export class TensorLakeClient {
     throw new Error(`Sandbox ${sandboxId} did not become running within ${timeoutMs}ms`)
   }
 
-  private connectSandbox(sandboxId: string): Sandbox {
-    return this.sdk.connect(sandboxId, getSandboxProxyUrl(sandboxId))
-  }
-
   async executeCommand(
     sandboxId: string,
     command: string,
     workingDir = '/tmp/workspace',
     timeoutMs = 120_000,
   ): Promise<ProcessResult> {
-    const sandbox = this.connectSandbox(sandboxId)
-    const result = await sandbox.run('sh', {
-      args: ['-c', command],
-      workingDir,
-      timeout: timeoutMs / 1000,
-    })
+    const result = await this.withSandbox(
+      sandboxId,
+      (sandbox) =>
+        sandbox.run('sh', {
+          args: ['-c', command],
+          workingDir,
+          timeout: timeoutMs / 1000,
+        }),
+      { retry: false },
+    )
     return {
       exitCode: result.exitCode ?? -1,
       stdout: result.stdout ?? '',
       stderr: result.stderr ?? '',
     }
   }
 
+  // Processes are started unnamed (non-managed) on purpose: the daemon keeps
+  // tracking them after exit or kill, so status and output stay queryable by PID.
+  async startBackgroundProcess(sandboxId: string, command: string, workingDir: string): Promise<number> {
+    const info = await this.withSandbox(
+      sandboxId,
+      (sandbox) =>
+        sandbox.startProcess('sh', {
+          args: ['-c', command],
+          workingDir,
+        }),
+      { retry: false },
+    )
+    return info.pid
+  }
+
+  async getProcessStatus(sandboxId: string, pid: number): Promise<ProcessStatusInfo> {
+    const info = await this.withSandbox(sandboxId, (sandbox) => sandbox.getProcess(pid))
+    return {
+      pid: info.pid,
+      status: info.status as unknown as string,
+      exitCode: info.exitCode,
+      signal: info.signal,
+      command: [info.command, ...(info.args ?? [])].join(' '),
+    }
+  }
+
+  async getProcessOutput(sandboxId: string, pid: number): Promise<string[]> {
+    const output = await this.withSandbox(sandboxId, (sandbox) => sandbox.getOutput(pid))
+    return output.lines
+  }
+
+  async killProcess(sandboxId: string, pid: number): Promise<void> {
+    await this.withSandbox(sandboxId, (sandbox) => sandbox.killProcess(pid))
+  }
+
   async readFile(sandboxId: string, path: string): Promise<Buffer> {
-    const sandbox = this.connectSandbox(sandboxId)
-    const data = await sandbox.readFile(path)
+    const data = await this.withSandbox(sandboxId, (sandbox) => sandbox.readFile(path))
     return Buffer.from(data)
   }
 
   async writeFile(sandboxId: string, path: string, content: Buffer): Promise<void> {
-    const sandbox = this.connectSandbox(sandboxId)
-    await sandbox.writeFile(path, content)
+    await this.withSandbox(sandboxId, (sandbox) => sandbox.writeFile(path, content))
   }
 
   async listDirectory(sandboxId: string, path: string): Promise<DirectoryEntry[]> {
-    const sandbox = this.connectSandbox(sandboxId)
-    const response = await sandbox.listDirectory(path)
+    const response = await this.withSandbox(sandboxId, (sandbox) => sandbox.listDirectory(path))
     return response.entries.map((e) => ({
       name: e.name,
       is_dir: e.isDir,