refactor(proxy,adapters): unify MCP proxy onto StitchToolClient and isolate tool catalog subpath (#371)

Author: davideast

packages/sdk/src/adk-adapter.ts

@@ -12,11 +12,24 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { FunctionTool } from "@google/adk";
+import type { FunctionTool as FunctionToolType } from "@google/adk";
 import type { Schema } from "@google/genai";
 import { toolDefinitions } from "../generated/src/tool-definitions.js";
 import { getOrCreateClient } from "./singleton.js";
 
+// @google/adk is an OPTIONAL peer dependency: it must not be required to
+// install or import the core SDK. Guarded dynamic import gives consumers an
+// actionable error instead of a bare ERR_MODULE_NOT_FOUND.
+let FunctionTool: typeof FunctionToolType;
+try {
+  ({ FunctionTool } = await import("@google/adk"));
+} catch {
+  throw new Error(
+    `"@google/stitch-sdk/adk" requires the optional peer dependency "@google/adk", ` +
+      `which is not installed. Install it with:\n\n  npm install @google/adk\n`,
+  );
+}
+
 /**
  * Recursively cleans and flattens a JSON Schema to make it compatible with the Google ADK/Gemini API.
  * It resolves internal `#/$defs/` references directly into the object, and removes
@@ -29,6 +42,11 @@ import { getOrCreateClient } from "./singleton.js";
 function cleanSchema(schema: any): any {
   if (!schema || typeof schema !== "object") return schema;
   const defs = schema.$defs || {};
+  // Defs currently being resolved. A self-recursive $def cannot be
+  // expressed in the Gemini schema dialect — break the cycle with {}
+  // instead of embedding an in-progress object by reference (which
+  // produced circular output that exploded on JSON serialization).
+  const inProgress = new Set<string>();
 
   function stripAndResolve(node: any, seen = new Map()): any {
     if (!node || typeof node !== "object") return node;
@@ -50,7 +68,12 @@ function cleanSchema(schema: any): any {
     ) {
       const defName = node.$ref.replace("#/$defs/", "");
       if (defs[defName]) {
+        if (inProgress.has(defName)) {
+          return {}; // recursive $def — cycle broken
+        }
+        inProgress.add(defName);
         const target = stripAndResolve(defs[defName], seen);
+        inProgress.delete(defName);
         const resolved = { ...target };
         for (const [k, v] of Object.entries(node)) {
           if (
@@ -108,7 +131,20 @@ function cleanSchema(schema: any): any {
 export function stitchAdkTools(options?: {
   apiKey?: string;
   include?: string[];
-}): FunctionTool<Schema>[] {
+}): FunctionToolType<Schema>[] {
+  // A misspelled tool name silently vanishing from an agent's toolbox
+  // is undebuggable — validate loudly [V1_PLAN §3.6].
+  if (options?.include) {
+    const known = new Set(toolDefinitions.map((t) => t.name));
+    const unknown = options.include.filter((name) => !known.has(name));
+    if (unknown.length > 0) {
+      throw new Error(
+        `Unknown Stitch tool name(s) in include filter: ${unknown.join(", ")}.\n` +
+          `Available tools: ${[...known].join(", ")}`,
+      );
+    }
+  }
+
   const client = getOrCreateClient(options);
 
   const filtered = options?.include

packages/sdk/src/proxy/client.ts

@@ -13,100 +13,34 @@
 // limitations under the License.
 
 import { StitchProxyConfig } from "../spec/proxy.js";
-import { buildAuthHeaders } from "../auth.js";
+import type { StitchToolClient } from "../client.js";
 import type { Tool } from "@modelcontextprotocol/sdk/types.js";
-import { repairToolSchemas } from "../schema-repair.js";
 
 /**
  * Shared state for proxy handlers.
+ *
+ * `client` is the ONE MCP stack in the repo (D9): a real StitchToolClient
+ * over the MCP SDK's StreamableHTTPClientTransport. The caller constructs
+ * and injects it (see core.ts) — this module never builds transports or
+ * speaks JSON-RPC itself.
  */
 export interface ProxyContext {
   config: StitchProxyConfig;
+  client: StitchToolClient;
   remoteTools: Tool[];
 }
 
 /**
- * Forward a JSON-RPC request to Stitch.
- */
-export async function forwardToStitch(
-  config: StitchProxyConfig,
-  method: string,
-  params?: unknown,
-): Promise<unknown> {
-  const request = {
-    jsonrpc: "2.0",
-    method,
-    params: params ?? {},
-    id: Date.now(),
-  };
-
-  let response: Response;
-  try {
-    response = await fetch(config.url, {
-      method: "POST",
-      headers: {
-        "Content-Type": "application/json",
-        Accept: "application/json",
-        ...buildAuthHeaders(config),
-      },
-      body: JSON.stringify(request),
-    });
-  } catch (err: any) {
-    throw new Error(`Network failure connecting to Stitch API: ${err.message}`);
-  }
-
-  if (!response.ok) {
-    const errorText = await response.text();
-    throw new Error(`Stitch API error (${response.status}): ${errorText}`);
-  }
-
-  const result = (await response.json()) as {
-    error?: { message: string };
-    result?: unknown;
-  };
-
-  if (result.error) {
-    throw new Error(`Stitch RPC error: ${result.error.message}`);
-  }
-
-  return result.result;
-}
-
-/**
- * Initialize connection to Stitch and fetch tools.
+ * Initialize the upstream Stitch connection and fetch tools.
+ *
+ * The MCP SDK handles the initialize handshake (protocol version,
+ * session id, awaited notifications/initialized) — the hand-rolled
+ * JSON-RPC machinery that used to live here is gone.
  */
 export async function initializeStitchConnection(
   ctx: ProxyContext,
 ): Promise<void> {
-  // Send initialize request
-  await forwardToStitch(ctx.config, "initialize", {
-    protocolVersion: ctx.config.protocolVersion || "2024-11-05",
-    capabilities: {},
-    clientInfo: {
-      name: ctx.config.name,
-      version: ctx.config.version,
-    },
-  });
-
-  // Send initialized notification (fire and forget)
-  fetch(ctx.config.url, {
-    method: "POST",
-    headers: {
-      "Content-Type": "application/json",
-      Accept: "application/json",
-      ...buildAuthHeaders(ctx.config),
-    },
-    body: JSON.stringify({
-      jsonrpc: "2.0",
-      method: "notifications/initialized",
-    }),
-  }).catch((err) => {
-    console.error(
-      "[stitch-proxy] Failed to send initialized notification:",
-      err,
-    );
-  });
-
+  await ctx.client.connect();
   await refreshTools(ctx);
   console.error(
     `[stitch-proxy] Connected to Stitch, discovered ${ctx.remoteTools.length} tools`,
@@ -116,15 +50,12 @@ export async function initializeStitchConnection(
 /**
  * Refresh the cached tools list from Stitch.
  *
- * Applies schema repair to inject missing $defs before the tools are
- * re-served to MCP clients whose AJV validators would otherwise crash
- * on unresolved $ref targets.
+ * Stores schemas RAW, exactly as served. Repair (injecting missing
+ * $defs) is a SERVING concern applied where schemas are consumed —
+ * see handlers/listTools.ts — never at capture: the tools-manifest
+ * is the pipeline's source of truth and must not be coupled to the
+ * repair heuristics.
  */
 export async function refreshTools(ctx: ProxyContext): Promise<void> {
-  const toolsResult = (await forwardToStitch(ctx.config, "tools/list", {})) as {
-    tools: Tool[];
-  };
-  const tools = toolsResult.tools || [];
-  repairToolSchemas(tools);
-  ctx.remoteTools = tools;
+  ctx.remoteTools = (await ctx.client.listToolsRaw()).tools;
 }

packages/sdk/src/proxy/core.ts

@@ -20,13 +20,16 @@ import {
   StitchProxyConfig,
   StitchProxySpec,
 } from "../spec/proxy.js";
+import { StitchToolClient } from "../client.js";
 import { ProxyContext, initializeStitchConnection } from "./client.js";
 import { registerListToolsHandler } from "./handlers/listTools.js";
 import { registerCallToolHandler } from "./handlers/callTool.js";
 
 /**
  * A proxy server that forwards MCP requests to Stitch.
- * Bypasses SDK transport layer to handle specific auth and JSON-RPC forwarding.
+ * Runs on the SAME StitchToolClient / MCP SDK transport as the rest of the
+ * SDK (D9 — one MCP stack); it does not bypass the transport or speak
+ * hand-rolled JSON-RPC. It adds Stitch auth + virtual-tool routing on top.
  */
 export class StitchProxy implements StitchProxySpec {
   private config: StitchProxyConfig;
@@ -41,7 +44,12 @@ export class StitchProxy implements StitchProxySpec {
         inputConfig?.quotaProjectId ||
         process.env.STITCH_PROJECT_ID ||
         process.env.GOOGLE_CLOUD_PROJECT,
-      url: inputConfig?.url || process.env.STITCH_MCP_URL,
+      // Precedence: explicit config > STITCH_BASE_URL > STITCH_MCP_URL
+      // (STITCH_MCP_URL stays accepted for the proxy binary).
+      url:
+        inputConfig?.url ||
+        process.env.STITCH_BASE_URL ||
+        process.env.STITCH_MCP_URL,
       name: inputConfig?.name,
       version: inputConfig?.version,
     };
@@ -67,9 +75,17 @@ export class StitchProxy implements StitchProxySpec {
       },
     );
 
-    // Shared context for handlers
+    // Shared context for handlers. The proxy runs on the real
+    // StitchToolClient (one MCP stack, D9) — map proxy config onto the
+    // client's config shape.
     this.ctx = {
       config: this.config,
+      client: new StitchToolClient({
+        apiKey: this.config.apiKey,
+        accessToken: this.config.accessToken,
+        projectId: this.config.quotaProjectId,
+        baseUrl: this.config.url,
+      }),
       remoteTools: [] as Tool[],
     };
 
@@ -90,5 +106,6 @@ export class StitchProxy implements StitchProxySpec {
 
   async close(): Promise<void> {
     await this.server.close();
+    await this.ctx.client.close();
   }
 }

packages/sdk/src/proxy/handlers/callTool.ts

@@ -15,7 +15,6 @@
 import { CallToolRequestSchema } from "@modelcontextprotocol/sdk/types.js";
 import type { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import type { ProxyContext } from "../client.js";
-import { forwardToStitch } from "../client.js";
 import { isVirtualTool, handleVirtualTool } from "../virtual-tools.js";
 
 /**
@@ -50,10 +49,19 @@ export function registerCallToolHandler(
     }
 
     try {
-      const result = await forwardToStitch(ctx.config, "tools/call", {
-        name,
-        arguments: args,
-      });
+      // Forward the RAW envelope verbatim — the downstream MCP client
+      // owns error semantics (isError stays an envelope, never a throw).
+      const result = await ctx.client.callToolRaw(name, args ?? {});
+      if (result == null) {
+        // A handler returning undefined is a protocol violation downstream;
+        // surface it as a proper MCP error envelope instead.
+        return {
+          isError: true,
+          content: [
+            { type: "text", text: `Upstream returned no result for ${name}` },
+          ],
+        };
+      }
       return result as { content: Array<{ type: string; text: string }> };
     } catch (err) {
       const errorMessage = err instanceof Error ? err.message : String(err);

packages/sdk/src/proxy/handlers/listTools.ts

@@ -16,9 +16,10 @@ import { ListToolsRequestSchema } from "@modelcontextprotocol/sdk/types.js";
 import type { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import type { ProxyContext } from "../client.js";
 import { refreshTools } from "../client.js";
-import { downloadAssetsTool } from "../virtual-tools.js";
+import { virtualTools } from "../virtual-tools.js";
+import { repairToolSchemas } from "../../schema-repair.js";
 
-const PROXY_VIRTUAL_TOOLS = [downloadAssetsTool].map((t) => ({
+const PROXY_VIRTUAL_TOOLS = virtualTools.map((t) => ({
   name: t.name,
   description: t.description,
   inputSchema: t.inputSchema,
@@ -44,6 +45,27 @@ export function registerListToolsHandler(
         );
       }
     }
-    return { tools: [...ctx.remoteTools, ...PROXY_VIRTUAL_TOOLS] };
+    // ctx.remoteTools holds RAW schemas (capture truth). Repair a copy at
+    // serving time: downstream MCP clients' AJV validators crash on
+    // unresolved $ref targets the backend sometimes omits.
+    const served = structuredClone(ctx.remoteTools);
+    repairToolSchemas(served);
+
+    // Collision detection: if the server ever ships a tool with a virtual
+    // tool's name, the virtual tool wins the callTool route (isVirtualTool
+    // is checked first) — so don't serve a duplicate listing, and say so
+    // loudly instead of silently shadowing.
+    const virtualNames = new Set(PROXY_VIRTUAL_TOOLS.map((t) => t.name));
+    const deduped = served.filter((t) => {
+      if (virtualNames.has(t.name)) {
+        console.warn(
+          `[stitch-proxy] Remote tool "${t.name}" collides with a local ` +
+            `virtual tool and is shadowed. Rename one of them.`,
+        );
+        return false;
+      }
+      return true;
+    });
+    return { tools: [...deduped, ...PROXY_VIRTUAL_TOOLS] };
   });
 }

packages/sdk/src/proxy/index.ts

@@ -14,9 +14,5 @@
 
 // Proxy barrel exports
 export { StitchProxy } from "./core.js";
-export {
-  forwardToStitch,
-  initializeStitchConnection,
-  refreshTools,
-} from "./client.js";
+export { initializeStitchConnection, refreshTools } from "./client.js";
 export type { ProxyContext } from "./client.js";