feat(toolbox-core): add client headers (#23)

* feat: add client headers

* lint

* lint

* resolve comments

* disable any lint warning on a line

* Only intercept request if it is sent to the toolbox server

* feat(toolbox-core): Add helper methods for retrieving Google ID Tokens (#32)

* working code

* fix tests

* lint

* Revert "feat(toolbox-core): Add helper methods for retrieving Google ID Token…" (#39)

This reverts commit f747e60.

* remove addHeaders method

* fix tests. Still need to update to use # for private vars

* fix tests

* minor cleanup

* lint

* do not use interceptors

* Update client.ts

* lint

Author: twishabansal

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

@@ -14,32 +14,62 @@
 
 import {ToolboxTool} from './tool.js';
 import axios from 'axios';
-import {type AxiosInstance, type AxiosResponse} from 'axios';
+import {
+  type AxiosInstance,
+  type AxiosRequestConfig,
+  type AxiosResponse,
+} from 'axios';
 import {ZodManifestSchema, createZodSchemaFromParams} from './protocol.js';
 import {logApiError} from './errorUtils.js';
 import {ZodError} from 'zod';
-import {BoundParams, BoundValue} from './utils.js';
+import {BoundParams, BoundValue, resolveValue} from './utils.js';
 
 type Manifest = import('zod').infer<typeof ZodManifestSchema>;
 type ToolSchemaFromManifest = Manifest['tools'][string];
 
+// Types for dynamic headers
+export type HeaderFunction = () => string | Promise<string>;
+export type ClientHeaderProvider = string | HeaderFunction;
+export type ClientHeadersConfig = Record<string, ClientHeaderProvider>;
+
 /**
  * An asynchronous client for interacting with a Toolbox service.
- * Manages an Axios Client Session, if not provided.
  */
 class ToolboxClient {
   #baseUrl: string;
   #session: AxiosInstance;
+  #clientHeaders: ClientHeadersConfig;
 
   /**
    * Initializes the ToolboxClient.
    * @param {string} url - The base URL for the Toolbox service API (e.g., "http://localhost:5000").
    * @param {AxiosInstance} [session] - Optional Axios instance for making HTTP
    * requests. If not provided, a new one will be created.
+   * @param {ClientHeadersConfig} [clientHeaders] - Optional initial headers to
+   * be included in each request.
    */
-  constructor(url: string, session?: AxiosInstance) {
+  constructor(
+    url: string,
+    session?: AxiosInstance | null,
+    clientHeaders?: ClientHeadersConfig | null
+  ) {
     this.#baseUrl = url;
     this.#session = session || axios.create({baseURL: this.#baseUrl});
+    this.#clientHeaders = clientHeaders || {};
+  }
+
+  /**
+   * Resolves client headers from their provider functions.
+   * @returns {Promise<Record<string, string>>} A promise that resolves to the resolved headers.
+   */
+  async #resolveClientHeaders(): Promise<Record<string, string>> {
+    const resolvedEntries = await Promise.all(
+      Object.entries(this.#clientHeaders).map(async ([key, value]) => {
+        const resolved = await resolveValue(value);
+        return [key, String(resolved)];
+      })
+    );
+    return Object.fromEntries(resolvedEntries);
   }
 
   /**
@@ -51,7 +81,9 @@ class ToolboxClient {
   async #fetchAndParseManifest(apiPath: string): Promise<Manifest> {
     const url = `${this.#baseUrl}${apiPath}`;
     try {
-      const response: AxiosResponse = await this.#session.get(url);
+      const headers = await this.#resolveClientHeaders();
+      const config: AxiosRequestConfig = {headers};
+      const response: AxiosResponse = await this.#session.get(url, config);
       const responseData = response.data;
 
       try {
@@ -114,7 +146,8 @@ class ToolboxClient {
       toolName,
       toolSchema.description,
       paramZodSchema,
-      boundParams
+      applicableBoundParams,
+      this.#clientHeaders
     );
     return {tool, usedBoundKeys};
   }
@@ -125,8 +158,8 @@ class ToolboxClient {
    * returns a callable (`ToolboxTool`) that can be used to invoke the
    * tool remotely.
    *
-   * @param {BoundParams} [boundParams] - Optional parameters to pre-bind to the tool.
    * @param {string} name - The unique name or identifier of the tool to load.
+   * @param {BoundParams} [boundParams] - Optional parameters to pre-bind to the tool.
    * @returns {Promise<ReturnType<typeof ToolboxTool>>} A promise that resolves
    * to a ToolboxTool function, ready for execution.
    * @throws {Error} If the tool is not found in the manifest, the manifest structure is invalid,
@@ -157,7 +190,9 @@ class ToolboxClient {
 
       if (unusedBound.length > 0) {
         throw new Error(
-          `Validation failed for tool '${name}': unused bound parameters: ${unusedBound.join(', ')}.`
+          `Validation failed for tool '${name}': unused bound parameters: ${unusedBound.join(
+            ', '
+          )}.`
         );
       }
       return tool;
@@ -205,7 +240,9 @@ class ToolboxClient {
       throw new Error(
         `Validation failed for toolset '${
           name || 'default'
-        }': unused bound parameters could not be applied to any tool: ${unusedBound.join(', ')}.`
+        }': unused bound parameters could not be applied to any tool: ${unusedBound.join(
+          ', '
+        )}.`
       );
     }
     return tools;

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

@@ -13,9 +13,10 @@
 // limitations under the License.
 
 import {ZodObject, ZodError, ZodRawShape} from 'zod';
-import {AxiosInstance, AxiosResponse} from 'axios';
+import {AxiosInstance, AxiosRequestConfig, AxiosResponse} from 'axios';
 import {logApiError} from './errorUtils.js';
 import {BoundParams, BoundValue, resolveValue} from './utils.js';
+import {ClientHeadersConfig} from './client.js';
 
 /**
  * Creates a callable tool function representing a specific tool on a remote
@@ -27,6 +28,7 @@ import {BoundParams, BoundValue, resolveValue} from './utils.js';
  * @param {string} description - A description of the remote tool.
  * @param {ZodObject<any>} paramSchema - The Zod schema for validating the tool's parameters.
  * @param {BoundParams} [boundParams] - Optional parameters to pre-bind to the tool.
+ * @param {ClientHeadersConfig} [clientHeaders] - Optional client-specific headers.
  * @returns {CallableTool & CallableToolProperties} An async function that, when
  * called, invokes the tool with the provided arguments. Validates arguments
  * against the tool's signature, then sends them
@@ -39,7 +41,8 @@ function ToolboxTool(
   name: string,
   description: string,
   paramSchema: ZodObject<ZodRawShape>,
-  boundParams: BoundParams = {}
+  boundParams: BoundParams = {},
+  clientHeaders: ClientHeadersConfig = {}
 ) {
   const toolUrl = `${baseUrl}/api/tool/${name}/invoke`;
   const boundKeys = Object.keys(boundParams);
@@ -59,23 +62,42 @@ function ToolboxTool(
           e => `${e.path.join('.') || 'payload'}: ${e.message}`
         );
         throw new Error(
-          `Argument validation failed for tool "${name}":\n - ${errorMessages.join('\n - ')}`
+          `Argument validation failed for tool "${name}":\n - ${errorMessages.join(
+            '\n - '
+          )}`
         );
       }
       throw new Error(`Argument validation failed: ${String(error)}`);
     }
 
     // Resolve any bound parameters that are functions.
-    const resolvedEntries = await Promise.all(
+    const resolvedBoundEntries = await Promise.all(
       Object.entries(boundParams).map(async ([key, value]) => {
         const resolved = await resolveValue(value);
         return [key, resolved];
       })
     );
-    const resolvedBoundParams = Object.fromEntries(resolvedEntries);
+    const resolvedBoundParams = Object.fromEntries(resolvedBoundEntries);
+
+    // Resolve client headers
+    const resolvedHeaderEntries = await Promise.all(
+      Object.entries(clientHeaders).map(async ([key, value]) => {
+        const resolved = await resolveValue(value);
+        return [key, resolved];
+      })
+    );
+    const resolvedHeaders = Object.fromEntries(resolvedHeaderEntries);
+
     const payload = {...validatedUserArgs, ...resolvedBoundParams};
     try {
-      const response: AxiosResponse = await session.post(toolUrl, payload);
+      const config: AxiosRequestConfig = {
+        headers: resolvedHeaders,
+      };
+      const response: AxiosResponse = await session.post(
+        toolUrl,
+        payload,
+        config
+      );
       return response.data;
     } catch (error) {
       logApiError(`Error posting data to ${toolUrl}:`, error);
@@ -119,7 +141,8 @@ function ToolboxTool(
       this.toolName,
       this.description,
       this.params,
-      newBoundParams
+      newBoundParams,
+      clientHeaders
     );
   };