@ziggs-ai/agent-sdk 0.1.11 → 0.1.12

package/src/adapters/McpAdapter.ts

@@ -37,6 +37,32 @@ export type McpTransportConfig =
       baseUrl?: string;
       /** Agent to impersonate (X-Agent-Id) — defaults to the host agentId. */
       agentId?: string;
+    }
+  | {
+      /**
+       * Per-user brokered connection (ZIG-771): the agent holds NO
+       * connectionId/grantId. It speaks StreamableHTTP to
+       * `/connections/mcp/relay` and asserts the invoking end-user via
+       * `x-on-behalf-of-user`; the gateway resolves that user's own connection
+       * + grant and injects their token. One MCP client is created per invoking
+       * user so each keeps an isolated upstream session. Tool schemas cannot be
+       * discovered at boot (there is no connection yet) — pass `toolSchema`.
+       */
+      type: 'ziggs-per-user';
+      /** Operator key — defaults to ZIGGS_OPERATOR_KEY. */
+      operatorKey?: string;
+      /** Ziggs backend base URL — defaults to BACKEND_URL / ZIGGS_BACKEND_URL. */
+      baseUrl?: string;
+      /** Agent to impersonate (X-Agent-Id) — defaults to the host agentId. */
+      agentId?: string;
+      /**
+       * The MCP server URL this agent fronts. Not used for gateway calls (the
+       * gateway derives the server from the resolved connection), but needed to
+       * open a connection request when the invoking user hasn't connected yet
+       * (ZIG-771). Without it, a NO_USER_CONNECTION surfaces as an error rather
+       * than an auto-created consent prompt.
+       */
+      serverUrl?: string;
     };
 
 /**
@@ -123,6 +149,43 @@ export function resolveZiggsGatewayTarget(
   };
 }
 
+/**
+ * Resolve the per-user gateway URL + headers (ZIG-771). No connectionId/grantId —
+ * the gateway resolves the invoking user's own connection + grant from
+ * `x-on-behalf-of-user` + the impersonated agent.
+ */
+export function resolvePerUserGatewayTarget(
+  config: Extract<McpTransportConfig, { type: 'ziggs-per-user' }>,
+  agentId: string,
+  userId: string,
+  onBehalfOfAgreementId?: string,
+): { url: URL; headers: Record<string, string> } {
+  const operatorKey = config.operatorKey || process.env.ZIGGS_OPERATOR_KEY;
+  if (!operatorKey) {
+    throw new Error(
+      'ziggs-per-user MCP transport requires an operatorKey (or ZIGGS_OPERATOR_KEY)',
+    );
+  }
+  const baseUrl = (
+    config.baseUrl ||
+    process.env.BACKEND_URL ||
+    process.env.ZIGGS_BACKEND_URL ||
+    'http://localhost:3000'
+  ).replace(/\/$/, '');
+  return {
+    url: new URL(`${baseUrl}/connections/mcp/relay`),
+    headers: {
+      authorization: `Bearer ${operatorKey}`,
+      'x-agent-id': agentId,
+      'x-on-behalf-of-user': userId,
+      // ZIG-831: when this call runs under a delegation, userId is the delegating
+      // agent (senderId), not the end-user who owns the connection. The gateway
+      // uses this agreement to resolve the delegation root's user.
+      ...(onBehalfOfAgreementId ? { 'x-on-behalf-of-agreement': onBehalfOfAgreementId } : {}),
+    },
+  };
+}
+
 export interface CreateMcpAgentOptions extends Omit<AgentHostOptions, 'tools' | 'workflow' | 'states' | 'initial'> {
   /** MCP server connection — stdio (local process) or HTTP (remote). */
   transport: McpTransportConfig;
@@ -138,6 +201,26 @@ export interface CreateMcpAgentOptions extends Omit<AgentHostOptions, 'tools' |
    * See ZIG-759; Phase 2 will replace truncation with query-aware selection.
    */
   maxTools?: number;
+  /**
+   * Restrict the bound tool set to only tools whose name matches one of these
+   * case-insensitive substrings, applied BEFORE the maxTools cap. Lets a large
+   * MCP server (e.g. PostHog's 146 tools) be narrowed to the handful this agent
+   * actually needs — cutting per-turn cost and stopping the model from thrashing
+   * across (or never binding) the tool it should call. Empty/undefined = bind
+   * all (unchanged behavior). ZIG-759 Phase 2 = name-based selection.
+   */
+  allowedTools?: string[];
+  /**
+   * Pre-discovered MCP tool schemas. REQUIRED for `type: 'ziggs-per-user'`
+   * (ZIG-771) — there is no boot connection to call `tools/list` on, so the
+   * schema is baked at provision time and passed here. Ignored for the other
+   * transports, which discover tools live at boot.
+   */
+  toolSchema?: Array<{
+    name: string;
+    description?: string;
+    inputSchema?: Record<string, unknown>;
+  }>;
 }
 
 /** Provider `tools` array cap (OpenAI's is 128). */
@@ -177,6 +260,11 @@ export function buildTransport(config: McpTransportConfig, defaults: { agentId?:
       fetch: wrapGatewayFetch(),
     });
   }
+  if (config.type === 'ziggs-per-user') {
+    // Per-user transports are built lazily per invoking user in createMcpAgent
+    // (each needs a different x-on-behalf-of-user header), not here.
+    throw new Error('ziggs-per-user transport is created per invoking user, not via buildTransport');
+  }
   return new StreamableHTTPClientTransport(
     new URL(config.url),
     config.headers && Object.keys(config.headers).length > 0
@@ -204,7 +292,17 @@ function mcpContentToResult(content: McpContent[]): unknown {
   }
 }
 
-function mapMcpTools(client: Client, tools: McpTool[]): ToolDefinition[] {
+/**
+ * Calls one MCP tool. The single-client path closes over a fixed client; the
+ * per-user path (ZIG-771) uses `context.senderId` to pick that user's client.
+ */
+type McpToolCaller = (
+  toolName: string,
+  args: Record<string, unknown>,
+  context: Record<string, unknown>,
+) => Promise<{ isError?: boolean; content: McpContent[] }>;
+
+function mapMcpTools(callTool: McpToolCaller, tools: McpTool[]): ToolDefinition[] {
   return tools.map(tool => ({
     schema: {
       type: 'function' as const,
@@ -214,11 +312,11 @@ function mapMcpTools(client: Client, tools: McpTool[]): ToolDefinition[] {
         parameters: tool.inputSchema ?? { type: 'object', properties: {} },
       },
     },
-    handler: async (args: Record<string, unknown>) => {
-      const result = await client.callTool({ name: tool.name, arguments: args }) as {
-        isError?: boolean;
-        content: McpContent[];
-      };
+    handler: async (
+      args: Record<string, unknown>,
+      context: Record<string, unknown>,
+    ) => {
+      const result = await callTool(tool.name, args, context);
       if (result.isError) {
         throw new Error(mcpContentToResult(result.content) as string);
       }
@@ -279,6 +377,90 @@ function buildDefaultWorkflow(description: string, mcpTools: McpTool[]): Workflo
   };
 }
 
+/**
+ * Max chars for a single tool description bound to the LLM. Verbose MCP servers
+ * embed whole manuals in their descriptions (PostHog's `execute-sql` alone is
+ * ~30K chars of HogQL docs; its 146 tools total ~54K tokens of descriptions) and
+ * that full block rides every turn. Keep the lead — what the tool does + how to
+ * call it — and drop the embedded reference; agents fetch detail via dedicated
+ * schema/help tools at call time. Tunable; big enough to preserve real guidance.
+ * ZIG-759.
+ */
+export const MAX_TOOL_DESC_CHARS = 2000;
+
+/** Truncate one description at a sentence/newline boundary near the cap. */
+export function trimDescription(desc: string | undefined): string | undefined {
+  if (!desc || desc.length <= MAX_TOOL_DESC_CHARS) return desc;
+  const slice = desc.slice(0, MAX_TOOL_DESC_CHARS);
+  const boundary = Math.max(slice.lastIndexOf('\n'), slice.lastIndexOf('. '));
+  const head = boundary > MAX_TOOL_DESC_CHARS * 0.6 ? slice.slice(0, boundary + 1) : slice;
+  return `${head.trimEnd()} […description truncated; use the server's schema/help tools for full detail]`;
+}
+
+/** Trim oversized descriptions across the tool set (cost control, ZIG-759). */
+function trimToolDescriptions(tools: McpTool[], agentId: string | undefined): McpTool[] {
+  let trimmedCount = 0;
+  let savedChars = 0;
+  const out = tools.map((t) => {
+    const nextDesc = trimDescription(t.description);
+    if (nextDesc === t.description) return t;
+    trimmedCount += 1;
+    savedChars += (t.description?.length ?? 0) - (nextDesc?.length ?? 0);
+    return { ...t, description: nextDesc };
+  });
+  if (trimmedCount > 0) {
+    console.warn(
+      `[McpAdapter] ${agentId}: trimmed ${trimmedCount} oversized tool description(s) to ${MAX_TOOL_DESC_CHARS} chars, saving ~${Math.round(savedChars / 4 / 1000)}K tokens/turn. (ZIG-759)`,
+    );
+  }
+  return out;
+}
+
+/**
+ * Keep only tools whose name matches one of the allowlist substrings
+ * (case-insensitive). No allowlist → pass through unchanged. Applied before the
+ * maxTools cap so a targeted agent binds exactly the tools it needs. ZIG-759.
+ */
+function filterToolsByAllowlist(
+  tools: McpTool[],
+  allowedTools: string[] | undefined,
+  agentId: string | undefined,
+): McpTool[] {
+  if (!allowedTools || allowedTools.length === 0) return tools;
+  const patterns = allowedTools.map((p) => p.toLowerCase());
+  const kept = tools.filter((t) =>
+    patterns.some((p) => t.name.toLowerCase().includes(p)),
+  );
+  console.warn(
+    `[McpAdapter] ${agentId}: tool allowlist active — bound ${kept.length}/${tools.length} tools matching [${allowedTools.join(', ')}]${kept.length === 0 ? ' ⚠️ NONE MATCHED — check patterns against the server\'s tool names' : ': ' + kept.map((t) => t.name).join(', ')}.`,
+  );
+  return kept;
+}
+
+/**
+ * Cap the tool set bound to the LLM at the provider limit (OpenAI 128) —
+ * truncate and log rather than crash — and trim oversized descriptions so the
+ * bound schema doesn't ship a manual every turn. ZIG-759 Phase 1. Shared by the
+ * boot-time (single-client) and per-user paths. When `allowedTools` is set the
+ * server's tools are narrowed by name first, so the cap rarely bites.
+ */
+function capTools(
+  tools: McpTool[],
+  agentId: string | undefined,
+  maxTools: number,
+  extraToolCount: number,
+  allowedTools?: string[],
+): McpTool[] {
+  const allowed = filterToolsByAllowlist(tools, allowedTools, agentId);
+  const trimmed = trimToolDescriptions(allowed, agentId);
+  const budget = mcpToolBudget(maxTools, extraToolCount);
+  if (trimmed.length <= budget) return trimmed;
+  console.warn(
+    `[McpAdapter] ${agentId}: MCP server exposes ${trimmed.length} tools, over the ${maxTools}-tool LLM cap; binding the first ${budget}, dropping ${trimmed.length - budget}. (ZIG-759 Phase 2 will select per query.)`,
+  );
+  return trimmed.slice(0, budget);
+}
+
 /**
  * Wraps any MCP server as a Ziggs AgentHost.
  *
@@ -313,27 +495,175 @@ function buildDefaultWorkflow(description: string, mcpTools: McpTool[]): Workflo
  *   });
  */
 export async function createMcpAgent(options: CreateMcpAgentOptions): Promise<McpAgentHandle> {
-  const { transport: transportConfig, extraTools = [], workflow, maxTools = DEFAULT_MAX_TOOLS, ...hostOptions } = options;
+  const { transport: transportConfig, extraTools = [], workflow, maxTools = DEFAULT_MAX_TOOLS, allowedTools, toolSchema, ...hostOptions } = options;
+
+  // ── Per-user path (ZIG-771): no boot connection; one MCP client per invoking
+  // user, each carrying that user's on-behalf-of header for an isolated session.
+  if (transportConfig.type === 'ziggs-per-user') {
+    const agentId = transportConfig.agentId || hostOptions.agentId;
+    if (!agentId) {
+      throw new Error('ziggs-per-user MCP transport requires an agentId to impersonate');
+    }
+    if (!toolSchema || toolSchema.length === 0) {
+      throw new Error('ziggs-per-user MCP transport requires toolSchema (no boot connection to discover tools)');
+    }
+    const tools = capTools(toolSchema as McpTool[], agentId, maxTools, extraTools.length, allowedTools);
 
+    const clients = new Map<string, Promise<Client>>();
+    const clientFor = (
+      userId: string,
+      onBehalfOfAgreementId?: string,
+    ): Promise<Client> => {
+      // Key by user + agreement: the gateway may resolve a different end-user
+      // per agreement (delegation), so sessions must not be shared across them.
+      const cacheKey = `${userId}|${onBehalfOfAgreementId ?? ''}`;
+      let p = clients.get(cacheKey);
+      if (!p) {
+        p = (async () => {
+          const target = resolvePerUserGatewayTarget(transportConfig, agentId, userId, onBehalfOfAgreementId);
+          const transport = new StreamableHTTPClientTransport(target.url, {
+            requestInit: { headers: target.headers },
+            fetch: wrapGatewayFetch(),
+          });
+          const c = new Client({ name: 'ziggs-mcp-agent', version: '1.0.0' });
+          await c.connect(transport);
+          return c;
+        })().catch((err: unknown) => {
+          // Don't cache a failed connect (e.g. NO_USER_CONNECTION before the
+          // user consents) — the next turn must be able to retry.
+          clients.delete(cacheKey);
+          throw err;
+        });
+        clients.set(cacheKey, p);
+      }
+      return p;
+    };
+
+    // When the invoking user hasn't connected this server, open a connection
+    // request targeting THEM (X-On-Behalf-Of-User) so the consent prompt lands
+    // in their inbox. The backend dedups pending requests, so retrying every
+    // turn is safe. ZIG-771.
+    const backendBase = (
+      transportConfig.baseUrl ||
+      process.env.BACKEND_URL ||
+      process.env.ZIGGS_BACKEND_URL ||
+      'http://localhost:3000'
+    ).replace(/\/$/, '');
+    // App URL for the user-facing connect link (ZIG-790): api.ziggsai.com → ziggsai.com.
+    const appUrl = backendBase.includes('localhost')
+      ? 'http://localhost:1234'
+      : backendBase.replace(/\/\/api\./, '//');
+    const serverLabel = transportConfig.serverUrl
+      ? new URL(transportConfig.serverUrl).host
+      : 'the external service';
+    // Open a connection request for this user; returns its id so the caller can
+    // deep-link the user straight to it on the Connections page (ZIG-790).
+    const triggerConnectionRequest = async (userId: string): Promise<string | null> => {
+      if (!transportConfig.serverUrl) return null;
+      const operatorKey = transportConfig.operatorKey || process.env.ZIGGS_OPERATOR_KEY;
+      const res = await fetch(`${backendBase}/connections/mcp/requests`, {
+        method: 'POST',
+        headers: {
+          'content-type': 'application/json',
+          authorization: `Bearer ${operatorKey}`,
+          'x-agent-id': agentId,
+          'x-on-behalf-of-user': userId,
+        },
+        body: JSON.stringify({
+          serverUrl: transportConfig.serverUrl,
+          tools: tools.map(t => t.name),
+        }),
+      });
+      const body = (await res.json().catch(() => null)) as { request?: { requestId?: string } } | null;
+      return body?.request?.requestId ?? null;
+    };
+
+    const caller: McpToolCaller = async (name, args, context) => {
+      const userId = typeof context.senderId === 'string' ? context.senderId : undefined;
+      if (!userId) {
+        throw new Error('per-user MCP tool call requires an invoking user (context.senderId is missing)');
+      }
+      // ZIG-831: for a delegated wake, senderId is the delegating agent, not the
+      // connection-owning end-user; pass the agreement so the gateway resolves
+      // the delegation root's user.
+      const onBehalfOfAgreementId =
+        typeof context.onBehalfOfAgreementId === 'string' ? context.onBehalfOfAgreementId : undefined;
+      // ZIG-832 (B): per-user resolution is intermittently flaky even when the
+      // connection is valid and resolvable — the same on-behalf call fails with
+      // NO_USER_CONNECTION on one wake and succeeds on the next. clientFor never
+      // caches a failed connect, so each attempt re-resolves from scratch. Retry
+      // a few times (with backoff) before surfacing the connect prompt, so a
+      // transient miss no longer fails the whole fetch step.
+      const MAX_ATTEMPTS = 3;
+      let lastErr: unknown;
+      for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) {
+        try {
+          const c = await clientFor(userId, onBehalfOfAgreementId);
+          return (await c.callTool({ name, arguments: args })) as {
+            isError?: boolean;
+            content: McpContent[];
+          };
+        } catch (err) {
+          lastErr = err;
+          const resolutionMiss =
+            err instanceof ZiggsGatewayError && err.code === 'NO_USER_CONNECTION';
+          const transient = resolutionMiss || !(err instanceof ZiggsGatewayError);
+          if (attempt < MAX_ATTEMPTS && transient) {
+            console.warn(
+              `[McpAdapter] per-user ${name} attempt ${attempt}/${MAX_ATTEMPTS} failed (${(err as Error).message}); retrying with a fresh connection`,
+            );
+            await new Promise((r) => setTimeout(r, 500 * attempt));
+            continue;
+          }
+          break;
+        }
+      }
+      if (
+        lastErr instanceof ZiggsGatewayError &&
+        lastErr.code === 'NO_USER_CONNECTION' &&
+        transportConfig.serverUrl
+      ) {
+        const requestId = await triggerConnectionRequest(userId).catch(() => null);
+        const link = `${appUrl}/app/settings/connections${requestId ? `?request=${encodeURIComponent(requestId)}` : ''}`;
+        return {
+          content: [
+            {
+              type: 'text',
+              text: `ACTION NEEDED: the user has not connected ${serverLabel}. Ask them to approve the connection here: ${link} — then retry.`,
+            },
+          ],
+        };
+      }
+      throw lastErr;
+    };
+
+    const host = new AgentHost({
+      ...hostOptions,
+      workflow: workflow ?? buildDefaultWorkflow(hostOptions.description ?? '', tools),
+      tools: [...mapMcpTools(caller, tools), ...extraTools],
+    });
+    return {
+      host,
+      closeMcp: async () => {
+        await Promise.all(
+          [...clients.values()].map(p => p.then(c => c.close()).catch(() => {})),
+        );
+      },
+      toolNames: tools.map(t => t.name),
+    };
+  }
+
+  // ── Single-client path (stdio / http / brokered ziggs): discover tools at boot.
   const transport = buildTransport(transportConfig, { agentId: hostOptions.agentId });
   const client = new Client({ name: 'ziggs-mcp-agent', version: '1.0.0' });
-
   await client.connect(transport);
 
-  const listed = (await client.listTools()).tools as McpTool[];
-  // Cap the tool set bound to the LLM: providers reject a `tools` array longer
-  // than their limit (OpenAI 128). Truncate — never crash — and log what dropped
-  // so a large server (e.g. PostHog: 146 tools) still answers. ZIG-759 Phase 1.
-  const budget = mcpToolBudget(maxTools, extraTools.length);
-  let rawTools = listed;
-  if (listed.length > budget) {
-    rawTools = listed.slice(0, budget);
-    console.warn(`[McpAdapter] ${hostOptions.agentId}: MCP server exposes ${listed.length} tools, over the ${maxTools}-tool LLM cap; binding the first ${rawTools.length}, dropping ${listed.length - rawTools.length}. (ZIG-759 Phase 2 will select per query.)`);
-  }
-  const mcpTools = mapMcpTools(client, rawTools);
+  const rawTools = capTools((await client.listTools()).tools as McpTool[], hostOptions.agentId, maxTools, extraTools.length, allowedTools);
+  const caller: McpToolCaller = (name, args) =>
+    client.callTool({ name, arguments: args }) as Promise<{ isError?: boolean; content: McpContent[] }>;
+  const mcpTools = mapMcpTools(caller, rawTools);
 
   const resolvedWorkflow = workflow ?? buildDefaultWorkflow(hostOptions.description ?? '', rawTools);
-
   const host = new AgentHost({
     ...hostOptions,
     workflow: resolvedWorkflow,

package/src/adapters/OpenAIAdapter.ts

@@ -167,21 +167,28 @@ export class OpenAIAdapter {
   async chatMessagesStream(
     messages: OpenAI.Chat.ChatCompletionMessageParam[],
     onChunk: (text: string) => void,
-  ): Promise<string> {
+  ): Promise<{ text: string; usage: OpenAI.CompletionUsage | undefined }> {
     const stream = await this.client.chat.completions.create({
       model: this.model,
       messages,
       stream: true,
+      // Without this, a streamed completion never returns a usage block — the
+      // streamed reply (often the largest output) would bill tokens that are
+      // counted nowhere (ZIG-813). The usage arrives on a final chunk whose
+      // `choices` is empty.
+      stream_options: { include_usage: true },
       max_completion_tokens: maxCompletionTokensForModel(this.model),
     });
     let text = '';
+    let usage: OpenAI.CompletionUsage | undefined;
     for await (const chunk of stream) {
+      if (chunk.usage) usage = chunk.usage;
       const delta = chunk.choices[0]?.delta?.content;
       if (delta) {
         text += delta;
         onChunk(delta);
       }
     }
-    return text;
+    return { text, usage };
   }
 }

package/src/agent/Agent.ts

@@ -42,8 +42,16 @@ export interface TickOutput {
   parked: boolean;
   /** True when a state-timeout outcome matched no transition this tick — the host must NOT re-arm. */
   stateTimeoutUnhandled: boolean;
-  /** Accumulated LLM call count and token usage for this tick. */
-  metrics: { llmCalls: number; tokens: { total: number; input: number; output: number } };
+  /**
+   * Accumulated LLM call count and token usage for this tick.
+   * `cacheRead` / `cacheCreation` are populated by the Claude brain (its usage
+   * exposes cache traffic, which is the bulk of real cost — ZIG-759/813); the
+   * FSM path leaves them undefined. `total` includes cache tokens when present.
+   */
+  metrics: {
+    llmCalls: number;
+    tokens: { total: number; input: number; output: number; cacheRead?: number; cacheCreation?: number };
+  };
 }
 
 /**