langsmith 0.5.9 → 0.5.11

package/dist/experimental/sandbox/ws_execute.cjs

@@ -0,0 +1,350 @@
+"use strict";
+/**
+ * WebSocket-based command execution for long-running commands.
+ *
+ * Uses the `ws` npm package (optional peer dependency).
+ * Install with: npm install ws
+ */
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.WSStreamControl = void 0;
+exports.buildWsUrl = buildWsUrl;
+exports.buildAuthHeaders = buildAuthHeaders;
+exports.raiseForWsError = raiseForWsError;
+exports.runWsStream = runWsStream;
+exports.reconnectWsStream = reconnectWsStream;
+const errors_js_1 = require("./errors.cjs");
+async function ensureWs() {
+    try {
+        const ws = await import("ws");
+        return { WebSocket: ws.default || ws.WebSocket || ws };
+    }
+    catch {
+        throw new Error("WebSocket-based execution requires the 'ws' package. " +
+            "Install it with: npm install ws");
+    }
+}
+// =============================================================================
+// URL and Auth Helpers
+// =============================================================================
+/**
+ * Convert a dataplane HTTP URL to a WebSocket URL for /execute/ws.
+ */
+function buildWsUrl(dataplaneUrl) {
+    const wsUrl = dataplaneUrl
+        .replace("https://", "wss://")
+        .replace("http://", "ws://");
+    return `${wsUrl}/execute/ws`;
+}
+/**
+ * Build auth headers for the WebSocket upgrade request.
+ */
+function buildAuthHeaders(apiKey) {
+    if (apiKey) {
+        return { "X-Api-Key": apiKey };
+    }
+    return {};
+}
+// =============================================================================
+// Stream Control
+// =============================================================================
+/**
+ * Control interface for an active WebSocket stream.
+ *
+ * Created before the async generator starts, bound to the WebSocket once
+ * the connection opens. The CommandHandle holds a reference to this
+ * object to send kill/input messages.
+ */
+class WSStreamControl {
+    constructor() {
+        Object.defineProperty(this, "_ws", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: null
+        });
+        Object.defineProperty(this, "_closed", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: false
+        });
+        Object.defineProperty(this, "_killed", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: false
+        });
+    }
+    /** Bind to the active WebSocket. Called inside the generator. */
+    _bind(ws) {
+        this._ws = ws;
+    }
+    /** Mark as closed. Called when the generator exits. */
+    _unbind() {
+        this._closed = true;
+        this._ws = null;
+    }
+    /** True if kill() has been called on this stream. */
+    get killed() {
+        return this._killed;
+    }
+    /** Send a kill message to abort the running command. */
+    sendKill() {
+        this._killed = true;
+        if (this._ws && !this._closed && this._ws.readyState === 1) {
+            this._ws.send(JSON.stringify({ type: "kill" }));
+        }
+    }
+    /** Send stdin data to the running command. */
+    sendInput(data) {
+        if (this._ws && !this._closed && this._ws.readyState === 1) {
+            this._ws.send(JSON.stringify({ type: "input", data }));
+        }
+    }
+}
+exports.WSStreamControl = WSStreamControl;
+// =============================================================================
+// Error Handling
+// =============================================================================
+/**
+ * Raise the appropriate exception from a server error message.
+ */
+function raiseForWsError(msg, commandId = "") {
+    const errorType = msg.error_type ?? "CommandError";
+    const errorMsg = msg.error ?? "Unknown error";
+    if (errorType === "CommandTimeout") {
+        throw new errors_js_1.LangSmithCommandTimeoutError(errorMsg);
+    }
+    if (errorType === "CommandNotFound") {
+        throw new errors_js_1.LangSmithSandboxOperationError(commandId ? `Command not found: ${commandId}` : errorMsg, commandId ? "reconnect" : "command", errorType);
+    }
+    if (errorType === "SessionExpired") {
+        throw new errors_js_1.LangSmithSandboxOperationError(commandId ? `Session expired: ${commandId}` : errorMsg, commandId ? "reconnect" : "command", errorType);
+    }
+    throw new errors_js_1.LangSmithSandboxOperationError(errorMsg, commandId ? "reconnect" : "command", errorType);
+}
+// =============================================================================
+// WebSocket Stream Helpers
+// =============================================================================
+/**
+ * Create a ws WebSocket connection and return a promise that resolves when open
+ * or rejects on error.
+ */
+async function connectWs(url, headers) {
+    const { WebSocket: WS } = await ensureWs();
+    return new Promise((resolve, reject) => {
+        const ws = new WS(url, { headers });
+        ws.on("open", () => {
+            ws.removeAllListeners("error");
+            resolve(ws);
+        });
+        ws.on("error", (err) => {
+            ws.removeAllListeners("open");
+            reject(new errors_js_1.LangSmithSandboxConnectionError(`Failed to connect to sandbox WebSocket: ${err.message}`));
+        });
+    });
+}
+/**
+ * Read messages from a ws WebSocket as an async iterable.
+ *
+ * Yields parsed WsMessage objects. Handles close events and errors,
+ * mapping them to appropriate exceptions.
+ */
+async function* readWsMessages(ws) {
+    // Buffer incoming messages so the consumer can process them at its own pace
+    const messageQueue = [];
+    let resolve = null;
+    let error = null;
+    let done = false;
+    const onMessage = (data) => {
+        const raw = typeof data === "string" ? data : data.toString();
+        const msg = JSON.parse(raw);
+        messageQueue.push(msg);
+        if (resolve) {
+            const r = resolve;
+            resolve = null;
+            r();
+        }
+    };
+    const onClose = (code, reason) => {
+        done = true;
+        if (code === 1001) {
+            error = new errors_js_1.LangSmithSandboxServerReloadError("Server is reloading, reconnect to resume");
+        }
+        else if (code !== 1000) {
+            error = new errors_js_1.LangSmithSandboxConnectionError(`WebSocket connection closed unexpectedly (code: ${code}, reason: ${reason.toString()})`);
+        }
+        if (resolve) {
+            const r = resolve;
+            resolve = null;
+            r();
+        }
+    };
+    const onError = (err) => {
+        done = true;
+        if (!error) {
+            error = new errors_js_1.LangSmithSandboxConnectionError(`WebSocket connection error: ${err.message}`);
+        }
+        if (resolve) {
+            const r = resolve;
+            resolve = null;
+            r();
+        }
+    };
+    ws.on("message", onMessage);
+    ws.on("close", onClose);
+    ws.on("error", onError);
+    try {
+        while (true) {
+            // Drain buffered messages first
+            while (messageQueue.length > 0) {
+                yield messageQueue.shift();
+            }
+            // If done and queue is empty, we're finished
+            if (done) {
+                if (error) {
+                    throw error;
+                }
+                return;
+            }
+            // Wait for next message or close/error
+            await new Promise((r) => {
+                resolve = r;
+            });
+        }
+    }
+    finally {
+        ws.removeListener("message", onMessage);
+        ws.removeListener("close", onClose);
+        ws.removeListener("error", onError);
+    }
+}
+// =============================================================================
+// Async Stream Functions
+// =============================================================================
+/**
+ * Execute a command over WebSocket, yielding raw message dicts.
+ *
+ * Returns a tuple of [async_message_iterator, control]. The control object
+ * provides sendKill() and sendInput() methods for the CommandHandle.
+ *
+ * The iterator yields WsMessage objects with a "type" field:
+ * - { type: "started", command_id: "...", pid: N }
+ * - { type: "stdout", data: "...", offset: N }
+ * - { type: "stderr", data: "...", offset: N }
+ * - { type: "exit", exit_code: N }
+ *
+ * If onStdout/onStderr callbacks are provided, they are invoked as
+ * data arrives in addition to yielding the messages.
+ */
+async function runWsStream(dataplaneUrl, apiKey, command, options = {}) {
+    const { timeout = 60, env, cwd, shell = "/bin/bash", onStdout, onStderr, commandId, idleTimeout = 300, killOnDisconnect = false, ttlSeconds = 600, pty, } = options;
+    const wsUrl = buildWsUrl(dataplaneUrl);
+    const headers = buildAuthHeaders(apiKey);
+    const control = new WSStreamControl();
+    async function* stream() {
+        let ws;
+        try {
+            ws = await connectWs(wsUrl, headers);
+            control._bind(ws);
+            // Send execute request
+            const payload = {
+                type: "execute",
+                command,
+                timeout_seconds: timeout,
+                shell,
+                idle_timeout_seconds: idleTimeout,
+                kill_on_disconnect: killOnDisconnect,
+                ttl_seconds: ttlSeconds,
+            };
+            if (env)
+                payload.env = env;
+            if (cwd)
+                payload.cwd = cwd;
+            if (commandId)
+                payload.command_id = commandId;
+            if (pty)
+                payload.pty = true;
+            ws.send(JSON.stringify(payload));
+            // Read messages until exit or error
+            for await (const msg of readWsMessages(ws)) {
+                const msgType = msg.type;
+                if (msgType === "started") {
+                    yield msg;
+                }
+                else if (msgType === "stdout") {
+                    if (onStdout)
+                        onStdout(msg.data);
+                    yield msg;
+                }
+                else if (msgType === "stderr") {
+                    if (onStderr)
+                        onStderr(msg.data);
+                    yield msg;
+                }
+                else if (msgType === "exit") {
+                    yield msg;
+                    return;
+                }
+                else if (msgType === "error") {
+                    raiseForWsError(msg);
+                }
+            }
+        }
+        finally {
+            control._unbind();
+            if (ws && ws.readyState === 1) {
+                ws.close();
+            }
+        }
+    }
+    return [stream(), control];
+}
+/**
+ * Reconnect to an existing command over WebSocket.
+ *
+ * Returns a tuple of [async_message_iterator, control], same as runWsStream.
+ * The iterator yields stdout, stderr, exit, and error messages.
+ * No 'started' message is sent on reconnection.
+ */
+async function reconnectWsStream(dataplaneUrl, apiKey, commandId, options = {}) {
+    const { stdoutOffset = 0, stderrOffset = 0 } = options;
+    const wsUrl = buildWsUrl(dataplaneUrl);
+    const headers = buildAuthHeaders(apiKey);
+    const control = new WSStreamControl();
+    async function* stream() {
+        let ws;
+        try {
+            ws = await connectWs(wsUrl, headers);
+            control._bind(ws);
+            // Send reconnect request
+            ws.send(JSON.stringify({
+                type: "reconnect",
+                command_id: commandId,
+                stdout_offset: stdoutOffset,
+                stderr_offset: stderrOffset,
+            }));
+            // Read messages until exit or error
+            for await (const msg of readWsMessages(ws)) {
+                const msgType = msg.type;
+                if (msgType === "stdout" || msgType === "stderr") {
+                    yield msg;
+                }
+                else if (msgType === "exit") {
+                    yield msg;
+                    return;
+                }
+                else if (msgType === "error") {
+                    raiseForWsError(msg, commandId);
+                }
+            }
+        }
+        finally {
+            control._unbind();
+            if (ws && ws.readyState === 1) {
+                ws.close();
+            }
+        }
+    }
+    return [stream(), control];
+}

package/dist/experimental/sandbox/ws_execute.d.ts

@@ -0,0 +1,70 @@
+/**
+ * WebSocket-based command execution for long-running commands.
+ *
+ * Uses the `ws` npm package (optional peer dependency).
+ * Install with: npm install ws
+ */
+import type { WsMessage, WsRunOptions } from "./types.js";
+type WsWebSocket = any;
+/**
+ * Convert a dataplane HTTP URL to a WebSocket URL for /execute/ws.
+ */
+export declare function buildWsUrl(dataplaneUrl: string): string;
+/**
+ * Build auth headers for the WebSocket upgrade request.
+ */
+export declare function buildAuthHeaders(apiKey: string | undefined): Record<string, string>;
+/**
+ * Control interface for an active WebSocket stream.
+ *
+ * Created before the async generator starts, bound to the WebSocket once
+ * the connection opens. The CommandHandle holds a reference to this
+ * object to send kill/input messages.
+ */
+export declare class WSStreamControl {
+    private _ws;
+    private _closed;
+    private _killed;
+    /** Bind to the active WebSocket. Called inside the generator. */
+    _bind(ws: WsWebSocket): void;
+    /** Mark as closed. Called when the generator exits. */
+    _unbind(): void;
+    /** True if kill() has been called on this stream. */
+    get killed(): boolean;
+    /** Send a kill message to abort the running command. */
+    sendKill(): void;
+    /** Send stdin data to the running command. */
+    sendInput(data: string): void;
+}
+/**
+ * Raise the appropriate exception from a server error message.
+ */
+export declare function raiseForWsError(msg: WsMessage, commandId?: string): never;
+/**
+ * Execute a command over WebSocket, yielding raw message dicts.
+ *
+ * Returns a tuple of [async_message_iterator, control]. The control object
+ * provides sendKill() and sendInput() methods for the CommandHandle.
+ *
+ * The iterator yields WsMessage objects with a "type" field:
+ * - { type: "started", command_id: "...", pid: N }
+ * - { type: "stdout", data: "...", offset: N }
+ * - { type: "stderr", data: "...", offset: N }
+ * - { type: "exit", exit_code: N }
+ *
+ * If onStdout/onStderr callbacks are provided, they are invoked as
+ * data arrives in addition to yielding the messages.
+ */
+export declare function runWsStream(dataplaneUrl: string, apiKey: string | undefined, command: string, options?: WsRunOptions): Promise<[AsyncIterableIterator<WsMessage>, WSStreamControl]>;
+/**
+ * Reconnect to an existing command over WebSocket.
+ *
+ * Returns a tuple of [async_message_iterator, control], same as runWsStream.
+ * The iterator yields stdout, stderr, exit, and error messages.
+ * No 'started' message is sent on reconnection.
+ */
+export declare function reconnectWsStream(dataplaneUrl: string, apiKey: string | undefined, commandId: string, options?: {
+    stdoutOffset?: number;
+    stderrOffset?: number;
+}): Promise<[AsyncIterableIterator<WsMessage>, WSStreamControl]>;
+export {};