@mastra/client-js 0.0.0-fix-fetching-workflow-snapshots-20250625000954 → 0.0.0-fix-traces-pagination-plus-share-for-cloud-20250717083008

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,7 @@
1
1
  import {
2
2
  parsePartialJson,
3
3
  processDataStream,
4
+ processTextStream,
4
5
  type JSONValue,
5
6
  type ReasoningUIPart,
6
7
  type TextUIPart,
@@ -14,6 +15,7 @@ import type { JSONSchema7 } from 'json-schema';
14
15
  import { ZodSchema } from 'zod';
15
16
  import { zodToJsonSchema } from '../utils/zod-to-json-schema';
16
17
  import { processClientTools } from '../utils/process-client-tools';
18
+ import { v4 as uuid } from '@lukeed/uuid';
17
19
 
18
20
  import type {
19
21
  GenerateParams,
@@ -144,12 +146,18 @@ export class Agent extends BaseResource {
144
146
  });
145
147
 
146
148
  if (response.finishReason === 'tool-calls') {
147
- for (const toolCall of (
149
+ const toolCalls = (
148
150
  response as unknown as {
149
151
  toolCalls: { toolName: string; args: any; toolCallId: string }[];
150
152
  messages: CoreMessage[];
151
153
  }
152
- ).toolCalls) {
154
+ ).toolCalls;
155
+
156
+ if (!toolCalls || !Array.isArray(toolCalls)) {
157
+ return response;
158
+ }
159
+
160
+ for (const toolCall of toolCalls) {
153
161
  const clientTool = params.clientTools?.[toolCall.toolName] as Tool;
154
162
 
155
163
  if (clientTool && clientTool.execute) {
@@ -198,6 +206,7 @@ export class Agent extends BaseResource {
198
206
  onFinish,
199
207
  getCurrentDate = () => new Date(),
200
208
  lastMessage,
209
+ streamProtocol,
201
210
  }: {
202
211
  stream: ReadableStream<Uint8Array>;
203
212
  update: (options: { message: UIMessage; data: JSONValue[] | undefined; replaceLastMessage: boolean }) => void;
@@ -206,6 +215,7 @@ export class Agent extends BaseResource {
206
215
  generateId?: () => string;
207
216
  getCurrentDate?: () => Date;
208
217
  lastMessage: UIMessage | undefined;
218
+ streamProtocol: 'text' | 'data';
209
219
  }) {
210
220
  const replaceLastMessage = lastMessage?.role === 'assistant';
211
221
  let step = replaceLastMessage
@@ -219,7 +229,7 @@ export class Agent extends BaseResource {
219
229
  const message: UIMessage = replaceLastMessage
220
230
  ? structuredClone(lastMessage)
221
231
  : {
222
- id: crypto.randomUUID(),
232
+ id: uuid(),
223
233
  createdAt: getCurrentDate(),
224
234
  role: 'assistant',
225
235
  content: '',
@@ -279,7 +289,7 @@ export class Agent extends BaseResource {
279
289
  // changes. This is why we need to add a revision id to ensure that the message
280
290
  // is updated with SWR (without it, the changes get stuck in SWR and are not
281
291
  // forwarded to rendering):
282
- revisionId: crypto.randomUUID(),
292
+ revisionId: uuid(),
283
293
  } as UIMessage;
284
294
 
285
295
  update({
@@ -289,263 +299,419 @@ export class Agent extends BaseResource {
289
299
  });
290
300
  }
291
301
 
292
- await processDataStream({
293
- stream,
294
- onTextPart(value) {
295
- if (currentTextPart == null) {
296
- currentTextPart = {
297
- type: 'text',
298
- text: value,
299
- };
300
- message.parts.push(currentTextPart);
301
- } else {
302
- currentTextPart.text += value;
303
- }
302
+ if (streamProtocol === 'text') {
303
+ await processTextStream({
304
+ stream,
305
+ onTextPart(value) {
306
+ message.content += value;
307
+ execUpdate();
308
+ },
309
+ });
304
310
 
305
- message.content += value;
306
- execUpdate();
307
- },
308
- onReasoningPart(value) {
309
- if (currentReasoningTextDetail == null) {
310
- currentReasoningTextDetail = { type: 'text', text: value };
311
- if (currentReasoningPart != null) {
312
- currentReasoningPart.details.push(currentReasoningTextDetail);
311
+ onFinish?.({ message, finishReason, usage });
312
+ } else {
313
+ await processDataStream({
314
+ stream,
315
+ onTextPart(value) {
316
+ if (currentTextPart == null) {
317
+ currentTextPart = {
318
+ type: 'text',
319
+ text: value,
320
+ };
321
+ message.parts.push(currentTextPart);
322
+ } else {
323
+ currentTextPart.text += value;
313
324
  }
314
- } else {
315
- currentReasoningTextDetail.text += value;
316
- }
317
325
 
318
- if (currentReasoningPart == null) {
319
- currentReasoningPart = {
320
- type: 'reasoning',
321
- reasoning: value,
322
- details: [currentReasoningTextDetail],
323
- };
324
- message.parts.push(currentReasoningPart);
325
- } else {
326
- currentReasoningPart.reasoning += value;
327
- }
326
+ message.content += value;
327
+ execUpdate();
328
+ },
329
+ onReasoningPart(value) {
330
+ if (currentReasoningTextDetail == null) {
331
+ currentReasoningTextDetail = { type: 'text', text: value };
332
+ if (currentReasoningPart != null) {
333
+ currentReasoningPart.details.push(currentReasoningTextDetail);
334
+ }
335
+ } else {
336
+ currentReasoningTextDetail.text += value;
337
+ }
328
338
 
329
- message.reasoning = (message.reasoning ?? '') + value;
339
+ if (currentReasoningPart == null) {
340
+ currentReasoningPart = {
341
+ type: 'reasoning',
342
+ reasoning: value,
343
+ details: [currentReasoningTextDetail],
344
+ };
345
+ message.parts.push(currentReasoningPart);
346
+ } else {
347
+ currentReasoningPart.reasoning += value;
348
+ }
330
349
 
331
- execUpdate();
332
- },
333
- onReasoningSignaturePart(value) {
334
- if (currentReasoningTextDetail != null) {
335
- currentReasoningTextDetail.signature = value.signature;
336
- }
337
- },
338
- onRedactedReasoningPart(value) {
339
- if (currentReasoningPart == null) {
340
- currentReasoningPart = {
341
- type: 'reasoning',
342
- reasoning: '',
343
- details: [],
344
- };
345
- message.parts.push(currentReasoningPart);
346
- }
350
+ message.reasoning = (message.reasoning ?? '') + value;
347
351
 
348
- currentReasoningPart.details.push({
349
- type: 'redacted',
350
- data: value.data,
351
- });
352
+ execUpdate();
353
+ },
354
+ onReasoningSignaturePart(value) {
355
+ if (currentReasoningTextDetail != null) {
356
+ currentReasoningTextDetail.signature = value.signature;
357
+ }
358
+ },
359
+ onRedactedReasoningPart(value) {
360
+ if (currentReasoningPart == null) {
361
+ currentReasoningPart = {
362
+ type: 'reasoning',
363
+ reasoning: '',
364
+ details: [],
365
+ };
366
+ message.parts.push(currentReasoningPart);
367
+ }
352
368
 
353
- currentReasoningTextDetail = undefined;
369
+ currentReasoningPart.details.push({
370
+ type: 'redacted',
371
+ data: value.data,
372
+ });
354
373
 
355
- execUpdate();
356
- },
357
- onFilePart(value) {
358
- message.parts.push({
359
- type: 'file',
360
- mimeType: value.mimeType,
361
- data: value.data,
362
- });
374
+ currentReasoningTextDetail = undefined;
363
375
 
364
- execUpdate();
365
- },
366
- onSourcePart(value) {
367
- message.parts.push({
368
- type: 'source',
369
- source: value,
370
- });
376
+ execUpdate();
377
+ },
378
+ onFilePart(value) {
379
+ message.parts.push({
380
+ type: 'file',
381
+ mimeType: value.mimeType,
382
+ data: value.data,
383
+ });
371
384
 
372
- execUpdate();
373
- },
374
- onToolCallStreamingStartPart(value) {
375
- if (message.toolInvocations == null) {
376
- message.toolInvocations = [];
377
- }
385
+ execUpdate();
386
+ },
387
+ onSourcePart(value) {
388
+ message.parts.push({
389
+ type: 'source',
390
+ source: value,
391
+ });
378
392
 
379
- // add the partial tool call to the map
380
- partialToolCalls[value.toolCallId] = {
381
- text: '',
382
- step,
383
- toolName: value.toolName,
384
- index: message.toolInvocations.length,
385
- };
393
+ execUpdate();
394
+ },
395
+ onToolCallStreamingStartPart(value) {
396
+ if (message.toolInvocations == null) {
397
+ message.toolInvocations = [];
398
+ }
399
+
400
+ // add the partial tool call to the map
401
+ partialToolCalls[value.toolCallId] = {
402
+ text: '',
403
+ step,
404
+ toolName: value.toolName,
405
+ index: message.toolInvocations.length,
406
+ };
386
407
 
387
- const invocation = {
388
- state: 'partial-call',
389
- step,
390
- toolCallId: value.toolCallId,
391
- toolName: value.toolName,
392
- args: undefined,
393
- } as const;
408
+ const invocation = {
409
+ state: 'partial-call',
410
+ step,
411
+ toolCallId: value.toolCallId,
412
+ toolName: value.toolName,
413
+ args: undefined,
414
+ } as const;
394
415
 
395
- message.toolInvocations.push(invocation);
416
+ message.toolInvocations.push(invocation);
396
417
 
397
- updateToolInvocationPart(value.toolCallId, invocation);
418
+ updateToolInvocationPart(value.toolCallId, invocation);
398
419
 
399
- execUpdate();
400
- },
401
- onToolCallDeltaPart(value) {
402
- const partialToolCall = partialToolCalls[value.toolCallId];
420
+ execUpdate();
421
+ },
422
+ onToolCallDeltaPart(value) {
423
+ const partialToolCall = partialToolCalls[value.toolCallId];
403
424
 
404
- partialToolCall!.text += value.argsTextDelta;
425
+ partialToolCall!.text += value.argsTextDelta;
405
426
 
406
- const { value: partialArgs } = parsePartialJson(partialToolCall!.text);
427
+ const { value: partialArgs } = parsePartialJson(partialToolCall!.text);
407
428
 
408
- const invocation = {
409
- state: 'partial-call',
410
- step: partialToolCall!.step,
411
- toolCallId: value.toolCallId,
412
- toolName: partialToolCall!.toolName,
413
- args: partialArgs,
414
- } as const;
429
+ const invocation = {
430
+ state: 'partial-call',
431
+ step: partialToolCall!.step,
432
+ toolCallId: value.toolCallId,
433
+ toolName: partialToolCall!.toolName,
434
+ args: partialArgs,
435
+ } as const;
415
436
 
416
- message.toolInvocations![partialToolCall!.index] = invocation;
437
+ message.toolInvocations![partialToolCall!.index] = invocation;
417
438
 
418
- updateToolInvocationPart(value.toolCallId, invocation);
439
+ updateToolInvocationPart(value.toolCallId, invocation);
419
440
 
420
- execUpdate();
421
- },
422
- async onToolCallPart(value) {
423
- const invocation = {
424
- state: 'call',
425
- step,
426
- ...value,
427
- } as const;
428
-
429
- if (partialToolCalls[value.toolCallId] != null) {
430
- // change the partial tool call to a full tool call
431
- message.toolInvocations![partialToolCalls[value.toolCallId]!.index] = invocation;
432
- } else {
433
- if (message.toolInvocations == null) {
434
- message.toolInvocations = [];
441
+ execUpdate();
442
+ },
443
+ async onToolCallPart(value) {
444
+ const invocation = {
445
+ state: 'call',
446
+ step,
447
+ ...value,
448
+ } as const;
449
+
450
+ if (partialToolCalls[value.toolCallId] != null) {
451
+ // change the partial tool call to a full tool call
452
+ message.toolInvocations![partialToolCalls[value.toolCallId]!.index] = invocation;
453
+ } else {
454
+ if (message.toolInvocations == null) {
455
+ message.toolInvocations = [];
456
+ }
457
+
458
+ message.toolInvocations.push(invocation);
435
459
  }
436
460
 
437
- message.toolInvocations.push(invocation);
438
- }
461
+ updateToolInvocationPart(value.toolCallId, invocation);
439
462
 
440
- updateToolInvocationPart(value.toolCallId, invocation);
463
+ execUpdate();
441
464
 
442
- execUpdate();
465
+ // invoke the onToolCall callback if it exists. This is blocking.
466
+ // In the future we should make this non-blocking, which
467
+ // requires additional state management for error handling etc.
468
+ if (onToolCall) {
469
+ const result = await onToolCall({ toolCall: value });
470
+ if (result != null) {
471
+ const invocation = {
472
+ state: 'result',
473
+ step,
474
+ ...value,
475
+ result,
476
+ } as const;
443
477
 
444
- // invoke the onToolCall callback if it exists. This is blocking.
445
- // In the future we should make this non-blocking, which
446
- // requires additional state management for error handling etc.
447
- if (onToolCall) {
448
- const result = await onToolCall({ toolCall: value });
449
- if (result != null) {
450
- const invocation = {
451
- state: 'result',
452
- step,
453
- ...value,
454
- result,
455
- } as const;
478
+ // store the result in the tool invocation
479
+ message.toolInvocations![message.toolInvocations!.length - 1] = invocation;
456
480
 
457
- // store the result in the tool invocation
458
- message.toolInvocations![message.toolInvocations!.length - 1] = invocation;
481
+ updateToolInvocationPart(value.toolCallId, invocation);
459
482
 
460
- updateToolInvocationPart(value.toolCallId, invocation);
483
+ execUpdate();
484
+ }
485
+ }
486
+ },
487
+ onToolResultPart(value) {
488
+ const toolInvocations = message.toolInvocations;
461
489
 
462
- execUpdate();
490
+ if (toolInvocations == null) {
491
+ throw new Error('tool_result must be preceded by a tool_call');
463
492
  }
464
- }
465
- },
466
- onToolResultPart(value) {
467
- const toolInvocations = message.toolInvocations;
468
493
 
469
- if (toolInvocations == null) {
470
- throw new Error('tool_result must be preceded by a tool_call');
471
- }
494
+ // find if there is any tool invocation with the same toolCallId
495
+ // and replace it with the result
496
+ const toolInvocationIndex = toolInvocations.findIndex(
497
+ invocation => invocation.toolCallId === value.toolCallId,
498
+ );
472
499
 
473
- // find if there is any tool invocation with the same toolCallId
474
- // and replace it with the result
475
- const toolInvocationIndex = toolInvocations.findIndex(invocation => invocation.toolCallId === value.toolCallId);
500
+ if (toolInvocationIndex === -1) {
501
+ throw new Error('tool_result must be preceded by a tool_call with the same toolCallId');
502
+ }
476
503
 
477
- if (toolInvocationIndex === -1) {
478
- throw new Error('tool_result must be preceded by a tool_call with the same toolCallId');
479
- }
504
+ const invocation = {
505
+ ...toolInvocations[toolInvocationIndex],
506
+ state: 'result' as const,
507
+ ...value,
508
+ } as const;
480
509
 
481
- const invocation = {
482
- ...toolInvocations[toolInvocationIndex],
483
- state: 'result' as const,
484
- ...value,
485
- } as const;
510
+ toolInvocations[toolInvocationIndex] = invocation as ToolInvocation;
486
511
 
487
- toolInvocations[toolInvocationIndex] = invocation as ToolInvocation;
512
+ updateToolInvocationPart(value.toolCallId, invocation as ToolInvocation);
488
513
 
489
- updateToolInvocationPart(value.toolCallId, invocation as ToolInvocation);
514
+ execUpdate();
515
+ },
516
+ onDataPart(value) {
517
+ data.push(...value);
518
+ execUpdate();
519
+ },
520
+ onMessageAnnotationsPart(value) {
521
+ if (messageAnnotations == null) {
522
+ messageAnnotations = [...value];
523
+ } else {
524
+ messageAnnotations.push(...value);
525
+ }
490
526
 
491
- execUpdate();
492
- },
493
- onDataPart(value) {
494
- data.push(...value);
495
- execUpdate();
496
- },
497
- onMessageAnnotationsPart(value) {
498
- if (messageAnnotations == null) {
499
- messageAnnotations = [...value];
500
- } else {
501
- messageAnnotations.push(...value);
502
- }
527
+ execUpdate();
528
+ },
529
+ onFinishStepPart(value) {
530
+ step += 1;
503
531
 
504
- execUpdate();
505
- },
506
- onFinishStepPart(value) {
507
- step += 1;
532
+ // reset the current text and reasoning parts
533
+ currentTextPart = value.isContinued ? currentTextPart : undefined;
534
+ currentReasoningPart = undefined;
535
+ currentReasoningTextDetail = undefined;
536
+ },
537
+ onStartStepPart(value) {
538
+ // keep message id stable when we are updating an existing message:
539
+ if (!replaceLastMessage) {
540
+ message.id = value.messageId;
541
+ }
508
542
 
509
- // reset the current text and reasoning parts
510
- currentTextPart = value.isContinued ? currentTextPart : undefined;
511
- currentReasoningPart = undefined;
512
- currentReasoningTextDetail = undefined;
513
- },
514
- onStartStepPart(value) {
515
- // keep message id stable when we are updating an existing message:
516
- if (!replaceLastMessage) {
517
- message.id = value.messageId;
518
- }
543
+ // add a step boundary part to the message
544
+ message.parts.push({ type: 'step-start' });
545
+ execUpdate();
546
+ },
547
+ onFinishMessagePart(value) {
548
+ finishReason = value.finishReason;
549
+ if (value.usage != null) {
550
+ // usage = calculateLanguageModelUsage(value.usage);
551
+ usage = value.usage;
552
+ }
553
+ },
554
+ onErrorPart(error) {
555
+ throw new Error(error);
556
+ },
557
+ });
519
558
 
520
- // add a step boundary part to the message
521
- message.parts.push({ type: 'step-start' });
522
- execUpdate();
523
- },
524
- onFinishMessagePart(value) {
525
- finishReason = value.finishReason;
526
- if (value.usage != null) {
527
- // usage = calculateLanguageModelUsage(value.usage);
528
- usage = value.usage;
529
- }
530
- },
531
- onErrorPart(error) {
532
- throw new Error(error);
533
- },
559
+ onFinish?.({ message, finishReason, usage });
560
+ }
561
+ }
562
+
563
+ /**
564
+ * Processes the stream response and handles tool calls
565
+ */
566
+ private async processStreamResponse(processedParams: any, writable: WritableStream<Uint8Array>) {
567
+ const response: Response & {
568
+ processDataStream: (options?: Omit<Parameters<typeof processDataStream>[0], 'stream'>) => Promise<void>;
569
+ } = await this.request(`/api/agents/${this.agentId}/stream`, {
570
+ method: 'POST',
571
+ body: processedParams,
572
+ stream: true,
534
573
  });
535
574
 
536
- onFinish?.({ message, finishReason, usage });
575
+ if (!response.body) {
576
+ throw new Error('No response body');
577
+ }
578
+
579
+ try {
580
+ const streamProtocol = processedParams.output ? 'text' : 'data';
581
+ let toolCalls: ToolInvocation[] = [];
582
+ let finishReasonToolCalls = false;
583
+ let messages: UIMessage[] = [];
584
+ let hasProcessedToolCalls = false;
585
+
586
+ // Use tee() to split the stream into two branches
587
+ const [streamForWritable, streamForProcessing] = response.body.tee();
588
+
589
+ // Pipe one branch to the writable stream
590
+ streamForWritable
591
+ .pipeTo(writable, {
592
+ preventClose: true,
593
+ })
594
+ .catch(error => {
595
+ console.error('Error piping to writable stream:', error);
596
+ });
597
+
598
+ // Process the other branch for chat response handling
599
+ this.processChatResponse({
600
+ stream: streamForProcessing,
601
+ update: ({ message }) => {
602
+ messages.push(message);
603
+ },
604
+ onFinish: async ({ finishReason, message }) => {
605
+ if (finishReason === 'tool-calls') {
606
+ const toolCall = [...(message?.parts ?? [])]
607
+ .reverse()
608
+ .find(part => part.type === 'tool-invocation')?.toolInvocation;
609
+ if (toolCall) {
610
+ toolCalls.push(toolCall);
611
+ }
612
+
613
+ // Handle tool calls if needed
614
+ for (const toolCall of toolCalls) {
615
+ const clientTool = processedParams.clientTools?.[toolCall.toolName] as Tool;
616
+ if (clientTool && clientTool.execute) {
617
+ const result = await clientTool.execute(
618
+ {
619
+ context: toolCall?.args,
620
+ runId: processedParams.runId,
621
+ resourceId: processedParams.resourceId,
622
+ threadId: processedParams.threadId,
623
+ runtimeContext: processedParams.runtimeContext as RuntimeContext,
624
+ },
625
+ {
626
+ messages: (response as unknown as { messages: CoreMessage[] }).messages,
627
+ toolCallId: toolCall?.toolCallId,
628
+ },
629
+ );
630
+
631
+ const lastMessage: UIMessage = JSON.parse(JSON.stringify(messages[messages.length - 1]));
632
+
633
+ const toolInvocationPart = lastMessage?.parts?.find(
634
+ part => part.type === 'tool-invocation' && part.toolInvocation?.toolCallId === toolCall.toolCallId,
635
+ ) as ToolInvocationUIPart | undefined;
636
+
637
+ if (toolInvocationPart) {
638
+ toolInvocationPart.toolInvocation = {
639
+ ...toolInvocationPart.toolInvocation,
640
+ state: 'result',
641
+ result,
642
+ };
643
+ }
644
+
645
+ const toolInvocation = lastMessage?.toolInvocations?.find(
646
+ toolInvocation => toolInvocation.toolCallId === toolCall.toolCallId,
647
+ ) as ToolInvocation | undefined;
648
+
649
+ if (toolInvocation) {
650
+ toolInvocation.state = 'result';
651
+ // @ts-ignore
652
+ toolInvocation.result = result;
653
+ }
654
+
655
+ // write the tool result part to the stream
656
+ const writer = writable.getWriter();
657
+
658
+ try {
659
+ await writer.write(
660
+ new TextEncoder().encode(
661
+ 'a:' +
662
+ JSON.stringify({
663
+ toolCallId: toolCall.toolCallId,
664
+ result,
665
+ }) +
666
+ '\n',
667
+ ),
668
+ );
669
+ } finally {
670
+ writer.releaseLock();
671
+ }
672
+
673
+ // Convert messages to the correct format for the recursive call
674
+ const originalMessages = processedParams.messages;
675
+ const messageArray = Array.isArray(originalMessages) ? originalMessages : [originalMessages];
676
+
677
+ // Recursively call stream with updated messages
678
+ this.processStreamResponse(
679
+ {
680
+ ...processedParams,
681
+ messages: [...messageArray, ...messages, lastMessage],
682
+ },
683
+ writable,
684
+ );
685
+ }
686
+ }
687
+ } else {
688
+ setTimeout(() => {
689
+ if (!writable.locked) {
690
+ writable.close();
691
+ }
692
+ }, 0);
693
+ }
694
+ },
695
+ lastMessage: undefined,
696
+ streamProtocol,
697
+ });
698
+ } catch (error) {
699
+ console.error('Error processing stream response:', error);
700
+ }
701
+ return response;
537
702
  }
538
703
 
539
704
  /**
540
705
  * Streams a response from the agent
541
706
  * @param params - Stream parameters including prompt
542
- * @returns Promise containing the enhanced Response object with processDataStream method
707
+ * @returns Promise containing the enhanced Response object with processDataStream and processTextStream methods
543
708
  */
544
709
  async stream<T extends JSONSchema7 | ZodSchema | undefined = undefined>(
545
710
  params: StreamParams<T>,
546
711
  ): Promise<
547
712
  Response & {
548
713
  processDataStream: (options?: Omit<Parameters<typeof processDataStream>[0], 'stream'>) => Promise<void>;
714
+ processTextStream: (options?: Omit<Parameters<typeof processTextStream>[0], 'stream'>) => Promise<void>;
549
715
  }
550
716
  > {
551
717
  const processedParams = {
@@ -558,7 +724,6 @@ export class Agent extends BaseResource {
558
724
 
559
725
  // Create a readable stream that will handle the response processing
560
726
  const { readable, writable } = new TransformStream<Uint8Array, Uint8Array>();
561
-
562
727
  // Start processing the response in the background
563
728
  const response = await this.processStreamResponse(processedParams, writable);
564
729
 
@@ -569,6 +734,7 @@ export class Agent extends BaseResource {
569
734
  headers: response.headers,
570
735
  }) as Response & {
571
736
  processDataStream: (options?: Omit<Parameters<typeof processDataStream>[0], 'stream'>) => Promise<void>;
737
+ processTextStream: (options?: Omit<Parameters<typeof processTextStream>[0], 'stream'>) => Promise<void>;
572
738
  };
573
739
 
574
740
  // Add the processDataStream method to the response
@@ -579,116 +745,15 @@ export class Agent extends BaseResource {
579
745
  });
580
746
  };
581
747
 
582
- return streamResponse;
583
- }
584
-
585
- /**
586
- * Processes the stream response and handles tool calls
587
- */
588
- private async processStreamResponse(processedParams: any, writable: WritableStream<Uint8Array>) {
589
- const response: Response & {
590
- processDataStream: (options?: Omit<Parameters<typeof processDataStream>[0], 'stream'>) => Promise<void>;
591
- } = await this.request(`/api/agents/${this.agentId}/stream`, {
592
- method: 'POST',
593
- body: processedParams,
594
- stream: true,
595
- });
596
-
597
- if (!response.body) {
598
- throw new Error('No response body');
599
- }
600
-
601
- try {
602
- let toolCalls: ToolInvocation[] = [];
603
- let finishReasonToolCalls = false;
604
- let messages: UIMessage[] = [];
605
- let hasProcessedToolCalls = false;
606
-
607
- response.clone().body!.pipeTo(writable);
608
-
609
- await this.processChatResponse({
610
- stream: response.clone().body!,
611
- update: ({ message }) => {
612
- messages.push(message);
613
- },
614
- onFinish: ({ finishReason, message }) => {
615
- if (finishReason === 'tool-calls') {
616
- finishReasonToolCalls = true;
617
- const toolCall = [...(message?.parts ?? [])]
618
- .reverse()
619
- .find(part => part.type === 'tool-invocation')?.toolInvocation;
620
- if (toolCall) {
621
- toolCalls.push(toolCall);
622
- }
623
- }
624
- },
625
- lastMessage: undefined,
748
+ //Add the processTextStream method to the response
749
+ streamResponse.processTextStream = async options => {
750
+ await processTextStream({
751
+ stream: streamResponse.body as ReadableStream<Uint8Array>,
752
+ onTextPart: options?.onTextPart ?? (() => {}),
626
753
  });
754
+ };
627
755
 
628
- // Handle tool calls if needed
629
- if (finishReasonToolCalls && !hasProcessedToolCalls) {
630
- hasProcessedToolCalls = true;
631
-
632
- for (const toolCall of toolCalls) {
633
- const clientTool = processedParams.clientTools?.[toolCall.toolName] as Tool;
634
- if (clientTool && clientTool.execute) {
635
- const result = await clientTool.execute(
636
- {
637
- context: toolCall?.args,
638
- runId: processedParams.runId,
639
- resourceId: processedParams.resourceId,
640
- threadId: processedParams.threadId,
641
- runtimeContext: processedParams.runtimeContext as RuntimeContext,
642
- },
643
- {
644
- messages: (response as unknown as { messages: CoreMessage[] }).messages,
645
- toolCallId: toolCall?.toolCallId,
646
- },
647
- );
648
-
649
- const lastMessage: UIMessage = JSON.parse(JSON.stringify(messages[messages.length - 1]));
650
-
651
- const toolInvocationPart = lastMessage?.parts?.find(
652
- part => part.type === 'tool-invocation' && part.toolInvocation?.toolCallId === toolCall.toolCallId,
653
- ) as ToolInvocationUIPart | undefined;
654
-
655
- if (toolInvocationPart) {
656
- toolInvocationPart.toolInvocation = {
657
- ...toolInvocationPart.toolInvocation,
658
- state: 'result',
659
- result,
660
- };
661
- }
662
-
663
- const toolInvocation = lastMessage?.toolInvocations?.find(
664
- toolInvocation => toolInvocation.toolCallId === toolCall.toolCallId,
665
- ) as ToolInvocation | undefined;
666
-
667
- if (toolInvocation) {
668
- toolInvocation.state = 'result';
669
- // @ts-ignore
670
- toolInvocation.result = result;
671
- }
672
-
673
- // Convert messages to the correct format for the recursive call
674
- const originalMessages = processedParams.messages;
675
- const messageArray = Array.isArray(originalMessages) ? originalMessages : [originalMessages];
676
-
677
- // Recursively call stream with updated messages
678
- this.processStreamResponse(
679
- {
680
- ...processedParams,
681
- messages: [...messageArray, ...messages, lastMessage],
682
- },
683
- writable,
684
- );
685
- }
686
- }
687
- }
688
- } catch (error) {
689
- console.error('Error processing stream response:', error);
690
- }
691
- return response;
756
+ return streamResponse;
692
757
  }
693
758
 
694
759
  /**