[ai] Fix "start" and "finish" chunk events in DurableAgent (#378)

* [ai] Remove "start" and "finish" chunk events

* Wire through sendStart and sendFinish properly

* changeset

---------

Co-authored-by: Pranay Prakash <[email protected]>

Author: TooTallNate

.changeset/cuddly-lobsters-build.md

@@ -0,0 +1,5 @@
+---
+"@workflow/ai": patch
+---
+
+DurableAgent#stream now sends `start` and `finish` chunks properly at the start and end

packages/ai/src/agent/do-stream-step.ts

@@ -21,7 +21,10 @@ export async function doStreamStep(
   conversationPrompt: LanguageModelV2Prompt,
   modelInit: string | (() => Promise<LanguageModelV2>),
   writable: WritableStream<UIMessageChunk>,
-  tools?: LanguageModelV2CallOptions['tools']
+  tools?: LanguageModelV2CallOptions['tools'],
+  options?: {
+    sendStart?: boolean;
+  }
 ) {
   'use step';
 
@@ -65,9 +68,11 @@ export async function doStreamStep(
     .pipeThrough(
       new TransformStream<LanguageModelV2StreamPart, UIMessageChunk>({
         start: (controller) => {
-          controller.enqueue({
-            type: 'start',
-          });
+          if (options?.sendStart) {
+            controller.enqueue({
+              type: 'start',
+            });
+          }
           controller.enqueue({
             type: 'start-step',
           });
@@ -76,9 +81,6 @@ export async function doStreamStep(
           controller.enqueue({
             type: 'finish-step',
           });
-          controller.enqueue({
-            type: 'finish',
-          });
         },
         transform: async (part, controller) => {
           const partType = part.type;

packages/ai/src/agent/durable-agent.ts

@@ -63,6 +63,18 @@ export interface DurableAgentStreamOptions {
    */
   preventClose?: boolean;
 
+  /**
+   * If true, sends a 'start' chunk at the beginning of the stream.
+   * Defaults to true.
+   */
+  sendStart?: boolean;
+
+  /**
+   * If true, sends a 'finish' chunk at the end of the stream.
+   * Defaults to true.
+   */
+  sendFinish?: boolean;
+
   /**
    * Condition for stopping the generation when there are tool results in the last step.
    * When the condition is an array, any of the conditions can be met to stop the generation.
@@ -133,6 +145,7 @@ export class DurableAgent {
       writable: options.writable,
       prompt: modelPrompt,
       stopConditions: options.stopWhen,
+      sendStart: options.sendStart ?? true,
     });
 
     let result = await iterator.next();
@@ -147,16 +160,37 @@ export class DurableAgent {
       result = await iterator.next(toolResults);
     }
 
-    if (!options.preventClose) {
-      await closeStream(options.writable);
+    const sendFinish = options.sendFinish ?? true;
+    const preventClose = options.preventClose ?? false;
+
+    // Only call closeStream if there's something to do
+    if (sendFinish || !preventClose) {
+      await closeStream(options.writable, preventClose, sendFinish);
     }
   }
 }
 
-async function closeStream(writable: WritableStream<UIMessageChunk>) {
+async function closeStream(
+  writable: WritableStream<UIMessageChunk>,
+  preventClose?: boolean,
+  sendFinish?: boolean
+) {
   'use step';
 
-  await writable.close();
+  // Conditionally write the finish chunk
+  if (sendFinish) {
+    const writer = writable.getWriter();
+    try {
+      await writer.write({ type: 'finish' });
+    } finally {
+      writer.releaseLock();
+    }
+  }
+
+  // Conditionally close the stream
+  if (!preventClose) {
+    await writable.close();
+  }
 }
 
 async function executeTool(

packages/ai/src/agent/stream-text-iterator.ts

@@ -15,12 +15,14 @@ export async function* streamTextIterator({
   writable,
   model,
   stopConditions,
+  sendStart = true,
 }: {
   prompt: LanguageModelV2Prompt;
   tools: ToolSet;
   writable: WritableStream<UIMessageChunk>;
   model: string | (() => Promise<LanguageModelV2>);
   stopConditions?: ModelStopCondition[] | ModelStopCondition;
+  sendStart?: boolean;
 }): AsyncGenerator<
   LanguageModelV2ToolCall[],
   void,
@@ -30,13 +32,19 @@ export async function* streamTextIterator({
 
   const steps: StepResult<any>[] = [];
   let done = false;
+  let isFirstIteration = true;
+
   while (!done) {
     const { toolCalls, finish, step } = await doStreamStep(
       conversationPrompt,
       model,
       writable,
-      toolsToModelTools(tools)
+      toolsToModelTools(tools),
+      {
+        sendStart: sendStart && isFirstIteration,
+      }
     );
+    isFirstIteration = false;
     steps.push(step);
 
     if (finish?.finishReason === 'tool-calls') {