Add middleware API

Author: mme

typescript-sdk/integrations/mastra/src/mastra.ts

@@ -59,7 +59,7 @@ export class MastraAgent extends AbstractAgent {
     this.runtimeContext = runtimeContext ?? new RuntimeContext();
   }
 
-  protected run(input: RunAgentInput): Observable<BaseEvent> {
+  public run(input: RunAgentInput): Observable<BaseEvent> {
     let messageId = randomUUID();
 
     return new Observable<BaseEvent>((subscriber) => {

typescript-sdk/integrations/middleware-starter/src/index.ts

@@ -2,7 +2,7 @@ import { AbstractAgent, BaseEvent, EventType, RunAgentInput } from "@ag-ui/clien
 import { Observable } from "rxjs";
 
 export class MiddlewareStarterAgent extends AbstractAgent {
-  protected run(input: RunAgentInput): Observable<BaseEvent> {
+  public run(input: RunAgentInput): Observable<BaseEvent> {
     const messageId = Date.now().toString();
     return new Observable<BaseEvent>((observer) => {
       observer.next({

typescript-sdk/integrations/vercel-ai-sdk/src/index.ts

@@ -55,7 +55,7 @@ export class VercelAISDKAgent extends AbstractAgent {
     this.toolChoice = toolChoice ?? "auto";
   }
 
-  protected run(input: RunAgentInput): Observable<BaseEvent> {
+  public run(input: RunAgentInput): Observable<BaseEvent> {
     const finalMessages: Message[] = input.messages;
 
     return new Observable<ProcessedEvent>((subscriber) => {

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

@@ -13,6 +13,7 @@ import { LegacyRuntimeProtocolEvent } from "@/legacy/types";
 import { lastValueFrom } from "rxjs";
 import { transformChunks } from "@/chunks";
 import { AgentStateMutation, AgentSubscriber, runSubscribersWithMutation } from "./subscriber";
+import { Middleware, MiddlewareFunction, FunctionMiddleware } from "@/middleware";
 
 export interface RunAgentResult {
   result: any;
@@ -27,6 +28,7 @@ export abstract class AbstractAgent {
   public state: State;
   public debug: boolean = false;
   public subscribers: AgentSubscriber[] = [];
+  private middlewares: Middleware[] = [];
 
   constructor({
     agentId,
@@ -53,7 +55,15 @@ export abstract class AbstractAgent {
     };
   }
 
-  protected abstract run(input: RunAgentInput): Observable<BaseEvent>;
+  public use(...middlewares: (Middleware | MiddlewareFunction)[]): this {
+    const normalizedMiddlewares = middlewares.map(m =>
+      typeof m === 'function' ? new FunctionMiddleware(m) : m
+    );
+    this.middlewares.push(...normalizedMiddlewares);
+    return this;
+  }
+
+  public abstract run(input: RunAgentInput): Observable<BaseEvent>;
 
   public async runAgent(
     parameters?: RunAgentParameters,
@@ -77,7 +87,21 @@ export abstract class AbstractAgent {
     await this.onInitialize(input, subscribers);
 
     const pipeline = pipe(
-      () => this.run(input),
+      () => {
+        // Build middleware chain using reduceRight
+        if (this.middlewares.length === 0) {
+          return this.run(input);
+        }
+
+        const chainedAgent = this.middlewares.reduceRight(
+          (nextAgent: AbstractAgent, middleware) => ({
+            run: (i: RunAgentInput) => middleware.run(i, nextAgent)
+          } as AbstractAgent),
+          this // Original agent is the final 'next'
+        );
+
+        return chainedAgent.run(input);
+      },
       transformChunks(this.debug),
       verifyEvents(this.debug),
       (source$) => this.apply(input, source$, subscribers),
@@ -416,7 +440,23 @@ export abstract class AbstractAgent {
     this.agentId = this.agentId ?? uuidv4();
     const input = this.prepareRunAgentInput(config);
 
-    return this.run(input).pipe(
+    // Build middleware chain for legacy bridge
+    const runObservable = (() => {
+      if (this.middlewares.length === 0) {
+        return this.run(input);
+      }
+
+      const chainedAgent = this.middlewares.reduceRight(
+        (nextAgent: AbstractAgent, middleware) => ({
+          run: (i: RunAgentInput) => middleware.run(i, nextAgent)
+        } as AbstractAgent),
+        this
+      );
+
+      return chainedAgent.run(input);
+    })();
+
+    return runObservable.pipe(
       transformChunks(this.debug),
       verifyEvents(this.debug),
       convertToLegacyEvents(this.threadId, input.runId, this.agentId),