@nocobase/plugin-ai 2.1.0-alpha.13 → 2.1.0-alpha.14

package/dist/server/ai-employees/ai-employee.js

@@ -36,11 +36,11 @@ var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__ge
 var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
 var ai_employee_exports = {};
 __export(ai_employee_exports, {
-  AIEmployee: () => AIEmployee
+  AIEmployee: () => AIEmployee,
+  ChatStreamProtocol: () => ChatStreamProtocol
 });
 module.exports = __toCommonJS(ai_employee_exports);
 var import_database = require("@nocobase/database");
-var import_provider = require("../llm-providers/provider");
 var import_utils = require("../utils");
 var import_prompts = require("./prompts");
 var import_lodash = __toESM(require("lodash"));
@@ -54,20 +54,32 @@ var import_langgraph = require("@langchain/langgraph");
 var import_stream = require("@langchain/core/utils/stream");
 var import_utils2 = require("./utils");
 var import_base = require("@langchain/core/callbacks/base");
+var import_shared = require("../../ai/tools/sub-agents/shared");
 class AIEmployee {
+  sessionId;
+  from = "main-agent";
   employee;
   aiChatConversation;
   skillSettings;
   plugin;
   db;
-  sessionId;
   ctx;
   systemMessage;
+  protocol;
   webSearch;
   model;
   legacy;
-  protocol;
-  constructor(ctx, employee, sessionId, systemMessage, skillSettings, webSearch, model, legacy) {
+  constructor({
+    ctx,
+    employee,
+    sessionId,
+    systemMessage,
+    skillSettings,
+    webSearch,
+    model,
+    legacy,
+    from = "main-agent"
+  }) {
     this.employee = employee;
     this.ctx = ctx;
     this.plugin = ctx.app.pm.get("ai");
@@ -78,10 +90,19 @@ class AIEmployee {
     this.skillSettings = skillSettings;
     this.model = model;
     this.legacy = legacy;
+    this.from = from;
     const builtInManager = this.plugin.builtInManager;
     builtInManager.setupBuiltInInfo(ctx, this.employee);
     this.webSearch = webSearch;
-    this.protocol = ChatStreamProtocol.create(ctx);
+    this.protocol = ChatStreamProtocol.fromContext(ctx);
+  }
+  async getFormatMessages(userMessages) {
+    const { provider } = await this.getLLMService();
+    const { messages } = await this.aiChatConversation.getChatContext({
+      userMessages,
+      formatMessages: (messages2) => this.formatMessages({ messages: messages2, provider })
+    });
+    return messages;
   }
   // === Chat flow ===
   buildState(messages) {
@@ -186,7 +207,9 @@ class AIEmployee {
   async invoke({
     messageId,
     userMessages = [],
-    userDecisions
+    userDecisions,
+    writer,
+    context
   }) {
     try {
       const { provider, chatContext, config, state } = await this.buildChatContext({
@@ -194,15 +217,21 @@ class AIEmployee {
         userMessages,
         userDecisions
       });
-      const { threadId } = await this.getCurrentThread();
       const invokeConfig = {
-        configurable: { thread_id: threadId },
-        context: { ctx: this.ctx },
+        context: { ctx: this.ctx, decisions: chatContext.decisions, ...context },
         recursionLimit: 100,
+        writer,
         ...config
       };
+      if (this.from === "main-agent") {
+        const { threadId } = await this.getCurrentThread();
+        invokeConfig.configurable = { thread_id: threadId };
+      }
       return await this.agentInvoke(provider, chatContext, invokeConfig, state);
     } catch (err) {
+      if (err.name === "GraphInterrupt") {
+        throw err;
+      }
       this.ctx.log.error(err);
       throw err;
     }
@@ -250,17 +279,24 @@ class AIEmployee {
     middleware
   }) {
     const model = provider.createModel();
-    const toolDefinitions = (tools == null ? void 0 : tools.map(import_provider.ToolDefinition.from("ToolsEntry"))) ?? [];
-    const allTools = provider.resolveTools(toolDefinitions);
-    const checkpointer = new import_checkpoints.SequelizeCollectionSaver(() => this.ctx.app.mainDataSource);
-    return (0, import_langchain.createAgent)({ model, tools: allTools, middleware, systemPrompt, checkpointer });
+    const allTools = provider.resolveTools((tools == null ? void 0 : tools.map(import_utils.buildTool)) ?? []);
+    if (this.from === "main-agent") {
+      const checkpointer = new import_checkpoints.SequelizeCollectionSaver(() => this.ctx.app.mainDataSource);
+      return (0, import_langchain.createAgent)({ model, tools: allTools, middleware, systemPrompt, checkpointer });
+    } else {
+      return (0, import_langchain.createAgent)({ model, tools: allTools, middleware, systemPrompt });
+    }
   }
   getAgentInput(context, state) {
-    var _a;
-    if ((_a = context.decisions) == null ? void 0 : _a.length) {
+    var _a, _b;
+    if ((_b = (_a = context.decisions) == null ? void 0 : _a.decisions) == null ? void 0 : _b.length) {
       return new import_langgraph.Command({
-        resume: {
-          decisions: context.decisions
+        resume: context.decisions.interruptId ? {
+          [context.decisions.interruptId]: {
+            decisions: context.decisions.decisions
+          }
+        } : {
+          decisions: context.decisions.decisions
         }
       });
     }
@@ -273,13 +309,16 @@ class AIEmployee {
     const { systemPrompt, tools, middleware } = context;
     const agent = await this.createAgent({ provider, systemPrompt, tools, middleware });
     const input = this.getAgentInput(context, state);
-    return agent.stream(input, config);
+    if (this.from === "sub-agent") {
+      delete config.configurable;
+    }
+    return agent.stream(input, this.withRunMetadata(config));
   }
   async agentInvoke(provider, context, config, state) {
     const { systemPrompt, tools, middleware } = context;
     const agent = await this.createAgent({ provider, systemPrompt, tools, middleware });
     const input = this.getAgentInput(context, state);
-    return agent.invoke(input, config);
+    return agent.invoke(input, this.withRunMetadata(config));
   }
   async prepareChatStream({
     chatContext,
@@ -297,8 +336,8 @@ class AIEmployee {
         {
           signal,
           streamMode: ["updates", "messages", "custom"],
-          configurable: { thread_id: threadId },
-          context: { ctx: this.ctx },
+          configurable: this.from === "main-agent" ? { thread_id: threadId } : void 0,
+          context: { ctx: this.ctx, decisions: chatContext.decisions },
           recursionLimit: 100,
           ...config
         },
@@ -311,8 +350,8 @@ class AIEmployee {
     }
   }
   async processChatStream(stream, options) {
-    var _a, _b, _c, _d, _e, _f, _g, _h, _i, _j, _k, _l, _m, _n, _o, _p, _q, _r, _s, _t;
-    let toolCalls;
+    var _a, _b, _c, _d, _e, _f, _g, _h, _i, _j, _k, _l, _m, _n, _o, _p, _q, _r, _s;
+    const aiMessageIdMap = /* @__PURE__ */ new Map();
     const { signal, providerName, model, provider, responseMetadata, allowEmpty = false } = options;
     let isReasoning = false;
     let gathered;
@@ -329,58 +368,69 @@ class AIEmployee {
         }
         await this.aiChatConversation.withTransaction(async (conversation, transaction) => {
           const result = await conversation.addMessages(values);
-          if (toolCalls == null ? void 0 : toolCalls.length) {
-            await this.initToolCall(transaction, result.messageId, toolCalls);
-          }
         });
       }
     });
     try {
-      this.protocol.startStream();
+      const aiEmployeeConversation = {
+        sessionId: this.sessionId,
+        from: this.from,
+        username: this.employee.username
+      };
+      this.protocol.with(aiEmployeeConversation).startStream();
       for await (const [mode, chunks] of stream) {
         if (mode === "messages") {
-          const [chunk] = chunks;
+          const [chunk, metadata] = chunks;
+          const { currentConversation } = metadata;
           if (chunk.type === "ai") {
             gathered = gathered !== void 0 ? (0, import_stream.concat)(gathered, chunk) : chunk;
             if (chunk.content) {
               if (isReasoning) {
                 isReasoning = false;
-                this.protocol.stopReasoning();
+                this.protocol.with(currentConversation).stopReasoning();
               }
               const parsedContent = provider.parseResponseChunk(chunk.content);
               if (parsedContent) {
-                this.protocol.content(parsedContent);
+                this.protocol.with(currentConversation).content(parsedContent);
               }
             }
             if ((_a = chunk.tool_call_chunks) == null ? void 0 : _a.length) {
-              this.protocol.toolCallChunks(chunk.tool_call_chunks);
+              this.protocol.with(currentConversation).toolCallChunks(chunk.tool_call_chunks);
             }
             const webSearch = provider.parseWebSearchAction(chunk);
             if (webSearch == null ? void 0 : webSearch.length) {
-              this.protocol.webSearch(webSearch);
+              this.protocol.with(currentConversation).webSearch(webSearch);
             }
             const reasoningContent = provider.parseReasoningContent(chunk);
             if (reasoningContent) {
               isReasoning = true;
-              this.protocol.reasoning(reasoningContent);
+              this.protocol.with(currentConversation).reasoning(reasoningContent);
             }
           }
         } else if (mode === "updates") {
           if ("__interrupt__" in chunks) {
+            const interruptId = chunks.__interrupt__[0].id;
             const interruptActions = this.toInterruptActions(chunks.__interrupt__[0].value);
             if (interruptActions.size) {
-              for (const toolCall of toolCalls ?? []) {
-                const interruptAction = interruptActions.get(toolCall.name);
-                if (!interruptAction) {
+              const toolsMap = await this.getToolsMap();
+              for (const interruptAction of interruptActions.values()) {
+                if (!interruptAction.currentConversation || !interruptAction.toolCall) {
+                  this.logger.warn("currentConversation or toolCall not exist in __interrupt__", interruptAction);
+                  continue;
+                }
+                const { sessionId, from, username } = interruptAction.currentConversation;
+                const { id: toolCallId, name: toolCallName } = interruptAction.toolCall;
+                const messageId = aiMessageIdMap.get(sessionId);
+                if (!messageId) {
                   continue;
                 }
-                await this.updateToolCallInterrupted(toolCall.messageId, toolCall.id, interruptAction);
-                this.protocol.toolCallStatus({
+                await this.updateToolCallInterrupted(sessionId, messageId, toolCallId, interruptId, interruptAction);
+                this.protocol.with(interruptAction.currentConversation).toolCallStatus({
                   toolCall: {
-                    messageId: toolCall.messageId,
-                    id: toolCall.id,
-                    name: toolCall.name,
-                    willInterrupt: toolCall.willInterrupt
+                    messageId,
+                    id: toolCallId,
+                    name: toolCallName,
+                    willInterrupt: this.shouldInterruptToolCall(toolsMap.get(toolCallName))
                   },
                   invokeStatus: "interrupted",
                   interruptAction
@@ -389,7 +439,9 @@ class AIEmployee {
             }
           }
         } else if (mode === "custom") {
+          const { currentConversation } = chunks;
           if (chunks.action === "AfterAIMessageSaved") {
+            aiMessageIdMap.set(currentConversation.sessionId, chunks.body.messageId);
             const data = responseMetadata.get(chunks.body.id);
             if (data) {
               const savedMessage = await this.aiMessagesModel.findOne({
@@ -417,32 +469,31 @@ class AIEmployee {
               }
             }
           } else if (chunks.action === "initToolCalls") {
-            toolCalls = ((_b = chunks.body) == null ? void 0 : _b.toolCalls) ?? [];
-            this.protocol.toolCalls(chunks.body);
+            this.protocol.with(currentConversation).toolCalls(chunks.body);
           } else if (chunks.action === "beforeToolCall") {
             const toolsMap = await this.getToolsMap();
-            const willInterrupt = this.shouldInterruptToolCall(toolsMap.get((_d = (_c = chunks.body) == null ? void 0 : _c.toolCall) == null ? void 0 : _d.name));
-            this.protocol.toolCallStatus({
+            const willInterrupt = this.shouldInterruptToolCall(toolsMap.get((_c = (_b = chunks.body) == null ? void 0 : _b.toolCall) == null ? void 0 : _c.name));
+            this.protocol.with(currentConversation).toolCallStatus({
               toolCall: {
-                messageId: (_f = (_e = chunks.body) == null ? void 0 : _e.toolCall) == null ? void 0 : _f.messageId,
-                id: (_h = (_g = chunks.body) == null ? void 0 : _g.toolCall) == null ? void 0 : _h.id,
-                name: (_j = (_i = chunks.body) == null ? void 0 : _i.toolCall) == null ? void 0 : _j.name,
+                messageId: (_e = (_d = chunks.body) == null ? void 0 : _d.toolCall) == null ? void 0 : _e.messageId,
+                id: (_g = (_f = chunks.body) == null ? void 0 : _f.toolCall) == null ? void 0 : _g.id,
+                name: (_i = (_h = chunks.body) == null ? void 0 : _h.toolCall) == null ? void 0 : _i.name,
                 willInterrupt
               },
               invokeStatus: "pending"
             });
           } else if (chunks.action === "afterToolCall") {
             const toolsMap = await this.getToolsMap();
-            const willInterrupt = this.shouldInterruptToolCall(toolsMap.get((_l = (_k = chunks.body) == null ? void 0 : _k.toolCall) == null ? void 0 : _l.name));
-            this.protocol.toolCallStatus({
+            const willInterrupt = this.shouldInterruptToolCall(toolsMap.get((_k = (_j = chunks.body) == null ? void 0 : _j.toolCall) == null ? void 0 : _k.name));
+            this.protocol.with(currentConversation).toolCallStatus({
               toolCall: {
-                messageId: (_n = (_m = chunks.body) == null ? void 0 : _m.toolCall) == null ? void 0 : _n.messageId,
-                id: (_p = (_o = chunks.body) == null ? void 0 : _o.toolCall) == null ? void 0 : _p.id,
-                name: (_r = (_q = chunks.body) == null ? void 0 : _q.toolCall) == null ? void 0 : _r.name,
+                messageId: (_m = (_l = chunks.body) == null ? void 0 : _l.toolCall) == null ? void 0 : _m.messageId,
+                id: (_o = (_n = chunks.body) == null ? void 0 : _n.toolCall) == null ? void 0 : _o.id,
+                name: (_q = (_p = chunks.body) == null ? void 0 : _p.toolCall) == null ? void 0 : _q.name,
                 willInterrupt
               },
               invokeStatus: "done",
-              status: (_t = (_s = chunks.body) == null ? void 0 : _s.toolCallResult) == null ? void 0 : _t.status
+              status: (_s = (_r = chunks.body) == null ? void 0 : _r.toolCallResult) == null ? void 0 : _s.status
             });
           } else if (chunks.action === "beforeSendToolMessage") {
             const { messageId, messages } = chunks.body ?? {};
@@ -455,7 +506,7 @@ class AIEmployee {
               for (const { metadata } of messages) {
                 const tools = toolsMap.get(metadata.toolName);
                 const toolCallResult = toolCallResultMap.get(metadata.toolCallId);
-                this.protocol.toolCallStatus({
+                this.protocol.with(currentConversation).toolCallStatus({
                   toolCall: {
                     messageId,
                     id: metadata.toolCallId,
@@ -467,7 +518,9 @@ class AIEmployee {
                 });
               }
             }
-            this.protocol.newMessage();
+            this.protocol.with(currentConversation).newMessage();
+          } else if (chunks.action === "afterSubAgentInvoke") {
+            this.protocol.with(currentConversation).subAgentCompleted();
           }
         }
       }
@@ -475,7 +528,7 @@ class AIEmployee {
         this.sendErrorResponse("Empty message");
         return;
       }
-      this.protocol.endStream();
+      this.protocol.with(aiEmployeeConversation).endStream();
     } catch (err) {
       this.ctx.log.error(err);
       if (err.name === "GraphRecursionError") {
@@ -484,7 +537,9 @@ class AIEmployee {
         this.sendErrorResponse(provider.parseResponseError(err));
       }
     } finally {
-      this.ctx.res.end();
+      if (this.from === "main-agent") {
+        this.ctx.res.end();
+      }
     }
   }
   // === Prompts & knowledge base ===
@@ -566,6 +621,7 @@ ${workContextBackground.join("\n")}`;
       }
     }
     const availableSkills = await this.getAvailableSkills();
+    const availableAIEmployees = await this.getAvailableAIEmployees();
     const systemPrompt = (0, import_prompts.getSystemPrompt)({
       aiEmployee: {
         nickname: this.employee.nickname,
@@ -581,7 +637,8 @@ ${workContextBackground.join("\n")}`;
         locale: this.ctx.getCurrentLocale() || "en-US"
       },
       knowledgeBase,
-      availableSkills
+      availableSkills,
+      availableAIEmployees
     });
     const { important } = this.ctx.action.params.values || {};
     if (important === "GraphRecursionError") {
@@ -706,23 +763,54 @@ If information is missing, clearly state it in the summary.</Important>`;
       transaction
     });
   }
-  async updateToolCallInterrupted(messageId, toolCallId, interruptAction) {
-    const [updated] = await this.aiToolMessagesModel.update(
-      {
-        invokeStatus: "interrupted",
-        interruptActionOrder: interruptAction.order,
-        interruptAction
-      },
-      {
+  async updateToolCallInterrupted(sessionId, messageId, toolCallId, interruptId, interruptAction) {
+    return await this.db.sequelize.transaction(async (transaction) => {
+      const [updated] = await this.aiToolMessagesModel.update(
+        {
+          invokeStatus: "interrupted",
+          interruptActionOrder: interruptAction.order,
+          interruptAction
+        },
+        {
+          where: {
+            sessionId,
+            messageId,
+            toolCallId,
+            invokeStatus: "init"
+          },
+          transaction
+        }
+      );
+      if (!updated) {
+        return updated;
+      }
+      const message = await this.aiMessagesModel.findOne({
         where: {
-          sessionId: this.sessionId,
           messageId,
-          toolCallId,
-          invokeStatus: "init"
-        }
+          sessionId
+        },
+        transaction
+      });
+      if (!message) {
+        return updated;
       }
-    );
-    return updated;
+      await this.aiMessagesModel.update(
+        {
+          metadata: {
+            ...message.get("metadata") ?? {},
+            interruptId
+          }
+        },
+        {
+          where: {
+            messageId,
+            sessionId
+          },
+          transaction
+        }
+      );
+      return updated;
+    });
   }
   async updateToolCallPending(messageId, toolCallId) {
     const [updated] = await this.aiToolMessagesModel.update(
@@ -784,7 +872,6 @@ If information is missing, clearly state it in the summary.</Important>`;
     var _a;
     return (_a = await this.aiToolMessagesModel.findOne({
       where: {
-        sessionId: this.sessionId,
         messageId,
         toolCallId
       }
@@ -793,7 +880,6 @@ If information is missing, clearly state it in the summary.</Important>`;
   async getToolCallResultMap(messageId, toolCallIds) {
     const list = (await this.aiToolMessagesModel.findAll({
       where: {
-        sessionId: this.sessionId,
         messageId,
         toolCallId: {
           [import_database.Op.in]: toolCallIds
@@ -802,19 +888,6 @@ If information is missing, clearly state it in the summary.</Important>`;
     })).map((it) => it.toJSON());
     return new Map(list.map((it) => [it.toolCallId, it]));
   }
-  async getUserDecisions(messageId) {
-    const allInterruptedToolCall = await this.aiToolMessagesModel.findAll({
-      where: {
-        messageId,
-        interruptActionOrder: { [import_database.Op.not]: null }
-      },
-      order: [["interruptActionOrder", "ASC"]]
-    });
-    if (!allInterruptedToolCall.every((t) => t.invokeStatus === "waiting")) {
-      return [];
-    }
-    return allInterruptedToolCall.map((item) => item.userDecision);
-  }
   async cancelToolCall() {
     var _a;
     let messageId;
@@ -829,7 +902,6 @@ If information is missing, clearly state it in the summary.</Important>`;
     }
     const toolMessages = (await this.aiToolMessagesModel.findAll({
       where: {
-        sessionId: this.sessionId,
         messageId,
         invokeStatus: {
           [import_database.Op.ne]: "confirmed"
@@ -1068,10 +1140,20 @@ If information is missing, clearly state it in the summary.</Important>`;
     const actionRequestsMap = new Map(actionRequests.map((x) => [x.name, x]));
     const reviewConfigsMap = new Map(reviewConfigs.map((x) => [x.actionName, x]));
     for (const [name, actionRequest] of actionRequestsMap.entries()) {
+      const payload = actionRequest.description ? JSON.parse(actionRequest.description) : null;
       result.set(name, {
         order: order++,
         description: actionRequest.description,
-        allowedDecisions: (_a = reviewConfigsMap.get(name)) == null ? void 0 : _a.allowedDecisions
+        allowedDecisions: (_a = reviewConfigsMap.get(name)) == null ? void 0 : _a.allowedDecisions,
+        toolCall: {
+          id: payload.toolCallId,
+          name: payload.toolCallName
+        },
+        currentConversation: {
+          sessionId: payload.sessionId,
+          from: payload.from,
+          username: payload.username
+        }
       });
     }
     return result;
@@ -1098,11 +1180,16 @@ If information is missing, clearly state it in the summary.</Important>`;
   async getAvailableSkills() {
     var _a, _b;
     const { skillsManager } = this.plugin.ai;
+    const aIEmployeeTools = await this.getAIEmployeeTools();
+    const getSkill = aIEmployeeTools.find((it) => it.definition.name === "getSkill");
+    if (!getSkill) {
+      return [];
+    }
     const generalSkills = await skillsManager.listSkills({ scope: "GENERAL" });
     const specifiedSkillNames = ((_a = this.employee.skillSettings) == null ? void 0 : _a.skills) ?? [];
     const specifiedSkills = specifiedSkillNames.length ? await skillsManager.getSkills(specifiedSkillNames) : [];
     const skillFilter = ((_b = this.skillSettings) == null ? void 0 : _b.skills) ?? [];
-    return import_lodash.default.uniqBy([...generalSkills || [], ...specifiedSkills || []], "name").filter(
+    return import_lodash.default.uniqBy([...specifiedSkills || [], ...generalSkills || []], "name").filter(
       (it) => skillFilter.length === 0 || skillFilter.includes(it.name)
     );
   }
@@ -1167,6 +1254,12 @@ If information is missing, clearly state it in the summary.</Important>`;
     }
     return result;
   }
+  async getAvailableAIEmployees() {
+    const availableAIEmployees = (await (0, import_shared.listAccessibleAIEmployees)(this.ctx)).map(
+      (employee) => (0, import_shared.serializeEmployeeSummary)(this.ctx, employee)
+    );
+    return availableAIEmployees;
+  }
   getMiddleware(options) {
     const { providerName, model, tools, baseToolNames, messageId, agentThread } = options;
     return [
@@ -1206,6 +1299,19 @@ If information is missing, clearly state it in the summary.</Important>`;
   listTools(filter) {
     return this.toolsManager.listTools(filter);
   }
+  withRunMetadata(config) {
+    return {
+      ...config,
+      metadata: {
+        ...(config == null ? void 0 : config.metadata) ?? {},
+        currentConversation: {
+          sessionId: this.sessionId,
+          from: this.from,
+          username: this.employee.get("username")
+        }
+      }
+    };
+  }
   get toolsManager() {
     return this.ctx.app.aiManager.toolsManager;
   }
@@ -1250,8 +1356,8 @@ class AgentThread {
   }
 }
 class ChatStreamProtocol {
-  constructor(ctx) {
-    this.ctx = ctx;
+  constructor(streamConsumer) {
+    this.streamConsumer = streamConsumer;
   }
   _statistics = {
     sent: 0,
@@ -1262,71 +1368,79 @@ class ChatStreamProtocol {
       this._statistics.sent = 0;
     }
   };
-  static create(ctx) {
-    return new ChatStreamProtocol(ctx);
+  static fromContext(ctx) {
+    return new ChatStreamProtocol(ctx.res);
   }
-  startStream() {
-    this._statistics.reset();
-    this.write({ type: "stream_start" });
-  }
-  endStream() {
-    this.write({ type: "stream_end" });
-  }
-  newMessage(content) {
-    this.write({ type: "new_message", body: content });
-  }
-  content(content) {
-    this.write({ type: "content", body: content });
-  }
-  webSearch(content) {
-    this.write({ type: "web_search", body: content });
-  }
-  reasoning(content) {
-    this.write({ type: "reasoning", body: content });
-  }
-  stopReasoning() {
-    this.write({
-      type: "reasoning",
-      body: {
-        status: "stop",
-        content: ""
-      }
-    });
-  }
-  toolCallChunks(content) {
-    this.write({ type: "tool_call_chunks", body: content });
-  }
-  toolCalls(content) {
-    this.write({ type: "tool_calls", body: content });
-  }
-  toolCallStatus({
-    toolCall,
-    invokeStatus,
-    status,
-    interruptAction
-  }) {
-    this.write({
-      type: "tool_call_status",
-      body: {
+  with(conversation) {
+    const write = ({ type, body }) => {
+      const { sessionId, from, username } = conversation;
+      const data = `data: ${JSON.stringify({ sessionId, from, username, type, body })}
+
+`;
+      this.streamConsumer.write(data);
+      this._statistics.addSent(data.length);
+    };
+    return {
+      startStream: () => {
+        this._statistics.reset();
+        write({ type: "stream_start" });
+      },
+      endStream: () => {
+        write({ type: "stream_end" });
+      },
+      subAgentCompleted: () => {
+        write({ type: "sub_agent_completed" });
+      },
+      newMessage: (content) => {
+        write({ type: "new_message", body: content });
+      },
+      content: (content) => {
+        write({ type: "content", body: content });
+      },
+      webSearch: (content) => {
+        write({ type: "web_search", body: content });
+      },
+      reasoning: (content) => {
+        write({ type: "reasoning", body: content });
+      },
+      stopReasoning: () => {
+        write({
+          type: "reasoning",
+          body: {
+            status: "stop",
+            content: ""
+          }
+        });
+      },
+      toolCallChunks: (content) => {
+        write({ type: "tool_call_chunks", body: content });
+      },
+      toolCalls: (content) => {
+        write({ type: "tool_calls", body: content });
+      },
+      toolCallStatus: ({
         toolCall,
         invokeStatus,
         status,
         interruptAction
+      }) => {
+        write({
+          type: "tool_call_status",
+          body: {
+            toolCall,
+            invokeStatus,
+            status,
+            interruptAction
+          }
+        });
       }
-    });
+    };
   }
   get statistics() {
     return {
       sent: this._statistics.sent
     };
   }
-  write({ type, body }) {
-    const data = `data: ${JSON.stringify({ type, body })}
-
-`;
-    this.ctx.res.write(data);
-    this._statistics.addSent(data.length);
-  }
 }
 class ResponseMetadataCollector extends import_base.BaseCallbackHandler {
   constructor(llmProvider, responseMetadata) {
@@ -1347,5 +1461,6 @@ class ResponseMetadataCollector extends import_base.BaseCallbackHandler {
 }
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
-  AIEmployee
+  AIEmployee,
+  ChatStreamProtocol
 });