@automattic/agenttic-client 0.1.9 → 0.1.11

package/dist/index.js

@@ -113,8 +113,8 @@ function createTextPart(text) {
     text
   };
 }
-function createSendMessageRequest(params, method = "message/send") {
-  return {
+function createSendMessageRequest(params, method = "message/send", tokenStreaming = false) {
+  const request = {
     jsonrpc: "2.0",
     id: createRequestId(),
     method,
@@ -123,6 +123,10 @@ function createSendMessageRequest(params, method = "message/send") {
       ...params
     }
   };
+  if (tokenStreaming && method === "message/stream") {
+    request.tokenStreaming = true;
+  }
+  return request;
 }
 function extractTextFromMessage(message) {
   if (!message || !message.parts || !Array.isArray(message.parts)) {
@@ -218,6 +222,9 @@ function createToolResultMessage(toolResults, historyDataParts = []) {
     }
   };
 }
+function createAbortController() {
+  return new AbortController();
+}
 
 // src/client/utils/internal/messages.ts
 async function enhanceMessageWithTools(message, toolProvider) {
@@ -312,10 +319,14 @@ function parseStreamChunk(chunk, buffer = "") {
   const nextBuffer = currentStreamData.substring(lastCompleteEventEnd);
   return { events, nextBuffer };
 }
-async function* parseSSEStream(stream) {
+async function* parseSSEStream(stream, options = {}) {
+  const { supportDeltas = false } = options;
   const reader = stream.getReader();
   const decoder = new TextDecoder();
   let buffer = "";
+  const accumulator = new DeltaAccumulator();
+  let currentTaskId = null;
+  let lastStatus = null;
   try {
     while (true) {
       const { done, value } = await reader.read();
@@ -325,13 +336,50 @@ async function* parseSSEStream(stream) {
       const chunk = decoder.decode(value, { stream: true });
       const { events, nextBuffer } = parseStreamChunk(chunk, buffer);
       if (events && Array.isArray(events)) {
-        for (const event of events) {
+        for (let i = 0; i < events.length; i++) {
+          const event = events[i];
+          if (i > 0 && event.method === "message/delta" && typeof requestAnimationFrame !== "undefined") {
+            await new Promise((resolve) => {
+              requestAnimationFrame(() => resolve(void 0));
+            });
+          }
           if (event.error) {
             throw new Error(
               `Streaming error: ${event.error.message}`
             );
           }
-          if (event.result && event.result.status) {
+          if (supportDeltas && event.method === "message/delta" && event.params?.delta) {
+            const delta = event.params.delta;
+            try {
+              if (delta.deltaType === "content") {
+                accumulator.processContentDelta(
+                  delta.content
+                );
+                if (!currentTaskId && event.params.id) {
+                  currentTaskId = event.params.id;
+                }
+                if (currentTaskId) {
+                  const currentMessage = accumulator.getCurrentMessage();
+                  yield {
+                    id: currentTaskId,
+                    status: {
+                      state: "working",
+                      message: currentMessage
+                    },
+                    final: false,
+                    text: accumulator.getTextContent()
+                  };
+                }
+              }
+            } catch (error) {
+              logger("Failed to process delta: %o", error);
+            }
+          } else if (event.result && event.result.status) {
+            currentTaskId = event.result.id;
+            lastStatus = event.result.status;
+            if (accumulator.getTextContent() || accumulator.getCurrentMessage().parts.length > 0) {
+              accumulator.reset();
+            }
             const update = {
               id: event.result.id,
               status: event.result.status,
@@ -344,6 +392,22 @@ async function* parseSSEStream(stream) {
               )
             };
             yield update;
+          } else if (event.id && event.result) {
+            currentTaskId = event.result.id;
+            if (event.result.status) {
+              const update = {
+                id: event.result.id,
+                status: event.result.status,
+                final: event.result.status.state === "completed" || event.result.status.state === "failed" || event.result.status.state === "canceled",
+                text: extractTextFromMessage(
+                  event.result.status?.message || {
+                    role: "agent",
+                    parts: []
+                  }
+                )
+              };
+              yield update;
+            }
           }
         }
       }
@@ -353,6 +417,103 @@ async function* parseSSEStream(stream) {
     reader.releaseLock();
   }
 }
+var DeltaAccumulator = class {
+  textContent = "";
+  toolCalls = /* @__PURE__ */ new Map();
+  /**
+   * Process a simple content delta (server's actual format)
+   * @param content - The text content to append
+   */
+  processContentDelta(content) {
+    this.textContent += content;
+  }
+  /**
+   * Process a delta message and accumulate the content (original format)
+   * @param delta - The delta message to process
+   */
+  processDelta(delta) {
+    switch (delta.type) {
+      case "content":
+        this.textContent += delta.content;
+        break;
+      case "tool_name":
+        if (!this.toolCalls.has(delta.toolCallIndex)) {
+          this.toolCalls.set(delta.toolCallIndex, {
+            toolCallId: delta.toolCallId,
+            toolName: "",
+            argumentFragments: []
+          });
+        }
+        const toolCall = this.toolCalls.get(delta.toolCallIndex);
+        toolCall.toolName += delta.content;
+        break;
+      case "tool_argument":
+        if (!this.toolCalls.has(delta.toolCallIndex)) {
+          this.toolCalls.set(delta.toolCallIndex, {
+            toolCallId: delta.toolCallId,
+            toolName: "",
+            argumentFragments: []
+          });
+        }
+        const call = this.toolCalls.get(delta.toolCallIndex);
+        call.argumentFragments.push(delta.content);
+        break;
+    }
+  }
+  /**
+   * Get the accumulated text content
+   */
+  getTextContent() {
+    return this.textContent;
+  }
+  /**
+   * Get the current accumulated message
+   * @param role - The role for the message (default: 'agent')
+   */
+  getCurrentMessage(role = "agent") {
+    const parts = [];
+    if (this.textContent) {
+      parts.push({
+        type: "text",
+        text: this.textContent
+      });
+    }
+    for (const [index, toolCall] of this.toolCalls) {
+      if (toolCall.toolName) {
+        const argumentsStr = toolCall.argumentFragments.join("");
+        let args = {};
+        if (argumentsStr) {
+          try {
+            args = JSON.parse(argumentsStr);
+          } catch {
+            args = { _raw: argumentsStr };
+          }
+        }
+        parts.push({
+          type: "data",
+          data: {
+            toolCallId: toolCall.toolCallId,
+            toolId: toolCall.toolName,
+            arguments: args
+          }
+        });
+      }
+    }
+    return {
+      role,
+      parts,
+      kind: "message",
+      messageId: generateMessageId()
+    };
+  }
+  /**
+   * Reset the accumulator
+   */
+  reset() {
+    this.textContent = "";
+    this.toolCalls.clear();
+  }
+};
 
 // src/client/utils/internal/errors.ts
 function handleRequestError(error, timeoutId, operation = "request") {
@@ -417,6 +578,32 @@ async function getHeaders(authProvider, isStreaming = false) {
   }
   return baseHeaders;
 }
+function combineSignals(signal1, signal2) {
+  if (!signal2) {
+    return signal1;
+  }
+  const controller = new AbortController();
+  const handleAbort = (signal) => {
+    if (!controller.signal.aborted) {
+      controller.abort(signal.reason);
+    }
+  };
+  if (signal1.aborted) {
+    controller.abort(signal1.reason);
+  } else {
+    signal1.addEventListener("abort", () => handleAbort(signal1), {
+      once: true
+    });
+  }
+  if (signal2.aborted) {
+    controller.abort(signal2.reason);
+  } else {
+    signal2.addEventListener("abort", () => handleAbort(signal2), {
+      once: true
+    });
+  }
+  return controller.signal;
+}
 function createFetchOptions(headers, body, signal) {
   return {
     method: "POST",
@@ -428,7 +615,7 @@ function createFetchOptions(headers, body, signal) {
 async function prepareRequest(params, config, options, toolProvider, contextProvider, defaultSessionId) {
   const { message, sessionId, taskId, metadata } = params;
   const { agentId, agentUrl, authProvider, proxy } = config;
-  const { isStreaming = false } = options;
+  const { isStreaming = false, enableTokenStreaming = false } = options;
   const effectiveSessionId = sessionId || defaultSessionId;
   const fullAgentUrl = constructAgentUrl(agentUrl, agentId);
   const enhancedMessage = await enhanceMessage(
@@ -443,7 +630,9 @@ async function prepareRequest(params, config, options, toolProvider, contextProv
       message: enhancedMessage,
       metadata
     },
-    isStreaming ? "message/stream" : "message/send"
+    isStreaming ? "message/stream" : "message/send",
+    enableTokenStreaming && isStreaming
+    // Only enable token streaming if using SSE
   );
   const headers = await getHeaders(authProvider, isStreaming);
   logRequest("POST", fullAgentUrl, headers, request);
@@ -455,24 +644,26 @@ async function prepareRequest(params, config, options, toolProvider, contextProv
     fullAgentUrl
   };
 }
-async function executeRequest(preparedRequest, config) {
+async function executeRequest(preparedRequest, config, options = {}) {
   const { request, headers, fullAgentUrl } = preparedRequest;
   const { timeout } = config;
-  const { timeoutId, controller } = createTimeoutHandler(
+  const { abortSignal: externalSignal } = options;
+  const { timeoutId, controller: timeoutController } = createTimeoutHandler(
     timeout,
     "request"
   );
+  const signal = externalSignal ? combineSignals(timeoutController.signal, externalSignal) : timeoutController.signal;
   try {
-    const options = createFetchOptions(
+    const fetchOptions = createFetchOptions(
       headers,
       JSON.stringify(request),
-      controller.signal
+      signal
     );
     logger("Making request to %s with options: %O", fullAgentUrl, {
-      method: options.method,
-      headers: options.headers
+      method: fetchOptions.method,
+      headers: fetchOptions.headers
     });
-    const response = await fetch(fullAgentUrl, options);
+    const response = await fetch(fullAgentUrl, fetchOptions);
     clearTimeout(timeoutId);
     validateHttpResponse(response, "request");
     const data = await response.json();
@@ -490,21 +681,31 @@ async function executeRequest(preparedRequest, config) {
 async function* executeStreamingRequest(preparedRequest, config, options) {
   const { request, headers, fullAgentUrl } = preparedRequest;
   const {} = config;
-  const { streamingTimeout = 6e4 } = options;
+  const {
+    streamingTimeout = 6e4,
+    abortSignal: externalSignal,
+    enableTokenStreaming = false
+  } = options;
   const { timeoutId, controller } = createTimeoutHandler(
     streamingTimeout,
     "streaming request"
   );
+  const signal = externalSignal ? combineSignals(controller.signal, externalSignal) : controller.signal;
   try {
-    const fetchOptions = createFetchOptions(
-      headers,
-      JSON.stringify(request),
-      controller.signal
-    );
+    const requestBody = JSON.stringify(request);
+    const fetchOptions = createFetchOptions(headers, requestBody, signal);
     const response = await fetch(fullAgentUrl, fetchOptions);
     clearTimeout(timeoutId);
     validateStreamingResponse(response, "streaming request");
-    yield* parseSSEStream(response.body);
+    if (!response.body) {
+      throw new Error(
+        "Response body is null - server may not support streaming"
+      );
+    }
+    const supportDeltas = enableTokenStreaming && request.tokenStreaming === true;
+    yield* parseSSEStream(response.body, {
+      supportDeltas
+    });
   } catch (error) {
     handleRequestError(error, timeoutId, "streaming request");
   }
@@ -512,6 +713,7 @@ async function* executeStreamingRequest(preparedRequest, config, options) {
 
 // src/client/index.ts
 var DEFAULT_TIMEOUT = 12e4;
+var toolResultPromises = /* @__PURE__ */ new Map();
 async function hasMatchingToolCallbacks(toolProvider, message) {
   if (!toolProvider || !message || !toolProvider.getAvailableTools) {
     return false;
@@ -535,6 +737,32 @@ async function hasMatchingToolCallbacks(toolProvider, message) {
   }
   return false;
 }
+function clearToolResultPromises() {
+  toolResultPromises.clear();
+}
+function updateToolResultsWithResolvedPromises(toolResults) {
+  return toolResults.map((toolResult) => {
+    const toolCallId = toolResult.data.toolCallId;
+    const promiseEntry = toolResultPromises.get(toolCallId);
+    if (promiseEntry && promiseEntry.resolvedValue !== null) {
+      const resolvedValue = promiseEntry.resolvedValue;
+      if (resolvedValue.error) {
+        return createToolResultDataPart(
+          toolCallId,
+          toolResult.data.toolId,
+          void 0,
+          resolvedValue.error
+        );
+      }
+      return createToolResultDataPart(
+        toolCallId,
+        toolResult.data.toolId,
+        resolvedValue
+      );
+    }
+    return toolResult;
+  });
+}
 async function executeToolCallBatch(toolCalls, toolProvider, messageId) {
   const results = [];
   const agentMessages = [];
@@ -595,7 +823,7 @@ function conversationHistoryToDataParts(conversationHistory) {
   }
   return historyParts;
 }
-async function continueTask(taskId, message, requestConfig, toolProvider, contextProvider, sessionId) {
+async function continueTask(taskId, message, requestConfig, toolProvider, contextProvider, sessionId, abortSignal) {
   const continueParams = {
     message,
     taskId,
@@ -605,30 +833,37 @@ async function continueTask(taskId, message, requestConfig, toolProvider, contex
   const preparedRequest = await prepareRequest(
     continueParams,
     requestConfig,
-    { isStreaming: false },
+    { isStreaming: false, abortSignal },
     toolProvider,
     contextProvider,
     sessionId
   );
-  return await executeRequest(preparedRequest, requestConfig);
+  return await executeRequest(preparedRequest, requestConfig, {
+    abortSignal
+  });
 }
-async function continueTaskStreamed(taskId, message, requestConfig, toolProvider, contextProvider, sessionId) {
+async function continueTaskStreamed(taskId, message, requestConfig, toolProvider, contextProvider, sessionId, abortSignal, requestOptions) {
   const continueParams = {
     message,
     taskId,
     sessionId: void 0
     // Use task's session
   };
+  const options = requestOptions || { isStreaming: true };
   const preparedRequest = await prepareRequest(
     continueParams,
     requestConfig,
-    { isStreaming: true },
+    {
+      ...options,
+      abortSignal
+    },
     toolProvider,
     contextProvider,
     sessionId
   );
   const stream = executeStreamingRequest(preparedRequest, requestConfig, {
-    isStreaming: true
+    ...options,
+    abortSignal
   });
   return processAgentResponseStream(
     stream,
@@ -638,11 +873,14 @@ async function continueTaskStreamed(taskId, message, requestConfig, toolProvider
     sessionId,
     true,
     // withHistory
-    []
+    [],
     // newConversationParts - empty for continueTask
+    abortSignal,
+    options
+    // Pass through the same request options
   );
 }
-async function* processAgentResponseStream(stream, toolProvider, contextProvider, requestConfig, sessionId, withHistory = true, newConversationParts = []) {
+async function* processAgentResponseStream(stream, toolProvider, contextProvider, requestConfig, sessionId, withHistory = true, newConversationParts = [], abortSignal, requestOptions) {
   for await (const update of stream) {
     yield update;
     if (update.status.state === "running" && update.status.message && toolProvider && await hasMatchingToolCallbacks(
@@ -712,6 +950,28 @@ async function* processAgentResponseStream(stream, toolProvider, contextProvider
                 createAgentTextMessage(agentMessage)
               );
             }
+            if (result.result instanceof Promise) {
+              const promise = result.result;
+              const promiseEntry = {
+                promise,
+                resolvedValue: null
+              };
+              toolResultPromises.set(
+                toolCallId,
+                promiseEntry
+              );
+              promise.then((resolvedValue) => {
+                promiseEntry.resolvedValue = resolvedValue;
+              }).catch((error) => {
+                console.error(
+                  `Promise rejected for tool call ${toolCallId}:`,
+                  error
+                );
+                promiseEntry.resolvedValue = {
+                  error: error instanceof Error ? error.message : String(error)
+                };
+              });
+            }
             const toolResult = createToolResultDataPart(
               toolCallId,
               toolId,
@@ -752,7 +1012,9 @@ async function* processAgentResponseStream(stream, toolProvider, contextProvider
             requestConfig,
             toolProvider,
             contextProvider,
-            sessionId
+            sessionId,
+            abortSignal,
+            requestOptions
           );
           let continuedTaskUpdate = null;
           for await (const streamUpdate of continuedTaskStream) {
@@ -837,7 +1099,9 @@ async function* processAgentResponseStream(stream, toolProvider, contextProvider
                   requestConfig,
                   toolProvider,
                   contextProvider,
-                  sessionId
+                  sessionId,
+                  abortSignal,
+                  requestOptions
                 );
                 let moreFinalTask = null;
                 for await (const streamUpdate of moreTaskStream) {
@@ -930,7 +1194,8 @@ function createClient(config) {
     defaultSessionId,
     timeout = DEFAULT_TIMEOUT,
     toolProvider,
-    contextProvider
+    contextProvider,
+    enableStreaming = false
   } = config;
   const requestConfig = {
     agentId,
@@ -940,21 +1205,22 @@ function createClient(config) {
   };
   return {
     async sendMessage(params) {
-      const { withHistory = true } = params;
+      const { withHistory = true, abortSignal } = params;
       const sessionId = params.sessionId || defaultSessionId || void 0;
       const newConversationParts = [];
       newConversationParts.push(params.message);
       const preparedRequest = await prepareRequest(
         params,
         requestConfig,
-        { isStreaming: false },
+        { isStreaming: false, abortSignal },
         toolProvider,
         contextProvider,
         sessionId
       );
       let currentTask = await executeRequest(
         preparedRequest,
-        requestConfig
+        requestConfig,
+        { abortSignal }
       );
       const allToolParts = [];
       const agentMessages = [];
@@ -1015,7 +1281,8 @@ function createClient(config) {
             requestConfig,
             toolProvider,
             contextProvider,
-            sessionId
+            sessionId,
+            abortSignal
           );
         } else {
           break;
@@ -1061,14 +1328,26 @@ function createClient(config) {
       };
     },
     async *sendMessageStream(params) {
-      const { withHistory = true } = params;
+      const {
+        withHistory = true,
+        abortSignal,
+        enableStreaming: requestStreaming
+      } = params;
       const sessionId = params.sessionId || defaultSessionId || void 0;
+      const useTokenStreaming = requestStreaming ?? enableStreaming;
       const newConversationParts = [];
       newConversationParts.push(params.message);
       const preparedRequest = await prepareRequest(
         params,
         requestConfig,
-        { isStreaming: true, streamingTimeout: timeout },
+        {
+          isStreaming: true,
+          // Always use message/stream endpoint for SSE
+          enableTokenStreaming: useTokenStreaming,
+          // Optional token-by-token streaming
+          streamingTimeout: timeout,
+          abortSignal
+        },
         toolProvider,
         contextProvider,
         sessionId
@@ -1078,7 +1357,10 @@ function createClient(config) {
         requestConfig,
         {
           isStreaming: true,
-          streamingTimeout: timeout
+          enableTokenStreaming: useTokenStreaming,
+          // Token streaming is optional
+          streamingTimeout: timeout,
+          abortSignal
         }
       );
       yield* processAgentResponseStream(
@@ -1088,7 +1370,13 @@ function createClient(config) {
         requestConfig,
         sessionId,
         withHistory,
-        newConversationParts
+        newConversationParts,
+        abortSignal,
+        {
+          isStreaming: true,
+          enableTokenStreaming: useTokenStreaming,
+          streamingTimeout: timeout
+        }
       );
     },
     async continueTask(taskId, userInput, sessionId) {
@@ -1378,6 +1666,31 @@ function extractToolResultsFromMessage(message) {
 }
 
 // src/react/agentManager.ts
+async function resolvePromisesInConversationHistory(conversationHistory) {
+  const resolvedHistory = [];
+  for (const message of conversationHistory) {
+    if (message.parts && Array.isArray(message.parts)) {
+      const hasToolResults = message.parts.some(
+        (part) => part.type === "data" && "toolCallId" in part.data && "result" in part.data
+      );
+      if (hasToolResults) {
+        const updatedParts = updateToolResultsWithResolvedPromises(
+          message.parts
+        );
+        resolvedHistory.push({
+          ...message,
+          parts: updatedParts
+        });
+      } else {
+        resolvedHistory.push(message);
+      }
+    } else {
+      resolvedHistory.push(message);
+    }
+  }
+  clearToolResultPromises();
+  return resolvedHistory;
+}
 function createAgentManager() {
   const agents = /* @__PURE__ */ new Map();
   async function persistConversationHistory(key, messages) {
@@ -1506,9 +1819,20 @@ function createAgentManager() {
         ...managedAgent.conversationHistory
       ];
       let currentToolCallIds = [];
+      const resolvedConversationHistory = await resolvePromisesInConversationHistory(
+        currentConversationHistory
+      );
+      managedAgent.conversationHistory = resolvedConversationHistory;
+      currentConversationHistory = resolvedConversationHistory;
+      if (withHistory) {
+        await persistConversationHistory(
+          key,
+          resolvedConversationHistory
+        );
+      }
       const messageObj = options.message || createTextMessageWithHistory(
         message,
-        currentConversationHistory
+        resolvedConversationHistory
       );
       const userMessage = createTextMessage(message);
       currentConversationHistory = [
@@ -1632,7 +1956,7 @@ function getAgentManager() {
 }
 
 // src/utils/createMessageRenderer.tsx
-import Markdown from "react-markdown";
+import Markdown from "streamdown";
 
 // src/markdown-extensions/index.ts
 import remarkGfm from "remark-gfm";
@@ -1699,30 +2023,40 @@ var createFeedbackActions = (config) => {
       }
       const feedbackState = feedbackByMessageId[message.id];
       const actions = [];
-      if (feedbackState !== "down") {
-        actions.push({
-          id: "feedback-up",
-          icon: config.icons.up,
-          label: "Good response",
-          onClick: async () => {
+      actions.push({
+        id: "feedback-up",
+        icon: config.icons.up,
+        label: "Good response",
+        onClick: async () => {
+          if (feedbackState === "up") {
+            delete feedbackByMessageId[message.id];
+          } else {
             await handleFeedback(message.id, "up");
-          },
-          tooltip: "This response was helpful",
-          disabled: feedbackState === "up"
-        });
-      }
-      if (feedbackState !== "up") {
-        actions.push({
-          id: "feedback-down",
-          icon: config.icons.down,
-          label: "Bad response",
-          onClick: async () => {
+          }
+          notifyListeners();
+        },
+        tooltip: "This response was helpful",
+        pressed: feedbackState === "up",
+        disabled: feedbackState === "down"
+        // Disable if other is selected
+      });
+      actions.push({
+        id: "feedback-down",
+        icon: config.icons.down,
+        label: "Bad response",
+        onClick: async () => {
+          if (feedbackState === "down") {
+            delete feedbackByMessageId[message.id];
+          } else {
             await handleFeedback(message.id, "down");
-          },
-          tooltip: "This response was not helpful",
-          disabled: feedbackState === "down"
-        });
-      }
+          }
+          notifyListeners();
+        },
+        tooltip: "This response was not helpful",
+        pressed: feedbackState === "down",
+        disabled: feedbackState === "up"
+        // Disable if other is selected
+      });
       return actions;
     },
     clearFeedback: (messageId) => {
@@ -1800,7 +2134,9 @@ function resolveActionsForMessage(message, registrations) {
     icon: action.icon,
     onClick: action.onClick,
     tooltip: action.tooltip,
-    disabled: action.disabled || false
+    disabled: action.disabled || false,
+    pressed: action.pressed,
+    showLabel: action.showLabel
   }));
   return filteredActions;
 }
@@ -1904,7 +2240,8 @@ function useAgentChat(config) {
     error: isValidConfig ? null : "Invalid agent configuration",
     suggestions: [],
     markdownComponents: {},
-    markdownExtensions: {}
+    markdownExtensions: {},
+    currentAbortController: null
   });
   const {
     registerMessageActions,
@@ -1932,7 +2269,8 @@ function useAgentChat(config) {
             sessionId: agentConfig.sessionId,
             contextProvider: config.contextProvider || createNoOpContextProvider(),
             toolProvider: config.toolProvider || createNoOpToolProvider(),
-            authProvider: config.authProvider
+            authProvider: config.authProvider,
+            enableStreaming: config.enableStreaming
           });
           const clientHistory = agentManager.getConversationHistory(agentKey);
           setState((prev) => {
@@ -1961,6 +2299,7 @@ function useAgentChat(config) {
     config.contextProvider,
     config.toolProvider,
     config.authProvider,
+    config.enableStreaming,
     isValidConfig
   ]);
   const onSubmit = useCallback3(
@@ -1979,51 +2318,143 @@ function useAgentChat(config) {
         archived: false,
         showIcon: false
       };
+      const abortController = new AbortController();
       setState((prev) => ({
         ...prev,
         uiMessages: [...prev.uiMessages, userMessage],
         isProcessing: true,
-        error: null
+        error: null,
+        currentAbortController: abortController
       }));
       try {
-        let lastUpdate = null;
+        let streamingMessageId = null;
+        let finalMessageAdded = false;
         for await (const update of agentManager.sendMessageStream(
           agentKey,
-          message
+          message,
+          {
+            abortSignal: abortController.signal
+          }
         )) {
-          lastUpdate = update;
-        }
-        const updatedClientHistory = agentManager.getConversationHistory(agentKey);
-        setState((prev) => {
-          const transformedClientMessages = updatedClientHistory.map(
-            (msg) => transformClientMessageToUI(
-              msg,
+          if (!update.final && update.text) {
+            if (!streamingMessageId) {
+              streamingMessageId = `agent-streaming-${Date.now()}`;
+              const streamingMessage = {
+                id: streamingMessageId,
+                role: "agent",
+                content: [
+                  { type: "text", text: update.text }
+                ],
+                timestamp: Date.now(),
+                archived: false,
+                showIcon: true,
+                icon: "assistant"
+              };
+              setState((prev) => ({
+                ...prev,
+                uiMessages: [
+                  ...prev.uiMessages,
+                  streamingMessage
+                ]
+              }));
+            } else {
+              setState((prev) => ({
+                ...prev,
+                uiMessages: prev.uiMessages.map(
+                  (msg) => msg.id === streamingMessageId ? {
+                    ...msg,
+                    content: [
+                      {
+                        type: "text",
+                        text: update.text
+                      }
+                    ]
+                  } : msg
+                )
+              }));
+            }
+          }
+          if (update.final && update.status?.message && streamingMessageId) {
+            finalMessageAdded = true;
+            const currentStreamingId = streamingMessageId;
+            const finalMessage = transformClientMessageToUI(
+              update.status.message,
               registrationsRef.current
-            )
-          ).filter((msg) => msg !== null);
-          const clientMessageIds = new Set(
-            updatedClientHistory.map((msg) => msg.messageId)
-          );
-          const uiOnlyMessages = prev.uiMessages.filter(
-            (msg) => !clientMessageIds.has(msg.id) && msg.content[0]?.type === "component"
-          );
-          const mergedUIMessages = sortUIMessagesByTime([
-            ...transformedClientMessages,
-            ...uiOnlyMessages
-          ]);
-          return {
-            ...prev,
-            clientMessages: updatedClientHistory,
-            uiMessages: mergedUIMessages,
-            isProcessing: false
-          };
-        });
+            );
+            if (finalMessage) {
+              setState((prev) => {
+                const updatedMessages = prev.uiMessages.map(
+                  (msg) => msg.id === currentStreamingId ? finalMessage : msg
+                );
+                const updatedClientHistory = agentManager.getConversationHistory(
+                  agentKey
+                );
+                return {
+                  ...prev,
+                  clientMessages: updatedClientHistory,
+                  uiMessages: updatedMessages,
+                  isProcessing: false,
+                  currentAbortController: null
+                };
+              });
+            }
+            streamingMessageId = null;
+          }
+        }
+        if (!finalMessageAdded) {
+          const updatedClientHistory = agentManager.getConversationHistory(agentKey);
+          setState((prev) => {
+            let filteredMessages = prev.uiMessages;
+            if (streamingMessageId) {
+              filteredMessages = prev.uiMessages.filter(
+                (msg) => msg.id !== streamingMessageId
+              );
+            }
+            const transformedClientMessages = updatedClientHistory.map(
+              (msg) => transformClientMessageToUI(
+                msg,
+                registrationsRef.current
+              )
+            ).filter(
+              (msg) => msg !== null
+            );
+            const clientMessageIds = new Set(
+              updatedClientHistory.map((msg) => msg.messageId)
+            );
+            const uiOnlyMessages = filteredMessages.filter(
+              (msg) => !clientMessageIds.has(msg.id) && msg.content[0]?.type === "component"
+            );
+            const mergedUIMessages = sortUIMessagesByTime([
+              ...transformedClientMessages,
+              ...uiOnlyMessages
+            ]);
+            return {
+              ...prev,
+              clientMessages: updatedClientHistory,
+              uiMessages: mergedUIMessages,
+              isProcessing: false,
+              currentAbortController: null
+            };
+          });
+        }
       } catch (error) {
+        if (error instanceof Error && error.name === "AbortError") {
+          console.log("Request was aborted by user");
+          setState((prev) => ({
+            ...prev,
+            isProcessing: false,
+            error: null,
+            // Don't show error for user-initiated abort
+            currentAbortController: null
+          }));
+          return;
+        }
         const errorMessage = error instanceof Error ? error.message : "Failed to send message";
         setState((prev) => ({
           ...prev,
           isProcessing: false,
-          error: errorMessage
+          error: errorMessage,
+          currentAbortController: null
         }));
         throw error;
       }
@@ -2142,6 +2573,14 @@ function useAgentChat(config) {
     }),
     [state.markdownComponents, state.markdownExtensions]
   );
+  const abortCurrentRequest = useCallback3(() => {
+    setState((currentState) => {
+      if (currentState.currentAbortController) {
+        currentState.currentAbortController.abort();
+      }
+      return currentState;
+    });
+  }, []);
   return {
     // AgentUI props
     messages: state.uiMessages,
@@ -2161,7 +2600,9 @@ function useAgentChat(config) {
     clearAllMessageActions,
     createFeedbackActions: createFeedbackActions2,
     // Tool integration
-    addMessage
+    addMessage,
+    // Abort control
+    abortCurrentRequest
   };
 }
 
@@ -2266,6 +2707,7 @@ export {
   BarChart,
   ChartBlock,
   LineChart,
+  createAbortController,
   createClient,
   createFeedbackActions,
   createJetpackAuthProvider,