@ai-sdk/langchain 3.0.0-beta.17 → 3.0.0-beta.178

package/src/utils.ts

@@ -2,26 +2,30 @@ import {
   AIMessage,
   HumanMessage,
   ToolMessage,
-  BaseMessage,
   AIMessageChunk,
-  BaseMessageChunk,
-  ToolCallChunk,
   type ContentBlock,
   type ToolCall,
+  type BaseMessage,
+  type BaseMessageChunk,
+  type ToolCallChunk,
 } from '@langchain/core/messages';
-import {
-  type UIMessageChunk,
-  type ToolResultPart,
-  type AssistantContent,
-  type UserContent,
+import type {
+  UIMessageChunk,
+  ToolResultPart,
+  AssistantContent,
+  UserContent,
+  ProviderMetadata,
+  JSONValue,
 } from 'ai';
 
-import {
-  type LangGraphEventState,
-  type ReasoningContentBlock,
-  type ThinkingContentBlock,
-  type GPT5ReasoningOutput,
-  type ImageGenerationOutput,
+import type {
+  LangGraphEventState,
+  LangGraphMessageSeen,
+  ReasoningContentBlock,
+  ThinkingContentBlock,
+  GPT5ReasoningOutput,
+  ImageGenerationOutput,
+  NormalizedCitation,
 } from './types';
 
 /**
@@ -114,8 +118,8 @@ function getDefaultFilename(
   mediaType: string,
   prefix: string = 'file',
 ): string {
-  const ext = mediaType.split('/')[1] || 'bin';
-  return `${prefix}.${ext}`;
+  const fileExtension = mediaType.split('/')[1] || 'bin';
+  return `${prefix}.${fileExtension}`;
 }
 
 /**
@@ -220,13 +224,53 @@ export function convertUserContent(content: UserContent): HumanMessage {
         });
       }
     } else if (part.type === 'file') {
-      const filePart = part as {
+      const rawFilePart = part as {
         type: 'file';
-        data: string | Uint8Array | URL | ArrayBuffer;
+        data:
+          | string
+          | Uint8Array
+          | URL
+          | ArrayBuffer
+          | { type: 'data'; data: string | Uint8Array | ArrayBuffer }
+          | { type: 'url'; url: URL }
+          | { type: 'reference'; reference: Record<string, string> }
+          | { type: 'text'; text: string };
         mediaType: string;
         filename?: string;
       };
 
+      // Normalize tagged data shape into the legacy bare value this code expects.
+      const normalizedData: string | Uint8Array | URL | ArrayBuffer = (() => {
+        const data = rawFilePart.data;
+        if (
+          typeof data === 'object' &&
+          data !== null &&
+          !(data instanceof URL) &&
+          !(data instanceof Uint8Array) &&
+          !(data instanceof ArrayBuffer) &&
+          'type' in data
+        ) {
+          switch (data.type) {
+            case 'data':
+              return data.data;
+            case 'url':
+              return data.url;
+            case 'text':
+              return data.text;
+            default:
+              return '';
+          }
+        }
+        return data as string | Uint8Array | URL | ArrayBuffer;
+      })();
+
+      const filePart = {
+        type: 'file' as const,
+        data: normalizedData,
+        mediaType: rawFilePart.mediaType,
+        filename: rawFilePart.filename,
+      };
+
       /**
        * Check if this is an image file - if so, use OpenAI's image_url format
        */
@@ -393,6 +437,7 @@ export function processModelChunk(
     /** Track the ID used for text-start to ensure text-end uses the same ID */
     textMessageId?: string | null;
     emittedImages?: Set<string>;
+    emittedSourceIds?: Set<string>;
   },
   controller: ReadableStreamDefaultController<UIMessageChunk>,
 ): void {
@@ -403,6 +448,10 @@ export function processModelChunk(
     state.emittedImages = new Set<string>();
   }
 
+  if (!state.emittedSourceIds) {
+    state.emittedSourceIds = new Set<string>();
+  }
+
   /**
    * Get the message ID from the chunk if available
    */
@@ -506,6 +555,16 @@ export function processModelChunk(
       id: state.textMessageId ?? state.messageId,
     });
   }
+
+  const citations = extractCitationsFromContentBlocks(chunk);
+  if (citations.length > 0) {
+    emitSourceChunks(
+      citations,
+      state.messageId,
+      state.emittedSourceIds,
+      controller,
+    );
+  }
 }
 
 /**
@@ -562,7 +621,8 @@ export function getMessageId(msg: unknown): string | undefined {
 /**
  * Checks if a message is an AI message chunk (works for both class instances and plain objects).
  * For class instances, only AIMessageChunk is matched (not AIMessage).
- * For plain objects from RemoteGraph API, matches type === 'ai'.
+ * For plain objects from RemoteGraph API, matches type === 'ai' (TypeScript langchain-core)
+ * or type === 'AIMessageChunk' (Python langchain-core).
  * For serialized LangChain messages, matches type === 'constructor' with AIMessageChunk in id path.
  *
  * @param msg - The message to check.
@@ -579,18 +639,25 @@ export function isAIMessageChunk(
    * Plain object from RemoteGraph API (not a LangChain class instance)
    */
   if (isPlainMessageObject(msg)) {
-    const obj = msg as Record<string, unknown>;
+    const messageRecord = msg as Record<string, unknown>;
     /**
-     * Direct type === 'ai' (RemoteGraph format)
+     * Direct type from RemoteGraph format. TypeScript langchain-core emits
+     * type === 'ai'; Python langchain-core emits type === 'AIMessageChunk'.
      */
-    if ('type' in obj && obj.type === 'ai') return true;
+    if (
+      'type' in messageRecord &&
+      (messageRecord.type === 'ai' || messageRecord.type === 'AIMessageChunk')
+    ) {
+      return true;
+    }
     /**
      * Serialized LangChain message format: { lc: 1, type: "constructor", id: ["...", "AIMessageChunk"], kwargs: {...} }
      */
     if (
-      obj.type === 'constructor' &&
-      Array.isArray(obj.id) &&
-      (obj.id.includes('AIMessageChunk') || obj.id.includes('AIMessage'))
+      messageRecord.type === 'constructor' &&
+      Array.isArray(messageRecord.id) &&
+      (messageRecord.id.includes('AIMessageChunk') ||
+        messageRecord.id.includes('AIMessage'))
     ) {
       return true;
     }
@@ -612,18 +679,18 @@ export function isToolMessageType(
    * Plain object from RemoteGraph API (not a LangChain class instance)
    */
   if (isPlainMessageObject(msg)) {
-    const obj = msg as Record<string, unknown>;
+    const messageRecord = msg as Record<string, unknown>;
     /**
      * Direct type === 'tool' (RemoteGraph format)
      */
-    if ('type' in obj && obj.type === 'tool') return true;
+    if ('type' in messageRecord && messageRecord.type === 'tool') return true;
     /**
      * Serialized LangChain message format
      */
     if (
-      obj.type === 'constructor' &&
-      Array.isArray(obj.id) &&
-      obj.id.includes('ToolMessage')
+      messageRecord.type === 'constructor' &&
+      Array.isArray(messageRecord.id) &&
+      messageRecord.id.includes('ToolMessage')
     ) {
       return true;
     }
@@ -916,6 +983,152 @@ export function extractReasoningFromValuesMessage(
   return undefined;
 }
 
+export function isCitationContentBlock(
+  obj: unknown,
+): obj is ContentBlock.Citation {
+  return (
+    obj != null &&
+    typeof obj === 'object' &&
+    'type' in obj &&
+    (obj as { type: unknown }).type === 'citation'
+  );
+}
+
+/**
+ * LangChain Core standardizes citations as `{ type: 'citation', ... }` entries in
+ * the `annotations` array of a text content block. The text extractors elsewhere
+ * in this file only read the `text` field, so without this the citations would be
+ * dropped instead of surfaced as AI SDK source parts.
+ *
+ * Handles both class instances and serialized LangChain message objects.
+ */
+export function extractCitationsFromContentBlocks(
+  msg: unknown,
+): NormalizedCitation[] {
+  if (msg == null || typeof msg !== 'object') return [];
+
+  const msgObj = msg as Record<string, unknown>;
+  const kwargs =
+    msgObj.kwargs && typeof msgObj.kwargs === 'object'
+      ? (msgObj.kwargs as Record<string, unknown>)
+      : msgObj;
+
+  // `contentBlocks` is the normalized view derived from `content`; reading both
+  // would double-count annotations, so prefer it and only fall back to `content`.
+  const contentBlocks = (kwargs as { contentBlocks?: unknown }).contentBlocks;
+  const content = (kwargs as { content?: unknown }).content;
+  const blockSources: unknown[] = Array.isArray(contentBlocks)
+    ? contentBlocks
+    : Array.isArray(content)
+      ? content
+      : [];
+
+  const citations: NormalizedCitation[] = [];
+
+  for (const block of blockSources) {
+    if (block == null || typeof block !== 'object') continue;
+
+    const annotations = (block as { annotations?: unknown }).annotations;
+    if (!Array.isArray(annotations)) continue;
+
+    for (const annotation of annotations) {
+      if (!isCitationContentBlock(annotation)) continue;
+
+      citations.push({
+        url: annotation.url,
+        title: annotation.title,
+        source: annotation.source,
+        citedText: annotation.citedText,
+        startIndex: annotation.startIndex,
+        endIndex: annotation.endIndex,
+      });
+    }
+  }
+
+  return citations;
+}
+
+function buildSourceProviderMetadata(
+  citation: NormalizedCitation,
+): ProviderMetadata | undefined {
+  const langchain: Record<string, JSONValue> = {};
+
+  if (typeof citation.citedText === 'string') {
+    langchain.citedText = citation.citedText;
+  }
+  if (typeof citation.startIndex === 'number') {
+    langchain.startIndex = citation.startIndex;
+  }
+  if (typeof citation.endIndex === 'number') {
+    langchain.endIndex = citation.endIndex;
+  }
+  if (typeof citation.source === 'string') {
+    langchain.source = citation.source;
+  }
+
+  if (Object.keys(langchain).length === 0) {
+    return undefined;
+  }
+
+  return { langchain };
+}
+
+/**
+ * Emits AI SDK source chunks (`source-url` / `source-document`) for the given
+ * citations.
+ *
+ * Citations with a `url` become `source-url` parts keyed by that url; the rest
+ * become `source-document` parts. Citation metadata that the source parts cannot
+ * represent natively (`citedText`, `startIndex`, `endIndex`, `source`) is preserved
+ * under `providerMetadata.langchain`.
+ *
+ * `emittedSourceIds` dedupes across the LangGraph messages and values events, which
+ * can each surface the same citation. URL-less citations are keyed by their content
+ * rather than position, so a differing citation subset/order between those two events
+ * does not cause an id collision.
+ */
+export function emitSourceChunks(
+  citations: NormalizedCitation[],
+  messageId: string,
+  emittedSourceIds: Set<string>,
+  controller: ReadableStreamDefaultController<UIMessageChunk>,
+): void {
+  for (const citation of citations) {
+    if (citation.url) {
+      if (emittedSourceIds.has(citation.url)) continue;
+      emittedSourceIds.add(citation.url);
+
+      const providerMetadata = buildSourceProviderMetadata(citation);
+      controller.enqueue({
+        type: 'source-url',
+        sourceId: citation.url,
+        url: citation.url,
+        ...(citation.title ? { title: citation.title } : {}),
+        ...(providerMetadata ? { providerMetadata } : {}),
+      });
+      continue;
+    }
+
+    // A citation with no url and no human-readable label carries no information a
+    // UI could render, so skip it rather than emit a placeholder source.
+    const title = citation.title ?? citation.source;
+    if (!title) continue;
+
+    const sourceId = `${messageId}:${title}:${citation.citedText ?? ''}:${citation.startIndex ?? ''}:${citation.endIndex ?? ''}`;
+    if (emittedSourceIds.has(sourceId)) continue;
+    emittedSourceIds.add(sourceId);
+
+    const providerMetadata = buildSourceProviderMetadata(citation);
+    controller.enqueue({
+      type: 'source-document',
+      sourceId,
+      mediaType: 'text/plain',
+      title,
+      ...(providerMetadata ? { providerMetadata } : {}),
+    });
+  }
+}
+
 /**
  * Checks if an object is an image generation output
  *
@@ -950,6 +1163,50 @@ export function extractImageOutputs(
   return toolOutputs.filter(isImageGenerationOutput);
 }
 
+function formatToolError(error: unknown): string {
+  if (error instanceof Error) return error.message;
+  if (typeof error === 'string') return error;
+
+  try {
+    const serialized = JSON.stringify(error);
+    return serialized ?? String(error);
+  } catch {
+    return String(error);
+  }
+}
+
+/**
+ * Returns per-message bookkeeping, creating it on first use.
+ * Map keys keep remote-controlled message IDs like "__proto__" from touching Object.prototype.
+ */
+function getOrCreateMessageSeen(
+  messageSeen: LangGraphEventState['messageSeen'],
+  msgId: string,
+): LangGraphMessageSeen {
+  let seen = messageSeen.get(msgId);
+  if (!seen) {
+    seen = {};
+    messageSeen.set(msgId, seen);
+  }
+  return seen;
+}
+
+/**
+ * Returns the per-index tool call metadata for a message, creating it on first use.
+ * Later streamed chunks use this to recover tool call IDs/names from earlier chunks.
+ */
+function getOrCreateToolCallInfoByIndex(
+  toolCallInfoByIndex: LangGraphEventState['toolCallInfoByIndex'],
+  msgId: string,
+): Map<number, { id: string; name: string }> {
+  let toolCallInfo = toolCallInfoByIndex.get(msgId);
+  if (!toolCallInfo) {
+    toolCallInfo = new Map();
+    toolCallInfoByIndex.set(msgId, toolCallInfo);
+  }
+  return toolCallInfo;
+}
+
 /**
  * Processes a LangGraph event and emits UI message chunks.
  *
@@ -971,6 +1228,7 @@ export function processLangGraphEvent(
     messageReasoningIds,
     toolCallInfoByIndex,
     emittedToolCallsByKey,
+    emittedToolInputs,
   } = state;
   const [type, data] = parseLangGraphEvent(event);
 
@@ -1021,7 +1279,10 @@ export function processLangGraphEvent(
       if (!msgId) return;
 
       /**
-       * Track LangGraph step changes and emit start-step/finish-step events
+       * Track LangGraph step changes and emit start-step/finish-step events.
+       * Before emitting finish-step, close any open text/reasoning parts so
+       * the client does not receive orphaned deltas after its
+       * activeReasoningParts / activeTextParts have been cleared.
        */
       const langgraphStep =
         typeof metadata?.langgraph_step === 'number'
@@ -1029,6 +1290,17 @@ export function processLangGraphEvent(
           : null;
       if (langgraphStep !== null && langgraphStep !== state.currentStep) {
         if (state.currentStep !== null) {
+          for (const [id, seen] of messageSeen) {
+            if (seen.text) {
+              controller.enqueue({ type: 'text-end', id });
+            }
+            if (seen.reasoning) {
+              controller.enqueue({ type: 'reasoning-end', id });
+            }
+            messageSeen.delete(id);
+            messageConcat.delete(id);
+            messageReasoningIds.delete(id);
+          }
           controller.enqueue({ type: 'finish-step' });
         }
         controller.enqueue({ type: 'start-step' });
@@ -1040,17 +1312,19 @@ export function processLangGraphEvent(
        * Note: Only works for actual class instances, not serialized messages
        */
       if (AIMessageChunk.isInstance(msg)) {
-        if (messageConcat[msgId]) {
-          messageConcat[msgId] = messageConcat[msgId].concat(
-            msg,
-          ) as AIMessageChunk;
+        const existingMessage = messageConcat.get(msgId);
+        if (existingMessage) {
+          messageConcat.set(
+            msgId,
+            existingMessage.concat(msg) as AIMessageChunk,
+          );
         } else {
-          messageConcat[msgId] = msg;
+          messageConcat.set(msgId, msg);
         }
       }
 
       if (isAIMessageChunk(msg)) {
-        const concatChunk = messageConcat[msgId];
+        const concatChunk = messageConcat.get(msgId);
 
         /**
          * Handle image generation outputs from additional_kwargs.tool_outputs
@@ -1103,29 +1377,34 @@ export function processLangGraphEvent(
           | undefined;
         if (toolCallChunks?.length) {
           for (const toolCallChunk of toolCallChunks) {
-            const idx = toolCallChunk.index ?? 0;
+            const toolCallIndex = toolCallChunk.index ?? 0;
 
             /**
              * If this chunk has an id, store it for future lookups by index
              */
             if (toolCallChunk.id) {
-              toolCallInfoByIndex[msgId] ??= {};
-              toolCallInfoByIndex[msgId][idx] = {
-                id: toolCallChunk.id,
-                name:
-                  toolCallChunk.name ||
-                  concatChunk?.tool_call_chunks?.[idx]?.name ||
-                  'unknown',
-              };
+              getOrCreateToolCallInfoByIndex(toolCallInfoByIndex, msgId).set(
+                toolCallIndex,
+                {
+                  id: toolCallChunk.id,
+                  name:
+                    toolCallChunk.name ||
+                    concatChunk?.tool_call_chunks?.[toolCallIndex]?.name ||
+                    'unknown',
+                },
+              );
             }
 
             /**
              * Get the tool call ID from the chunk, stored info, or accumulated chunks
              */
+            const storedToolCallInfo = toolCallInfoByIndex
+              .get(msgId)
+              ?.get(toolCallIndex);
             const toolCallId =
               toolCallChunk.id ||
-              toolCallInfoByIndex[msgId]?.[idx]?.id ||
-              concatChunk?.tool_call_chunks?.[idx]?.id;
+              storedToolCallInfo?.id ||
+              concatChunk?.tool_call_chunks?.[toolCallIndex]?.id;
 
             /**
              * Skip if we don't have a proper tool call ID - we'll handle it in values
@@ -1136,8 +1415,8 @@ export function processLangGraphEvent(
 
             const toolName =
               toolCallChunk.name ||
-              toolCallInfoByIndex[msgId]?.[idx]?.name ||
-              concatChunk?.tool_call_chunks?.[idx]?.name ||
+              storedToolCallInfo?.name ||
+              concatChunk?.tool_call_chunks?.[toolCallIndex]?.name ||
               'unknown';
 
             /**
@@ -1145,18 +1424,21 @@ export function processLangGraphEvent(
              * (even if args is empty - the first chunk often has empty args)
              * Set dynamic: true to enable HITL approval requests
              */
-            if (!messageSeen[msgId]?.tool?.[toolCallId]) {
-              controller.enqueue({
-                type: 'tool-input-start',
-                toolCallId: toolCallId,
-                toolName: toolName,
-                dynamic: true,
-              });
+            const seen = messageSeen.get(msgId);
+            if (!seen?.tool?.has(toolCallId)) {
+              const updatedSeen = getOrCreateMessageSeen(messageSeen, msgId);
+              updatedSeen.tool ??= new Set();
+              updatedSeen.tool.add(toolCallId);
 
-              messageSeen[msgId] ??= {};
-              messageSeen[msgId].tool ??= {};
-              messageSeen[msgId].tool![toolCallId] = true;
-              emittedToolCalls.add(toolCallId);
+              if (!emittedToolCalls.has(toolCallId)) {
+                emittedToolCalls.add(toolCallId);
+                controller.enqueue({
+                  type: 'tool-input-start',
+                  toolCallId: toolCallId,
+                  toolName: toolName,
+                  dynamic: true,
+                });
+              }
             }
 
             /**
@@ -1187,8 +1469,8 @@ export function processLangGraphEvent(
         // Capture reasoning ID when we first see it (even if no content yet)
         const chunkReasoningId = extractReasoningId(msg);
         if (chunkReasoningId) {
-          if (!messageReasoningIds[msgId]) {
-            messageReasoningIds[msgId] = chunkReasoningId;
+          if (!messageReasoningIds.has(msgId)) {
+            messageReasoningIds.set(msgId, chunkReasoningId);
           }
           // Immediately mark as emitted to prevent values from duplicating
           // This must happen as soon as we see the ID, before content arrives
@@ -1199,12 +1481,12 @@ export function processLangGraphEvent(
         if (reasoning) {
           // Use stored reasoning ID, or current chunk's ID, or fall back to message ID
           const reasoningId =
-            messageReasoningIds[msgId] ?? chunkReasoningId ?? msgId;
+            messageReasoningIds.get(msgId) ?? chunkReasoningId ?? msgId;
 
-          if (!messageSeen[msgId]?.reasoning) {
+          const seen = messageSeen.get(msgId);
+          if (!seen?.reasoning) {
             controller.enqueue({ type: 'reasoning-start', id: msgId });
-            messageSeen[msgId] ??= {};
-            messageSeen[msgId].reasoning = true;
+            getOrCreateMessageSeen(messageSeen, msgId).reasoning = true;
           }
 
           // Streaming chunks have delta text, emit directly without slicing
@@ -1222,10 +1504,10 @@ export function processLangGraphEvent(
          */
         const text = getMessageText(msg);
         if (text) {
-          if (!messageSeen[msgId]?.text) {
+          const seen = messageSeen.get(msgId);
+          if (!seen?.text) {
             controller.enqueue({ type: 'text-start', id: msgId });
-            messageSeen[msgId] ??= {};
-            messageSeen[msgId].text = true;
+            getOrCreateMessageSeen(messageSeen, msgId).text = true;
           }
 
           controller.enqueue({
@@ -1234,6 +1516,16 @@ export function processLangGraphEvent(
             id: msgId,
           });
         }
+
+        const citations = extractCitationsFromContentBlocks(msg);
+        if (citations.length > 0) {
+          emitSourceChunks(
+            citations,
+            msgId,
+            state.emittedSourceIds,
+            controller,
+          );
+        }
       } else if (isToolMessageType(msg)) {
         // Handle both direct properties and serialized messages (kwargs)
         const msgObj = msg as unknown as Record<string, unknown>;
@@ -1272,31 +1564,122 @@ export function processLangGraphEvent(
       return;
     }
 
+    case 'tools': {
+      if (data == null || typeof data !== 'object' || Array.isArray(data)) {
+        return;
+      }
+
+      const payload = data as {
+        event?: unknown;
+        name?: unknown;
+        input?: unknown;
+        data?: unknown;
+        output?: unknown;
+        error?: unknown;
+        toolCallId?: unknown;
+      };
+      const toolCallId =
+        typeof payload.toolCallId === 'string' ? payload.toolCallId : undefined;
+      const toolName =
+        typeof payload.name === 'string' ? payload.name : 'unknown';
+
+      if (!toolCallId) return;
+
+      const ensureToolInputLifecycle = () => {
+        if (!emittedToolCalls.has(toolCallId)) {
+          emittedToolCalls.add(toolCallId);
+          controller.enqueue({
+            type: 'tool-input-start',
+            toolCallId,
+            toolName,
+            dynamic: true,
+          });
+        }
+
+        if (!emittedToolInputs.has(toolCallId)) {
+          emittedToolInputs.add(toolCallId);
+          controller.enqueue({
+            type: 'tool-input-available',
+            toolCallId,
+            toolName,
+            input: payload.input,
+            dynamic: true,
+          });
+        }
+      };
+
+      switch (payload.event) {
+        case 'on_tool_start': {
+          const toolCallKey = `${toolName}:${JSON.stringify(payload.input)}`;
+          emittedToolCallsByKey.set(toolCallKey, toolCallId);
+
+          ensureToolInputLifecycle();
+          break;
+        }
+
+        case 'on_tool_event': {
+          ensureToolInputLifecycle();
+          controller.enqueue({
+            type: 'tool-output-available',
+            toolCallId,
+            output: payload.data,
+            preliminary: true,
+          });
+          break;
+        }
+
+        case 'on_tool_end': {
+          ensureToolInputLifecycle();
+          controller.enqueue({
+            type: 'tool-output-available',
+            toolCallId,
+            output: payload.output,
+          });
+          break;
+        }
+
+        case 'on_tool_error': {
+          ensureToolInputLifecycle();
+          controller.enqueue({
+            type: 'tool-output-error',
+            toolCallId,
+            errorText: formatToolError(payload.error),
+          });
+          break;
+        }
+      }
+
+      return;
+    }
+
     case 'values': {
       /**
        * Finalize all pending message chunks
        */
-      for (const [id, seen] of Object.entries(messageSeen)) {
+      for (const [id, seen] of messageSeen) {
         if (seen.text) controller.enqueue({ type: 'text-end', id });
         if (seen.tool) {
-          for (const [toolCallId, toolCallSeen] of Object.entries(seen.tool)) {
-            const concatMsg = messageConcat[id];
+          for (const toolCallId of seen.tool) {
+            const concatMsg = messageConcat.get(id);
             const toolCall = concatMsg?.tool_calls?.find(
               call => call.id === toolCallId,
             );
 
-            if (toolCallSeen && toolCall) {
+            if (toolCall) {
               emittedToolCalls.add(toolCallId);
               // Store mapping for HITL interrupt lookup
               const toolCallKey = `${toolCall.name}:${JSON.stringify(toolCall.args)}`;
               emittedToolCallsByKey.set(toolCallKey, toolCallId);
-              controller.enqueue({
-                type: 'tool-input-available',
-                toolCallId,
-                toolName: toolCall.name,
-                input: toolCall.args,
-                dynamic: true,
-              });
+              if (!emittedToolInputs.has(toolCallId)) {
+                emittedToolInputs.add(toolCallId);
+                controller.enqueue({
+                  type: 'tool-input-available',
+                  toolCallId,
+                  toolName: toolCall.name,
+                  input: toolCall.args,
+                  dynamic: true,
+                });
+              }
             }
           }
         }
@@ -1305,9 +1688,9 @@ export function processLangGraphEvent(
           controller.enqueue({ type: 'reasoning-end', id });
         }
 
-        delete messageSeen[id];
-        delete messageConcat[id];
-        delete messageReasoningIds[id];
+        messageSeen.delete(id);
+        messageConcat.delete(id);
+        messageReasoningIds.delete(id);
       }
 
       /**
@@ -1367,21 +1750,25 @@ export function processLangGraphEvent(
               /**
                * For plain objects from RemoteGraph API or serialized LangChain messages
                */
-              const obj = msg as Record<string, unknown>;
+              const messageRecord = msg as Record<string, unknown>;
 
               /**
                * Determine the data source (handle both direct and serialized formats)
                */
               const isSerializedFormat =
-                obj.type === 'constructor' &&
-                Array.isArray(obj.id) &&
-                ((obj.id as string[]).includes('AIMessageChunk') ||
-                  (obj.id as string[]).includes('AIMessage'));
+                messageRecord.type === 'constructor' &&
+                Array.isArray(messageRecord.id) &&
+                ((messageRecord.id as string[]).includes('AIMessageChunk') ||
+                  (messageRecord.id as string[]).includes('AIMessage'));
               const dataSource = isSerializedFormat
-                ? (obj.kwargs as Record<string, unknown>)
-                : obj;
-
-              if (obj.type === 'ai' || isSerializedFormat) {
+                ? (messageRecord.kwargs as Record<string, unknown>)
+                : messageRecord;
+
+              if (
+                messageRecord.type === 'ai' ||
+                messageRecord.type === 'AIMessageChunk' ||
+                isSerializedFormat
+              ) {
                 /**
                  * Try tool_calls first (normalized format)
                  */
@@ -1454,6 +1841,7 @@ export function processLangGraphEvent(
                     toolName: toolCall.name,
                     dynamic: true,
                   });
+                  emittedToolInputs.add(toolCall.id);
                   controller.enqueue({
                     type: 'tool-input-available',
                     toolCallId: toolCall.id,
@@ -1461,6 +1849,11 @@ export function processLangGraphEvent(
                     input: toolCall.args,
                     dynamic: true,
                   });
+                } else if (toolCall.id && emittedToolCalls.has(toolCall.id)) {
+                  // Register key mapping for tool calls already emitted via messages mode
+                  // so that __interrupt__ handling can match them by key
+                  const toolCallKey = `${toolCall.name}:${JSON.stringify(toolCall.args)}`;
+                  emittedToolCallsByKey.set(toolCallKey, toolCall.id);
                 }
               }
             }
@@ -1476,7 +1869,7 @@ export function processLangGraphEvent(
              *    AND tool_calls. We skip those to avoid duplicate reasoning entries.)
              */
             const reasoningId = extractReasoningId(msg);
-            const wasStreamedThisRequest = !!messageSeen[msgId];
+            const wasStreamedThisRequest = messageSeen.has(msgId);
             const hasToolCalls = toolCalls && toolCalls.length > 0;
 
             /**
@@ -1509,6 +1902,16 @@ export function processLangGraphEvent(
                 emittedReasoningIds.add(reasoningId);
               }
             }
+
+            const valuesCitations = extractCitationsFromContentBlocks(msg);
+            if (valuesCitations.length > 0) {
+              emitSourceChunks(
+                valuesCitations,
+                msgId,
+                state.emittedSourceIds,
+                controller,
+              );
+            }
           }
         }
       }
@@ -1572,6 +1975,7 @@ export function processLangGraphEvent(
                   toolName,
                   dynamic: true,
                 });
+                emittedToolInputs.add(toolCallId);
                 controller.enqueue({
                   type: 'tool-input-available',
                   toolCallId,