@mastra/client-js 0.10.15-alpha.1 → 0.10.15-alpha.2

package/src/resources/agent.ts

@@ -1,7 +1,6 @@
 import {
   parsePartialJson,
   processDataStream,
-  processTextStream,
   type JSONValue,
   type ReasoningUIPart,
   type TextUIPart,
@@ -10,7 +9,8 @@ import {
   type UIMessage,
   type UseChatOptions,
 } from '@ai-sdk/ui-utils';
-import { Tool, type CoreMessage, type GenerateReturn } from '@mastra/core';
+import { Tool, type CoreMessage } from '@mastra/core';
+import { type GenerateReturn } from '@mastra/core/llm';
 import type { JSONSchema7 } from 'json-schema';
 import { ZodSchema } from 'zod';
 import { zodToJsonSchema } from '../utils/zod-to-json-schema';
@@ -118,18 +118,19 @@ export class Agent extends BaseResource {
    * @param params - Generation parameters including prompt
    * @returns Promise containing the generated response
    */
-  async generate<T extends JSONSchema7 | ZodSchema | undefined = undefined>(
-    params: GenerateParams<T> & { output?: never; experimental_output?: never },
-  ): Promise<GenerateReturn<T>>;
-  async generate<T extends JSONSchema7 | ZodSchema | undefined = undefined>(
-    params: GenerateParams<T> & { output: T; experimental_output?: never },
-  ): Promise<GenerateReturn<T>>;
-  async generate<T extends JSONSchema7 | ZodSchema | undefined = undefined>(
-    params: GenerateParams<T> & { output?: never; experimental_output: T },
-  ): Promise<GenerateReturn<T>>;
-  async generate<T extends JSONSchema7 | ZodSchema | undefined = undefined>(
-    params: GenerateParams<T>,
-  ): Promise<GenerateReturn<T>> {
+  async generate(
+    params: GenerateParams<undefined> & { output?: never; experimental_output?: never },
+  ): Promise<GenerateReturn<any, undefined, undefined>>;
+  async generate<Output extends JSONSchema7 | ZodSchema>(
+    params: GenerateParams<Output> & { output: Output; experimental_output?: never },
+  ): Promise<GenerateReturn<any, Output, undefined>>;
+  async generate<StructuredOutput extends JSONSchema7 | ZodSchema>(
+    params: GenerateParams<StructuredOutput> & { output?: never; experimental_output: StructuredOutput },
+  ): Promise<GenerateReturn<any, undefined, StructuredOutput>>;
+  async generate<
+    Output extends JSONSchema7 | ZodSchema | undefined = undefined,
+    StructuredOutput extends JSONSchema7 | ZodSchema | undefined = undefined,
+  >(params: GenerateParams<Output>): Promise<GenerateReturn<any, Output, StructuredOutput>> {
     const processedParams = {
       ...params,
       output: params.output ? zodToJsonSchema(params.output) : undefined,
@@ -140,10 +141,13 @@ export class Agent extends BaseResource {
 
     const { runId, resourceId, threadId, runtimeContext } = processedParams as GenerateParams;
 
-    const response: GenerateReturn<T> = await this.request(`/api/agents/${this.agentId}/generate`, {
-      method: 'POST',
-      body: processedParams,
-    });
+    const response: GenerateReturn<any, Output, StructuredOutput> = await this.request(
+      `/api/agents/${this.agentId}/generate`,
+      {
+        method: 'POST',
+        body: processedParams,
+      },
+    );
 
     if (response.finishReason === 'tool-calls') {
       const toolCalls = (
@@ -206,7 +210,6 @@ export class Agent extends BaseResource {
     onFinish,
     getCurrentDate = () => new Date(),
     lastMessage,
-    streamProtocol,
   }: {
     stream: ReadableStream<Uint8Array>;
     update: (options: { message: UIMessage; data: JSONValue[] | undefined; replaceLastMessage: boolean }) => void;
@@ -215,7 +218,6 @@ export class Agent extends BaseResource {
     generateId?: () => string;
     getCurrentDate?: () => Date;
     lastMessage: UIMessage | undefined;
-    streamProtocol: 'text' | 'data';
   }) {
     const replaceLastMessage = lastMessage?.role === 'assistant';
     let step = replaceLastMessage
@@ -299,265 +301,297 @@ export class Agent extends BaseResource {
       });
     }
 
-    if (streamProtocol === 'text') {
-      await processTextStream({
-        stream,
-        onTextPart(value) {
-          message.content += value;
-          execUpdate();
-        },
-      });
+    await processDataStream({
+      stream,
+      onTextPart(value) {
+        if (currentTextPart == null) {
+          currentTextPart = {
+            type: 'text',
+            text: value,
+          };
+          message.parts.push(currentTextPart);
+        } else {
+          currentTextPart.text += value;
+        }
 
-      onFinish?.({ message, finishReason, usage });
-    } else {
-      await processDataStream({
-        stream,
-        onTextPart(value) {
-          if (currentTextPart == null) {
-            currentTextPart = {
-              type: 'text',
-              text: value,
-            };
-            message.parts.push(currentTextPart);
-          } else {
-            currentTextPart.text += value;
+        message.content += value;
+        execUpdate();
+      },
+      onReasoningPart(value) {
+        if (currentReasoningTextDetail == null) {
+          currentReasoningTextDetail = { type: 'text', text: value };
+          if (currentReasoningPart != null) {
+            currentReasoningPart.details.push(currentReasoningTextDetail);
           }
+        } else {
+          currentReasoningTextDetail.text += value;
+        }
 
-          message.content += value;
-          execUpdate();
-        },
-        onReasoningPart(value) {
-          if (currentReasoningTextDetail == null) {
-            currentReasoningTextDetail = { type: 'text', text: value };
-            if (currentReasoningPart != null) {
-              currentReasoningPart.details.push(currentReasoningTextDetail);
-            }
-          } else {
-            currentReasoningTextDetail.text += value;
-          }
+        if (currentReasoningPart == null) {
+          currentReasoningPart = {
+            type: 'reasoning',
+            reasoning: value,
+            details: [currentReasoningTextDetail],
+          };
+          message.parts.push(currentReasoningPart);
+        } else {
+          currentReasoningPart.reasoning += value;
+        }
 
-          if (currentReasoningPart == null) {
-            currentReasoningPart = {
-              type: 'reasoning',
-              reasoning: value,
-              details: [currentReasoningTextDetail],
-            };
-            message.parts.push(currentReasoningPart);
-          } else {
-            currentReasoningPart.reasoning += value;
-          }
+        message.reasoning = (message.reasoning ?? '') + value;
 
-          message.reasoning = (message.reasoning ?? '') + value;
+        execUpdate();
+      },
+      onReasoningSignaturePart(value) {
+        if (currentReasoningTextDetail != null) {
+          currentReasoningTextDetail.signature = value.signature;
+        }
+      },
+      onRedactedReasoningPart(value) {
+        if (currentReasoningPart == null) {
+          currentReasoningPart = {
+            type: 'reasoning',
+            reasoning: '',
+            details: [],
+          };
+          message.parts.push(currentReasoningPart);
+        }
 
-          execUpdate();
-        },
-        onReasoningSignaturePart(value) {
-          if (currentReasoningTextDetail != null) {
-            currentReasoningTextDetail.signature = value.signature;
-          }
-        },
-        onRedactedReasoningPart(value) {
-          if (currentReasoningPart == null) {
-            currentReasoningPart = {
-              type: 'reasoning',
-              reasoning: '',
-              details: [],
-            };
-            message.parts.push(currentReasoningPart);
-          }
+        currentReasoningPart.details.push({
+          type: 'redacted',
+          data: value.data,
+        });
 
-          currentReasoningPart.details.push({
-            type: 'redacted',
-            data: value.data,
-          });
+        currentReasoningTextDetail = undefined;
 
-          currentReasoningTextDetail = undefined;
+        execUpdate();
+      },
+      onFilePart(value) {
+        message.parts.push({
+          type: 'file',
+          mimeType: value.mimeType,
+          data: value.data,
+        });
 
-          execUpdate();
-        },
-        onFilePart(value) {
-          message.parts.push({
-            type: 'file',
-            mimeType: value.mimeType,
-            data: value.data,
-          });
+        execUpdate();
+      },
+      onSourcePart(value) {
+        message.parts.push({
+          type: 'source',
+          source: value,
+        });
 
-          execUpdate();
-        },
-        onSourcePart(value) {
-          message.parts.push({
-            type: 'source',
-            source: value,
-          });
+        execUpdate();
+      },
+      onToolCallStreamingStartPart(value) {
+        if (message.toolInvocations == null) {
+          message.toolInvocations = [];
+        }
 
-          execUpdate();
-        },
-        onToolCallStreamingStartPart(value) {
-          if (message.toolInvocations == null) {
-            message.toolInvocations = [];
-          }
+        // add the partial tool call to the map
+        partialToolCalls[value.toolCallId] = {
+          text: '',
+          step,
+          toolName: value.toolName,
+          index: message.toolInvocations.length,
+        };
 
-          // add the partial tool call to the map
-          partialToolCalls[value.toolCallId] = {
-            text: '',
-            step,
-            toolName: value.toolName,
-            index: message.toolInvocations.length,
-          };
+        const invocation = {
+          state: 'partial-call',
+          step,
+          toolCallId: value.toolCallId,
+          toolName: value.toolName,
+          args: undefined,
+        } as const;
 
-          const invocation = {
-            state: 'partial-call',
-            step,
-            toolCallId: value.toolCallId,
-            toolName: value.toolName,
-            args: undefined,
-          } as const;
+        message.toolInvocations.push(invocation);
 
-          message.toolInvocations.push(invocation);
+        updateToolInvocationPart(value.toolCallId, invocation);
 
-          updateToolInvocationPart(value.toolCallId, invocation);
+        execUpdate();
+      },
+      onToolCallDeltaPart(value) {
+        const partialToolCall = partialToolCalls[value.toolCallId];
 
-          execUpdate();
-        },
-        onToolCallDeltaPart(value) {
-          const partialToolCall = partialToolCalls[value.toolCallId];
+        partialToolCall!.text += value.argsTextDelta;
 
-          partialToolCall!.text += value.argsTextDelta;
+        const { value: partialArgs } = parsePartialJson(partialToolCall!.text);
 
-          const { value: partialArgs } = parsePartialJson(partialToolCall!.text);
+        const invocation = {
+          state: 'partial-call',
+          step: partialToolCall!.step,
+          toolCallId: value.toolCallId,
+          toolName: partialToolCall!.toolName,
+          args: partialArgs,
+        } as const;
 
-          const invocation = {
-            state: 'partial-call',
-            step: partialToolCall!.step,
-            toolCallId: value.toolCallId,
-            toolName: partialToolCall!.toolName,
-            args: partialArgs,
-          } as const;
+        message.toolInvocations![partialToolCall!.index] = invocation;
 
-          message.toolInvocations![partialToolCall!.index] = invocation;
+        updateToolInvocationPart(value.toolCallId, invocation);
 
-          updateToolInvocationPart(value.toolCallId, invocation);
+        execUpdate();
+      },
+      async onToolCallPart(value) {
+        const invocation = {
+          state: 'call',
+          step,
+          ...value,
+        } as const;
+
+        if (partialToolCalls[value.toolCallId] != null) {
+          // change the partial tool call to a full tool call
+          message.toolInvocations![partialToolCalls[value.toolCallId]!.index] = invocation;
+        } else {
+          if (message.toolInvocations == null) {
+            message.toolInvocations = [];
+          }
 
-          execUpdate();
-        },
-        async onToolCallPart(value) {
-          const invocation = {
-            state: 'call',
-            step,
-            ...value,
-          } as const;
-
-          if (partialToolCalls[value.toolCallId] != null) {
-            // change the partial tool call to a full tool call
-            message.toolInvocations![partialToolCalls[value.toolCallId]!.index] = invocation;
-          } else {
-            if (message.toolInvocations == null) {
-              message.toolInvocations = [];
-            }
+          message.toolInvocations.push(invocation);
+        }
+
+        updateToolInvocationPart(value.toolCallId, invocation);
+
+        execUpdate();
+
+        // invoke the onToolCall callback if it exists. This is blocking.
+        // In the future we should make this non-blocking, which
+        // requires additional state management for error handling etc.
+        if (onToolCall) {
+          const result = await onToolCall({ toolCall: value });
+          if (result != null) {
+            const invocation = {
+              state: 'result',
+              step,
+              ...value,
+              result,
+            } as const;
+
+            // store the result in the tool invocation
+            message.toolInvocations![message.toolInvocations!.length - 1] = invocation;
 
-            message.toolInvocations.push(invocation);
+            updateToolInvocationPart(value.toolCallId, invocation);
+
+            execUpdate();
           }
+        }
+      },
+      onToolResultPart(value) {
+        const toolInvocations = message.toolInvocations;
 
-          updateToolInvocationPart(value.toolCallId, invocation);
+        if (toolInvocations == null) {
+          throw new Error('tool_result must be preceded by a tool_call');
+        }
 
-          execUpdate();
+        // find if there is any tool invocation with the same toolCallId
+        // and replace it with the result
+        const toolInvocationIndex = toolInvocations.findIndex(invocation => invocation.toolCallId === value.toolCallId);
 
-          // invoke the onToolCall callback if it exists. This is blocking.
-          // In the future we should make this non-blocking, which
-          // requires additional state management for error handling etc.
-          if (onToolCall) {
-            const result = await onToolCall({ toolCall: value });
-            if (result != null) {
-              const invocation = {
-                state: 'result',
-                step,
-                ...value,
-                result,
-              } as const;
+        if (toolInvocationIndex === -1) {
+          throw new Error('tool_result must be preceded by a tool_call with the same toolCallId');
+        }
 
-              // store the result in the tool invocation
-              message.toolInvocations![message.toolInvocations!.length - 1] = invocation;
+        const invocation = {
+          ...toolInvocations[toolInvocationIndex],
+          state: 'result' as const,
+          ...value,
+        } as const;
 
-              updateToolInvocationPart(value.toolCallId, invocation);
+        toolInvocations[toolInvocationIndex] = invocation as ToolInvocation;
 
-              execUpdate();
-            }
-          }
-        },
-        onToolResultPart(value) {
-          const toolInvocations = message.toolInvocations;
+        updateToolInvocationPart(value.toolCallId, invocation as ToolInvocation);
 
-          if (toolInvocations == null) {
-            throw new Error('tool_result must be preceded by a tool_call');
-          }
+        execUpdate();
+      },
+      onDataPart(value) {
+        data.push(...value);
+        execUpdate();
+      },
+      onMessageAnnotationsPart(value) {
+        if (messageAnnotations == null) {
+          messageAnnotations = [...value];
+        } else {
+          messageAnnotations.push(...value);
+        }
 
-          // find if there is any tool invocation with the same toolCallId
-          // and replace it with the result
-          const toolInvocationIndex = toolInvocations.findIndex(
-            invocation => invocation.toolCallId === value.toolCallId,
-          );
+        execUpdate();
+      },
+      onFinishStepPart(value) {
+        step += 1;
 
-          if (toolInvocationIndex === -1) {
-            throw new Error('tool_result must be preceded by a tool_call with the same toolCallId');
-          }
+        // reset the current text and reasoning parts
+        currentTextPart = value.isContinued ? currentTextPart : undefined;
+        currentReasoningPart = undefined;
+        currentReasoningTextDetail = undefined;
+      },
+      onStartStepPart(value) {
+        // keep message id stable when we are updating an existing message:
+        if (!replaceLastMessage) {
+          message.id = value.messageId;
+        }
 
-          const invocation = {
-            ...toolInvocations[toolInvocationIndex],
-            state: 'result' as const,
-            ...value,
-          } as const;
+        // add a step boundary part to the message
+        message.parts.push({ type: 'step-start' });
+        execUpdate();
+      },
+      onFinishMessagePart(value) {
+        finishReason = value.finishReason;
+        if (value.usage != null) {
+          // usage = calculateLanguageModelUsage(value.usage);
+          usage = value.usage;
+        }
+      },
+      onErrorPart(error) {
+        throw new Error(error);
+      },
+    });
 
-          toolInvocations[toolInvocationIndex] = invocation as ToolInvocation;
+    onFinish?.({ message, finishReason, usage });
+  }
 
-          updateToolInvocationPart(value.toolCallId, invocation as ToolInvocation);
+  /**
+   * Streams a response from the agent
+   * @param params - Stream parameters including prompt
+   * @returns Promise containing the enhanced Response object with processDataStream method
+   */
+  async stream<T extends JSONSchema7 | ZodSchema | undefined = undefined>(
+    params: StreamParams<T>,
+  ): Promise<
+    Response & {
+      processDataStream: (options?: Omit<Parameters<typeof processDataStream>[0], 'stream'>) => Promise<void>;
+    }
+  > {
+    const processedParams = {
+      ...params,
+      output: params.output ? zodToJsonSchema(params.output) : undefined,
+      experimental_output: params.experimental_output ? zodToJsonSchema(params.experimental_output) : undefined,
+      runtimeContext: parseClientRuntimeContext(params.runtimeContext),
+      clientTools: processClientTools(params.clientTools),
+    };
 
-          execUpdate();
-        },
-        onDataPart(value) {
-          data.push(...value);
-          execUpdate();
-        },
-        onMessageAnnotationsPart(value) {
-          if (messageAnnotations == null) {
-            messageAnnotations = [...value];
-          } else {
-            messageAnnotations.push(...value);
-          }
+    // Create a readable stream that will handle the response processing
+    const { readable, writable } = new TransformStream<Uint8Array, Uint8Array>();
 
-          execUpdate();
-        },
-        onFinishStepPart(value) {
-          step += 1;
+    // Start processing the response in the background
+    const response = await this.processStreamResponse(processedParams, writable);
 
-          // reset the current text and reasoning parts
-          currentTextPart = value.isContinued ? currentTextPart : undefined;
-          currentReasoningPart = undefined;
-          currentReasoningTextDetail = undefined;
-        },
-        onStartStepPart(value) {
-          // keep message id stable when we are updating an existing message:
-          if (!replaceLastMessage) {
-            message.id = value.messageId;
-          }
+    // Create a new response with the readable stream
+    const streamResponse = new Response(readable, {
+      status: response.status,
+      statusText: response.statusText,
+      headers: response.headers,
+    }) as Response & {
+      processDataStream: (options?: Omit<Parameters<typeof processDataStream>[0], 'stream'>) => Promise<void>;
+    };
 
-          // add a step boundary part to the message
-          message.parts.push({ type: 'step-start' });
-          execUpdate();
-        },
-        onFinishMessagePart(value) {
-          finishReason = value.finishReason;
-          if (value.usage != null) {
-            // usage = calculateLanguageModelUsage(value.usage);
-            usage = value.usage;
-          }
-        },
-        onErrorPart(error) {
-          throw new Error(error);
-        },
+    // Add the processDataStream method to the response
+    streamResponse.processDataStream = async (options = {}) => {
+      await processDataStream({
+        stream: streamResponse.body as ReadableStream<Uint8Array>,
+        ...options,
       });
+    };
 
-      onFinish?.({ message, finishReason, usage });
-    }
+    return streamResponse;
   }
 
   /**
@@ -577,7 +611,6 @@ export class Agent extends BaseResource {
     }
 
     try {
-      const streamProtocol = processedParams.output ? 'text' : 'data';
       let toolCalls: ToolInvocation[] = [];
       let finishReasonToolCalls = false;
       let messages: UIMessage[] = [];
@@ -686,14 +719,11 @@ export class Agent extends BaseResource {
             }
           } else {
             setTimeout(() => {
-              if (!writable.locked) {
-                writable.close();
-              }
+              writable.close();
             }, 0);
           }
         },
         lastMessage: undefined,
-        streamProtocol,
       });
     } catch (error) {
       console.error('Error processing stream response:', error);
@@ -701,61 +731,6 @@ export class Agent extends BaseResource {
     return response;
   }
 
-  /**
-   * Streams a response from the agent
-   * @param params - Stream parameters including prompt
-   * @returns Promise containing the enhanced Response object with processDataStream and processTextStream methods
-   */
-  async stream<T extends JSONSchema7 | ZodSchema | undefined = undefined>(
-    params: StreamParams<T>,
-  ): Promise<
-    Response & {
-      processDataStream: (options?: Omit<Parameters<typeof processDataStream>[0], 'stream'>) => Promise<void>;
-      processTextStream: (options?: Omit<Parameters<typeof processTextStream>[0], 'stream'>) => Promise<void>;
-    }
-  > {
-    const processedParams = {
-      ...params,
-      output: params.output ? zodToJsonSchema(params.output) : undefined,
-      experimental_output: params.experimental_output ? zodToJsonSchema(params.experimental_output) : undefined,
-      runtimeContext: parseClientRuntimeContext(params.runtimeContext),
-      clientTools: processClientTools(params.clientTools),
-    };
-
-    // Create a readable stream that will handle the response processing
-    const { readable, writable } = new TransformStream<Uint8Array, Uint8Array>();
-    // Start processing the response in the background
-    const response = await this.processStreamResponse(processedParams, writable);
-
-    // Create a new response with the readable stream
-    const streamResponse = new Response(readable, {
-      status: response.status,
-      statusText: response.statusText,
-      headers: response.headers,
-    }) as Response & {
-      processDataStream: (options?: Omit<Parameters<typeof processDataStream>[0], 'stream'>) => Promise<void>;
-      processTextStream: (options?: Omit<Parameters<typeof processTextStream>[0], 'stream'>) => Promise<void>;
-    };
-
-    // Add the processDataStream method to the response
-    streamResponse.processDataStream = async (options = {}) => {
-      await processDataStream({
-        stream: streamResponse.body as ReadableStream<Uint8Array>,
-        ...options,
-      });
-    };
-
-    //Add the processTextStream method to the response
-    streamResponse.processTextStream = async options => {
-      await processTextStream({
-        stream: streamResponse.body as ReadableStream<Uint8Array>,
-        onTextPart: options?.onTextPart ?? (() => {}),
-      });
-    };
-
-    return streamResponse;
-  }
-
   /**
    * Gets details about a specific tool available to the agent
    * @param toolId - ID of the tool to retrieve