@granular-software/sdk 0.2.0 → 0.2.1

package/dist/index.mjs

@@ -1,7 +1,11 @@
-import WebSocket from 'ws';
 import * as Automerge from '@automerge/automerge';
 
 // src/ws-client.ts
+var GlobalWebSocket = void 0;
+if (typeof globalThis !== "undefined" && globalThis.WebSocket) {
+  GlobalWebSocket = globalThis.WebSocket;
+}
+var READY_STATE_OPEN = 1;
 var WSClient = class {
   ws = null;
   url;
@@ -16,7 +20,9 @@ var WSClient = class {
   syncState = Automerge.initSyncState();
   reconnectTimer = null;
   isExplicitlyDisconnected = false;
+  options;
   constructor(options) {
+    this.options = options;
     this.url = options.url;
     this.sessionId = options.sessionId;
     this.token = options.token;
@@ -27,34 +33,68 @@ var WSClient = class {
    */
   async connect() {
     this.isExplicitlyDisconnected = false;
+    const WebSocketClass = this.options.WebSocketCtor || GlobalWebSocket;
+    if (!WebSocketClass) {
+      throw new Error('No WebSocket implementation found. If using Node.js, please install "ws" and pass the constructor to the SDK options: { WebSocketCtor: WebSocket }.');
+    }
     return new Promise((resolve, reject) => {
       try {
         const wsUrl = new URL(this.url);
         wsUrl.searchParams.set("sessionId", this.sessionId);
         wsUrl.searchParams.set("token", this.token);
-        this.ws = new WebSocket(wsUrl.toString());
-        this.ws.on("open", () => {
-          this.emit("open", {});
-          resolve();
-        });
-        this.ws.on("message", (data) => {
-          try {
-            const message = JSON.parse(data.toString());
-            this.handleMessage(message);
-          } catch (error) {
-            console.error("[Granular] Failed to parse message:", error);
-          }
-        });
-        this.ws.on("error", (error) => {
-          this.emit("error", error);
-          if (this.ws?.readyState !== WebSocket.OPEN) {
-            reject(error);
-          }
-        });
-        this.ws.on("close", () => {
-          this.emit("close", {});
-          this.handleDisconnect();
-        });
+        this.ws = new WebSocketClass(wsUrl.toString());
+        if (!this.ws) throw new Error("Failed to create WebSocket");
+        const socket = this.ws;
+        if (typeof socket.on === "function") {
+          socket.on("open", () => {
+            this.emit("open", {});
+            resolve();
+          });
+          socket.on("message", (data) => {
+            try {
+              const message = JSON.parse(data.toString());
+              this.handleMessage(message);
+            } catch (error) {
+              console.error("[Granular] Failed to parse message:", error);
+            }
+          });
+          socket.on("error", (error) => {
+            this.emit("error", error);
+            if (socket.readyState !== READY_STATE_OPEN) {
+              reject(error);
+            }
+          });
+          socket.on("close", () => {
+            this.emit("close", {});
+            this.handleDisconnect();
+          });
+        } else {
+          this.ws.onopen = () => {
+            this.emit("open", {});
+            resolve();
+          };
+          this.ws.onmessage = (event) => {
+            try {
+              const data = event.data;
+              const message = JSON.parse(data.toString());
+              this.handleMessage(message);
+            } catch (error) {
+              console.error("[Granular] Failed to parse message:", error);
+            }
+          };
+          this.ws.onerror = (event) => {
+            const error = new Error("WebSocket error");
+            error.event = event;
+            this.emit("error", error);
+            if (this.ws?.readyState !== READY_STATE_OPEN) {
+              reject(error);
+            }
+          };
+          this.ws.onclose = () => {
+            this.emit("close", {});
+            this.handleDisconnect();
+          };
+        }
       } catch (error) {
         reject(error);
       }
@@ -74,8 +114,8 @@ var WSClient = class {
     console.log("[Granular DEBUG] Received message:", JSON.stringify(message).slice(0, 500));
     if ("type" in message && message.type === "sync") {
       const syncMessage = message;
+      let bytes;
       try {
-        let bytes;
         const payload = syncMessage.message || syncMessage.data;
         if (typeof payload === "string") {
           const binaryString = atob(payload);
@@ -91,6 +131,7 @@ var WSClient = class {
         } else {
           return;
         }
+        console.log("[Granular DEBUG] Applying sync bytes:", bytes.length);
         const [newDoc, newSyncState] = Automerge.receiveSyncMessage(
           this.doc,
           this.syncState,
@@ -98,8 +139,49 @@ var WSClient = class {
         );
         this.doc = newDoc;
         this.syncState = newSyncState;
+        const docAny = this.doc;
+        if (docAny.catalog) {
+          console.log("[Granular DEBUG] Doc catalog sync applied. Keys in catalog:", Object.keys(docAny.catalog || {}));
+          console.log("[Granular DEBUG] RawToolCatalogs:", Object.keys(docAny.catalog.rawToolCatalogs || {}));
+        } else {
+          console.log("[Granular DEBUG] Doc synced but no catalog yet. Keys in doc:", Object.keys(docAny));
+        }
+        this.emit("sync", this.doc);
+      } catch (e) {
+        try {
+          console.log("[Granular DEBUG] receiveSyncMessage failed, trying applyChanges...");
+          const [newDoc] = Automerge.applyChanges(this.doc, [bytes]);
+          this.doc = newDoc;
+          this.emit("sync", this.doc);
+          console.log("[Granular DEBUG] applyChanges succeeded. Doc:", JSON.stringify(Automerge.toJS(this.doc)));
+        } catch (applyError) {
+          console.warn("[Granular] Failed to apply sync message (both sync & applyChanges)", e, applyError);
+        }
+      }
+      return;
+    }
+    if ("type" in message && message.type === "snapshot") {
+      const snapshotMessage = message;
+      try {
+        const bytes = new Uint8Array(snapshotMessage.data);
+        console.log("[Granular DEBUG] Loading snapshot bytes:", bytes.length);
+        this.doc = Automerge.load(bytes);
+        this.emit("sync", this.doc);
+        console.log("[Granular DEBUG] Snapshot loaded. Doc:", JSON.stringify(Automerge.toJS(this.doc)));
+      } catch (e) {
+        console.warn("[Granular] Failed to load snapshot message", e);
+      }
+      return;
+    }
+    if ("type" in message && message.type === "change") {
+      const changeMessage = message;
+      try {
+        const bytes = new Uint8Array(changeMessage.data);
+        const [newDoc] = Automerge.applyChanges(this.doc, [bytes]);
+        this.doc = newDoc;
         this.emit("sync", this.doc);
       } catch (e) {
+        console.warn("[Granular] Failed to apply change message", e);
       }
       return;
     }
@@ -137,7 +219,7 @@ var WSClient = class {
    * @throws {Error} If connection is closed or timeout occurs
    */
   async call(method, params) {
-    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
+    if (!this.ws || this.ws.readyState !== READY_STATE_OPEN) {
       throw new Error("WebSocket not connected");
     }
     const id = `rpc-${this.nextRpcId++}`;
@@ -266,6 +348,10 @@ var Session = class {
   /** Tracks which tools are instance methods (className set, not static) */
   instanceTools = /* @__PURE__ */ new Set();
   currentDomainRevision = null;
+  /** Local effect registry: name → full ToolWithHandler */
+  effects = /* @__PURE__ */ new Map();
+  /** Last known tools for diffing */
+  lastKnownTools = /* @__PURE__ */ new Map();
   constructor(client, clientId) {
     this.client = client;
     this.clientId = clientId || `client_${Date.now()}`;
@@ -381,6 +467,85 @@ var Session = class {
       rejected: result.rejected
     };
   }
+  /**
+   * Publish a single effect (tool) to the sandbox.
+   *
+   * Adds the effect to the local registry and re-publishes the
+   * full tool catalog to the server. If an effect with the same
+   * name already exists, it is replaced.
+   *
+   * @param effect - The effect (tool) to publish
+   * @returns PublishToolsResult with domainRevision
+   *
+   * @example
+   * ```typescript
+   * await env.publishEffect({
+   *   name: 'get_bio',
+   *   description: 'Get biography of an author',
+   *   className: 'author',
+   *   inputSchema: { type: 'object', properties: { detailed: { type: 'boolean' } } },
+   *   handler: async (id, params) => ({ bio: `Bio of ${id}` }),
+   * });
+   * ```
+   */
+  async publishEffect(effect) {
+    this.effects.set(effect.name, effect);
+    return this._syncEffects();
+  }
+  /**
+   * Publish multiple effects (tools) at once.
+   *
+   * Adds all effects to the local registry and re-publishes the
+   * full tool catalog in a single RPC call.
+   *
+   * @param effects - Array of effects to publish
+   * @returns PublishToolsResult with domainRevision
+   *
+   * @example
+   * ```typescript
+   * await env.publishEffects([
+   *   { name: 'get_bio', description: '...', inputSchema: {}, handler: async (id) => ({}) },
+   *   { name: 'search', description: '...', inputSchema: {}, handler: async (params) => ({}) },
+   * ]);
+   * ```
+   */
+  async publishEffects(effects) {
+    for (const effect of effects) {
+      this.effects.set(effect.name, effect);
+    }
+    return this._syncEffects();
+  }
+  /**
+   * Remove an effect by name and re-publish the remaining catalog.
+   *
+   * @param name - The name of the effect to remove
+   * @returns PublishToolsResult with domainRevision
+   */
+  async unpublishEffect(name) {
+    this.effects.delete(name);
+    this.toolHandlers.delete(name);
+    this.instanceTools.delete(name);
+    return this._syncEffects();
+  }
+  /**
+   * Remove all effects and publish an empty catalog.
+   *
+   * @returns PublishToolsResult with domainRevision
+   */
+  async unpublishAllEffects() {
+    this.effects.clear();
+    this.toolHandlers.clear();
+    this.instanceTools.clear();
+    return this._syncEffects();
+  }
+  /**
+   * Internal: re-publish the full effect catalog to the server.
+   * Called after any mutation to the effects Map.
+   */
+  async _syncEffects() {
+    const allEffects = Array.from(this.effects.values());
+    return this.publishTools(allEffects);
+  }
   /**
    * Submit a job to execute code in the sandbox.
    *
@@ -398,9 +563,9 @@ var Session = class {
    * execute locally and return the result to the sandbox.
    */
   async submitJob(code, domainRevision) {
-    const revision = domainRevision || this.currentDomainRevision;
+    const revision = domainRevision || this.currentDomainRevision || this.client.doc?.domain?.active || void 0;
     if (!revision) {
-      throw new Error("No domain revision available. Call publishTools() first.");
+      throw new Error("No domain revision available. Call publishTools() or ensure schema is activated.");
     }
     const result = await this.client.call("job.submit", {
       domainRevision: revision,
@@ -431,6 +596,72 @@ var Session = class {
   async answerPrompt(promptId, answer) {
     await this.client.call("prompt.answer", { promptId, value: answer });
   }
+  /**
+   * Get the current list of available tools.
+   * Consolidates tools from all connected clients.
+   */
+  getTools() {
+    const doc = this.client.doc;
+    const toolMap = /* @__PURE__ */ new Map();
+    const domainPkg = doc.domain?.packages?.domain;
+    if (domainPkg?.tools && Array.isArray(domainPkg.tools)) {
+      for (const tool of domainPkg.tools) {
+        if (!tool?.name) continue;
+        toolMap.set(tool.name, {
+          name: tool.name,
+          description: tool.description,
+          inputSchema: tool.inputSchema,
+          outputSchema: tool.outputSchema,
+          className: tool.className || void 0,
+          static: tool.static || false,
+          ready: false,
+          publishedAt: void 0
+        });
+      }
+    }
+    const catalogs = doc.catalog?.rawToolCatalogs || {};
+    for (const [clientId, catalog] of Object.entries(catalogs)) {
+      const cat = catalog;
+      if (!cat.tools) continue;
+      for (const tool of cat.tools) {
+        if (!tool?.name) continue;
+        const existing = toolMap.get(tool.name);
+        if (existing?.publishedAt && cat.publishedAt && existing.publishedAt > cat.publishedAt) continue;
+        const isLocal = clientId === this.clientId;
+        const ready = isLocal ? this.effects.has(tool.name) || this.toolHandlers.has(tool.name) || this.instanceTools.has(tool.name) : true;
+        toolMap.set(tool.name, {
+          name: tool.name,
+          description: tool.description,
+          inputSchema: tool.inputSchema,
+          outputSchema: tool.outputSchema,
+          className: tool.className || existing?.className,
+          static: tool.static || existing?.static || false,
+          clientId,
+          ready,
+          publishedAt: cat.publishedAt
+        });
+      }
+    }
+    return Array.from(toolMap.values());
+  }
+  /**
+   * Subscribe to tool changes (added, removed, updated).
+   * @param callback - Function called with change events
+   * @returns Unsubscribe function
+   */
+  onToolsChanged(callback) {
+    const handler = (data) => callback(data);
+    if (!this.eventListeners.has("tools:changed")) {
+      this.eventListeners.set("tools:changed", []);
+    }
+    this.eventListeners.get("tools:changed").push(handler);
+    return () => {
+      const listeners = this.eventListeners.get("tools:changed");
+      if (listeners) {
+        this.eventListeners.set("tools:changed", listeners.filter((h) => h !== handler));
+      }
+    };
+  }
   /**
    * Get the current domain state and available tools
    */
@@ -639,7 +870,10 @@ import { ${allImports} } from "./sandbox-tools";
     });
   }
   setupEventHandlers() {
-    this.client.on("sync", (doc) => this.emit("sync", doc));
+    this.client.on("sync", (doc) => {
+      this.emit("sync", doc);
+      this.checkForToolChanges();
+    });
     this.client.on("prompt", (prompt) => this.emit("prompt", prompt));
     this.client.on("disconnect", () => this.emit("disconnect", {}));
     this.client.on("job.status", (data) => {
@@ -658,6 +892,33 @@ import { ${allImports} } from "./sandbox-tools";
       handlers.forEach((h) => h(data));
     }
   }
+  /**
+   * Check for changes in the tool catalog and emit 'tools:changed' if needed
+   */
+  checkForToolChanges() {
+    const currentTools = this.getTools();
+    const currentMap = new Map(currentTools.map((t) => [t.name, t]));
+    const added = [];
+    const removed = [];
+    for (const [name, tool] of currentMap) {
+      if (!this.lastKnownTools.has(name)) {
+        added.push(name);
+      }
+    }
+    for (const name of this.lastKnownTools.keys()) {
+      if (!currentMap.has(name)) {
+        removed.push(name);
+      }
+    }
+    if (added.length > 0 || removed.length > 0) {
+      this.lastKnownTools = currentMap;
+      this.emit("tools:changed", {
+        tools: currentTools,
+        added,
+        removed
+      });
+    }
+  }
 };
 var JobImplementation = class {
   id;
@@ -731,6 +992,18 @@ var JobImplementation = class {
         this._rejectResult(jobData.error || new Error("Job failed"));
       }
     });
+    this.client.on("tool.call.start", (data) => {
+      const d = data;
+      if (d.jobId === id) {
+        this.emit("toolCallStart", { callId: d.callId, toolName: d.toolName, input: d.input, timestamp: d.timestamp });
+      }
+    });
+    this.client.on("tool.call.end", (data) => {
+      const d = data;
+      if (d.jobId === id) {
+        this.emit("toolCallEnd", { callId: d.callId, toolName: d.toolName, result: d.result, error: d.error, durationMs: d.durationMs, timestamp: d.timestamp });
+      }
+    });
   }
   get result() {
     return this._resultPromise;
@@ -1445,11 +1718,57 @@ var Environment = class _Environment extends Session {
   async publishTools(tools, revision = "1.0.0") {
     return super.publishTools(tools, revision);
   }
+  /**
+   * Publish a single effect (tool) incrementally.
+   *
+   * Adds the effect to the local registry and re-publishes the
+   * full tool catalog to the server. If an effect with the same
+   * name already exists, it is replaced.
+   *
+   * @example
+   * ```typescript
+   * await env.publishEffect({
+   *   name: 'get_bio',
+   *   description: 'Get biography',
+   *   inputSchema: { type: 'object', properties: { detailed: { type: 'boolean' } } },
+   *   handler: async (params) => ({ bio: 'Hello' }),
+   * });
+   * ```
+   */
+  async publishEffect(effect) {
+    return super.publishEffect(effect);
+  }
+  /**
+   * Publish multiple effects (tools) at once.
+   *
+   * Adds all effects to the local registry and re-publishes the
+   * full tool catalog in a single RPC call.
+   */
+  async publishEffects(effects) {
+    return super.publishEffects(effects);
+  }
+  /**
+   * Remove an effect by name and re-publish the remaining catalog.
+   */
+  async unpublishEffect(name) {
+    return super.unpublishEffect(name);
+  }
+  /**
+   * Remove all effects and publish an empty catalog.
+   */
+  async unpublishAllEffects() {
+    return super.unpublishAllEffects();
+  }
 };
 var Granular = class {
   apiKey;
   apiUrl;
   httpUrl;
+  WebSocketCtor;
+  /** Sandbox-level effect registry: sandboxId → (toolName → ToolWithHandler) */
+  sandboxEffects = /* @__PURE__ */ new Map();
+  /** Active environments tracker: sandboxId → Environment[] */
+  activeEnvironments = /* @__PURE__ */ new Map();
   /**
    * Create a new Granular client
    * @param options - Client configuration
@@ -1457,6 +1776,7 @@ var Granular = class {
   constructor(options) {
     this.apiKey = options.apiKey;
     this.apiUrl = options.apiUrl || "wss://api.granular.dev/v2/ws";
+    this.WebSocketCtor = options.WebSocketCtor;
     this.httpUrl = this.apiUrl.replace("wss://", "https://").replace("/ws", "");
   }
   /**
@@ -1527,7 +1847,7 @@ var Granular = class {
    * ```
    */
   async connect(options) {
-    const clientId = `client_${Date.now()}`;
+    const clientId = options.clientId || `client_${Date.now()}`;
     const sandbox = await this.findOrCreateSandbox(options.sandbox);
     for (const profileName of options.user.permissions) {
       const profileId = await this.ensurePermissionProfile(sandbox.sandboxId, profileName);
@@ -1540,14 +1860,104 @@ var Granular = class {
     const client = new WSClient({
       url: this.apiUrl,
       sessionId: envData.environmentId,
-      token: this.apiKey
+      token: this.apiKey,
+      WebSocketCtor: this.WebSocketCtor
     });
     await client.connect();
     const graphqlEndpoint = `${this.httpUrl}/orchestrator/graphql`;
     const environment = new Environment(client, envData, clientId, this.apiKey, graphqlEndpoint);
+    if (!this.activeEnvironments.has(sandbox.sandboxId)) {
+      this.activeEnvironments.set(sandbox.sandboxId, []);
+    }
+    this.activeEnvironments.get(sandbox.sandboxId).push(environment);
+    environment.on("disconnect", () => {
+      const list = this.activeEnvironments.get(sandbox.sandboxId);
+      if (list) {
+        this.activeEnvironments.set(sandbox.sandboxId, list.filter((e) => e !== environment));
+      }
+    });
     await environment.hello();
+    const effects = this.sandboxEffects.get(sandbox.sandboxId);
+    if (effects && effects.size > 0) {
+      const effectsList = Array.from(effects.values());
+      console.log(`[Granular] Auto-publishing ${effectsList.length} effects for sandbox ${sandbox.sandboxId}`);
+      await environment.publishEffects(effectsList);
+    }
     return environment;
   }
+  // ── Sandbox-Level Effects ──
+  /**
+   * Register an effect (tool) for a specific sandbox.
+   * 
+   * The effect will be automatically published to:
+   * 1. Any currently active environments for this sandbox
+   * 2. Any new environments created/connected for this sandbox
+   * 
+   * @param sandboxNameOrId - The name or ID of the sandbox
+   * @param effect - The tool definition and handler
+   */
+  async registerEffect(sandboxNameOrId, effect) {
+    const sandbox = await this.findOrCreateSandbox(sandboxNameOrId);
+    const sandboxId = sandbox.sandboxId;
+    if (!this.sandboxEffects.has(sandboxId)) {
+      this.sandboxEffects.set(sandboxId, /* @__PURE__ */ new Map());
+    }
+    this.sandboxEffects.get(sandboxId).set(effect.name, effect);
+    const activeEnvs = this.activeEnvironments.get(sandboxId) || [];
+    for (const env of activeEnvs) {
+      await env.publishEffect(effect);
+    }
+  }
+  /**
+   * Register multiple effects (tools) for a specific sandbox.
+   * 
+   * batch version of `registerEffect`.
+   */
+  async registerEffects(sandboxNameOrId, effects) {
+    const sandbox = await this.findOrCreateSandbox(sandboxNameOrId);
+    const sandboxId = sandbox.sandboxId;
+    if (!this.sandboxEffects.has(sandboxId)) {
+      this.sandboxEffects.set(sandboxId, /* @__PURE__ */ new Map());
+    }
+    const map = this.sandboxEffects.get(sandboxId);
+    for (const effect of effects) {
+      map.set(effect.name, effect);
+    }
+    const activeEnvs = this.activeEnvironments.get(sandboxId) || [];
+    for (const env of activeEnvs) {
+      await env.publishEffects(effects);
+    }
+  }
+  /**
+   * Unregister an effect from a sandbox.
+   * 
+   * Removes it from the local registry and unpublishes it from
+   * all active environments.
+   */
+  async unregisterEffect(sandboxNameOrId, name) {
+    const sandbox = await this.findOrCreateSandbox(sandboxNameOrId);
+    const sandboxId = sandbox.sandboxId;
+    const map = this.sandboxEffects.get(sandboxId);
+    if (map) {
+      map.delete(name);
+    }
+    const activeEnvs = this.activeEnvironments.get(sandboxId) || [];
+    for (const env of activeEnvs) {
+      await env.unpublishEffect(name);
+    }
+  }
+  /**
+   * Unregister all effects for a sandbox.
+   */
+  async unregisterAllEffects(sandboxNameOrId) {
+    const sandbox = await this.findOrCreateSandbox(sandboxNameOrId);
+    const sandboxId = sandbox.sandboxId;
+    this.sandboxEffects.delete(sandboxId);
+    const activeEnvs = this.activeEnvironments.get(sandboxId) || [];
+    for (const env of activeEnvs) {
+      await env.unpublishAllEffects();
+    }
+  }
   /**
    * Find a sandbox by name or create it if it doesn't exist
    */