update subscriber interface

Author: mme

typescript-sdk/packages/client/src/agent/agent.ts

@@ -14,6 +14,10 @@ import { lastValueFrom } from "rxjs";
 import { transformChunks } from "@/chunks";
 import { AgentStateMutation, RunAgentSubscriber, runSubscribersWithMutation } from "./subscriber";
 
+export interface RunAgentResult {
+  result: any;
+}
+
 export abstract class AbstractAgent {
   public agentId?: string;
   public description: string;
@@ -53,7 +57,7 @@ export abstract class AbstractAgent {
   public async runAgent(
     parameters?: RunAgentParameters,
     subscriber?: RunAgentSubscriber,
-  ): Promise<any> {
+  ): Promise<RunAgentResult> {
     this.agentId = this.agentId ?? uuidv4();
     const input = this.prepareRunAgentInput(parameters);
     let result: any = undefined;
@@ -83,7 +87,7 @@ export abstract class AbstractAgent {
       }),
     );
 
-    return lastValueFrom(pipeline(of(null))).then(() => result);
+    return lastValueFrom(pipeline(of(null))).then(() => ({ result }));
   }
 
   public abortRun() {}
@@ -108,6 +112,7 @@ export abstract class AbstractAgent {
           subscribers.forEach((subscriber) => {
             subscriber.onMessagesChanged?.({
               messages: this.messages,
+              state: this.state,
               agent: this,
               input,
             });
@@ -118,6 +123,7 @@ export abstract class AbstractAgent {
           subscribers.forEach((subscriber) => {
             subscriber.onStateChanged?.({
               state: this.state,
+              messages: this.messages,
               agent: this,
               input,
             });
@@ -157,6 +163,7 @@ export abstract class AbstractAgent {
         subscribers.forEach((subscriber) => {
           subscriber.onMessagesChanged?.({
             messages: this.messages,
+            state: this.state,
             agent: this,
             input,
           });
@@ -168,6 +175,7 @@ export abstract class AbstractAgent {
         subscribers.forEach((subscriber) => {
           subscriber.onStateChanged?.({
             state: this.state,
+            messages: this.messages,
             agent: this,
             input,
           });
@@ -194,6 +202,7 @@ export abstract class AbstractAgent {
             subscribers.forEach((subscriber) => {
               subscriber.onMessagesChanged?.({
                 messages: this.messages,
+                state: this.state,
                 agent: this,
                 input,
               });
@@ -204,6 +213,7 @@ export abstract class AbstractAgent {
             subscribers.forEach((subscriber) => {
               subscriber.onStateChanged?.({
                 state: this.state,
+                messages: this.messages,
                 agent: this,
                 input,
               });
@@ -240,6 +250,7 @@ export abstract class AbstractAgent {
         subscribers.forEach((subscriber) => {
           subscriber.onMessagesChanged?.({
             messages: this.messages,
+            state: this.state,
             agent: this,
             input,
           });
@@ -250,6 +261,7 @@ export abstract class AbstractAgent {
         subscribers.forEach((subscriber) => {
           subscriber.onStateChanged?.({
             state: this.state,
+            messages: this.messages,
             agent: this,
             input,
           });

typescript-sdk/packages/client/src/agent/subscriber.ts

@@ -20,6 +20,7 @@ import {
   MessagesSnapshotEvent,
   RawEvent,
   CustomEvent,
+  ToolCall,
 } from "@ag-ui/core";
 import { AbstractAgent } from "./agent";
 import { structuredClone_ } from "@/utils";
@@ -131,8 +132,10 @@ export interface RunAgentSubscriber {
   ): MaybePromise<AgentStateMutation | void>;
 
   // State changes
-  onMessagesChanged?(params: Omit<RunAgentSubscriberParams, "state">): MaybePromise<void>;
-  onStateChanged?(params: Omit<RunAgentSubscriberParams, "messages">): MaybePromise<void>;
+  onMessagesChanged?(params: RunAgentSubscriberParams): MaybePromise<void>;
+  onStateChanged?(params: RunAgentSubscriberParams): MaybePromise<void>;
+  onNewMessage?(params: { message: Message } & RunAgentSubscriberParams): MaybePromise<void>;
+  onNewToolCall?(params: { toolCall: ToolCall } & RunAgentSubscriberParams): MaybePromise<void>;
 }
 
 export async function runSubscribersWithMutation(

typescript-sdk/packages/client/src/apply/default.ts

@@ -148,6 +148,18 @@ export const defaultApplyEvents = (
           );
           applyMutation(mutation);
 
+          await Promise.all(
+            subscribers.map((subscriber) => {
+              subscriber.onNewMessage?.({
+                message: messages[messages.length - 1],
+                messages,
+                state,
+                agent,
+                input,
+              });
+            }),
+          );
+
           return emitUpdates();
         }
 
@@ -287,6 +299,21 @@ export const defaultApplyEvents = (
           );
           applyMutation(mutation);
 
+          await Promise.all(
+            subscribers.map((subscriber) => {
+              subscriber.onNewToolCall?.({
+                toolCall:
+                  messages[messages.length - 1].toolCalls[
+                    messages[messages.length - 1].toolCalls.length - 1
+                  ],
+                messages,
+                state,
+                agent,
+                input,
+              });
+            }),
+          );
+
           return emitUpdates();
         }
 
@@ -319,6 +346,18 @@ export const defaultApplyEvents = (
 
             messages.push(toolMessage);
 
+            await Promise.all(
+              subscribers.map((subscriber) => {
+                subscriber.onNewMessage?.({
+                  message: toolMessage,
+                  messages,
+                  state,
+                  agent,
+                  input,
+                });
+              }),
+            );
+
             applyMutation({ messages });
           }