@databricks/appkit 0.41.7 → 0.43.0

package/dist/connectors/mcp/client.js.map

@@ -1 +1 @@
-{"version":3,"file":"client.js","names":[],"sources":["../../../src/connectors/mcp/client.ts"],"sourcesContent":["/**\n * Custom MCP over HTTP (Streamable) — not `@modelcontextprotocol/sdk`\n *\n * This module implements a tiny JSON-RPC 2.0 client on `fetch` for the subset\n * of MCP we need: `initialize`, `notifications/initialized`, `tools/list`,\n * `tools/call` over a single JSON request/response. We do not use the official\n * SDK because:\n *\n * - **Policy and auth are the product** — every outbound URL is checked with\n *   {@link McpHostPolicy} (allowlist, DNS, private/blocked IP ranges) before\n *   the first byte is sent, and workspace tokens are only forwarded when\n *   `forwardWorkspaceAuth` is true for that destination. A generic transport\n *   from the SDK would still need the same hooks; re-wrapping it would be\n *   about as much code, with a larger third-party surface to audit.\n * - **Narrow scope** — we only target Databricks-hosted MCP over Streamable\n *   HTTP, not stdio, full SSE sessions, or the rest of the protocol. A\n *   hand-rolled path keeps the call graph obvious in code review.\n * - **Zero extra runtime dependency** for this path, consistent with other\n *   small, security-sensitive AppKit pieces.\n *\n * Revisit if we add more transports, or if the SDK ships a first-class way to\n * inject our host policy and per-URL auth without fighting the default\n * transport.\n */\nimport type { AgentToolDefinition } from \"shared\";\nimport { createLogger } from \"../../logging/logger\";\nimport {\n  assertResolvedHostSafe,\n  checkMcpUrl,\n  type DnsLookup,\n  type McpHostPolicy,\n} from \"./host-policy\";\nimport type { McpEndpointConfig } from \"./types\";\n\nconst logger = createLogger(\"connector:mcp\");\n\n/**\n * Hard cap on the size of a single MCP response body, including SSE\n * frames bundled into one HTTP response. MCP `initialize` / `tools/list`\n * / `tools/call` responses are JSON-RPC payloads — single-digit kilobytes\n * in normal use. A response anywhere near this size signals either a\n * misbehaving server or an attempt to exhaust client memory; we'd rather\n * fail loudly than allocate unbounded buffers from a remote.\n */\nconst MCP_RESPONSE_BODY_LIMIT_BYTES = 1024 * 1024;\n\n/**\n * Read a fetch Response body into a string with a hard size cap. Aborts\n * and throws if the cumulative bytes read cross {@link\n * MCP_RESPONSE_BODY_LIMIT_BYTES}, so a remote server cannot keep\n * streaming data past the limit. Returns the empty string when the\n * response has no readable body.\n */\n/**\n * Empty-object fallback used when an MCP server ships a missing or\n * malformed `inputSchema`. Matches the shape downstream adapters expect\n * for a function tool that takes no arguments.\n */\nconst EMPTY_TOOL_PARAMETERS: AgentToolDefinition[\"parameters\"] = {\n  type: \"object\",\n  properties: {},\n};\n\n/**\n * Coerce a remote MCP server's reported `inputSchema` into the\n * JSONSchema7 shape AppKit's adapters expect for a function tool's\n * `parameters`. The MCP wire type is `Record<string, unknown>`, so a\n * misbehaving (or malicious) server could ship arbitrary JSON. We accept\n * only the standard `{ type: \"object\", properties: {...} }` shape and\n * fall back to an empty-parameters schema otherwise — the tool still\n * registers, it just can't accept arguments.\n */\nfunction coerceToolParameters(\n  inputSchema: Record<string, unknown> | undefined,\n): AgentToolDefinition[\"parameters\"] {\n  if (!inputSchema || typeof inputSchema !== \"object\") {\n    return EMPTY_TOOL_PARAMETERS;\n  }\n  const { type, properties } = inputSchema;\n  if (type !== \"object\") return EMPTY_TOOL_PARAMETERS;\n  if (\n    properties !== undefined &&\n    (typeof properties !== \"object\" ||\n      properties === null ||\n      Array.isArray(properties))\n  ) {\n    return EMPTY_TOOL_PARAMETERS;\n  }\n  return inputSchema as AgentToolDefinition[\"parameters\"];\n}\n\nasync function readResponseTextCapped(\n  response: Response,\n  maxBytes: number,\n  contextLabel: string,\n): Promise<string> {\n  if (!response.body) return \"\";\n  const reader = response.body.getReader();\n  const decoder = new TextDecoder(\"utf-8\");\n  let total = 0;\n  let out = \"\";\n  try {\n    while (true) {\n      const { done, value } = await reader.read();\n      if (done) break;\n      total += value.byteLength;\n      if (total > maxBytes) {\n        throw new Error(\n          `MCP ${contextLabel}: response body exceeded ${maxBytes} bytes — refusing to allocate unbounded buffer from a remote server.`,\n        );\n      }\n      out += decoder.decode(value, { stream: true });\n    }\n    out += decoder.decode();\n  } finally {\n    reader.releaseLock();\n  }\n  return out;\n}\n\ninterface JsonRpcRequest {\n  jsonrpc: \"2.0\";\n  id: number;\n  method: string;\n  params?: Record<string, unknown>;\n}\n\ninterface JsonRpcResponse {\n  jsonrpc: \"2.0\";\n  id: number;\n  result?: unknown;\n  error?: { code: number; message: string; data?: unknown };\n}\n\ninterface McpToolSchema {\n  name: string;\n  description?: string;\n  inputSchema?: Record<string, unknown>;\n}\n\ninterface McpToolCallResult {\n  content: Array<{ type: string; text?: string }>;\n  isError?: boolean;\n}\n\n/**\n * Per-endpoint outcome of {@link AppKitMcpClient.connectAll}. Callers (the\n * agents plugin in particular) use the split to warn at startup when some\n * MCP servers are unreachable without aborting boot for the rest.\n */\nexport interface McpConnectAllResult {\n  connected: string[];\n  failed: Array<{ name: string; error: Error }>;\n}\n\ninterface McpServerConnection {\n  config: McpEndpointConfig;\n  resolvedUrl: string;\n  /**\n   * Whether workspace auth (SP / OBO) may be forwarded to this endpoint's URL.\n   * Decided at `connect()` time via {@link McpHostPolicy} and cached for the\n   * lifetime of the connection.\n   */\n  forwardWorkspaceAuth: boolean;\n  tools: Map<string, McpToolSchema>;\n}\n\n/**\n * Lightweight MCP client for Databricks-hosted MCP servers.\n *\n * Uses raw fetch() with JSON-RPC 2.0 over HTTP — no @modelcontextprotocol/sdk\n * or LangChain dependency. Supports the Streamable HTTP transport only\n * (POST with JSON-RPC request, single JSON-RPC response). Implements exactly\n * four methods: `initialize`, `notifications/initialized`, `tools/list`,\n * `tools/call`. No prompts/resources/completion/sampling.\n *\n * All outbound URLs are gated by an {@link McpHostPolicy}: unallowlisted hosts\n * are rejected before the first byte is sent, and workspace credentials are\n * only forwarded to the same-origin workspace. See `mcp-host-policy.ts`.\n *\n * Rationale for hand-rolling JSON-RPC instead of `@modelcontextprotocol/sdk`:\n * see the file-level comment at the top of this module.\n */\nexport class AppKitMcpClient {\n  private connections = new Map<string, McpServerConnection>();\n  private sessionIds = new Map<string, string>();\n  private requestId = 0;\n  private closed = false;\n\n  constructor(\n    private workspaceHost: string,\n    private authenticate: () => Promise<Record<string, string>>,\n    private policy: McpHostPolicy,\n    private options: { dnsLookup?: DnsLookup; fetchImpl?: typeof fetch } = {},\n  ) {}\n\n  /**\n   * Connects every endpoint in parallel and returns a structured summary so\n   * callers can distinguish \"all connected\" from \"some failed\".\n   *\n   * Returning the result instead of throwing is deliberate: one\n   * misconfigured MCP server should not take down the entire agents plugin\n   * at boot, and the agents plugin uses the summary to warn at startup with\n   * the failed-endpoint names. Errors are also logged here so a caller\n   * that ignores the return still gets per-endpoint diagnostics.\n   *\n   * @returns `connected` lists the endpoint names that initialised\n   *   successfully; `failed` carries `{ name, error }` for the rest.\n   */\n  async connectAll(\n    endpoints: McpEndpointConfig[],\n  ): Promise<McpConnectAllResult> {\n    const results = await Promise.allSettled(\n      endpoints.map((ep) => this.connect(ep)),\n    );\n    const out: McpConnectAllResult = { connected: [], failed: [] };\n    for (let i = 0; i < results.length; i++) {\n      const r = results[i];\n      const name = endpoints[i].name;\n      if (r.status === \"fulfilled\") {\n        out.connected.push(name);\n      } else {\n        const error =\n          r.reason instanceof Error ? r.reason : new Error(String(r.reason));\n        logger.error(\"Failed to connect MCP server %s: %O\", name, error);\n        out.failed.push({ name, error });\n      }\n    }\n    return out;\n  }\n\n  private resolveUrl(endpoint: McpEndpointConfig): string {\n    if (\n      endpoint.url.startsWith(\"http://\") ||\n      endpoint.url.startsWith(\"https://\")\n    ) {\n      return endpoint.url;\n    }\n    return `${this.workspaceHost}${endpoint.url}`;\n  }\n\n  async connect(endpoint: McpEndpointConfig): Promise<void> {\n    const resolvedUrl = this.resolveUrl(endpoint);\n    const check = checkMcpUrl(resolvedUrl, this.policy);\n    if (!check.ok) {\n      throw new Error(\n        `MCP endpoint '${endpoint.name}' refused at connect: ${check.reason}`,\n      );\n    }\n    await assertResolvedHostSafe(\n      check.url.hostname,\n      this.policy,\n      this.options.dnsLookup,\n    );\n\n    logger.info(\n      \"Connecting to MCP server: %s at %s (forwardWorkspaceAuth=%s)\",\n      endpoint.name,\n      resolvedUrl,\n      check.forwardWorkspaceAuth,\n    );\n\n    const initResponse = await this.sendRpc(\n      resolvedUrl,\n      \"initialize\",\n      {\n        protocolVersion: \"2025-03-26\",\n        capabilities: {},\n        clientInfo: { name: \"appkit-agent\", version: \"0.1.0\" },\n      },\n      { forwardWorkspaceAuth: check.forwardWorkspaceAuth },\n    );\n\n    if (initResponse.sessionId) {\n      this.sessionIds.set(endpoint.name, initResponse.sessionId);\n    }\n    const sessionId = this.sessionIds.get(endpoint.name);\n\n    await this.sendNotification(resolvedUrl, \"notifications/initialized\", {\n      sessionId,\n      forwardWorkspaceAuth: check.forwardWorkspaceAuth,\n    });\n\n    const listResponse = await this.sendRpc(\n      resolvedUrl,\n      \"tools/list\",\n      {},\n      { sessionId, forwardWorkspaceAuth: check.forwardWorkspaceAuth },\n    );\n    const toolList =\n      (listResponse.result as { tools?: McpToolSchema[] })?.tools ?? [];\n\n    const tools = new Map<string, McpToolSchema>();\n    for (const tool of toolList) {\n      tools.set(tool.name, tool);\n    }\n\n    this.connections.set(endpoint.name, {\n      config: endpoint,\n      resolvedUrl,\n      forwardWorkspaceAuth: check.forwardWorkspaceAuth,\n      tools,\n    });\n    logger.info(\n      \"Connected to MCP server %s: %d tools available\",\n      endpoint.name,\n      tools.size,\n    );\n  }\n\n  getAllToolDefinitions(): AgentToolDefinition[] {\n    const defs: AgentToolDefinition[] = [];\n    for (const [serverName, conn] of this.connections) {\n      for (const [toolName, schema] of conn.tools) {\n        defs.push({\n          name: `mcp.${serverName}.${toolName}`,\n          description: schema.description ?? toolName,\n          parameters: coerceToolParameters(schema.inputSchema),\n        });\n      }\n    }\n    return defs;\n  }\n\n  /**\n   * Whether the named MCP server may receive workspace-scoped auth headers\n   * (e.g., an OBO bearer token from an end-user request). Callers should gate\n   * auth-forwarding decisions on this to prevent credential exfiltration to\n   * non-workspace hosts.\n   */\n  canForwardWorkspaceAuth(serverName: string): boolean {\n    return this.connections.get(serverName)?.forwardWorkspaceAuth ?? false;\n  }\n\n  async callTool(\n    qualifiedName: string,\n    args: unknown,\n    authHeaders?: Record<string, string>,\n    callerSignal?: AbortSignal,\n  ): Promise<string> {\n    const parts = qualifiedName.split(\".\");\n    if (parts.length < 3 || parts[0] !== \"mcp\") {\n      throw new Error(`Invalid MCP tool name: ${qualifiedName}`);\n    }\n    const serverName = parts[1];\n    const toolName = parts.slice(2).join(\".\");\n\n    const conn = this.connections.get(serverName);\n    if (!conn) {\n      throw new Error(`MCP server not connected: ${serverName}`);\n    }\n\n    const sessionId = this.sessionIds.get(serverName);\n    // authHeaders are caller-supplied credentials (typically the OBO token).\n    // Only honor them if the destination URL was admitted with\n    // forwardWorkspaceAuth=true at connect time.\n    const scopedAuthOverride = conn.forwardWorkspaceAuth\n      ? authHeaders\n      : undefined;\n\n    const rpcResult = await this.sendRpc(\n      conn.resolvedUrl,\n      \"tools/call\",\n      { name: toolName, arguments: args },\n      {\n        authOverride: scopedAuthOverride,\n        sessionId,\n        forwardWorkspaceAuth: conn.forwardWorkspaceAuth,\n        callerSignal,\n      },\n    );\n    const result = rpcResult.result as McpToolCallResult;\n\n    // `text` is optional on `McpToolCallResult.content[]` per the MCP\n    // spec; filtering only on `type === \"text\"` lets `c.text` be\n    // `undefined`, which `Array.join` would render as the literal\n    // string `\"undefined\"` and ship to the agent. Narrow on both\n    // fields so the joined string only contains real text.\n    const textContent = (result.content ?? []).filter(\n      (c): c is { type: \"text\"; text: string } =>\n        c.type === \"text\" && typeof c.text === \"string\",\n    );\n\n    if (result.isError) {\n      const errText = textContent.map((c) => c.text).join(\"\\n\");\n      throw new Error(errText || \"MCP tool call failed\");\n    }\n\n    return textContent.map((c) => c.text).join(\"\\n\");\n  }\n\n  async close(): Promise<void> {\n    this.closed = true;\n    this.connections.clear();\n    this.sessionIds.clear();\n  }\n\n  private async sendRpc(\n    url: string,\n    method: string,\n    params?: Record<string, unknown>,\n    options?: {\n      authOverride?: Record<string, string>;\n      sessionId?: string;\n      forwardWorkspaceAuth?: boolean;\n      /**\n       * Optional external abort signal (typically the agent's stream signal).\n       * Composed with the built-in 30 s timeout so `/cancel` or agent-run\n       * shutdown immediately propagates to the MCP fetch rather than waiting\n       * for the remote server to respond.\n       */\n      callerSignal?: AbortSignal;\n    },\n  ): Promise<{ result: unknown; sessionId?: string }> {\n    if (this.closed) throw new Error(\"MCP client is closed\");\n\n    const request: JsonRpcRequest = {\n      jsonrpc: \"2.0\",\n      id: ++this.requestId,\n      method,\n      ...(params && { params }),\n    };\n\n    const authHeaders = await this.resolveAuthHeaders(options);\n    const headers: Record<string, string> = {\n      \"Content-Type\": \"application/json\",\n      Accept: \"application/json, text/event-stream\",\n      ...authHeaders,\n    };\n    if (options?.sessionId) {\n      headers[\"Mcp-Session-Id\"] = options.sessionId;\n    }\n\n    const fetchImpl = this.options.fetchImpl ?? fetch;\n    const signals: AbortSignal[] = [AbortSignal.timeout(30_000)];\n    if (options?.callerSignal) signals.push(options.callerSignal);\n    const response = await fetchImpl(url, {\n      method: \"POST\",\n      headers,\n      body: JSON.stringify(request),\n      signal: signals.length > 1 ? AbortSignal.any(signals) : signals[0],\n    });\n\n    if (!response.ok) {\n      throw new Error(\n        `MCP request to ${method} failed: ${response.status} ${response.statusText}`,\n      );\n    }\n\n    const contentType = response.headers.get(\"content-type\") ?? \"\";\n    // Always read the body via the capped helper so a misconfigured or\n    // malicious server can't exhaust client memory by streaming an\n    // unbounded payload. Applies to both SSE (`response.text()` would\n    // have buffered the whole stream) and plain JSON (`response.json()`\n    // does the same internally).\n    const bodyText = await readResponseTextCapped(\n      response,\n      MCP_RESPONSE_BODY_LIMIT_BYTES,\n      method,\n    );\n    let json: JsonRpcResponse;\n\n    if (contentType.includes(\"text/event-stream\")) {\n      const lastData = bodyText\n        .split(\"\\n\")\n        .filter((line) => line.startsWith(\"data: \"))\n        .map((line) => line.slice(6))\n        .pop();\n      if (!lastData) {\n        throw new Error(`MCP SSE response for ${method} contained no data`);\n      }\n      json = JSON.parse(lastData) as JsonRpcResponse;\n    } else {\n      if (bodyText.length === 0) {\n        throw new Error(`MCP response for ${method} had an empty body`);\n      }\n      json = JSON.parse(bodyText) as JsonRpcResponse;\n    }\n\n    if (json.error) {\n      throw new Error(`MCP error (${json.error.code}): ${json.error.message}`);\n    }\n\n    const sid = response.headers.get(\"mcp-session-id\") ?? undefined;\n    return { result: json.result, sessionId: sid };\n  }\n\n  private async sendNotification(\n    url: string,\n    method: string,\n    options?: {\n      sessionId?: string;\n      forwardWorkspaceAuth?: boolean;\n    },\n  ): Promise<void> {\n    if (this.closed) return;\n\n    const authHeaders = await this.resolveAuthHeaders(options);\n    const headers: Record<string, string> = {\n      \"Content-Type\": \"application/json\",\n      Accept: \"application/json, text/event-stream\",\n      ...authHeaders,\n    };\n    if (options?.sessionId) {\n      headers[\"Mcp-Session-Id\"] = options.sessionId;\n    }\n\n    const fetchImpl = this.options.fetchImpl ?? fetch;\n    // MCP notifications are fire-and-forget per spec — we don't throw on\n    // failure. But silently swallowing 4xx/5xx hides server-side\n    // rejections that would otherwise look like a successful connect()\n    // followed by mysterious tool-call failures. Surface the bad status\n    // via the logger so the dev sees it without breaking the protocol\n    // contract.\n    try {\n      const response = await fetchImpl(url, {\n        method: \"POST\",\n        headers,\n        body: JSON.stringify({ jsonrpc: \"2.0\", method }),\n        signal: AbortSignal.timeout(30_000),\n      });\n      if (!response.ok) {\n        logger.warn(\n          \"MCP notification %s to %s returned %d %s — the server may have rejected the request, but per MCP spec notifications are fire-and-forget and the connection is considered established.\",\n          method,\n          url,\n          response.status,\n          response.statusText,\n        );\n      }\n    } catch (err) {\n      logger.warn(\n        \"MCP notification %s to %s failed before a response was received: %O\",\n        method,\n        url,\n        err,\n      );\n    }\n  }\n\n  /**\n   * Return the auth headers to send on an outbound request. Workspace auth\n   * (SP or OBO) is only resolved when `forwardWorkspaceAuth` is true; for\n   * non-workspace hosts no bearer token is attached.\n   */\n  private async resolveAuthHeaders(options?: {\n    authOverride?: Record<string, string>;\n    forwardWorkspaceAuth?: boolean;\n  }): Promise<Record<string, string>> {\n    if (!options?.forwardWorkspaceAuth) return {};\n    if (options.authOverride) return options.authOverride;\n    return this.authenticate();\n  }\n}\n"],"mappings":";;;;AAkCA,MAAM,SAAS,aAAa,gBAAgB;;;;;;;;;AAU5C,MAAM,gCAAgC,OAAO;;;;;;;;;;;;;AAc7C,MAAM,wBAA2D;CAC/D,MAAM;CACN,YAAY,EAAE;CACf;;;;;;;;;;AAWD,SAAS,qBACP,aACmC;AACnC,KAAI,CAAC,eAAe,OAAO,gBAAgB,SACzC,QAAO;CAET,MAAM,EAAE,MAAM,eAAe;AAC7B,KAAI,SAAS,SAAU,QAAO;AAC9B,KACE,eAAe,WACd,OAAO,eAAe,YACrB,eAAe,QACf,MAAM,QAAQ,WAAW,EAE3B,QAAO;AAET,QAAO;;AAGT,eAAe,uBACb,UACA,UACA,cACiB;AACjB,KAAI,CAAC,SAAS,KAAM,QAAO;CAC3B,MAAM,SAAS,SAAS,KAAK,WAAW;CACxC,MAAM,UAAU,IAAI,YAAY,QAAQ;CACxC,IAAI,QAAQ;CACZ,IAAI,MAAM;AACV,KAAI;AACF,SAAO,MAAM;GACX,MAAM,EAAE,MAAM,UAAU,MAAM,OAAO,MAAM;AAC3C,OAAI,KAAM;AACV,YAAS,MAAM;AACf,OAAI,QAAQ,SACV,OAAM,IAAI,MACR,OAAO,aAAa,2BAA2B,SAAS,sEACzD;AAEH,UAAO,QAAQ,OAAO,OAAO,EAAE,QAAQ,MAAM,CAAC;;AAEhD,SAAO,QAAQ,QAAQ;WACf;AACR,SAAO,aAAa;;AAEtB,QAAO;;;;;;;;;;;;;;;;;;AAkET,IAAa,kBAAb,MAA6B;CAC3B,AAAQ,8BAAc,IAAI,KAAkC;CAC5D,AAAQ,6BAAa,IAAI,KAAqB;CAC9C,AAAQ,YAAY;CACpB,AAAQ,SAAS;CAEjB,YACE,AAAQ,eACR,AAAQ,cACR,AAAQ,QACR,AAAQ,UAA+D,EAAE,EACzE;EAJQ;EACA;EACA;EACA;;;;;;;;;;;;;;;CAgBV,MAAM,WACJ,WAC8B;EAC9B,MAAM,UAAU,MAAM,QAAQ,WAC5B,UAAU,KAAK,OAAO,KAAK,QAAQ,GAAG,CAAC,CACxC;EACD,MAAM,MAA2B;GAAE,WAAW,EAAE;GAAE,QAAQ,EAAE;GAAE;AAC9D,OAAK,IAAI,IAAI,GAAG,IAAI,QAAQ,QAAQ,KAAK;GACvC,MAAM,IAAI,QAAQ;GAClB,MAAM,OAAO,UAAU,GAAG;AAC1B,OAAI,EAAE,WAAW,YACf,KAAI,UAAU,KAAK,KAAK;QACnB;IACL,MAAM,QACJ,EAAE,kBAAkB,QAAQ,EAAE,SAAS,IAAI,MAAM,OAAO,EAAE,OAAO,CAAC;AACpE,WAAO,MAAM,uCAAuC,MAAM,MAAM;AAChE,QAAI,OAAO,KAAK;KAAE;KAAM;KAAO,CAAC;;;AAGpC,SAAO;;CAGT,AAAQ,WAAW,UAAqC;AACtD,MACE,SAAS,IAAI,WAAW,UAAU,IAClC,SAAS,IAAI,WAAW,WAAW,CAEnC,QAAO,SAAS;AAElB,SAAO,GAAG,KAAK,gBAAgB,SAAS;;CAG1C,MAAM,QAAQ,UAA4C;EACxD,MAAM,cAAc,KAAK,WAAW,SAAS;EAC7C,MAAM,QAAQ,YAAY,aAAa,KAAK,OAAO;AACnD,MAAI,CAAC,MAAM,GACT,OAAM,IAAI,MACR,iBAAiB,SAAS,KAAK,wBAAwB,MAAM,SAC9D;AAEH,QAAM,uBACJ,MAAM,IAAI,UACV,KAAK,QACL,KAAK,QAAQ,UACd;AAED,SAAO,KACL,gEACA,SAAS,MACT,aACA,MAAM,qBACP;EAED,MAAM,eAAe,MAAM,KAAK,QAC9B,aACA,cACA;GACE,iBAAiB;GACjB,cAAc,EAAE;GAChB,YAAY;IAAE,MAAM;IAAgB,SAAS;IAAS;GACvD,EACD,EAAE,sBAAsB,MAAM,sBAAsB,CACrD;AAED,MAAI,aAAa,UACf,MAAK,WAAW,IAAI,SAAS,MAAM,aAAa,UAAU;EAE5D,MAAM,YAAY,KAAK,WAAW,IAAI,SAAS,KAAK;AAEpD,QAAM,KAAK,iBAAiB,aAAa,6BAA6B;GACpE;GACA,sBAAsB,MAAM;GAC7B,CAAC;EAQF,MAAM,YANe,MAAM,KAAK,QAC9B,aACA,cACA,EAAE,EACF;GAAE;GAAW,sBAAsB,MAAM;GAAsB,CAChE,EAEe,QAAwC,SAAS,EAAE;EAEnE,MAAM,wBAAQ,IAAI,KAA4B;AAC9C,OAAK,MAAM,QAAQ,SACjB,OAAM,IAAI,KAAK,MAAM,KAAK;AAG5B,OAAK,YAAY,IAAI,SAAS,MAAM;GAClC,QAAQ;GACR;GACA,sBAAsB,MAAM;GAC5B;GACD,CAAC;AACF,SAAO,KACL,kDACA,SAAS,MACT,MAAM,KACP;;CAGH,wBAA+C;EAC7C,MAAM,OAA8B,EAAE;AACtC,OAAK,MAAM,CAAC,YAAY,SAAS,KAAK,YACpC,MAAK,MAAM,CAAC,UAAU,WAAW,KAAK,MACpC,MAAK,KAAK;GACR,MAAM,OAAO,WAAW,GAAG;GAC3B,aAAa,OAAO,eAAe;GACnC,YAAY,qBAAqB,OAAO,YAAY;GACrD,CAAC;AAGN,SAAO;;;;;;;;CAST,wBAAwB,YAA6B;AACnD,SAAO,KAAK,YAAY,IAAI,WAAW,EAAE,wBAAwB;;CAGnE,MAAM,SACJ,eACA,MACA,aACA,cACiB;EACjB,MAAM,QAAQ,cAAc,MAAM,IAAI;AACtC,MAAI,MAAM,SAAS,KAAK,MAAM,OAAO,MACnC,OAAM,IAAI,MAAM,0BAA0B,gBAAgB;EAE5D,MAAM,aAAa,MAAM;EACzB,MAAM,WAAW,MAAM,MAAM,EAAE,CAAC,KAAK,IAAI;EAEzC,MAAM,OAAO,KAAK,YAAY,IAAI,WAAW;AAC7C,MAAI,CAAC,KACH,OAAM,IAAI,MAAM,6BAA6B,aAAa;EAG5D,MAAM,YAAY,KAAK,WAAW,IAAI,WAAW;EAIjD,MAAM,qBAAqB,KAAK,uBAC5B,cACA;EAaJ,MAAM,UAXY,MAAM,KAAK,QAC3B,KAAK,aACL,cACA;GAAE,MAAM;GAAU,WAAW;GAAM,EACnC;GACE,cAAc;GACd;GACA,sBAAsB,KAAK;GAC3B;GACD,CACF,EACwB;EAOzB,MAAM,eAAe,OAAO,WAAW,EAAE,EAAE,QACxC,MACC,EAAE,SAAS,UAAU,OAAO,EAAE,SAAS,SAC1C;AAED,MAAI,OAAO,SAAS;GAClB,MAAM,UAAU,YAAY,KAAK,MAAM,EAAE,KAAK,CAAC,KAAK,KAAK;AACzD,SAAM,IAAI,MAAM,WAAW,uBAAuB;;AAGpD,SAAO,YAAY,KAAK,MAAM,EAAE,KAAK,CAAC,KAAK,KAAK;;CAGlD,MAAM,QAAuB;AAC3B,OAAK,SAAS;AACd,OAAK,YAAY,OAAO;AACxB,OAAK,WAAW,OAAO;;CAGzB,MAAc,QACZ,KACA,QACA,QACA,SAYkD;AAClD,MAAI,KAAK,OAAQ,OAAM,IAAI,MAAM,uBAAuB;EAExD,MAAM,UAA0B;GAC9B,SAAS;GACT,IAAI,EAAE,KAAK;GACX;GACA,GAAI,UAAU,EAAE,QAAQ;GACzB;EAED,MAAM,cAAc,MAAM,KAAK,mBAAmB,QAAQ;EAC1D,MAAM,UAAkC;GACtC,gBAAgB;GAChB,QAAQ;GACR,GAAG;GACJ;AACD,MAAI,SAAS,UACX,SAAQ,oBAAoB,QAAQ;EAGtC,MAAM,YAAY,KAAK,QAAQ,aAAa;EAC5C,MAAM,UAAyB,CAAC,YAAY,QAAQ,IAAO,CAAC;AAC5D,MAAI,SAAS,aAAc,SAAQ,KAAK,QAAQ,aAAa;EAC7D,MAAM,WAAW,MAAM,UAAU,KAAK;GACpC,QAAQ;GACR;GACA,MAAM,KAAK,UAAU,QAAQ;GAC7B,QAAQ,QAAQ,SAAS,IAAI,YAAY,IAAI,QAAQ,GAAG,QAAQ;GACjE,CAAC;AAEF,MAAI,CAAC,SAAS,GACZ,OAAM,IAAI,MACR,kBAAkB,OAAO,WAAW,SAAS,OAAO,GAAG,SAAS,aACjE;EAGH,MAAM,cAAc,SAAS,QAAQ,IAAI,eAAe,IAAI;EAM5D,MAAM,WAAW,MAAM,uBACrB,UACA,+BACA,OACD;EACD,IAAI;AAEJ,MAAI,YAAY,SAAS,oBAAoB,EAAE;GAC7C,MAAM,WAAW,SACd,MAAM,KAAK,CACX,QAAQ,SAAS,KAAK,WAAW,SAAS,CAAC,CAC3C,KAAK,SAAS,KAAK,MAAM,EAAE,CAAC,CAC5B,KAAK;AACR,OAAI,CAAC,SACH,OAAM,IAAI,MAAM,wBAAwB,OAAO,oBAAoB;AAErE,UAAO,KAAK,MAAM,SAAS;SACtB;AACL,OAAI,SAAS,WAAW,EACtB,OAAM,IAAI,MAAM,oBAAoB,OAAO,oBAAoB;AAEjE,UAAO,KAAK,MAAM,SAAS;;AAG7B,MAAI,KAAK,MACP,OAAM,IAAI,MAAM,cAAc,KAAK,MAAM,KAAK,KAAK,KAAK,MAAM,UAAU;EAG1E,MAAM,MAAM,SAAS,QAAQ,IAAI,iBAAiB,IAAI;AACtD,SAAO;GAAE,QAAQ,KAAK;GAAQ,WAAW;GAAK;;CAGhD,MAAc,iBACZ,KACA,QACA,SAIe;AACf,MAAI,KAAK,OAAQ;EAEjB,MAAM,cAAc,MAAM,KAAK,mBAAmB,QAAQ;EAC1D,MAAM,UAAkC;GACtC,gBAAgB;GAChB,QAAQ;GACR,GAAG;GACJ;AACD,MAAI,SAAS,UACX,SAAQ,oBAAoB,QAAQ;EAGtC,MAAM,YAAY,KAAK,QAAQ,aAAa;AAO5C,MAAI;GACF,MAAM,WAAW,MAAM,UAAU,KAAK;IACpC,QAAQ;IACR;IACA,MAAM,KAAK,UAAU;KAAE,SAAS;KAAO;KAAQ,CAAC;IAChD,QAAQ,YAAY,QAAQ,IAAO;IACpC,CAAC;AACF,OAAI,CAAC,SAAS,GACZ,QAAO,KACL,yLACA,QACA,KACA,SAAS,QACT,SAAS,WACV;WAEI,KAAK;AACZ,UAAO,KACL,uEACA,QACA,KACA,IACD;;;;;;;;CASL,MAAc,mBAAmB,SAGG;AAClC,MAAI,CAAC,SAAS,qBAAsB,QAAO,EAAE;AAC7C,MAAI,QAAQ,aAAc,QAAO,QAAQ;AACzC,SAAO,KAAK,cAAc"}
+{"version":3,"file":"client.js","names":[],"sources":["../../../src/connectors/mcp/client.ts"],"sourcesContent":["/**\n * Custom MCP over HTTP (Streamable) — not `@modelcontextprotocol/sdk`\n *\n * This module implements a tiny JSON-RPC 2.0 client on `fetch` for the subset\n * of MCP we need: `initialize`, `notifications/initialized`, `tools/list`,\n * `tools/call` over a single JSON request/response. We do not use the official\n * SDK because:\n *\n * - **Policy and auth are the product** — every outbound URL is checked with\n *   {@link McpHostPolicy} (allowlist, DNS, private/blocked IP ranges) before\n *   the first byte is sent, and workspace tokens are only forwarded when\n *   `forwardWorkspaceAuth` is true for that destination. A generic transport\n *   from the SDK would still need the same hooks; re-wrapping it would be\n *   about as much code, with a larger third-party surface to audit.\n * - **Narrow scope** — we only target Databricks-hosted MCP over Streamable\n *   HTTP, not stdio, full SSE sessions, or the rest of the protocol. A\n *   hand-rolled path keeps the call graph obvious in code review.\n * - **Zero extra runtime dependency** for this path, consistent with other\n *   small, security-sensitive AppKit pieces.\n *\n * Revisit if we add more transports, or if the SDK ships a first-class way to\n * inject our host policy and per-URL auth without fighting the default\n * transport.\n */\nimport type { AgentToolDefinition } from \"shared\";\nimport { APPKIT_USER_AGENT } from \"../../context/client-options\";\nimport { createLogger } from \"../../logging/logger\";\nimport {\n  assertResolvedHostSafe,\n  checkMcpUrl,\n  type DnsLookup,\n  type McpHostPolicy,\n} from \"./host-policy\";\nimport type { McpEndpointConfig } from \"./types\";\n\nconst logger = createLogger(\"connector:mcp\");\n\n/**\n * Hard cap on the size of a single MCP response body, including SSE\n * frames bundled into one HTTP response. MCP `initialize` / `tools/list`\n * / `tools/call` responses are JSON-RPC payloads — single-digit kilobytes\n * in normal use. A response anywhere near this size signals either a\n * misbehaving server or an attempt to exhaust client memory; we'd rather\n * fail loudly than allocate unbounded buffers from a remote.\n */\nconst MCP_RESPONSE_BODY_LIMIT_BYTES = 1024 * 1024;\n\n/**\n * Read a fetch Response body into a string with a hard size cap. Aborts\n * and throws if the cumulative bytes read cross {@link\n * MCP_RESPONSE_BODY_LIMIT_BYTES}, so a remote server cannot keep\n * streaming data past the limit. Returns the empty string when the\n * response has no readable body.\n */\n/**\n * Empty-object fallback used when an MCP server ships a missing or\n * malformed `inputSchema`. Matches the shape downstream adapters expect\n * for a function tool that takes no arguments.\n */\nconst EMPTY_TOOL_PARAMETERS: AgentToolDefinition[\"parameters\"] = {\n  type: \"object\",\n  properties: {},\n};\n\n/**\n * Coerce a remote MCP server's reported `inputSchema` into the\n * JSONSchema7 shape AppKit's adapters expect for a function tool's\n * `parameters`. The MCP wire type is `Record<string, unknown>`, so a\n * misbehaving (or malicious) server could ship arbitrary JSON. We accept\n * only the standard `{ type: \"object\", properties: {...} }` shape and\n * fall back to an empty-parameters schema otherwise — the tool still\n * registers, it just can't accept arguments.\n */\nfunction coerceToolParameters(\n  inputSchema: Record<string, unknown> | undefined,\n): AgentToolDefinition[\"parameters\"] {\n  if (!inputSchema || typeof inputSchema !== \"object\") {\n    return EMPTY_TOOL_PARAMETERS;\n  }\n  const { type, properties } = inputSchema;\n  if (type !== \"object\") return EMPTY_TOOL_PARAMETERS;\n  if (\n    properties !== undefined &&\n    (typeof properties !== \"object\" ||\n      properties === null ||\n      Array.isArray(properties))\n  ) {\n    return EMPTY_TOOL_PARAMETERS;\n  }\n  return inputSchema as AgentToolDefinition[\"parameters\"];\n}\n\nasync function readResponseTextCapped(\n  response: Response,\n  maxBytes: number,\n  contextLabel: string,\n): Promise<string> {\n  if (!response.body) return \"\";\n  const reader = response.body.getReader();\n  const decoder = new TextDecoder(\"utf-8\");\n  let total = 0;\n  let out = \"\";\n  try {\n    while (true) {\n      const { done, value } = await reader.read();\n      if (done) break;\n      total += value.byteLength;\n      if (total > maxBytes) {\n        throw new Error(\n          `MCP ${contextLabel}: response body exceeded ${maxBytes} bytes — refusing to allocate unbounded buffer from a remote server.`,\n        );\n      }\n      out += decoder.decode(value, { stream: true });\n    }\n    out += decoder.decode();\n  } finally {\n    reader.releaseLock();\n  }\n  return out;\n}\n\ninterface JsonRpcRequest {\n  jsonrpc: \"2.0\";\n  id: number;\n  method: string;\n  params?: Record<string, unknown>;\n}\n\ninterface JsonRpcResponse {\n  jsonrpc: \"2.0\";\n  id: number;\n  result?: unknown;\n  error?: { code: number; message: string; data?: unknown };\n}\n\ninterface McpToolSchema {\n  name: string;\n  description?: string;\n  inputSchema?: Record<string, unknown>;\n}\n\ninterface McpToolCallResult {\n  content: Array<{ type: string; text?: string }>;\n  isError?: boolean;\n}\n\n/**\n * Per-endpoint outcome of {@link AppKitMcpClient.connectAll}. Callers (the\n * agents plugin in particular) use the split to warn at startup when some\n * MCP servers are unreachable without aborting boot for the rest.\n */\nexport interface McpConnectAllResult {\n  connected: string[];\n  failed: Array<{ name: string; error: Error }>;\n}\n\ninterface McpServerConnection {\n  config: McpEndpointConfig;\n  resolvedUrl: string;\n  /**\n   * Whether workspace auth (SP / OBO) may be forwarded to this endpoint's URL.\n   * Decided at `connect()` time via {@link McpHostPolicy} and cached for the\n   * lifetime of the connection.\n   */\n  forwardWorkspaceAuth: boolean;\n  tools: Map<string, McpToolSchema>;\n}\n\n/**\n * Lightweight MCP client for Databricks-hosted MCP servers.\n *\n * Uses raw fetch() with JSON-RPC 2.0 over HTTP — no @modelcontextprotocol/sdk\n * or LangChain dependency. Supports the Streamable HTTP transport only\n * (POST with JSON-RPC request, single JSON-RPC response). Implements exactly\n * four methods: `initialize`, `notifications/initialized`, `tools/list`,\n * `tools/call`. No prompts/resources/completion/sampling.\n *\n * All outbound URLs are gated by an {@link McpHostPolicy}: unallowlisted hosts\n * are rejected before the first byte is sent, and workspace credentials are\n * only forwarded to the same-origin workspace. See `mcp-host-policy.ts`.\n *\n * Rationale for hand-rolling JSON-RPC instead of `@modelcontextprotocol/sdk`:\n * see the file-level comment at the top of this module.\n */\nexport class AppKitMcpClient {\n  private connections = new Map<string, McpServerConnection>();\n  private sessionIds = new Map<string, string>();\n  private requestId = 0;\n  private closed = false;\n\n  constructor(\n    private workspaceHost: string,\n    private authenticate: () => Promise<Record<string, string>>,\n    private policy: McpHostPolicy,\n    private options: { dnsLookup?: DnsLookup; fetchImpl?: typeof fetch } = {},\n  ) {}\n\n  /**\n   * Connects every endpoint in parallel and returns a structured summary so\n   * callers can distinguish \"all connected\" from \"some failed\".\n   *\n   * Returning the result instead of throwing is deliberate: one\n   * misconfigured MCP server should not take down the entire agents plugin\n   * at boot, and the agents plugin uses the summary to warn at startup with\n   * the failed-endpoint names. Errors are also logged here so a caller\n   * that ignores the return still gets per-endpoint diagnostics.\n   *\n   * @returns `connected` lists the endpoint names that initialised\n   *   successfully; `failed` carries `{ name, error }` for the rest.\n   */\n  async connectAll(\n    endpoints: McpEndpointConfig[],\n  ): Promise<McpConnectAllResult> {\n    const results = await Promise.allSettled(\n      endpoints.map((ep) => this.connect(ep)),\n    );\n    const out: McpConnectAllResult = { connected: [], failed: [] };\n    for (let i = 0; i < results.length; i++) {\n      const r = results[i];\n      const name = endpoints[i].name;\n      if (r.status === \"fulfilled\") {\n        out.connected.push(name);\n      } else {\n        const error =\n          r.reason instanceof Error ? r.reason : new Error(String(r.reason));\n        logger.error(\"Failed to connect MCP server %s: %O\", name, error);\n        out.failed.push({ name, error });\n      }\n    }\n    return out;\n  }\n\n  private resolveUrl(endpoint: McpEndpointConfig): string {\n    if (\n      endpoint.url.startsWith(\"http://\") ||\n      endpoint.url.startsWith(\"https://\")\n    ) {\n      return endpoint.url;\n    }\n    return `${this.workspaceHost}${endpoint.url}`;\n  }\n\n  async connect(endpoint: McpEndpointConfig): Promise<void> {\n    const resolvedUrl = this.resolveUrl(endpoint);\n    const check = checkMcpUrl(resolvedUrl, this.policy);\n    if (!check.ok) {\n      throw new Error(\n        `MCP endpoint '${endpoint.name}' refused at connect: ${check.reason}`,\n      );\n    }\n    await assertResolvedHostSafe(\n      check.url.hostname,\n      this.policy,\n      this.options.dnsLookup,\n    );\n\n    logger.info(\n      \"Connecting to MCP server: %s at %s (forwardWorkspaceAuth=%s)\",\n      endpoint.name,\n      resolvedUrl,\n      check.forwardWorkspaceAuth,\n    );\n\n    const initResponse = await this.sendRpc(\n      resolvedUrl,\n      \"initialize\",\n      {\n        protocolVersion: \"2025-03-26\",\n        capabilities: {},\n        clientInfo: { name: \"appkit-agent\", version: \"0.1.0\" },\n      },\n      { forwardWorkspaceAuth: check.forwardWorkspaceAuth },\n    );\n\n    if (initResponse.sessionId) {\n      this.sessionIds.set(endpoint.name, initResponse.sessionId);\n    }\n    const sessionId = this.sessionIds.get(endpoint.name);\n\n    await this.sendNotification(resolvedUrl, \"notifications/initialized\", {\n      sessionId,\n      forwardWorkspaceAuth: check.forwardWorkspaceAuth,\n    });\n\n    const listResponse = await this.sendRpc(\n      resolvedUrl,\n      \"tools/list\",\n      {},\n      { sessionId, forwardWorkspaceAuth: check.forwardWorkspaceAuth },\n    );\n    const toolList =\n      (listResponse.result as { tools?: McpToolSchema[] })?.tools ?? [];\n\n    const tools = new Map<string, McpToolSchema>();\n    for (const tool of toolList) {\n      tools.set(tool.name, tool);\n    }\n\n    this.connections.set(endpoint.name, {\n      config: endpoint,\n      resolvedUrl,\n      forwardWorkspaceAuth: check.forwardWorkspaceAuth,\n      tools,\n    });\n    logger.info(\n      \"Connected to MCP server %s: %d tools available\",\n      endpoint.name,\n      tools.size,\n    );\n  }\n\n  getAllToolDefinitions(): AgentToolDefinition[] {\n    const defs: AgentToolDefinition[] = [];\n    for (const [serverName, conn] of this.connections) {\n      for (const [toolName, schema] of conn.tools) {\n        defs.push({\n          name: `mcp.${serverName}.${toolName}`,\n          description: schema.description ?? toolName,\n          parameters: coerceToolParameters(schema.inputSchema),\n        });\n      }\n    }\n    return defs;\n  }\n\n  /**\n   * Whether the named MCP server may receive workspace-scoped auth headers\n   * (e.g., an OBO bearer token from an end-user request). Callers should gate\n   * auth-forwarding decisions on this to prevent credential exfiltration to\n   * non-workspace hosts.\n   */\n  canForwardWorkspaceAuth(serverName: string): boolean {\n    return this.connections.get(serverName)?.forwardWorkspaceAuth ?? false;\n  }\n\n  async callTool(\n    qualifiedName: string,\n    args: unknown,\n    authHeaders?: Record<string, string>,\n    callerSignal?: AbortSignal,\n  ): Promise<string> {\n    const parts = qualifiedName.split(\".\");\n    if (parts.length < 3 || parts[0] !== \"mcp\") {\n      throw new Error(`Invalid MCP tool name: ${qualifiedName}`);\n    }\n    const serverName = parts[1];\n    const toolName = parts.slice(2).join(\".\");\n\n    const conn = this.connections.get(serverName);\n    if (!conn) {\n      throw new Error(`MCP server not connected: ${serverName}`);\n    }\n\n    const sessionId = this.sessionIds.get(serverName);\n    // authHeaders are caller-supplied credentials (typically the OBO token).\n    // Only honor them if the destination URL was admitted with\n    // forwardWorkspaceAuth=true at connect time.\n    const scopedAuthOverride = conn.forwardWorkspaceAuth\n      ? authHeaders\n      : undefined;\n\n    const rpcResult = await this.sendRpc(\n      conn.resolvedUrl,\n      \"tools/call\",\n      { name: toolName, arguments: args },\n      {\n        authOverride: scopedAuthOverride,\n        sessionId,\n        forwardWorkspaceAuth: conn.forwardWorkspaceAuth,\n        callerSignal,\n      },\n    );\n    const result = rpcResult.result as McpToolCallResult;\n\n    // `text` is optional on `McpToolCallResult.content[]` per the MCP\n    // spec; filtering only on `type === \"text\"` lets `c.text` be\n    // `undefined`, which `Array.join` would render as the literal\n    // string `\"undefined\"` and ship to the agent. Narrow on both\n    // fields so the joined string only contains real text.\n    const textContent = (result.content ?? []).filter(\n      (c): c is { type: \"text\"; text: string } =>\n        c.type === \"text\" && typeof c.text === \"string\",\n    );\n\n    if (result.isError) {\n      const errText = textContent.map((c) => c.text).join(\"\\n\");\n      throw new Error(errText || \"MCP tool call failed\");\n    }\n\n    return textContent.map((c) => c.text).join(\"\\n\");\n  }\n\n  async close(): Promise<void> {\n    this.closed = true;\n    this.connections.clear();\n    this.sessionIds.clear();\n  }\n\n  private async sendRpc(\n    url: string,\n    method: string,\n    params?: Record<string, unknown>,\n    options?: {\n      authOverride?: Record<string, string>;\n      sessionId?: string;\n      forwardWorkspaceAuth?: boolean;\n      /**\n       * Optional external abort signal (typically the agent's stream signal).\n       * Composed with the built-in 30 s timeout so `/cancel` or agent-run\n       * shutdown immediately propagates to the MCP fetch rather than waiting\n       * for the remote server to respond.\n       */\n      callerSignal?: AbortSignal;\n    },\n  ): Promise<{ result: unknown; sessionId?: string }> {\n    if (this.closed) throw new Error(\"MCP client is closed\");\n\n    const request: JsonRpcRequest = {\n      jsonrpc: \"2.0\",\n      id: ++this.requestId,\n      method,\n      ...(params && { params }),\n    };\n\n    const authHeaders = await this.resolveAuthHeaders(options);\n    const headers: Record<string, string> = {\n      \"User-Agent\": APPKIT_USER_AGENT,\n      \"Content-Type\": \"application/json\",\n      Accept: \"application/json, text/event-stream\",\n      ...authHeaders,\n    };\n    if (options?.sessionId) {\n      headers[\"Mcp-Session-Id\"] = options.sessionId;\n    }\n\n    const fetchImpl = this.options.fetchImpl ?? fetch;\n    const signals: AbortSignal[] = [AbortSignal.timeout(30_000)];\n    if (options?.callerSignal) signals.push(options.callerSignal);\n    const response = await fetchImpl(url, {\n      method: \"POST\",\n      headers,\n      body: JSON.stringify(request),\n      signal: signals.length > 1 ? AbortSignal.any(signals) : signals[0],\n    });\n\n    if (!response.ok) {\n      throw new Error(\n        `MCP request to ${method} failed: ${response.status} ${response.statusText}`,\n      );\n    }\n\n    const contentType = response.headers.get(\"content-type\") ?? \"\";\n    // Always read the body via the capped helper so a misconfigured or\n    // malicious server can't exhaust client memory by streaming an\n    // unbounded payload. Applies to both SSE (`response.text()` would\n    // have buffered the whole stream) and plain JSON (`response.json()`\n    // does the same internally).\n    const bodyText = await readResponseTextCapped(\n      response,\n      MCP_RESPONSE_BODY_LIMIT_BYTES,\n      method,\n    );\n    let json: JsonRpcResponse;\n\n    if (contentType.includes(\"text/event-stream\")) {\n      const lastData = bodyText\n        .split(\"\\n\")\n        .filter((line) => line.startsWith(\"data: \"))\n        .map((line) => line.slice(6))\n        .pop();\n      if (!lastData) {\n        throw new Error(`MCP SSE response for ${method} contained no data`);\n      }\n      json = JSON.parse(lastData) as JsonRpcResponse;\n    } else {\n      if (bodyText.length === 0) {\n        throw new Error(`MCP response for ${method} had an empty body`);\n      }\n      json = JSON.parse(bodyText) as JsonRpcResponse;\n    }\n\n    if (json.error) {\n      throw new Error(`MCP error (${json.error.code}): ${json.error.message}`);\n    }\n\n    const sid = response.headers.get(\"mcp-session-id\") ?? undefined;\n    return { result: json.result, sessionId: sid };\n  }\n\n  private async sendNotification(\n    url: string,\n    method: string,\n    options?: {\n      sessionId?: string;\n      forwardWorkspaceAuth?: boolean;\n    },\n  ): Promise<void> {\n    if (this.closed) return;\n\n    const authHeaders = await this.resolveAuthHeaders(options);\n    const headers: Record<string, string> = {\n      \"User-Agent\": APPKIT_USER_AGENT,\n      \"Content-Type\": \"application/json\",\n      Accept: \"application/json, text/event-stream\",\n      ...authHeaders,\n    };\n    if (options?.sessionId) {\n      headers[\"Mcp-Session-Id\"] = options.sessionId;\n    }\n\n    const fetchImpl = this.options.fetchImpl ?? fetch;\n    // MCP notifications are fire-and-forget per spec — we don't throw on\n    // failure. But silently swallowing 4xx/5xx hides server-side\n    // rejections that would otherwise look like a successful connect()\n    // followed by mysterious tool-call failures. Surface the bad status\n    // via the logger so the dev sees it without breaking the protocol\n    // contract.\n    try {\n      const response = await fetchImpl(url, {\n        method: \"POST\",\n        headers,\n        body: JSON.stringify({ jsonrpc: \"2.0\", method }),\n        signal: AbortSignal.timeout(30_000),\n      });\n      if (!response.ok) {\n        logger.warn(\n          \"MCP notification %s to %s returned %d %s — the server may have rejected the request, but per MCP spec notifications are fire-and-forget and the connection is considered established.\",\n          method,\n          url,\n          response.status,\n          response.statusText,\n        );\n      }\n    } catch (err) {\n      logger.warn(\n        \"MCP notification %s to %s failed before a response was received: %O\",\n        method,\n        url,\n        err,\n      );\n    }\n  }\n\n  /**\n   * Return the auth headers to send on an outbound request. Workspace auth\n   * (SP or OBO) is only resolved when `forwardWorkspaceAuth` is true; for\n   * non-workspace hosts no bearer token is attached.\n   */\n  private async resolveAuthHeaders(options?: {\n    authOverride?: Record<string, string>;\n    forwardWorkspaceAuth?: boolean;\n  }): Promise<Record<string, string>> {\n    if (!options?.forwardWorkspaceAuth) return {};\n    if (options.authOverride) return options.authOverride;\n    return this.authenticate();\n  }\n}\n"],"mappings":";;;;;AAmCA,MAAM,SAAS,aAAa,gBAAgB;;;;;;;;;AAU5C,MAAM,gCAAgC,OAAO;;;;;;;;;;;;;AAc7C,MAAM,wBAA2D;CAC/D,MAAM;CACN,YAAY,EAAE;CACf;;;;;;;;;;AAWD,SAAS,qBACP,aACmC;AACnC,KAAI,CAAC,eAAe,OAAO,gBAAgB,SACzC,QAAO;CAET,MAAM,EAAE,MAAM,eAAe;AAC7B,KAAI,SAAS,SAAU,QAAO;AAC9B,KACE,eAAe,WACd,OAAO,eAAe,YACrB,eAAe,QACf,MAAM,QAAQ,WAAW,EAE3B,QAAO;AAET,QAAO;;AAGT,eAAe,uBACb,UACA,UACA,cACiB;AACjB,KAAI,CAAC,SAAS,KAAM,QAAO;CAC3B,MAAM,SAAS,SAAS,KAAK,WAAW;CACxC,MAAM,UAAU,IAAI,YAAY,QAAQ;CACxC,IAAI,QAAQ;CACZ,IAAI,MAAM;AACV,KAAI;AACF,SAAO,MAAM;GACX,MAAM,EAAE,MAAM,UAAU,MAAM,OAAO,MAAM;AAC3C,OAAI,KAAM;AACV,YAAS,MAAM;AACf,OAAI,QAAQ,SACV,OAAM,IAAI,MACR,OAAO,aAAa,2BAA2B,SAAS,sEACzD;AAEH,UAAO,QAAQ,OAAO,OAAO,EAAE,QAAQ,MAAM,CAAC;;AAEhD,SAAO,QAAQ,QAAQ;WACf;AACR,SAAO,aAAa;;AAEtB,QAAO;;;;;;;;;;;;;;;;;;AAkET,IAAa,kBAAb,MAA6B;CAC3B,AAAQ,8BAAc,IAAI,KAAkC;CAC5D,AAAQ,6BAAa,IAAI,KAAqB;CAC9C,AAAQ,YAAY;CACpB,AAAQ,SAAS;CAEjB,YACE,AAAQ,eACR,AAAQ,cACR,AAAQ,QACR,AAAQ,UAA+D,EAAE,EACzE;EAJQ;EACA;EACA;EACA;;;;;;;;;;;;;;;CAgBV,MAAM,WACJ,WAC8B;EAC9B,MAAM,UAAU,MAAM,QAAQ,WAC5B,UAAU,KAAK,OAAO,KAAK,QAAQ,GAAG,CAAC,CACxC;EACD,MAAM,MAA2B;GAAE,WAAW,EAAE;GAAE,QAAQ,EAAE;GAAE;AAC9D,OAAK,IAAI,IAAI,GAAG,IAAI,QAAQ,QAAQ,KAAK;GACvC,MAAM,IAAI,QAAQ;GAClB,MAAM,OAAO,UAAU,GAAG;AAC1B,OAAI,EAAE,WAAW,YACf,KAAI,UAAU,KAAK,KAAK;QACnB;IACL,MAAM,QACJ,EAAE,kBAAkB,QAAQ,EAAE,SAAS,IAAI,MAAM,OAAO,EAAE,OAAO,CAAC;AACpE,WAAO,MAAM,uCAAuC,MAAM,MAAM;AAChE,QAAI,OAAO,KAAK;KAAE;KAAM;KAAO,CAAC;;;AAGpC,SAAO;;CAGT,AAAQ,WAAW,UAAqC;AACtD,MACE,SAAS,IAAI,WAAW,UAAU,IAClC,SAAS,IAAI,WAAW,WAAW,CAEnC,QAAO,SAAS;AAElB,SAAO,GAAG,KAAK,gBAAgB,SAAS;;CAG1C,MAAM,QAAQ,UAA4C;EACxD,MAAM,cAAc,KAAK,WAAW,SAAS;EAC7C,MAAM,QAAQ,YAAY,aAAa,KAAK,OAAO;AACnD,MAAI,CAAC,MAAM,GACT,OAAM,IAAI,MACR,iBAAiB,SAAS,KAAK,wBAAwB,MAAM,SAC9D;AAEH,QAAM,uBACJ,MAAM,IAAI,UACV,KAAK,QACL,KAAK,QAAQ,UACd;AAED,SAAO,KACL,gEACA,SAAS,MACT,aACA,MAAM,qBACP;EAED,MAAM,eAAe,MAAM,KAAK,QAC9B,aACA,cACA;GACE,iBAAiB;GACjB,cAAc,EAAE;GAChB,YAAY;IAAE,MAAM;IAAgB,SAAS;IAAS;GACvD,EACD,EAAE,sBAAsB,MAAM,sBAAsB,CACrD;AAED,MAAI,aAAa,UACf,MAAK,WAAW,IAAI,SAAS,MAAM,aAAa,UAAU;EAE5D,MAAM,YAAY,KAAK,WAAW,IAAI,SAAS,KAAK;AAEpD,QAAM,KAAK,iBAAiB,aAAa,6BAA6B;GACpE;GACA,sBAAsB,MAAM;GAC7B,CAAC;EAQF,MAAM,YANe,MAAM,KAAK,QAC9B,aACA,cACA,EAAE,EACF;GAAE;GAAW,sBAAsB,MAAM;GAAsB,CAChE,EAEe,QAAwC,SAAS,EAAE;EAEnE,MAAM,wBAAQ,IAAI,KAA4B;AAC9C,OAAK,MAAM,QAAQ,SACjB,OAAM,IAAI,KAAK,MAAM,KAAK;AAG5B,OAAK,YAAY,IAAI,SAAS,MAAM;GAClC,QAAQ;GACR;GACA,sBAAsB,MAAM;GAC5B;GACD,CAAC;AACF,SAAO,KACL,kDACA,SAAS,MACT,MAAM,KACP;;CAGH,wBAA+C;EAC7C,MAAM,OAA8B,EAAE;AACtC,OAAK,MAAM,CAAC,YAAY,SAAS,KAAK,YACpC,MAAK,MAAM,CAAC,UAAU,WAAW,KAAK,MACpC,MAAK,KAAK;GACR,MAAM,OAAO,WAAW,GAAG;GAC3B,aAAa,OAAO,eAAe;GACnC,YAAY,qBAAqB,OAAO,YAAY;GACrD,CAAC;AAGN,SAAO;;;;;;;;CAST,wBAAwB,YAA6B;AACnD,SAAO,KAAK,YAAY,IAAI,WAAW,EAAE,wBAAwB;;CAGnE,MAAM,SACJ,eACA,MACA,aACA,cACiB;EACjB,MAAM,QAAQ,cAAc,MAAM,IAAI;AACtC,MAAI,MAAM,SAAS,KAAK,MAAM,OAAO,MACnC,OAAM,IAAI,MAAM,0BAA0B,gBAAgB;EAE5D,MAAM,aAAa,MAAM;EACzB,MAAM,WAAW,MAAM,MAAM,EAAE,CAAC,KAAK,IAAI;EAEzC,MAAM,OAAO,KAAK,YAAY,IAAI,WAAW;AAC7C,MAAI,CAAC,KACH,OAAM,IAAI,MAAM,6BAA6B,aAAa;EAG5D,MAAM,YAAY,KAAK,WAAW,IAAI,WAAW;EAIjD,MAAM,qBAAqB,KAAK,uBAC5B,cACA;EAaJ,MAAM,UAXY,MAAM,KAAK,QAC3B,KAAK,aACL,cACA;GAAE,MAAM;GAAU,WAAW;GAAM,EACnC;GACE,cAAc;GACd;GACA,sBAAsB,KAAK;GAC3B;GACD,CACF,EACwB;EAOzB,MAAM,eAAe,OAAO,WAAW,EAAE,EAAE,QACxC,MACC,EAAE,SAAS,UAAU,OAAO,EAAE,SAAS,SAC1C;AAED,MAAI,OAAO,SAAS;GAClB,MAAM,UAAU,YAAY,KAAK,MAAM,EAAE,KAAK,CAAC,KAAK,KAAK;AACzD,SAAM,IAAI,MAAM,WAAW,uBAAuB;;AAGpD,SAAO,YAAY,KAAK,MAAM,EAAE,KAAK,CAAC,KAAK,KAAK;;CAGlD,MAAM,QAAuB;AAC3B,OAAK,SAAS;AACd,OAAK,YAAY,OAAO;AACxB,OAAK,WAAW,OAAO;;CAGzB,MAAc,QACZ,KACA,QACA,QACA,SAYkD;AAClD,MAAI,KAAK,OAAQ,OAAM,IAAI,MAAM,uBAAuB;EAExD,MAAM,UAA0B;GAC9B,SAAS;GACT,IAAI,EAAE,KAAK;GACX;GACA,GAAI,UAAU,EAAE,QAAQ;GACzB;EAED,MAAM,cAAc,MAAM,KAAK,mBAAmB,QAAQ;EAC1D,MAAM,UAAkC;GACtC,cAAc;GACd,gBAAgB;GAChB,QAAQ;GACR,GAAG;GACJ;AACD,MAAI,SAAS,UACX,SAAQ,oBAAoB,QAAQ;EAGtC,MAAM,YAAY,KAAK,QAAQ,aAAa;EAC5C,MAAM,UAAyB,CAAC,YAAY,QAAQ,IAAO,CAAC;AAC5D,MAAI,SAAS,aAAc,SAAQ,KAAK,QAAQ,aAAa;EAC7D,MAAM,WAAW,MAAM,UAAU,KAAK;GACpC,QAAQ;GACR;GACA,MAAM,KAAK,UAAU,QAAQ;GAC7B,QAAQ,QAAQ,SAAS,IAAI,YAAY,IAAI,QAAQ,GAAG,QAAQ;GACjE,CAAC;AAEF,MAAI,CAAC,SAAS,GACZ,OAAM,IAAI,MACR,kBAAkB,OAAO,WAAW,SAAS,OAAO,GAAG,SAAS,aACjE;EAGH,MAAM,cAAc,SAAS,QAAQ,IAAI,eAAe,IAAI;EAM5D,MAAM,WAAW,MAAM,uBACrB,UACA,+BACA,OACD;EACD,IAAI;AAEJ,MAAI,YAAY,SAAS,oBAAoB,EAAE;GAC7C,MAAM,WAAW,SACd,MAAM,KAAK,CACX,QAAQ,SAAS,KAAK,WAAW,SAAS,CAAC,CAC3C,KAAK,SAAS,KAAK,MAAM,EAAE,CAAC,CAC5B,KAAK;AACR,OAAI,CAAC,SACH,OAAM,IAAI,MAAM,wBAAwB,OAAO,oBAAoB;AAErE,UAAO,KAAK,MAAM,SAAS;SACtB;AACL,OAAI,SAAS,WAAW,EACtB,OAAM,IAAI,MAAM,oBAAoB,OAAO,oBAAoB;AAEjE,UAAO,KAAK,MAAM,SAAS;;AAG7B,MAAI,KAAK,MACP,OAAM,IAAI,MAAM,cAAc,KAAK,MAAM,KAAK,KAAK,KAAK,MAAM,UAAU;EAG1E,MAAM,MAAM,SAAS,QAAQ,IAAI,iBAAiB,IAAI;AACtD,SAAO;GAAE,QAAQ,KAAK;GAAQ,WAAW;GAAK;;CAGhD,MAAc,iBACZ,KACA,QACA,SAIe;AACf,MAAI,KAAK,OAAQ;EAEjB,MAAM,cAAc,MAAM,KAAK,mBAAmB,QAAQ;EAC1D,MAAM,UAAkC;GACtC,cAAc;GACd,gBAAgB;GAChB,QAAQ;GACR,GAAG;GACJ;AACD,MAAI,SAAS,UACX,SAAQ,oBAAoB,QAAQ;EAGtC,MAAM,YAAY,KAAK,QAAQ,aAAa;AAO5C,MAAI;GACF,MAAM,WAAW,MAAM,UAAU,KAAK;IACpC,QAAQ;IACR;IACA,MAAM,KAAK,UAAU;KAAE,SAAS;KAAO;KAAQ,CAAC;IAChD,QAAQ,YAAY,QAAQ,IAAO;IACpC,CAAC;AACF,OAAI,CAAC,SAAS,GACZ,QAAO,KACL,yLACA,QACA,KACA,SAAS,QACT,SAAS,WACV;WAEI,KAAK;AACZ,UAAO,KACL,uEACA,QACA,KACA,IACD;;;;;;;;CASL,MAAc,mBAAmB,SAGG;AAClC,MAAI,CAAC,SAAS,qBAAsB,QAAO,EAAE;AAC7C,MAAI,QAAQ,aAAc,QAAO,QAAQ;AACzC,SAAO,KAAK,cAAc"}

package/dist/context/client-options.js

@@ -0,0 +1,28 @@
+import { name, version } from "../appkit/package.js";
+import { coerce } from "semver";
+
+//#region src/context/client-options.ts
+const normalizedVersion = coerce(version)?.version ?? version;
+/**
+* SDK client options that stamp every `apiClient.request()` with an AppKit
+* User-Agent (`@databricks/appkit/<version>`), so outbound Databricks traffic
+* is attributable to AppKit. Use this for every `WorkspaceClient` AppKit
+* constructs at runtime.
+*/
+function getClientOptions() {
+	return {
+		product: name,
+		productVersion: normalizedVersion,
+		...process.env.NODE_ENV === "development" && { userAgentExtra: { mode: "dev" } }
+	};
+}
+/**
+* Product/version User-Agent string matching the SDK stamp, for raw `fetch`
+* call sites that bypass the SDK's `apiClient` and have no client to derive it
+* from (e.g. the MCP connector).
+*/
+const APPKIT_USER_AGENT = `${name}/${normalizedVersion}`;
+
+//#endregion
+export { APPKIT_USER_AGENT, getClientOptions };
+//# sourceMappingURL=client-options.js.map

package/dist/context/client-options.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"client-options.js","names":["productVersion","productName"],"sources":["../../src/context/client-options.ts"],"sourcesContent":["import type { ClientOptions } from \"@databricks/sdk-experimental\";\nimport { coerce } from \"semver\";\nimport {\n  name as productName,\n  version as productVersion,\n} from \"../../package.json\";\n\nconst normalizedVersion = (coerce(productVersion)?.version ??\n  productVersion) as ClientOptions[\"productVersion\"];\n\n/**\n * SDK client options that stamp every `apiClient.request()` with an AppKit\n * User-Agent (`@databricks/appkit/<version>`), so outbound Databricks traffic\n * is attributable to AppKit. Use this for every `WorkspaceClient` AppKit\n * constructs at runtime.\n */\nexport function getClientOptions(): ClientOptions {\n  const isDev = process.env.NODE_ENV === \"development\";\n\n  return {\n    product: productName,\n    productVersion: normalizedVersion,\n    ...(isDev && { userAgentExtra: { mode: \"dev\" } }),\n  };\n}\n\n/**\n * Product/version User-Agent string matching the SDK stamp, for raw `fetch`\n * call sites that bypass the SDK's `apiClient` and have no client to derive it\n * from (e.g. the MCP connector).\n */\nexport const APPKIT_USER_AGENT = `${productName}/${normalizedVersion}`;\n"],"mappings":";;;;AAOA,MAAM,oBAAqB,OAAOA,QAAe,EAAE,WACjDA;;;;;;;AAQF,SAAgB,mBAAkC;AAGhD,QAAO;EACL,SAASC;EACT,gBAAgB;EAChB,GALY,QAAQ,IAAI,aAAa,iBAKxB,EAAE,gBAAgB,EAAE,MAAM,OAAO,EAAE;EACjD;;;;;;;AAQH,MAAa,oBAAoB,GAAGA,KAAY,GAAG"}

package/dist/context/service-context.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"service-context.d.ts","names":[],"sources":["../../src/context/service-context.ts"],"mappings":";;;;;;AAuBA;;UAAiB,mBAAA;EAEP;EAAR,MAAA,EAAQ,eAAA;EAMK;EAJb,aAAA;EAIoB;EAFpB,WAAA,GAAc,OAAA;EAJN;EAMR,WAAA,EAAa,OAAA;AAAA"}
+{"version":3,"file":"service-context.d.ts","names":[],"sources":["../../src/context/service-context.ts"],"mappings":";;;;;;AAmBA;;UAAiB,mBAAA;EAEP;EAAR,MAAA,EAAQ,eAAA;EAMK;EAJb,aAAA;EAIoB;EAFpB,WAAA,GAAc,OAAA;EAJN;EAMR,WAAA,EAAa,OAAA;AAAA"}

package/dist/context/service-context.js

@@ -2,20 +2,11 @@ import { AuthenticationError } from "../errors/authentication.js";
 import { ConfigurationError } from "../errors/configuration.js";
 import { InitializationError } from "../errors/initialization.js";
 import "../errors/index.js";
-import { name, version } from "../appkit/package.js";
+import { getClientOptions } from "./client-options.js";
 import { createHash } from "node:crypto";
 import { ConfigError, WorkspaceClient } from "@databricks/sdk-experimental";
-import { coerce } from "semver";
 
 //#region src/context/service-context.ts
-function getClientOptions() {
-	const isDev = process.env.NODE_ENV === "development";
-	return {
-		product: name,
-		productVersion: coerce(version)?.version ?? version,
-		...isDev && { userAgentExtra: { mode: "dev" } }
-	};
-}
 /**
 * ServiceContext is a singleton that manages the service principal's
 * WorkspaceClient and shared resources like warehouse/workspace IDs.
@@ -126,7 +117,8 @@ var ServiceContext = class ServiceContext {
 	}
 	static async getWarehouseId(client) {
 		if (process.env.DATABRICKS_WAREHOUSE_ID) return process.env.DATABRICKS_WAREHOUSE_ID;
-		if (process.env.NODE_ENV === "development") {
+		const agenticMode = process.env.DATABRICKS_APPS_AGENTIC_MODE === "true" || process.env.DATABRICKS_APPS_AGENTIC_MODE === "1";
+		if (process.env.NODE_ENV === "development" && !agenticMode) {
 			const response = await client.apiClient.request({
 				path: "/api/2.0/sql/warehouses",
 				method: "GET",

package/dist/context/service-context.js.map

@@ -1 +1 @@
-{"version":3,"file":"service-context.js","names":["productName","productVersion"],"sources":["../../src/context/service-context.ts"],"sourcesContent":["import { createHash } from \"node:crypto\";\nimport {\n  type ClientOptions,\n  ConfigError,\n  type sql,\n  WorkspaceClient,\n} from \"@databricks/sdk-experimental\";\nimport { coerce } from \"semver\";\nimport {\n  name as productName,\n  version as productVersion,\n} from \"../../package.json\";\nimport {\n  AuthenticationError,\n  ConfigurationError,\n  InitializationError,\n} from \"../errors\";\nimport type { UserContext } from \"./user-context\";\n\n/**\n * Service context holds the service principal client and shared resources.\n * This is initialized once at app startup and shared across all requests.\n */\nexport interface ServiceContextState {\n  /** WorkspaceClient authenticated as the service principal */\n  client: WorkspaceClient;\n  /** The service principal's user ID */\n  serviceUserId: string;\n  /** Promise that resolves to the warehouse ID (only present when a plugin requires `SQL_WAREHOUSE` resource) */\n  warehouseId?: Promise<string>;\n  /** Promise that resolves to the workspace ID */\n  workspaceId: Promise<string>;\n}\n\nfunction getClientOptions(): ClientOptions {\n  const isDev = process.env.NODE_ENV === \"development\";\n  const semver = coerce(productVersion);\n  const normalizedVersion = (semver?.version ??\n    productVersion) as ClientOptions[\"productVersion\"];\n\n  return {\n    product: productName,\n    productVersion: normalizedVersion,\n    ...(isDev && { userAgentExtra: { mode: \"dev\" } }),\n  };\n}\n\n/**\n * ServiceContext is a singleton that manages the service principal's\n * WorkspaceClient and shared resources like warehouse/workspace IDs.\n *\n * It's initialized once at app startup and provides the foundation\n * for both service principal and user context execution.\n */\nexport class ServiceContext {\n  private static instance: ServiceContextState | null = null;\n  private static initPromise: Promise<ServiceContextState> | null = null;\n\n  /**\n   * Initialize the service context. Should be called once at app startup.\n   * Safe to call multiple times - will return the same instance.\n   *\n   * @param options - Which shared resources to resolve (derived from plugin manifests).\n   * @param client - Optional pre-configured WorkspaceClient to use instead\n   *   of creating one from environment credentials.\n   */\n  static async initialize(\n    options?: { warehouseId?: boolean },\n    client?: WorkspaceClient,\n  ): Promise<ServiceContextState> {\n    if (ServiceContext.instance) {\n      return ServiceContext.instance;\n    }\n\n    if (ServiceContext.initPromise) {\n      return ServiceContext.initPromise;\n    }\n\n    ServiceContext.initPromise = ServiceContext.createContext(options, client);\n    ServiceContext.instance = await ServiceContext.initPromise;\n    return ServiceContext.instance;\n  }\n\n  /**\n   * Get the initialized service context.\n   * @throws Error if not initialized\n   */\n  static get(): ServiceContextState {\n    if (!ServiceContext.instance) {\n      throw InitializationError.notInitialized(\n        \"ServiceContext\",\n        \"Call ServiceContext.initialize() first\",\n      );\n    }\n    return ServiceContext.instance;\n  }\n\n  /**\n   * Check if the service context has been initialized.\n   */\n  static isInitialized(): boolean {\n    return ServiceContext.instance !== null;\n  }\n\n  /**\n   * Create a user context from request headers.\n   *\n   * @param token - The user's access token from x-forwarded-access-token header\n   * @param userId - The user's ID from x-forwarded-user header\n   * @param userName - Optional user name\n   * @throws Error if token is not provided\n   */\n  static createUserContext(\n    token: string,\n    userId: string,\n    userName?: string,\n    userEmail?: string,\n  ): UserContext {\n    if (!token) {\n      throw AuthenticationError.missingToken(\"user token\");\n    }\n\n    const host = process.env.DATABRICKS_HOST;\n    if (!host) {\n      throw ConfigurationError.missingEnvVar(\"DATABRICKS_HOST\");\n    }\n\n    const serviceCtx = ServiceContext.get();\n\n    // Create user client with the OAuth token from Databricks Apps\n    // Note: We use authType: \"pat\" because the token is passed as a Bearer token\n    // just like a PAT, even though it's technically an OAuth token\n    const userClient = new WorkspaceClient(\n      {\n        token,\n        host,\n        authType: \"pat\",\n      },\n      getClientOptions(),\n    );\n\n    const tokenFingerprint = createHash(\"sha256\")\n      .update(token)\n      .digest(\"hex\")\n      .slice(0, 16);\n\n    return {\n      client: userClient,\n      userId,\n      userName,\n      userEmail,\n      tokenFingerprint,\n      warehouseId: serviceCtx.warehouseId,\n      workspaceId: serviceCtx.workspaceId,\n      isUserContext: true,\n    };\n  }\n\n  /**\n   * Get the client options for WorkspaceClient.\n   * Exposed for testing purposes.\n   */\n  static getClientOptions(): ClientOptions {\n    return getClientOptions();\n  }\n\n  private static async createContext(\n    options?: { warehouseId?: boolean },\n    client?: WorkspaceClient,\n  ): Promise<ServiceContextState> {\n    try {\n      const wsClient = client ?? new WorkspaceClient({}, getClientOptions());\n\n      const [resolvedWorkspaceId, currentUser, resolvedWarehouseId] =\n        await Promise.all([\n          ServiceContext.getWorkspaceId(wsClient),\n          wsClient.currentUser.me(),\n          options?.warehouseId\n            ? ServiceContext.getWarehouseId(wsClient)\n            : Promise.resolve(undefined as string | undefined),\n        ]);\n\n      if (!currentUser.id) {\n        throw ConfigurationError.resourceNotFound(\"Service user ID\");\n      }\n\n      const warehouseId =\n        options?.warehouseId && resolvedWarehouseId !== undefined\n          ? Promise.resolve(resolvedWarehouseId)\n          : undefined;\n\n      return {\n        client: wsClient,\n        serviceUserId: currentUser.id,\n        warehouseId,\n        workspaceId: Promise.resolve(resolvedWorkspaceId),\n      };\n    } catch (e) {\n      if (e instanceof ConfigError) {\n        throw ConfigurationError.databricksAuthenticationSetupFailed(\n          e.baseMessage,\n          { cause: e },\n        );\n      }\n      throw e;\n    }\n  }\n\n  private static async getWorkspaceId(\n    client: WorkspaceClient,\n  ): Promise<string> {\n    if (process.env.DATABRICKS_WORKSPACE_ID) {\n      return process.env.DATABRICKS_WORKSPACE_ID;\n    }\n\n    const response = (await client.apiClient.request({\n      path: \"/api/2.0/preview/scim/v2/Me\",\n      method: \"GET\",\n      headers: new Headers(),\n      raw: false,\n      query: {},\n      responseHeaders: [\"x-databricks-org-id\"],\n    })) as { \"x-databricks-org-id\": string };\n\n    if (!response[\"x-databricks-org-id\"]) {\n      throw ConfigurationError.resourceNotFound(\"Workspace ID\");\n    }\n\n    return response[\"x-databricks-org-id\"];\n  }\n\n  private static async getWarehouseId(\n    client: WorkspaceClient,\n  ): Promise<string> {\n    if (process.env.DATABRICKS_WAREHOUSE_ID) {\n      return process.env.DATABRICKS_WAREHOUSE_ID;\n    }\n\n    if (process.env.NODE_ENV === \"development\") {\n      const response = (await client.apiClient.request({\n        path: \"/api/2.0/sql/warehouses\",\n        method: \"GET\",\n        headers: new Headers(),\n        raw: false,\n        query: {\n          skip_cannot_use: \"true\",\n        },\n      })) as { warehouses: sql.EndpointInfo[] };\n\n      const priorities: Record<sql.State, number> = {\n        RUNNING: 0,\n        STOPPED: 1,\n        STARTING: 2,\n        STOPPING: 3,\n        DELETED: 99,\n        DELETING: 99,\n      };\n\n      const warehouses = (response.warehouses || []).sort((a, b) => {\n        return (\n          priorities[a.state as sql.State] - priorities[b.state as sql.State]\n        );\n      });\n\n      if (response.warehouses.length === 0) {\n        throw ConfigurationError.resourceNotFound(\n          \"Warehouse ID\",\n          \"Please configure the DATABRICKS_WAREHOUSE_ID environment variable\",\n        );\n      }\n\n      const firstWarehouse = warehouses[0];\n      if (\n        firstWarehouse.state === \"DELETED\" ||\n        firstWarehouse.state === \"DELETING\" ||\n        !firstWarehouse.id\n      ) {\n        throw ConfigurationError.resourceNotFound(\n          \"Warehouse ID\",\n          \"Please configure the DATABRICKS_WAREHOUSE_ID environment variable\",\n        );\n      }\n\n      return firstWarehouse.id;\n    }\n\n    throw ConfigurationError.resourceNotFound(\n      \"Warehouse ID\",\n      \"Please configure the DATABRICKS_WAREHOUSE_ID environment variable\",\n    );\n  }\n\n  /**\n   * Reset the service context. Only for testing purposes.\n   */\n  static reset(): void {\n    ServiceContext.instance = null;\n    ServiceContext.initPromise = null;\n  }\n}\n"],"mappings":";;;;;;;;;;AAkCA,SAAS,mBAAkC;CACzC,MAAM,QAAQ,QAAQ,IAAI,aAAa;AAKvC,QAAO;EACL,SAASA;EACT,gBANa,OAAOC,QAAe,EACF,WACjCA;EAKA,GAAI,SAAS,EAAE,gBAAgB,EAAE,MAAM,OAAO,EAAE;EACjD;;;;;;;;;AAUH,IAAa,iBAAb,MAAa,eAAe;CAC1B,OAAe,WAAuC;CACtD,OAAe,cAAmD;;;;;;;;;CAUlE,aAAa,WACX,SACA,QAC8B;AAC9B,MAAI,eAAe,SACjB,QAAO,eAAe;AAGxB,MAAI,eAAe,YACjB,QAAO,eAAe;AAGxB,iBAAe,cAAc,eAAe,cAAc,SAAS,OAAO;AAC1E,iBAAe,WAAW,MAAM,eAAe;AAC/C,SAAO,eAAe;;;;;;CAOxB,OAAO,MAA2B;AAChC,MAAI,CAAC,eAAe,SAClB,OAAM,oBAAoB,eACxB,kBACA,yCACD;AAEH,SAAO,eAAe;;;;;CAMxB,OAAO,gBAAyB;AAC9B,SAAO,eAAe,aAAa;;;;;;;;;;CAWrC,OAAO,kBACL,OACA,QACA,UACA,WACa;AACb,MAAI,CAAC,MACH,OAAM,oBAAoB,aAAa,aAAa;EAGtD,MAAM,OAAO,QAAQ,IAAI;AACzB,MAAI,CAAC,KACH,OAAM,mBAAmB,cAAc,kBAAkB;EAG3D,MAAM,aAAa,eAAe,KAAK;AAmBvC,SAAO;GACL,QAfiB,IAAI,gBACrB;IACE;IACA;IACA,UAAU;IACX,EACD,kBAAkB,CACnB;GASC;GACA;GACA;GACA,kBAVuB,WAAW,SAAS,CAC1C,OAAO,MAAM,CACb,OAAO,MAAM,CACb,MAAM,GAAG,GAAG;GAQb,aAAa,WAAW;GACxB,aAAa,WAAW;GACxB,eAAe;GAChB;;;;;;CAOH,OAAO,mBAAkC;AACvC,SAAO,kBAAkB;;CAG3B,aAAqB,cACnB,SACA,QAC8B;AAC9B,MAAI;GACF,MAAM,WAAW,UAAU,IAAI,gBAAgB,EAAE,EAAE,kBAAkB,CAAC;GAEtE,MAAM,CAAC,qBAAqB,aAAa,uBACvC,MAAM,QAAQ,IAAI;IAChB,eAAe,eAAe,SAAS;IACvC,SAAS,YAAY,IAAI;IACzB,SAAS,cACL,eAAe,eAAe,SAAS,GACvC,QAAQ,QAAQ,OAAgC;IACrD,CAAC;AAEJ,OAAI,CAAC,YAAY,GACf,OAAM,mBAAmB,iBAAiB,kBAAkB;GAG9D,MAAM,cACJ,SAAS,eAAe,wBAAwB,SAC5C,QAAQ,QAAQ,oBAAoB,GACpC;AAEN,UAAO;IACL,QAAQ;IACR,eAAe,YAAY;IAC3B;IACA,aAAa,QAAQ,QAAQ,oBAAoB;IAClD;WACM,GAAG;AACV,OAAI,aAAa,YACf,OAAM,mBAAmB,oCACvB,EAAE,aACF,EAAE,OAAO,GAAG,CACb;AAEH,SAAM;;;CAIV,aAAqB,eACnB,QACiB;AACjB,MAAI,QAAQ,IAAI,wBACd,QAAO,QAAQ,IAAI;EAGrB,MAAM,WAAY,MAAM,OAAO,UAAU,QAAQ;GAC/C,MAAM;GACN,QAAQ;GACR,SAAS,IAAI,SAAS;GACtB,KAAK;GACL,OAAO,EAAE;GACT,iBAAiB,CAAC,sBAAsB;GACzC,CAAC;AAEF,MAAI,CAAC,SAAS,uBACZ,OAAM,mBAAmB,iBAAiB,eAAe;AAG3D,SAAO,SAAS;;CAGlB,aAAqB,eACnB,QACiB;AACjB,MAAI,QAAQ,IAAI,wBACd,QAAO,QAAQ,IAAI;AAGrB,MAAI,QAAQ,IAAI,aAAa,eAAe;GAC1C,MAAM,WAAY,MAAM,OAAO,UAAU,QAAQ;IAC/C,MAAM;IACN,QAAQ;IACR,SAAS,IAAI,SAAS;IACtB,KAAK;IACL,OAAO,EACL,iBAAiB,QAClB;IACF,CAAC;GAEF,MAAM,aAAwC;IAC5C,SAAS;IACT,SAAS;IACT,UAAU;IACV,UAAU;IACV,SAAS;IACT,UAAU;IACX;GAED,MAAM,cAAc,SAAS,cAAc,EAAE,EAAE,MAAM,GAAG,MAAM;AAC5D,WACE,WAAW,EAAE,SAAsB,WAAW,EAAE;KAElD;AAEF,OAAI,SAAS,WAAW,WAAW,EACjC,OAAM,mBAAmB,iBACvB,gBACA,oEACD;GAGH,MAAM,iBAAiB,WAAW;AAClC,OACE,eAAe,UAAU,aACzB,eAAe,UAAU,cACzB,CAAC,eAAe,GAEhB,OAAM,mBAAmB,iBACvB,gBACA,oEACD;AAGH,UAAO,eAAe;;AAGxB,QAAM,mBAAmB,iBACvB,gBACA,oEACD;;;;;CAMH,OAAO,QAAc;AACnB,iBAAe,WAAW;AAC1B,iBAAe,cAAc"}
+{"version":3,"file":"service-context.js","names":[],"sources":["../../src/context/service-context.ts"],"sourcesContent":["import { createHash } from \"node:crypto\";\nimport {\n  type ClientOptions,\n  ConfigError,\n  type sql,\n  WorkspaceClient,\n} from \"@databricks/sdk-experimental\";\nimport {\n  AuthenticationError,\n  ConfigurationError,\n  InitializationError,\n} from \"../errors\";\nimport { getClientOptions } from \"./client-options\";\nimport type { UserContext } from \"./user-context\";\n\n/**\n * Service context holds the service principal client and shared resources.\n * This is initialized once at app startup and shared across all requests.\n */\nexport interface ServiceContextState {\n  /** WorkspaceClient authenticated as the service principal */\n  client: WorkspaceClient;\n  /** The service principal's user ID */\n  serviceUserId: string;\n  /** Promise that resolves to the warehouse ID (only present when a plugin requires `SQL_WAREHOUSE` resource) */\n  warehouseId?: Promise<string>;\n  /** Promise that resolves to the workspace ID */\n  workspaceId: Promise<string>;\n}\n\n/**\n * ServiceContext is a singleton that manages the service principal's\n * WorkspaceClient and shared resources like warehouse/workspace IDs.\n *\n * It's initialized once at app startup and provides the foundation\n * for both service principal and user context execution.\n */\nexport class ServiceContext {\n  private static instance: ServiceContextState | null = null;\n  private static initPromise: Promise<ServiceContextState> | null = null;\n\n  /**\n   * Initialize the service context. Should be called once at app startup.\n   * Safe to call multiple times - will return the same instance.\n   *\n   * @param options - Which shared resources to resolve (derived from plugin manifests).\n   * @param client - Optional pre-configured WorkspaceClient to use instead\n   *   of creating one from environment credentials.\n   */\n  static async initialize(\n    options?: { warehouseId?: boolean },\n    client?: WorkspaceClient,\n  ): Promise<ServiceContextState> {\n    if (ServiceContext.instance) {\n      return ServiceContext.instance;\n    }\n\n    if (ServiceContext.initPromise) {\n      return ServiceContext.initPromise;\n    }\n\n    ServiceContext.initPromise = ServiceContext.createContext(options, client);\n    ServiceContext.instance = await ServiceContext.initPromise;\n    return ServiceContext.instance;\n  }\n\n  /**\n   * Get the initialized service context.\n   * @throws Error if not initialized\n   */\n  static get(): ServiceContextState {\n    if (!ServiceContext.instance) {\n      throw InitializationError.notInitialized(\n        \"ServiceContext\",\n        \"Call ServiceContext.initialize() first\",\n      );\n    }\n    return ServiceContext.instance;\n  }\n\n  /**\n   * Check if the service context has been initialized.\n   */\n  static isInitialized(): boolean {\n    return ServiceContext.instance !== null;\n  }\n\n  /**\n   * Create a user context from request headers.\n   *\n   * @param token - The user's access token from x-forwarded-access-token header\n   * @param userId - The user's ID from x-forwarded-user header\n   * @param userName - Optional user name\n   * @throws Error if token is not provided\n   */\n  static createUserContext(\n    token: string,\n    userId: string,\n    userName?: string,\n    userEmail?: string,\n  ): UserContext {\n    if (!token) {\n      throw AuthenticationError.missingToken(\"user token\");\n    }\n\n    const host = process.env.DATABRICKS_HOST;\n    if (!host) {\n      throw ConfigurationError.missingEnvVar(\"DATABRICKS_HOST\");\n    }\n\n    const serviceCtx = ServiceContext.get();\n\n    // Create user client with the OAuth token from Databricks Apps\n    // Note: We use authType: \"pat\" because the token is passed as a Bearer token\n    // just like a PAT, even though it's technically an OAuth token\n    const userClient = new WorkspaceClient(\n      {\n        token,\n        host,\n        authType: \"pat\",\n      },\n      getClientOptions(),\n    );\n\n    const tokenFingerprint = createHash(\"sha256\")\n      .update(token)\n      .digest(\"hex\")\n      .slice(0, 16);\n\n    return {\n      client: userClient,\n      userId,\n      userName,\n      userEmail,\n      tokenFingerprint,\n      warehouseId: serviceCtx.warehouseId,\n      workspaceId: serviceCtx.workspaceId,\n      isUserContext: true,\n    };\n  }\n\n  /**\n   * Get the client options for WorkspaceClient.\n   * Exposed for testing purposes.\n   */\n  static getClientOptions(): ClientOptions {\n    return getClientOptions();\n  }\n\n  private static async createContext(\n    options?: { warehouseId?: boolean },\n    client?: WorkspaceClient,\n  ): Promise<ServiceContextState> {\n    try {\n      const wsClient = client ?? new WorkspaceClient({}, getClientOptions());\n\n      const [resolvedWorkspaceId, currentUser, resolvedWarehouseId] =\n        await Promise.all([\n          ServiceContext.getWorkspaceId(wsClient),\n          wsClient.currentUser.me(),\n          options?.warehouseId\n            ? ServiceContext.getWarehouseId(wsClient)\n            : Promise.resolve(undefined as string | undefined),\n        ]);\n\n      if (!currentUser.id) {\n        throw ConfigurationError.resourceNotFound(\"Service user ID\");\n      }\n\n      const warehouseId =\n        options?.warehouseId && resolvedWarehouseId !== undefined\n          ? Promise.resolve(resolvedWarehouseId)\n          : undefined;\n\n      return {\n        client: wsClient,\n        serviceUserId: currentUser.id,\n        warehouseId,\n        workspaceId: Promise.resolve(resolvedWorkspaceId),\n      };\n    } catch (e) {\n      if (e instanceof ConfigError) {\n        throw ConfigurationError.databricksAuthenticationSetupFailed(\n          e.baseMessage,\n          { cause: e },\n        );\n      }\n      throw e;\n    }\n  }\n\n  private static async getWorkspaceId(\n    client: WorkspaceClient,\n  ): Promise<string> {\n    if (process.env.DATABRICKS_WORKSPACE_ID) {\n      return process.env.DATABRICKS_WORKSPACE_ID;\n    }\n\n    const response = (await client.apiClient.request({\n      path: \"/api/2.0/preview/scim/v2/Me\",\n      method: \"GET\",\n      headers: new Headers(),\n      raw: false,\n      query: {},\n      responseHeaders: [\"x-databricks-org-id\"],\n    })) as { \"x-databricks-org-id\": string };\n\n    if (!response[\"x-databricks-org-id\"]) {\n      throw ConfigurationError.resourceNotFound(\"Workspace ID\");\n    }\n\n    return response[\"x-databricks-org-id\"];\n  }\n\n  private static async getWarehouseId(\n    client: WorkspaceClient,\n  ): Promise<string> {\n    if (process.env.DATABRICKS_WAREHOUSE_ID) {\n      return process.env.DATABRICKS_WAREHOUSE_ID;\n    }\n\n    const agenticMode =\n      process.env.DATABRICKS_APPS_AGENTIC_MODE === \"true\" ||\n      process.env.DATABRICKS_APPS_AGENTIC_MODE === \"1\";\n\n    if (process.env.NODE_ENV === \"development\" && !agenticMode) {\n      const response = (await client.apiClient.request({\n        path: \"/api/2.0/sql/warehouses\",\n        method: \"GET\",\n        headers: new Headers(),\n        raw: false,\n        query: {\n          skip_cannot_use: \"true\",\n        },\n      })) as { warehouses: sql.EndpointInfo[] };\n\n      const priorities: Record<sql.State, number> = {\n        RUNNING: 0,\n        STOPPED: 1,\n        STARTING: 2,\n        STOPPING: 3,\n        DELETED: 99,\n        DELETING: 99,\n      };\n\n      const warehouses = (response.warehouses || []).sort((a, b) => {\n        return (\n          priorities[a.state as sql.State] - priorities[b.state as sql.State]\n        );\n      });\n\n      if (response.warehouses.length === 0) {\n        throw ConfigurationError.resourceNotFound(\n          \"Warehouse ID\",\n          \"Please configure the DATABRICKS_WAREHOUSE_ID environment variable\",\n        );\n      }\n\n      const firstWarehouse = warehouses[0];\n      if (\n        firstWarehouse.state === \"DELETED\" ||\n        firstWarehouse.state === \"DELETING\" ||\n        !firstWarehouse.id\n      ) {\n        throw ConfigurationError.resourceNotFound(\n          \"Warehouse ID\",\n          \"Please configure the DATABRICKS_WAREHOUSE_ID environment variable\",\n        );\n      }\n\n      return firstWarehouse.id;\n    }\n\n    throw ConfigurationError.resourceNotFound(\n      \"Warehouse ID\",\n      \"Please configure the DATABRICKS_WAREHOUSE_ID environment variable\",\n    );\n  }\n\n  /**\n   * Reset the service context. Only for testing purposes.\n   */\n  static reset(): void {\n    ServiceContext.instance = null;\n    ServiceContext.initPromise = null;\n  }\n}\n"],"mappings":";;;;;;;;;;;;;;;;AAqCA,IAAa,iBAAb,MAAa,eAAe;CAC1B,OAAe,WAAuC;CACtD,OAAe,cAAmD;;;;;;;;;CAUlE,aAAa,WACX,SACA,QAC8B;AAC9B,MAAI,eAAe,SACjB,QAAO,eAAe;AAGxB,MAAI,eAAe,YACjB,QAAO,eAAe;AAGxB,iBAAe,cAAc,eAAe,cAAc,SAAS,OAAO;AAC1E,iBAAe,WAAW,MAAM,eAAe;AAC/C,SAAO,eAAe;;;;;;CAOxB,OAAO,MAA2B;AAChC,MAAI,CAAC,eAAe,SAClB,OAAM,oBAAoB,eACxB,kBACA,yCACD;AAEH,SAAO,eAAe;;;;;CAMxB,OAAO,gBAAyB;AAC9B,SAAO,eAAe,aAAa;;;;;;;;;;CAWrC,OAAO,kBACL,OACA,QACA,UACA,WACa;AACb,MAAI,CAAC,MACH,OAAM,oBAAoB,aAAa,aAAa;EAGtD,MAAM,OAAO,QAAQ,IAAI;AACzB,MAAI,CAAC,KACH,OAAM,mBAAmB,cAAc,kBAAkB;EAG3D,MAAM,aAAa,eAAe,KAAK;AAmBvC,SAAO;GACL,QAfiB,IAAI,gBACrB;IACE;IACA;IACA,UAAU;IACX,EACD,kBAAkB,CACnB;GASC;GACA;GACA;GACA,kBAVuB,WAAW,SAAS,CAC1C,OAAO,MAAM,CACb,OAAO,MAAM,CACb,MAAM,GAAG,GAAG;GAQb,aAAa,WAAW;GACxB,aAAa,WAAW;GACxB,eAAe;GAChB;;;;;;CAOH,OAAO,mBAAkC;AACvC,SAAO,kBAAkB;;CAG3B,aAAqB,cACnB,SACA,QAC8B;AAC9B,MAAI;GACF,MAAM,WAAW,UAAU,IAAI,gBAAgB,EAAE,EAAE,kBAAkB,CAAC;GAEtE,MAAM,CAAC,qBAAqB,aAAa,uBACvC,MAAM,QAAQ,IAAI;IAChB,eAAe,eAAe,SAAS;IACvC,SAAS,YAAY,IAAI;IACzB,SAAS,cACL,eAAe,eAAe,SAAS,GACvC,QAAQ,QAAQ,OAAgC;IACrD,CAAC;AAEJ,OAAI,CAAC,YAAY,GACf,OAAM,mBAAmB,iBAAiB,kBAAkB;GAG9D,MAAM,cACJ,SAAS,eAAe,wBAAwB,SAC5C,QAAQ,QAAQ,oBAAoB,GACpC;AAEN,UAAO;IACL,QAAQ;IACR,eAAe,YAAY;IAC3B;IACA,aAAa,QAAQ,QAAQ,oBAAoB;IAClD;WACM,GAAG;AACV,OAAI,aAAa,YACf,OAAM,mBAAmB,oCACvB,EAAE,aACF,EAAE,OAAO,GAAG,CACb;AAEH,SAAM;;;CAIV,aAAqB,eACnB,QACiB;AACjB,MAAI,QAAQ,IAAI,wBACd,QAAO,QAAQ,IAAI;EAGrB,MAAM,WAAY,MAAM,OAAO,UAAU,QAAQ;GAC/C,MAAM;GACN,QAAQ;GACR,SAAS,IAAI,SAAS;GACtB,KAAK;GACL,OAAO,EAAE;GACT,iBAAiB,CAAC,sBAAsB;GACzC,CAAC;AAEF,MAAI,CAAC,SAAS,uBACZ,OAAM,mBAAmB,iBAAiB,eAAe;AAG3D,SAAO,SAAS;;CAGlB,aAAqB,eACnB,QACiB;AACjB,MAAI,QAAQ,IAAI,wBACd,QAAO,QAAQ,IAAI;EAGrB,MAAM,cACJ,QAAQ,IAAI,iCAAiC,UAC7C,QAAQ,IAAI,iCAAiC;AAE/C,MAAI,QAAQ,IAAI,aAAa,iBAAiB,CAAC,aAAa;GAC1D,MAAM,WAAY,MAAM,OAAO,UAAU,QAAQ;IAC/C,MAAM;IACN,QAAQ;IACR,SAAS,IAAI,SAAS;IACtB,KAAK;IACL,OAAO,EACL,iBAAiB,QAClB;IACF,CAAC;GAEF,MAAM,aAAwC;IAC5C,SAAS;IACT,SAAS;IACT,UAAU;IACV,UAAU;IACV,SAAS;IACT,UAAU;IACX;GAED,MAAM,cAAc,SAAS,cAAc,EAAE,EAAE,MAAM,GAAG,MAAM;AAC5D,WACE,WAAW,EAAE,SAAsB,WAAW,EAAE;KAElD;AAEF,OAAI,SAAS,WAAW,WAAW,EACjC,OAAM,mBAAmB,iBACvB,gBACA,oEACD;GAGH,MAAM,iBAAiB,WAAW;AAClC,OACE,eAAe,UAAU,aACzB,eAAe,UAAU,cACzB,CAAC,eAAe,GAEhB,OAAM,mBAAmB,iBACvB,gBACA,oEACD;AAGH,UAAO,eAAe;;AAGxB,QAAM,mBAAmB,iBACvB,gBACA,oEACD;;;;;CAMH,OAAO,QAAc;AACnB,iBAAe,WAAW;AAC1B,iBAAe,cAAc"}

package/dist/plugins/lakebase/lakebase.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"lakebase.d.ts","names":[],"sources":["../../../src/plugins/lakebase/lakebase.ts"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;;;AA+DA;;;;;;;;;;;;;;;;cAAa,cAAA,SAAuB,MAAA,YAAkB,YAAA;EA6O5B;EAAA,OA3OjB,QAAA,EAAuB,cAAA;EAAA,UAEZ,MAAA,EAAQ,eAAA;EAAA,QAClB,IAAA;EAAA,QACA,cAAA;;;;;;;;;EAUF,KAAA,CAAA,GAAK,OAAA;EAhBqD;;;;;;;;;;;;;;;;;;EAiE1D,KAAA,WAAgB,cAAA,OAAA,CACpB,IAAA,UACA,MAAA,eACC,OAAA,CAAQ,WAAA,CAAY,CAAA;EAmBT;;;;;;;;;;;;;;EAAA,QAAA,oBAAA;EAsHd;;;;EA/FA,qBAAA,CAAA;EAuGQ;;;;;EAAA,QA7EA,KAAA;cAEI,MAAA,EAAQ,eAAA;EAAA,QAYZ,cAAA;EA2CR,aAAA,CAAA,GAAiB,mBAAA;EAIX,gBAAA,CACJ,IAAA,UACA,IAAA,WACA,MAAA,GAAS,WAAA,GACR,OAAA;EAIH,OAAA,CAAQ,IAAA,GAJE,cAAA,GAIoD,MAAA,SAAA,YAAA;EAzIvC;;;;EAAA,QAiJf,gBAAA;;;;;;;;;;;;;EAqBR,OAAA,CAAA;UAGwB,YAAA;sBA5KF,cAAA,QAAc,IAAA,UACtB,MAAA,iBAEX,OAAA,CAAQ,WAAA,CAAY,CAAA;;;;;;;;;;iBA/EH,QAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;cAmQT,QAAA,EAAQ,QAAA,QAAA,cAAA,EAAA,eAAA"}
+{"version":3,"file":"lakebase.d.ts","names":[],"sources":["../../../src/plugins/lakebase/lakebase.ts"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;;;AAiEA;;;;;;;;;;;;;;;;cAAa,cAAA,SAAuB,MAAA,YAAkB,YAAA;EAkP5B;EAAA,OAhPjB,QAAA,EAAuB,cAAA;EAAA,UAEZ,MAAA,EAAQ,eAAA;EAAA,QAClB,IAAA;EAAA,QACA,cAAA;;;;;;;;;EAUF,KAAA,CAAA,GAAK,OAAA;EAhBqD;;;;;;;;;;;;;;;;;;EAsE1D,KAAA,WAAgB,cAAA,OAAA,CACpB,IAAA,UACA,MAAA,eACC,OAAA,CAAQ,WAAA,CAAY,CAAA;EAmBT;;;;;;;;;;;;;;EAAA,QAAA,oBAAA;EAsHd;;;;EA/FA,qBAAA,CAAA;EAuGQ;;;;;EAAA,QA7EA,KAAA;cAEI,MAAA,EAAQ,eAAA;EAAA,QAYZ,cAAA;EA2CR,aAAA,CAAA,GAAiB,mBAAA;EAIX,gBAAA,CACJ,IAAA,UACA,IAAA,WACA,MAAA,GAAS,WAAA,GACR,OAAA;EAIH,OAAA,CAAQ,IAAA,GAJE,cAAA,GAIoD,MAAA,SAAA,YAAA;EAzIvC;;;;EAAA,QAiJf,gBAAA;;;;;;;;;;;;;EAqBR,OAAA,CAAA;UAGwB,YAAA;sBA5KF,cAAA,QAAc,IAAA,UACtB,MAAA,iBAEX,OAAA,CAAQ,WAAA,CAAY,CAAA;;;;;;;;;;iBA1Ff,QAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;cA8QG,QAAA,EAAQ,QAAA,QAAA,cAAA,EAAA,eAAA"}

package/dist/plugins/lakebase/lakebase.js

@@ -1,5 +1,6 @@
 import { createLogger } from "../../logging/logger.js";
 import { createLakebasePoolManager } from "../../connectors/lakebase/pool-manager.js";
+import { getClientOptions } from "../../context/client-options.js";
 import { getUserContext } from "../../context/execution-context.js";
 import { RoutingPool } from "../../connectors/lakebase/routing-pool.js";
 import { createLakebasePool, getLakebaseOrmConfig, getLakebasePgConfig, getUsernameWithApiLookup } from "../../connectors/lakebase/index.js";
@@ -10,6 +11,7 @@ import { buildToolkitEntries } from "../../core/agent/build-toolkit.js";
 import { defineTool, executeFromRegistry, toolsFromRegistry } from "../../core/agent/tools/define-tool.js";
 import { assertReadOnlySql } from "../../core/agent/tools/sql-policy.js";
 import manifest_default from "./manifest.js";
+import { WorkspaceClient } from "@databricks/sdk-experimental";
 import { z } from "zod";
 
 //#region src/plugins/lakebase/lakebase.ts
@@ -61,7 +63,10 @@ var LakebasePlugin = class extends Plugin {
 	* context (set by `Plugin.asUser(req)` via AsyncLocalStorage).
 	*/
 	async setup() {
-		const poolConfig = this.config.pool;
+		const poolConfig = {
+			...this.config.pool,
+			workspaceClient: this.config.pool?.workspaceClient ?? new WorkspaceClient({}, getClientOptions())
+		};
 		const user = await getUsernameWithApiLookup(poolConfig);
 		const spPool = createLakebasePool({
 			...poolConfig,

package/dist/plugins/lakebase/lakebase.js.map

@@ -1 +1 @@
-{"version":3,"file":"lakebase.js","names":["manifest"],"sources":["../../../src/plugins/lakebase/lakebase.ts"],"sourcesContent":["import type { QueryResult, QueryResultRow } from \"pg\";\nimport type { AgentToolDefinition, ToolProvider } from \"shared\";\nimport { z } from \"zod\";\nimport {\n  createLakebasePool,\n  createLakebasePoolManager,\n  getLakebaseOrmConfig,\n  getLakebasePgConfig,\n  getUsernameWithApiLookup,\n  type LakebasePool,\n  type LakebasePoolManager,\n  RoutingPool,\n} from \"../../connectors/lakebase\";\nimport { getUserContext } from \"../../context/execution-context\";\nimport { buildToolkitEntries } from \"../../core/agent/build-toolkit\";\nimport {\n  defineTool,\n  executeFromRegistry,\n  toolsFromRegistry,\n} from \"../../core/agent/tools/define-tool\";\nimport { assertReadOnlySql } from \"../../core/agent/tools/sql-policy\";\nimport { createLogger } from \"../../logging/logger\";\nimport { Plugin, toPlugin } from \"../../plugin\";\nimport type { PluginManifest } from \"../../registry\";\nimport manifest from \"./manifest.json\";\nimport type { ILakebaseConfig } from \"./types\";\n\nconst logger = createLogger(\"lakebase\");\n\n/** Default pool settings for per-user OBO pools. */\nconst OBO_POOL_DEFAULTS = {\n  max: 3,\n  allowExitOnIdle: true,\n  idleTimeoutMillis: 30_000,\n};\n\n/**\n * AppKit plugin for Databricks Lakebase Autoscaling.\n *\n * Wraps `@databricks/lakebase` to provide a standard `pg.Pool` with automatic\n * OAuth token refresh, integrated with AppKit's logger and OpenTelemetry setup.\n *\n * Supports On-Behalf-Of (OBO) via `asUser(req)` — each user gets a separate\n * `pg.Pool` authenticated with their Databricks identity, enabling features\n * like Row-Level Security (RLS). Routing is handled transparently by\n * {@link RoutingPool}, which reads the execution context set by the base\n * class `asUser()`.\n *\n * @example\n * ```ts\n * import { createApp, lakebase, server } from \"@databricks/appkit\";\n *\n * const AppKit = await createApp({\n *   plugins: [server(), lakebase()],\n * });\n *\n * // Service principal query\n * const result = await AppKit.lakebase.query(\"SELECT * FROM users WHERE id = $1\", [userId]);\n *\n * // User-scoped query (per-user pool, RLS enforced)\n * const mine = await AppKit.lakebase.asUser(req).query(\"SELECT * FROM my_data\");\n * ```\n */\nexport class LakebasePlugin extends Plugin implements ToolProvider {\n  /** Plugin manifest declaring metadata and resource requirements */\n  static manifest = manifest as PluginManifest<\"lakebase\">;\n\n  protected declare config: ILakebaseConfig;\n  private pool: RoutingPool | null = null;\n  private oboPoolManager: LakebasePoolManager | null = null;\n\n  /**\n   * Initializes the Lakebase connection pool and OBO pool manager.\n   * Called automatically by AppKit during the plugin setup phase.\n   *\n   * Creates a {@link RoutingPool} that automatically routes queries to either\n   * the service-principal pool or a per-user pool based on the execution\n   * context (set by `Plugin.asUser(req)` via AsyncLocalStorage).\n   */\n  async setup() {\n    const poolConfig = this.config.pool;\n    const user = await getUsernameWithApiLookup(poolConfig);\n\n    const spPool = createLakebasePool({ ...poolConfig, user });\n    logger.info(\"Lakebase SP pool initialized\");\n\n    this.oboPoolManager = createLakebasePoolManager({\n      ...poolConfig,\n      ...OBO_POOL_DEFAULTS,\n    });\n    logger.info(\"Lakebase OBO pool manager initialized\");\n\n    const oboManager = this.oboPoolManager;\n    this.pool = new RoutingPool(spPool, (ctx) => {\n      if (!oboManager) throw new Error(\"OBO pool manager not initialized\");\n      // Lakebase OAuth roles use email as the postgres role when available\n      const userKey = ctx.userEmail ?? ctx.userId;\n      const isNew = !oboManager.hasPool(userKey);\n      const pool = oboManager.getPool(\n        userKey,\n        { workspaceClient: ctx.client, user: userKey },\n        ctx.tokenFingerprint,\n      );\n      if (isNew) {\n        logger.debug(\"Created OBO pool for user (total: %d)\", oboManager.size);\n      }\n      return pool;\n    });\n  }\n\n  /**\n   * Executes a parameterized SQL query against the Lakebase pool.\n   *\n   * When called inside `asUser(req)`, the query automatically routes to\n   * the per-user pool via {@link RoutingPool}.\n   *\n   * @param text - SQL query string, using `$1`, `$2`, ... placeholders\n   * @param values - Parameter values corresponding to placeholders\n   * @returns Query result with typed rows\n   *\n   * @example\n   * ```ts\n   * const result = await AppKit.lakebase.query<{ id: number; name: string }>(\n   *   \"SELECT id, name FROM users WHERE active = $1\",\n   *   [true],\n   * );\n   * ```\n   */\n  async query<T extends QueryResultRow = any>(\n    text: string,\n    values?: unknown[],\n  ): Promise<QueryResult<T>> {\n    // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup(), which AppKit always awaits before exposing the plugin API\n    return this.pool!.query<T>(text, values);\n  }\n\n  /**\n   * Execute a single statement inside a `BEGIN READ ONLY … ROLLBACK`\n   * transaction on a dedicated client.\n   *\n   * The three commands MUST share a connection — a naive\n   * `pool.query(\"BEGIN READ ONLY; <stmt>; ROLLBACK\")` batch cannot accept\n   * parameter values (PostgreSQL's Extended Query protocol rejects multi-\n   * statement prepared queries), which would silently break every\n   * parameterized query the agent tool issues.\n   *\n   * Returns the raw `rows` array for the user's statement. Side effects the\n   * statement may attempt (writes, writable-function side effects) are\n   * rejected by PostgreSQL under the read-only transaction posture.\n   */\n  private async runReadOnlyStatement(\n    text: string,\n    values?: unknown[],\n  ): Promise<unknown[]> {\n    // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup()\n    const client = await this.pool!.connect();\n    try {\n      await client.query(\"BEGIN READ ONLY\");\n      const result = await client.query(text, values);\n      return result.rows;\n    } finally {\n      try {\n        await client.query(\"ROLLBACK\");\n      } finally {\n        client.release();\n      }\n    }\n  }\n\n  /**\n   * Gracefully drains and closes all connection pools (SP + OBO).\n   * Called automatically by AppKit during shutdown.\n   */\n  abortActiveOperations(): void {\n    super.abortActiveOperations();\n    if (this.pool) {\n      logger.info(\"Closing Lakebase SP pool\");\n      this.pool.end().catch((err) => {\n        logger.error(\"Error closing Lakebase SP pool: %O\", err);\n      });\n      this.pool = null;\n    }\n    if (this.oboPoolManager) {\n      logger.info(\n        \"Closing all Lakebase OBO pools (%d)\",\n        this.oboPoolManager.size,\n      );\n      this.oboPoolManager.closeAll().catch((err) => {\n        logger.error(\"Error closing Lakebase OBO pools: %O\", err);\n      });\n      this.oboPoolManager = null;\n    }\n  }\n\n  /**\n   * Agent tool registry. Empty by default — the Lakebase plugin does NOT\n   * expose its SQL connection to LLM agents unless the developer explicitly\n   * opts in via `config.exposeAsAgentTool`. See {@link buildQueryTool}.\n   */\n  private tools: Record<string, ReturnType<typeof this.buildQueryTool>> = {};\n\n  constructor(config: ILakebaseConfig) {\n    super(config);\n    this.config = config;\n    if (config.exposeAsAgentTool) {\n      this.tools = { query: this.buildQueryTool(config.exposeAsAgentTool) };\n      logger.warn(\n        \"Lakebase agent tool is enabled (readOnly=%s). Every agent with access to this plugin can execute SQL against the Lakebase database as the requesting user's identity.\",\n        config.exposeAsAgentTool.readOnly !== false,\n      );\n    }\n  }\n\n  private buildQueryTool(\n    opt: NonNullable<ILakebaseConfig[\"exposeAsAgentTool\"]>,\n  ) {\n    const readOnly = opt.readOnly !== false;\n    return defineTool({\n      description: readOnly\n        ? \"Execute a read-only SQL query against the Lakebase PostgreSQL database. Only SELECT, WITH, SHOW, EXPLAIN, and DESCRIBE statements are accepted. Use $1, $2, etc. as placeholders and pass values separately.\"\n        : \"Execute a parameterized SQL statement against the Lakebase PostgreSQL database. Use $1, $2, etc. as placeholders and pass values separately. This tool can modify data; every invocation requires explicit human approval.\",\n      schema: z.object({\n        text: z\n          .string()\n          .describe(\n            \"SQL statement with $1, $2, ... placeholders for parameters\",\n          ),\n        values: z\n          .array(z.unknown())\n          .optional()\n          .describe(\"Parameter values corresponding to placeholders\"),\n      }),\n      annotations: {\n        effect: readOnly ? \"read\" : \"destructive\",\n        idempotent: false,\n        requiresUserContext: true,\n      },\n      execute: async (args, signal) => {\n        // Matches the files plugin pattern: the pg connection API\n        // doesn't accept AbortSignal in its current shape, so deeper\n        // mid-call cancellation needs a separate plumbing pass on the\n        // connector. This entry check still catches the common case —\n        // a tool dispatched after the user already cancelled the\n        // stream — and unwinds cleanly instead of running to\n        // completion against the SQL warehouse.\n        signal?.throwIfAborted();\n        if (readOnly) {\n          assertReadOnlySql(args.text);\n          return this.runReadOnlyStatement(args.text, args.values);\n        }\n        const result = await this.query(args.text, args.values);\n        return result.rows;\n      },\n    });\n  }\n\n  getAgentTools(): AgentToolDefinition[] {\n    return toolsFromRegistry(this.tools);\n  }\n\n  async executeAgentTool(\n    name: string,\n    args: unknown,\n    signal?: AbortSignal,\n  ): Promise<unknown> {\n    return executeFromRegistry(this.tools, name, args, signal);\n  }\n\n  toolkit(opts?: import(\"../../core/agent/types\").ToolkitOptions) {\n    return buildToolkitEntries(this.name, this.tools, opts);\n  }\n\n  /**\n   * Returns the pool config for the current execution context.\n   * Inside `asUser(req)`, returns user-scoped config; otherwise SP config.\n   */\n  private activePoolConfig() {\n    const ctx = getUserContext();\n    if (ctx) {\n      const user = ctx.userEmail ?? ctx.userId;\n      return { ...this.config.pool, workspaceClient: ctx.client, user };\n    }\n    return this.config.pool;\n  }\n\n  /**\n   * Returns the plugin's public API, accessible via `AppKit.lakebase`.\n   *\n   * - `pool` — The connection pool (routes to per-user pool when inside `asUser(req)`)\n   * - `query` — Convenience method for executing parameterized SQL queries\n   * - `getOrmConfig()` — Returns a config object compatible with Drizzle, TypeORM, Sequelize, etc.\n   *   Inside `asUser(req)`, returns user-scoped config.\n   * - `getPgConfig()` — Returns a `pg.PoolConfig` object for manual pool construction.\n   *   Inside `asUser(req)`, returns user-scoped config.\n   *\n   * Use `AppKit.lakebase.asUser(req)` to get the same API backed by a per-user pool.\n   */\n  exports() {\n    return {\n      // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup(), which AppKit always awaits before exposing the plugin API\n      pool: this.pool! as LakebasePool,\n      query: this.query.bind(this),\n      getOrmConfig: () => getLakebaseOrmConfig(this.activePoolConfig()),\n      getPgConfig: () => getLakebasePgConfig(this.activePoolConfig()),\n    };\n  }\n}\n\n/**\n * @internal\n */\nexport const lakebase = toPlugin(LakebasePlugin);\n"],"mappings":";;;;;;;;;;;;;;;AA2BA,MAAM,SAAS,aAAa,WAAW;;AAGvC,MAAM,oBAAoB;CACxB,KAAK;CACL,iBAAiB;CACjB,mBAAmB;CACpB;;;;;;;;;;;;;;;;;;;;;;;;;;;;AA6BD,IAAa,iBAAb,cAAoC,OAA+B;;CAEjE,OAAO,WAAWA;CAGlB,AAAQ,OAA2B;CACnC,AAAQ,iBAA6C;;;;;;;;;CAUrD,MAAM,QAAQ;EACZ,MAAM,aAAa,KAAK,OAAO;EAC/B,MAAM,OAAO,MAAM,yBAAyB,WAAW;EAEvD,MAAM,SAAS,mBAAmB;GAAE,GAAG;GAAY;GAAM,CAAC;AAC1D,SAAO,KAAK,+BAA+B;AAE3C,OAAK,iBAAiB,0BAA0B;GAC9C,GAAG;GACH,GAAG;GACJ,CAAC;AACF,SAAO,KAAK,wCAAwC;EAEpD,MAAM,aAAa,KAAK;AACxB,OAAK,OAAO,IAAI,YAAY,SAAS,QAAQ;AAC3C,OAAI,CAAC,WAAY,OAAM,IAAI,MAAM,mCAAmC;GAEpE,MAAM,UAAU,IAAI,aAAa,IAAI;GACrC,MAAM,QAAQ,CAAC,WAAW,QAAQ,QAAQ;GAC1C,MAAM,OAAO,WAAW,QACtB,SACA;IAAE,iBAAiB,IAAI;IAAQ,MAAM;IAAS,EAC9C,IAAI,iBACL;AACD,OAAI,MACF,QAAO,MAAM,yCAAyC,WAAW,KAAK;AAExE,UAAO;IACP;;;;;;;;;;;;;;;;;;;;CAqBJ,MAAM,MACJ,MACA,QACyB;AAEzB,SAAO,KAAK,KAAM,MAAS,MAAM,OAAO;;;;;;;;;;;;;;;;CAiB1C,MAAc,qBACZ,MACA,QACoB;EAEpB,MAAM,SAAS,MAAM,KAAK,KAAM,SAAS;AACzC,MAAI;AACF,SAAM,OAAO,MAAM,kBAAkB;AAErC,WADe,MAAM,OAAO,MAAM,MAAM,OAAO,EACjC;YACN;AACR,OAAI;AACF,UAAM,OAAO,MAAM,WAAW;aACtB;AACR,WAAO,SAAS;;;;;;;;CAStB,wBAA8B;AAC5B,QAAM,uBAAuB;AAC7B,MAAI,KAAK,MAAM;AACb,UAAO,KAAK,2BAA2B;AACvC,QAAK,KAAK,KAAK,CAAC,OAAO,QAAQ;AAC7B,WAAO,MAAM,sCAAsC,IAAI;KACvD;AACF,QAAK,OAAO;;AAEd,MAAI,KAAK,gBAAgB;AACvB,UAAO,KACL,uCACA,KAAK,eAAe,KACrB;AACD,QAAK,eAAe,UAAU,CAAC,OAAO,QAAQ;AAC5C,WAAO,MAAM,wCAAwC,IAAI;KACzD;AACF,QAAK,iBAAiB;;;;;;;;CAS1B,AAAQ,QAAgE,EAAE;CAE1E,YAAY,QAAyB;AACnC,QAAM,OAAO;AACb,OAAK,SAAS;AACd,MAAI,OAAO,mBAAmB;AAC5B,QAAK,QAAQ,EAAE,OAAO,KAAK,eAAe,OAAO,kBAAkB,EAAE;AACrE,UAAO,KACL,yKACA,OAAO,kBAAkB,aAAa,MACvC;;;CAIL,AAAQ,eACN,KACA;EACA,MAAM,WAAW,IAAI,aAAa;AAClC,SAAO,WAAW;GAChB,aAAa,WACT,iNACA;GACJ,QAAQ,EAAE,OAAO;IACf,MAAM,EACH,QAAQ,CACR,SACC,6DACD;IACH,QAAQ,EACL,MAAM,EAAE,SAAS,CAAC,CAClB,UAAU,CACV,SAAS,iDAAiD;IAC9D,CAAC;GACF,aAAa;IACX,QAAQ,WAAW,SAAS;IAC5B,YAAY;IACZ,qBAAqB;IACtB;GACD,SAAS,OAAO,MAAM,WAAW;AAQ/B,YAAQ,gBAAgB;AACxB,QAAI,UAAU;AACZ,uBAAkB,KAAK,KAAK;AAC5B,YAAO,KAAK,qBAAqB,KAAK,MAAM,KAAK,OAAO;;AAG1D,YADe,MAAM,KAAK,MAAM,KAAK,MAAM,KAAK,OAAO,EACzC;;GAEjB,CAAC;;CAGJ,gBAAuC;AACrC,SAAO,kBAAkB,KAAK,MAAM;;CAGtC,MAAM,iBACJ,MACA,MACA,QACkB;AAClB,SAAO,oBAAoB,KAAK,OAAO,MAAM,MAAM,OAAO;;CAG5D,QAAQ,MAAwD;AAC9D,SAAO,oBAAoB,KAAK,MAAM,KAAK,OAAO,KAAK;;;;;;CAOzD,AAAQ,mBAAmB;EACzB,MAAM,MAAM,gBAAgB;AAC5B,MAAI,KAAK;GACP,MAAM,OAAO,IAAI,aAAa,IAAI;AAClC,UAAO;IAAE,GAAG,KAAK,OAAO;IAAM,iBAAiB,IAAI;IAAQ;IAAM;;AAEnE,SAAO,KAAK,OAAO;;;;;;;;;;;;;;CAerB,UAAU;AACR,SAAO;GAEL,MAAM,KAAK;GACX,OAAO,KAAK,MAAM,KAAK,KAAK;GAC5B,oBAAoB,qBAAqB,KAAK,kBAAkB,CAAC;GACjE,mBAAmB,oBAAoB,KAAK,kBAAkB,CAAC;GAChE;;;;;;AAOL,MAAa,WAAW,SAAS,eAAe"}
+{"version":3,"file":"lakebase.js","names":["manifest"],"sources":["../../../src/plugins/lakebase/lakebase.ts"],"sourcesContent":["import { WorkspaceClient } from \"@databricks/sdk-experimental\";\nimport type { QueryResult, QueryResultRow } from \"pg\";\nimport type { AgentToolDefinition, ToolProvider } from \"shared\";\nimport { z } from \"zod\";\nimport {\n  createLakebasePool,\n  createLakebasePoolManager,\n  getLakebaseOrmConfig,\n  getLakebasePgConfig,\n  getUsernameWithApiLookup,\n  type LakebasePool,\n  type LakebasePoolManager,\n  RoutingPool,\n} from \"../../connectors/lakebase\";\nimport { getClientOptions } from \"../../context/client-options\";\nimport { getUserContext } from \"../../context/execution-context\";\nimport { buildToolkitEntries } from \"../../core/agent/build-toolkit\";\nimport {\n  defineTool,\n  executeFromRegistry,\n  toolsFromRegistry,\n} from \"../../core/agent/tools/define-tool\";\nimport { assertReadOnlySql } from \"../../core/agent/tools/sql-policy\";\nimport { createLogger } from \"../../logging/logger\";\nimport { Plugin, toPlugin } from \"../../plugin\";\nimport type { PluginManifest } from \"../../registry\";\nimport manifest from \"./manifest.json\";\nimport type { ILakebaseConfig } from \"./types\";\n\nconst logger = createLogger(\"lakebase\");\n\n/** Default pool settings for per-user OBO pools. */\nconst OBO_POOL_DEFAULTS = {\n  max: 3,\n  allowExitOnIdle: true,\n  idleTimeoutMillis: 30_000,\n};\n\n/**\n * AppKit plugin for Databricks Lakebase Autoscaling.\n *\n * Wraps `@databricks/lakebase` to provide a standard `pg.Pool` with automatic\n * OAuth token refresh, integrated with AppKit's logger and OpenTelemetry setup.\n *\n * Supports On-Behalf-Of (OBO) via `asUser(req)` — each user gets a separate\n * `pg.Pool` authenticated with their Databricks identity, enabling features\n * like Row-Level Security (RLS). Routing is handled transparently by\n * {@link RoutingPool}, which reads the execution context set by the base\n * class `asUser()`.\n *\n * @example\n * ```ts\n * import { createApp, lakebase, server } from \"@databricks/appkit\";\n *\n * const AppKit = await createApp({\n *   plugins: [server(), lakebase()],\n * });\n *\n * // Service principal query\n * const result = await AppKit.lakebase.query(\"SELECT * FROM users WHERE id = $1\", [userId]);\n *\n * // User-scoped query (per-user pool, RLS enforced)\n * const mine = await AppKit.lakebase.asUser(req).query(\"SELECT * FROM my_data\");\n * ```\n */\nexport class LakebasePlugin extends Plugin implements ToolProvider {\n  /** Plugin manifest declaring metadata and resource requirements */\n  static manifest = manifest as PluginManifest<\"lakebase\">;\n\n  protected declare config: ILakebaseConfig;\n  private pool: RoutingPool | null = null;\n  private oboPoolManager: LakebasePoolManager | null = null;\n\n  /**\n   * Initializes the Lakebase connection pool and OBO pool manager.\n   * Called automatically by AppKit during the plugin setup phase.\n   *\n   * Creates a {@link RoutingPool} that automatically routes queries to either\n   * the service-principal pool or a per-user pool based on the execution\n   * context (set by `Plugin.asUser(req)` via AsyncLocalStorage).\n   */\n  async setup() {\n    const poolConfig = {\n      ...this.config.pool,\n      workspaceClient:\n        this.config.pool?.workspaceClient ??\n        new WorkspaceClient({}, getClientOptions()),\n    };\n    const user = await getUsernameWithApiLookup(poolConfig);\n\n    const spPool = createLakebasePool({ ...poolConfig, user });\n    logger.info(\"Lakebase SP pool initialized\");\n\n    this.oboPoolManager = createLakebasePoolManager({\n      ...poolConfig,\n      ...OBO_POOL_DEFAULTS,\n    });\n    logger.info(\"Lakebase OBO pool manager initialized\");\n\n    const oboManager = this.oboPoolManager;\n    this.pool = new RoutingPool(spPool, (ctx) => {\n      if (!oboManager) throw new Error(\"OBO pool manager not initialized\");\n      // Lakebase OAuth roles use email as the postgres role when available\n      const userKey = ctx.userEmail ?? ctx.userId;\n      const isNew = !oboManager.hasPool(userKey);\n      const pool = oboManager.getPool(\n        userKey,\n        { workspaceClient: ctx.client, user: userKey },\n        ctx.tokenFingerprint,\n      );\n      if (isNew) {\n        logger.debug(\"Created OBO pool for user (total: %d)\", oboManager.size);\n      }\n      return pool;\n    });\n  }\n\n  /**\n   * Executes a parameterized SQL query against the Lakebase pool.\n   *\n   * When called inside `asUser(req)`, the query automatically routes to\n   * the per-user pool via {@link RoutingPool}.\n   *\n   * @param text - SQL query string, using `$1`, `$2`, ... placeholders\n   * @param values - Parameter values corresponding to placeholders\n   * @returns Query result with typed rows\n   *\n   * @example\n   * ```ts\n   * const result = await AppKit.lakebase.query<{ id: number; name: string }>(\n   *   \"SELECT id, name FROM users WHERE active = $1\",\n   *   [true],\n   * );\n   * ```\n   */\n  async query<T extends QueryResultRow = any>(\n    text: string,\n    values?: unknown[],\n  ): Promise<QueryResult<T>> {\n    // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup(), which AppKit always awaits before exposing the plugin API\n    return this.pool!.query<T>(text, values);\n  }\n\n  /**\n   * Execute a single statement inside a `BEGIN READ ONLY … ROLLBACK`\n   * transaction on a dedicated client.\n   *\n   * The three commands MUST share a connection — a naive\n   * `pool.query(\"BEGIN READ ONLY; <stmt>; ROLLBACK\")` batch cannot accept\n   * parameter values (PostgreSQL's Extended Query protocol rejects multi-\n   * statement prepared queries), which would silently break every\n   * parameterized query the agent tool issues.\n   *\n   * Returns the raw `rows` array for the user's statement. Side effects the\n   * statement may attempt (writes, writable-function side effects) are\n   * rejected by PostgreSQL under the read-only transaction posture.\n   */\n  private async runReadOnlyStatement(\n    text: string,\n    values?: unknown[],\n  ): Promise<unknown[]> {\n    // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup()\n    const client = await this.pool!.connect();\n    try {\n      await client.query(\"BEGIN READ ONLY\");\n      const result = await client.query(text, values);\n      return result.rows;\n    } finally {\n      try {\n        await client.query(\"ROLLBACK\");\n      } finally {\n        client.release();\n      }\n    }\n  }\n\n  /**\n   * Gracefully drains and closes all connection pools (SP + OBO).\n   * Called automatically by AppKit during shutdown.\n   */\n  abortActiveOperations(): void {\n    super.abortActiveOperations();\n    if (this.pool) {\n      logger.info(\"Closing Lakebase SP pool\");\n      this.pool.end().catch((err) => {\n        logger.error(\"Error closing Lakebase SP pool: %O\", err);\n      });\n      this.pool = null;\n    }\n    if (this.oboPoolManager) {\n      logger.info(\n        \"Closing all Lakebase OBO pools (%d)\",\n        this.oboPoolManager.size,\n      );\n      this.oboPoolManager.closeAll().catch((err) => {\n        logger.error(\"Error closing Lakebase OBO pools: %O\", err);\n      });\n      this.oboPoolManager = null;\n    }\n  }\n\n  /**\n   * Agent tool registry. Empty by default — the Lakebase plugin does NOT\n   * expose its SQL connection to LLM agents unless the developer explicitly\n   * opts in via `config.exposeAsAgentTool`. See {@link buildQueryTool}.\n   */\n  private tools: Record<string, ReturnType<typeof this.buildQueryTool>> = {};\n\n  constructor(config: ILakebaseConfig) {\n    super(config);\n    this.config = config;\n    if (config.exposeAsAgentTool) {\n      this.tools = { query: this.buildQueryTool(config.exposeAsAgentTool) };\n      logger.warn(\n        \"Lakebase agent tool is enabled (readOnly=%s). Every agent with access to this plugin can execute SQL against the Lakebase database as the requesting user's identity.\",\n        config.exposeAsAgentTool.readOnly !== false,\n      );\n    }\n  }\n\n  private buildQueryTool(\n    opt: NonNullable<ILakebaseConfig[\"exposeAsAgentTool\"]>,\n  ) {\n    const readOnly = opt.readOnly !== false;\n    return defineTool({\n      description: readOnly\n        ? \"Execute a read-only SQL query against the Lakebase PostgreSQL database. Only SELECT, WITH, SHOW, EXPLAIN, and DESCRIBE statements are accepted. Use $1, $2, etc. as placeholders and pass values separately.\"\n        : \"Execute a parameterized SQL statement against the Lakebase PostgreSQL database. Use $1, $2, etc. as placeholders and pass values separately. This tool can modify data; every invocation requires explicit human approval.\",\n      schema: z.object({\n        text: z\n          .string()\n          .describe(\n            \"SQL statement with $1, $2, ... placeholders for parameters\",\n          ),\n        values: z\n          .array(z.unknown())\n          .optional()\n          .describe(\"Parameter values corresponding to placeholders\"),\n      }),\n      annotations: {\n        effect: readOnly ? \"read\" : \"destructive\",\n        idempotent: false,\n        requiresUserContext: true,\n      },\n      execute: async (args, signal) => {\n        // Matches the files plugin pattern: the pg connection API\n        // doesn't accept AbortSignal in its current shape, so deeper\n        // mid-call cancellation needs a separate plumbing pass on the\n        // connector. This entry check still catches the common case —\n        // a tool dispatched after the user already cancelled the\n        // stream — and unwinds cleanly instead of running to\n        // completion against the SQL warehouse.\n        signal?.throwIfAborted();\n        if (readOnly) {\n          assertReadOnlySql(args.text);\n          return this.runReadOnlyStatement(args.text, args.values);\n        }\n        const result = await this.query(args.text, args.values);\n        return result.rows;\n      },\n    });\n  }\n\n  getAgentTools(): AgentToolDefinition[] {\n    return toolsFromRegistry(this.tools);\n  }\n\n  async executeAgentTool(\n    name: string,\n    args: unknown,\n    signal?: AbortSignal,\n  ): Promise<unknown> {\n    return executeFromRegistry(this.tools, name, args, signal);\n  }\n\n  toolkit(opts?: import(\"../../core/agent/types\").ToolkitOptions) {\n    return buildToolkitEntries(this.name, this.tools, opts);\n  }\n\n  /**\n   * Returns the pool config for the current execution context.\n   * Inside `asUser(req)`, returns user-scoped config; otherwise SP config.\n   */\n  private activePoolConfig() {\n    const ctx = getUserContext();\n    if (ctx) {\n      const user = ctx.userEmail ?? ctx.userId;\n      return { ...this.config.pool, workspaceClient: ctx.client, user };\n    }\n    return this.config.pool;\n  }\n\n  /**\n   * Returns the plugin's public API, accessible via `AppKit.lakebase`.\n   *\n   * - `pool` — The connection pool (routes to per-user pool when inside `asUser(req)`)\n   * - `query` — Convenience method for executing parameterized SQL queries\n   * - `getOrmConfig()` — Returns a config object compatible with Drizzle, TypeORM, Sequelize, etc.\n   *   Inside `asUser(req)`, returns user-scoped config.\n   * - `getPgConfig()` — Returns a `pg.PoolConfig` object for manual pool construction.\n   *   Inside `asUser(req)`, returns user-scoped config.\n   *\n   * Use `AppKit.lakebase.asUser(req)` to get the same API backed by a per-user pool.\n   */\n  exports() {\n    return {\n      // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup(), which AppKit always awaits before exposing the plugin API\n      pool: this.pool! as LakebasePool,\n      query: this.query.bind(this),\n      getOrmConfig: () => getLakebaseOrmConfig(this.activePoolConfig()),\n      getPgConfig: () => getLakebasePgConfig(this.activePoolConfig()),\n    };\n  }\n}\n\n/**\n * @internal\n */\nexport const lakebase = toPlugin(LakebasePlugin);\n"],"mappings":";;;;;;;;;;;;;;;;;AA6BA,MAAM,SAAS,aAAa,WAAW;;AAGvC,MAAM,oBAAoB;CACxB,KAAK;CACL,iBAAiB;CACjB,mBAAmB;CACpB;;;;;;;;;;;;;;;;;;;;;;;;;;;;AA6BD,IAAa,iBAAb,cAAoC,OAA+B;;CAEjE,OAAO,WAAWA;CAGlB,AAAQ,OAA2B;CACnC,AAAQ,iBAA6C;;;;;;;;;CAUrD,MAAM,QAAQ;EACZ,MAAM,aAAa;GACjB,GAAG,KAAK,OAAO;GACf,iBACE,KAAK,OAAO,MAAM,mBAClB,IAAI,gBAAgB,EAAE,EAAE,kBAAkB,CAAC;GAC9C;EACD,MAAM,OAAO,MAAM,yBAAyB,WAAW;EAEvD,MAAM,SAAS,mBAAmB;GAAE,GAAG;GAAY;GAAM,CAAC;AAC1D,SAAO,KAAK,+BAA+B;AAE3C,OAAK,iBAAiB,0BAA0B;GAC9C,GAAG;GACH,GAAG;GACJ,CAAC;AACF,SAAO,KAAK,wCAAwC;EAEpD,MAAM,aAAa,KAAK;AACxB,OAAK,OAAO,IAAI,YAAY,SAAS,QAAQ;AAC3C,OAAI,CAAC,WAAY,OAAM,IAAI,MAAM,mCAAmC;GAEpE,MAAM,UAAU,IAAI,aAAa,IAAI;GACrC,MAAM,QAAQ,CAAC,WAAW,QAAQ,QAAQ;GAC1C,MAAM,OAAO,WAAW,QACtB,SACA;IAAE,iBAAiB,IAAI;IAAQ,MAAM;IAAS,EAC9C,IAAI,iBACL;AACD,OAAI,MACF,QAAO,MAAM,yCAAyC,WAAW,KAAK;AAExE,UAAO;IACP;;;;;;;;;;;;;;;;;;;;CAqBJ,MAAM,MACJ,MACA,QACyB;AAEzB,SAAO,KAAK,KAAM,MAAS,MAAM,OAAO;;;;;;;;;;;;;;;;CAiB1C,MAAc,qBACZ,MACA,QACoB;EAEpB,MAAM,SAAS,MAAM,KAAK,KAAM,SAAS;AACzC,MAAI;AACF,SAAM,OAAO,MAAM,kBAAkB;AAErC,WADe,MAAM,OAAO,MAAM,MAAM,OAAO,EACjC;YACN;AACR,OAAI;AACF,UAAM,OAAO,MAAM,WAAW;aACtB;AACR,WAAO,SAAS;;;;;;;;CAStB,wBAA8B;AAC5B,QAAM,uBAAuB;AAC7B,MAAI,KAAK,MAAM;AACb,UAAO,KAAK,2BAA2B;AACvC,QAAK,KAAK,KAAK,CAAC,OAAO,QAAQ;AAC7B,WAAO,MAAM,sCAAsC,IAAI;KACvD;AACF,QAAK,OAAO;;AAEd,MAAI,KAAK,gBAAgB;AACvB,UAAO,KACL,uCACA,KAAK,eAAe,KACrB;AACD,QAAK,eAAe,UAAU,CAAC,OAAO,QAAQ;AAC5C,WAAO,MAAM,wCAAwC,IAAI;KACzD;AACF,QAAK,iBAAiB;;;;;;;;CAS1B,AAAQ,QAAgE,EAAE;CAE1E,YAAY,QAAyB;AACnC,QAAM,OAAO;AACb,OAAK,SAAS;AACd,MAAI,OAAO,mBAAmB;AAC5B,QAAK,QAAQ,EAAE,OAAO,KAAK,eAAe,OAAO,kBAAkB,EAAE;AACrE,UAAO,KACL,yKACA,OAAO,kBAAkB,aAAa,MACvC;;;CAIL,AAAQ,eACN,KACA;EACA,MAAM,WAAW,IAAI,aAAa;AAClC,SAAO,WAAW;GAChB,aAAa,WACT,iNACA;GACJ,QAAQ,EAAE,OAAO;IACf,MAAM,EACH,QAAQ,CACR,SACC,6DACD;IACH,QAAQ,EACL,MAAM,EAAE,SAAS,CAAC,CAClB,UAAU,CACV,SAAS,iDAAiD;IAC9D,CAAC;GACF,aAAa;IACX,QAAQ,WAAW,SAAS;IAC5B,YAAY;IACZ,qBAAqB;IACtB;GACD,SAAS,OAAO,MAAM,WAAW;AAQ/B,YAAQ,gBAAgB;AACxB,QAAI,UAAU;AACZ,uBAAkB,KAAK,KAAK;AAC5B,YAAO,KAAK,qBAAqB,KAAK,MAAM,KAAK,OAAO;;AAG1D,YADe,MAAM,KAAK,MAAM,KAAK,MAAM,KAAK,OAAO,EACzC;;GAEjB,CAAC;;CAGJ,gBAAuC;AACrC,SAAO,kBAAkB,KAAK,MAAM;;CAGtC,MAAM,iBACJ,MACA,MACA,QACkB;AAClB,SAAO,oBAAoB,KAAK,OAAO,MAAM,MAAM,OAAO;;CAG5D,QAAQ,MAAwD;AAC9D,SAAO,oBAAoB,KAAK,MAAM,KAAK,OAAO,KAAK;;;;;;CAOzD,AAAQ,mBAAmB;EACzB,MAAM,MAAM,gBAAgB;AAC5B,MAAI,KAAK;GACP,MAAM,OAAO,IAAI,aAAa,IAAI;AAClC,UAAO;IAAE,GAAG,KAAK,OAAO;IAAM,iBAAiB,IAAI;IAAQ;IAAM;;AAEnE,SAAO,KAAK,OAAO;;;;;;;;;;;;;;CAerB,UAAU;AACR,SAAO;GAEL,MAAM,KAAK;GACX,OAAO,KAAK,MAAM,KAAK,KAAK;GAC5B,oBAAoB,qBAAqB,KAAK,kBAAkB,CAAC;GACjE,mBAAmB,oBAAoB,KAAK,kBAAkB,CAAC;GAChE;;;;;;AAOL,MAAa,WAAW,SAAS,eAAe"}

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@databricks/appkit",
   "type": "module",
-  "version": "0.41.7",
+  "version": "0.43.0",
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "bin": {