feat(js/core): added middleware to reflection api

js/ai/src/index.ts

@@ -154,6 +154,8 @@ export {
   defineInterrupt,
   defineTool,
   interrupt,
+  respondTool,
+  restartTool,
   type InterruptConfig,
   type ToolAction,
   type ToolArgument,

js/ai/src/tool.ts

@@ -348,20 +348,11 @@ function implementTool<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(
         "The 'tool.reply' method is part of the 'interrupts' beta feature."
       );
     }
-    parseSchema(responseData, {
+    responseData = parseSchema(responseData, {
       jsonSchema: config.outputJsonSchema,
       schema: config.outputSchema,
     });
-    return {
-      toolResponse: stripUndefinedProps({
-        name: interrupt.toolRequest.name,
-        ref: interrupt.toolRequest.ref,
-        output: responseData,
-      }),
-      metadata: {
-        interruptResponse: options?.metadata || true,
-      },
-    };
+    return respondTool(interrupt, responseData, options);
   };
 
   (a as ToolAction<I, O>).restart = (interrupt, resumedMetadata, options) => {
@@ -379,19 +370,7 @@ function implementTool<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(
         jsonSchema: config.inputJsonSchema,
       });
     }
-    return {
-      toolRequest: stripUndefinedProps({
-        name: interrupt.toolRequest.name,
-        ref: interrupt.toolRequest.ref,
-        input: replaceInput || interrupt.toolRequest.input,
-      }),
-      metadata: stripUndefinedProps({
-        ...interrupt.metadata,
-        resumed: resumedMetadata || true,
-        // annotate the original input if replacing it
-        replacedInput: replaceInput ? interrupt.toolRequest.input : undefined,
-      }),
-    };
+    return restartTool(interrupt, resumedMetadata, { replaceInput });
   };
 }
 
@@ -408,6 +387,72 @@ export type InterruptConfig<
       ) => Record<string, any> | Promise<Record<string, any>>);
 };
 
+/**
+ * restartTool constructs a tool request corresponding to the provided interrupt tool request
+ * that will then re-trigger the tool after e.g. a user confirms. In contrast to ToolAction.restart,
+ * this is a standalone utility that does not require an active ToolAction instance.
+ *
+ * @param interrupt The interrupt tool request you want to restart.
+ * @param resumedMetadata The metadata you want to provide to the tool to aide in reprocessing. Defaults to `true` if none is supplied.
+ * @param options Additional options for restarting the tool.
+ *
+ * @beta
+ */
+export function restartTool(
+  interrupt: ToolRequestPart,
+  resumedMetadata?: any,
+  options?: {
+    /**
+     * Replace the existing input arguments to the tool with different ones.
+     **/
+    replaceInput?: any;
+  }
+): ToolRequestPart {
+  let replaceInput = options?.replaceInput;
+  return {
+    toolRequest: stripUndefinedProps({
+      name: interrupt.toolRequest.name,
+      ref: interrupt.toolRequest.ref,
+      input: replaceInput ?? interrupt.toolRequest.input,
+    }),
+    metadata: stripUndefinedProps({
+      ...interrupt.metadata,
+      resumed: resumedMetadata ?? true,
+      // annotate the original input if replacing it
+      replacedInput:
+        replaceInput !== undefined ? interrupt.toolRequest.input : undefined,
+    }),
+  };
+}
+
+/**
+ * respondTool constructs a tool response part corresponding to the provided interrupt tool request
+ * that bypasses normal tool execution and sends a manual output result. In contrast to ToolAction.respond,
+ * this is a standalone utility that does not require an active ToolAction instance.
+ *
+ * @param interrupt The interrupt tool request you are responding to.
+ * @param responseData The manual output result you want to send back to the model.
+ * @param options Additional options for responding to the tool.
+ *
+ * @beta
+ */
+export function respondTool(
+  interrupt: ToolRequestPart,
+  responseData: any,
+  options?: { metadata?: any }
+): ToolResponsePart {
+  return {
+    toolResponse: stripUndefinedProps({
+      name: interrupt.toolRequest.name,
+      ref: interrupt.toolRequest.ref,
+      output: responseData,
+    }),
+    metadata: stripUndefinedProps({
+      interruptResponse: options?.metadata ?? true,
+    }),
+  };
+}
+
 export function isToolRequest(part: Part): part is ToolRequestPart {
   return !!part.toolRequest;
 }

js/ai/tests/tool_test.ts

@@ -24,6 +24,8 @@ import {
   defineTool,
   isDynamicTool,
   isMultipartTool,
+  respondTool,
+  restartTool,
   tool,
 } from '../src/tool.js';
 
@@ -458,3 +460,66 @@ describe('defineTool', () => {
     assert.deepStrictEqual(result, { output: 'foo' });
   });
 });
+
+describe('respondTool', () => {
+  it('constructs a ToolResponsePart standalone', () => {
+    const interrupt = { toolRequest: { name: 'test', input: {} } };
+    assert.deepStrictEqual(respondTool(interrupt, 'output'), {
+      toolResponse: {
+        name: 'test',
+        output: 'output',
+      },
+      metadata: {
+        interruptResponse: true,
+      },
+    });
+  });
+
+  it('includes metadata standalone', () => {
+    const interrupt = { toolRequest: { name: 'test', input: {} } };
+    assert.deepStrictEqual(
+      respondTool(interrupt, 'output', { metadata: { extra: 'data' } }),
+      {
+        toolResponse: {
+          name: 'test',
+          output: 'output',
+        },
+        metadata: {
+          interruptResponse: { extra: 'data' },
+        },
+      }
+    );
+  });
+});
+
+describe('restartTool', () => {
+  it('constructs a ToolRequestPart standalone', () => {
+    const interrupt = { toolRequest: { name: 'test', input: {} } };
+    assert.deepStrictEqual(restartTool(interrupt), {
+      toolRequest: {
+        name: 'test',
+        input: {},
+      },
+      metadata: {
+        resumed: true,
+      },
+    });
+  });
+
+  it('allows replacing input standalone', () => {
+    const interrupt = { toolRequest: { name: 'test', input: {} } };
+    assert.deepStrictEqual(
+      restartTool(interrupt, undefined, { replaceInput: 'new' }),
+      {
+        toolRequest: {
+          name: 'test',
+          input: 'new',
+        },
+        metadata: {
+          resumed: true,
+          replacedInput: {},
+        },
+      }
+    );
+  });
+});

js/core/src/reflection-v2.ts

@@ -346,10 +346,28 @@ export class ReflectionServerV2 {
   private async handleListValues(request: JsonRpcRequest) {
     if (!request.id) return;
     const { type } = ReflectionListValuesParamsSchema.parse(request.params);
+    if (type !== 'defaultModel' && type !== 'middleware') {
+      this.sendError(
+        request.id,
+        -32602,
+        `'type' ${type} is not supported. Only 'defaultModel' and 'middleware' are supported`
+      );
+      return;
+    }
     const values = await this.registry.listValues(type);
+    const mappedValues: Record<string, any> = {};
+    for (const [key, value] of Object.entries(values)) {
+      mappedValues[key] =
+        value &&
+        typeof value === 'object' &&
+        'toJson' in value &&
+        typeof (value as any).toJson === 'function'
+          ? (value as any).toJson()
+          : value;
+    }
     this.sendResponse(
       request.id,
-      ReflectionListValuesResponseSchema.parse({ values })
+      ReflectionListValuesResponseSchema.parse({ values: mappedValues })
     );
   }
 
@@ -361,7 +379,7 @@ export class ReflectionServerV2 {
     const action = await this.registry.lookupAction(key);
 
     if (!action) {
-      this.sendError(request.id, 404, `action ${key} not found`);
+      this.sendError(request.id, -32602, `action ${key} not found`);
       return;
     }
 
@@ -476,7 +494,11 @@ export class ReflectionServerV2 {
       this.activeActions.delete(traceId);
       this.sendResponse(request.id, { message: 'Action cancelled' });
     } else {
-      this.sendError(request.id, 404, 'Action not found or already completed');
+      this.sendError(
+        request.id,
+        -32602,
+        'Action not found or already completed'
+      );
     }
   }
 

js/core/src/reflection.ts

@@ -179,16 +179,25 @@ export class ReflectionServer {
           response.status(400).send('Query parameter "type" is required.');
           return;
         }
-        if (type !== 'defaultModel') {
+        if (type !== 'defaultModel' && type !== 'middleware') {
           response
             .status(400)
             .send(
-              `'type' ${type} is not supported. Only 'defaultModel' is supported`
+              `'type' ${type} is not supported. Only 'defaultModel' and 'middleware' are supported`
             );
           return;
         }
         const values = await this.registry.listValues(type as string);
-        response.send(values);
+        const mappedValues: Record<string, any> = {};
+        for (const [key, value] of Object.entries(values)) {
+          mappedValues[key] =
+            value &&
+            (value as any).toJson &&
+            typeof (value as any).toJson === 'function'
+              ? (value as any).toJson()
+              : value;
+        }
+        response.send(mappedValues);
       } catch (err) {
         const { message, stack } = err as Error;
         next({ message, stack });