refactor: move transport logic to a ToolboxTransport class (#188)

* refactor: move transport logic to a ToolboxTransport class

* rename file

* update license

* basic lint

* lint

* log errors properly

* fix error handling

* add test cov

* lint

* small refactor

* fix consoleErrorSpy issue

Author: twishabansal

packages/toolbox-core/src/toolbox_core/client.ts

@@ -13,19 +13,14 @@
 // limitations under the License.
 
 import {ToolboxTool} from './tool.js';
-import axios from 'axios';
+import {AxiosInstance} from 'axios';
+import {ITransport} from './transport.types.js';
+import {ToolboxTransport} from './toolboxTransport.js';
 import {
-  type AxiosInstance,
-  type AxiosRequestConfig,
-  type AxiosResponse,
-} from 'axios';
-import {
-  ZodManifestSchema,
   createZodSchemaFromParams,
   ParameterSchema,
+  ZodManifestSchema,
 } from './protocol.js';
-import {logApiError} from './errorUtils.js';
-import {ZodError} from 'zod';
 import {BoundParams, identifyAuthRequirements, resolveValue} from './utils.js';
 import {AuthTokenGetters, RequiredAuthnParams} from './tool.js';
 
@@ -41,8 +36,7 @@ export type ClientHeadersConfig = Record<string, ClientHeaderProvider>;
  * An asynchronous client for interacting with a Toolbox service.
  */
 class ToolboxClient {
-  #baseUrl: string;
-  #session: AxiosInstance;
+  #transport: ITransport;
   #clientHeaders: ClientHeadersConfig;
 
   /**
@@ -58,8 +52,7 @@ class ToolboxClient {
     session?: AxiosInstance | null,
     clientHeaders?: ClientHeadersConfig | null,
   ) {
-    this.#baseUrl = url;
-    this.#session = session || axios.create({baseURL: this.#baseUrl});
+    this.#transport = new ToolboxTransport(url, session || undefined);
     this.#clientHeaders = clientHeaders || {};
   }
 
@@ -77,47 +70,6 @@ class ToolboxClient {
     return Object.fromEntries(resolvedEntries);
   }
 
-  /**
-   * Fetches and parses the manifest from a given API path.
-   * @param {string} apiPath - The API path to fetch the manifest from (e.g., "/api/tool/mytool").
-   * @returns {Promise<Manifest>} A promise that resolves to the parsed manifest.
-   * @throws {Error} If there's an error fetching data or if the manifest structure is invalid.
-   */
-  async #fetchAndParseManifest(apiPath: string): Promise<Manifest> {
-    const url = `${this.#baseUrl}${apiPath}`;
-    try {
-      const headers = await this.#resolveClientHeaders();
-      const config: AxiosRequestConfig = {headers};
-      const response: AxiosResponse = await this.#session.get(url, config);
-      const responseData = response.data;
-
-      try {
-        const manifest = ZodManifestSchema.parse(responseData);
-        return manifest;
-      } catch (validationError) {
-        let detailedMessage = `Invalid manifest structure received from ${url}: `;
-        if (validationError instanceof ZodError) {
-          const issueDetails = validationError.issues;
-          detailedMessage += JSON.stringify(issueDetails, null, 2);
-        } else if (validationError instanceof Error) {
-          detailedMessage += validationError.message;
-        } else {
-          detailedMessage += 'Unknown validation error.';
-        }
-        throw new Error(detailedMessage);
-      }
-    } catch (error) {
-      if (
-        error instanceof Error &&
-        error.message.startsWith('Invalid manifest structure received from')
-      ) {
-        throw error;
-      }
-      logApiError(`Error fetching data from ${url}:`, error);
-      throw error;
-    }
-  }
-
   /**
    * Creates a ToolboxTool instance from its schema.
    * @param {string} toolName - The name of the tool.
@@ -159,8 +111,7 @@ class ToolboxClient {
     const paramZodSchema = createZodSchemaFromParams(params);
 
     const tool = ToolboxTool(
-      this.#session,
-      this.#baseUrl,
+      this.#transport,
       toolName,
       toolSchema.description,
       paramZodSchema,
@@ -195,8 +146,8 @@ class ToolboxClient {
     authTokenGetters: AuthTokenGetters | null = {},
     boundParams: BoundParams | null = {},
   ): Promise<ReturnType<typeof ToolboxTool>> {
-    const apiPath = `/api/tool/${name}`;
-    const manifest = await this.#fetchAndParseManifest(apiPath);
+    const headers = await this.#resolveClientHeaders();
+    const manifest = await this.#transport.toolGet(name, headers);
 
     if (
       manifest.tools &&
@@ -240,7 +191,9 @@ class ToolboxClient {
       }
       return tool;
     } else {
-      throw new Error(`Tool "${name}" not found in manifest from ${apiPath}.`);
+      throw new Error(
+        `Tool "${name}" not found in manifest from ${this.#transport.baseUrl}/api/tool/${name}.`,
+      );
     }
   }
 
@@ -262,9 +215,9 @@ class ToolboxClient {
     strict = false,
   ): Promise<Array<ReturnType<typeof ToolboxTool>>> {
     const toolsetName = name || '';
-    const apiPath = `/api/toolset/${toolsetName}`;
+    const headers = await this.#resolveClientHeaders();
 
-    const manifest = await this.#fetchAndParseManifest(apiPath);
+    const manifest = await this.#transport.toolsList(toolsetName, headers);
     const tools: Array<ReturnType<typeof ToolboxTool>> = [];
 
     const overallUsedAuthKeys: Set<string> = new Set();

packages/toolbox-core/src/toolbox_core/index.ts

@@ -16,6 +16,9 @@
 
 export {ToolboxClient} from './client.js';
 
+// Export transport classes
+export {ToolboxTransport} from './toolboxTransport.js';
+
 // Export the main factory function and the core tool type
 export {ToolboxTool} from './tool.js';
 

packages/toolbox-core/src/toolbox_core/tool.ts

@@ -13,9 +13,8 @@
 // limitations under the License.
 
 import {ZodObject, ZodError, ZodRawShape} from 'zod';
-import {AxiosInstance, AxiosResponse} from 'axios';
 
-import {logApiError} from './errorUtils.js';
+import {ITransport} from './transport.types.js';
 import {
   BoundParams,
   BoundValue,
@@ -41,8 +40,7 @@ function getAuthHeaderName(authTokenName: string): string {
  * Creates a callable tool function representing a specific tool on a remote
  * Toolbox server.
  *
- * @param {AxiosInstance} session - The Axios session for making HTTP requests.
- * @param {string} baseUrl - The base URL of the Toolbox Server API.
+ * @param {ITransport} transport - The transport for making API requests.
  * @param {string} name - The name of the remote tool.
  * @param {string} description - A description of the remote tool.
  * @param {ZodObject<any>} paramSchema - The Zod schema for validating the tool's parameters.
@@ -55,8 +53,7 @@ function getAuthHeaderName(authTokenName: string): string {
  * called, invokes the tool with the provided arguments.
  */
 function ToolboxTool(
-  session: AxiosInstance,
-  baseUrl: string,
+  transport: ITransport,
   name: string,
   description: string,
   paramSchema: ZodObject<ZodRawShape>,
@@ -69,14 +66,13 @@ function ToolboxTool(
   if (
     (Object.keys(authTokenGetters).length > 0 ||
       Object.keys(clientHeaders).length > 0) &&
-    !baseUrl.startsWith('https://')
+    !transport.baseUrl.startsWith('https://')
   ) {
     console.warn(
       'Sending ID token over HTTP. User data may be exposed. Use HTTPS for secure communication.',
     );
   }
 
-  const toolUrl = `${baseUrl}/api/tool/${name}/invoke`;
   const boundKeys = Object.keys(boundParams);
   const userParamSchema = paramSchema.omit(
     Object.fromEntries(boundKeys.map(k => [k, true])),
@@ -159,19 +155,7 @@ function ToolboxTool(
       headers[getAuthHeaderName(authService)] = token;
     }
 
-    try {
-      const response: AxiosResponse = await session.post(
-        toolUrl,
-        filteredPayload,
-        {
-          headers,
-        },
-      );
-      return response.data.result;
-    } catch (error) {
-      logApiError(`Error posting data to ${toolUrl}:`, error);
-      throw error;
-    }
+    return await transport.toolInvoke(name, filteredPayload, headers);
   };
   callable.toolName = name;
   callable.description = description;
@@ -234,8 +218,7 @@ function ToolboxTool(
     }
 
     return ToolboxTool(
-      session,
-      baseUrl,
+      transport,
       this.toolName,
       this.description,
       this.params,
@@ -271,8 +254,7 @@ function ToolboxTool(
 
     const newBoundParams = {...this.boundParams, ...paramsToBind};
     return ToolboxTool(
-      session,
-      baseUrl,
+      transport,
       this.toolName,
       this.description,
       this.params,

packages/toolbox-core/src/toolbox_core/toolboxTransport.ts

@@ -0,0 +1,106 @@
+/**
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import axios, {AxiosInstance, AxiosResponse} from 'axios';
+import {ITransport} from './transport.types.js';
+import {ZodManifest, ZodManifestSchema} from './protocol.js';
+import {logApiError} from './errorUtils.js';
+
+/**
+ * Transport for the native Toolbox protocol.
+ */
+export class ToolboxTransport implements ITransport {
+  readonly #baseUrl: string;
+  #session: AxiosInstance;
+
+  constructor(baseUrl: string, session?: AxiosInstance) {
+    this.#baseUrl = baseUrl;
+    // If no axios session is provided, make our own
+    this.#session = session || axios.create({baseURL: this.baseUrl});
+  }
+
+  get baseUrl(): string {
+    return this.#baseUrl;
+  }
+
+  async #getManifest(
+    url: string,
+    headers?: Record<string, string>,
+  ): Promise<ZodManifest> {
+    /** Helper method to perform GET requests and parse the ManifestSchema. */
+    try {
+      const response: AxiosResponse = await this.#session.get(url, {headers});
+      return ZodManifestSchema.parse(response.data);
+    } catch (error) {
+      logApiError(`Error fetching data from ${url}:`, error);
+      throw error;
+    }
+  }
+
+  async toolGet(
+    toolName: string,
+    headers?: Record<string, string>,
+  ): Promise<ZodManifest> {
+    const url = `${this.#baseUrl}/api/tool/${toolName}`;
+    return await this.#getManifest(url, headers);
+  }
+
+  async toolsList(
+    toolsetName?: string,
+    headers?: Record<string, string>,
+  ): Promise<ZodManifest> {
+    const url = `${this.#baseUrl}/api/toolset/${toolsetName || ''}`;
+    return await this.#getManifest(url, headers);
+  }
+
+  async toolInvoke(
+    toolName: string,
+    arguments_: Record<string, unknown>,
+    headers: Record<string, string>,
+  ): Promise<string> {
+    // ID tokens contain sensitive user information (claims). Transmitting
+    // these over HTTP exposes the data to interception and unauthorized
+    // access. Always use HTTPS to ensure secure communication and protect
+    // user privacy.
+    if (
+      this.baseUrl.startsWith('http://') &&
+      headers &&
+      Object.keys(headers).length > 0
+    ) {
+      console.warn(
+        'Sending data token over HTTP. User data may be exposed. Use HTTPS for secure communication.',
+      );
+    }
+    const url = `${this.#baseUrl}/api/tool/${toolName}/invoke`;
+    try {
+      const response: AxiosResponse = await this.#session.post(
+        url,
+        arguments_,
+        {
+          headers,
+        },
+      );
+      const body = response.data;
+      if (body?.error) {
+        throw new Error(body.error);
+      }
+      return body.result;
+    } catch (error) {
+      logApiError(`Error posting data to ${url}:`, error);
+      throw error;
+    }
+  }
+}

packages/toolbox-core/src/toolbox_core/transport.types.ts

@@ -0,0 +1,53 @@
+/**
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import {ZodManifest} from './protocol.js';
+
+/**
+ * Defines the contract for a 'smart' transport that handles both
+ * protocol formatting and network communication.
+ */
+export interface ITransport {
+  /**
+   * The base URL for the transport.
+   */
+  readonly baseUrl: string;
+
+  /**
+   * Gets a single tool from the server.
+   */
+  toolGet(
+    toolName: string,
+    headers?: Record<string, string>,
+  ): Promise<ZodManifest>;
+
+  /**
+   * Lists available tools from the server.
+   */
+  toolsList(
+    toolsetName?: string,
+    headers?: Record<string, string>,
+  ): Promise<ZodManifest>;
+
+  /**
+   * Invokes a specific tool on the server.
+   */
+  toolInvoke(
+    toolName: string,
+    arguments_: Record<string, unknown>,
+    headers: Record<string, string>,
+  ): Promise<string>;
+}