@lleverage-ai/agent-sdk 0.0.12 → 0.0.14

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1 +1 @@
1
- {"version":3,"file":"agent.d.ts","sourceRoot":"","sources":["../src/agent.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AA4DH,OAAO,KAAK,EACV,KAAK,EACL,YAAY,EAwBb,MAAM,YAAY,CAAC;AA6pBpB;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA2CG;AACH,wBAAgB,WAAW,CAAC,OAAO,EAAE,YAAY,GAAG,KAAK,CA2zFxD"}
1
+ {"version":3,"file":"agent.d.ts","sourceRoot":"","sources":["../src/agent.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AA4DH,OAAO,KAAK,EACV,KAAK,EACL,YAAY,EAwBb,MAAM,YAAY,CAAC;AA6pBpB;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA2CG;AACH,wBAAgB,WAAW,CAAC,OAAO,EAAE,YAAY,GAAG,KAAK,CAq7FxD"}
package/dist/agent.js CHANGED
@@ -697,17 +697,10 @@ export function createAgent(options) {
697
697
  let hasProxiedTools = false;
698
698
  // Auto-created subagent definitions from delegated plugins
699
699
  const autoSubagents = [];
700
- // Count total plugin tools for threshold calculation and collect plugin skills.
701
- // Note: Function-based (streaming) tools are not counted since we don't know
702
- // their count until they're invoked with a streaming context.
700
+ // Collect plugin skills early so they're available for skill tool creation.
703
701
  // IMPORTANT: Plugin skills must be collected BEFORE createCoreTools() is called
704
702
  // so the skill tool includes them in progressive disclosure.
705
- let totalPluginToolCount = 0;
706
703
  for (const plugin of options.plugins ?? []) {
707
- if (plugin.tools && typeof plugin.tools !== "function") {
708
- totalPluginToolCount += Object.keys(plugin.tools).length;
709
- }
710
- // Collect plugin skills early so they're available for skill tool creation
711
704
  if (plugin.skills) {
712
705
  skills.push(...plugin.skills);
713
706
  }
@@ -740,10 +733,28 @@ export function createAgent(options) {
740
733
  // Note: Plugin skills are collected earlier (before createCoreTools) so
741
734
  // the skill tool can include them in progressive disclosure.
742
735
  for (const plugin of options.plugins ?? []) {
743
- // Handle tools via MCP manager for unified interface
744
- // Note: Function-based (streaming) tools are handled separately in
745
- // getActiveToolSetWithStreaming() and are not registered here
746
- if (plugin.tools && typeof plugin.tools !== "function") {
736
+ if (plugin.tools && typeof plugin.tools === "function") {
737
+ // Function-based (streaming) tools
738
+ const factory = plugin.tools;
739
+ const isDeferred = plugin.deferred === true || (pluginLoadingMode === "proxy" && plugin.deferred !== false);
740
+ if (isDeferred) {
741
+ // Deferred streaming plugin: register factory in MCP for discovery via search_tools + call_tool
742
+ mcpManager.registerStreamingPluginTools(plugin.name, factory, {
743
+ autoLoad: false,
744
+ });
745
+ hasProxiedTools = true;
746
+ }
747
+ else if (deferredLoadingActive && plugin.deferred !== false) {
748
+ // Deferred loading active: register but don't auto-load
749
+ mcpManager.registerStreamingPluginTools(plugin.name, factory, {
750
+ autoLoad: false,
751
+ });
752
+ }
753
+ // else: Eager non-deferred streaming plugins are handled directly in
754
+ // getActiveToolSetWithStreaming() — they're NOT registered in MCP.
755
+ }
756
+ else if (plugin.tools) {
757
+ // Static tools
747
758
  // Check if this plugin is deferred (proxy mode or per-plugin opt-in)
748
759
  const isDeferred = plugin.deferred === true || (pluginLoadingMode === "proxy" && plugin.deferred !== false);
749
760
  if (isDeferred) {
@@ -768,6 +779,8 @@ export function createAgent(options) {
768
779
  }
769
780
  }
770
781
  }
782
+ // Only tools registered in MCP are discoverable via search_tools.
783
+ const discoverablePluginToolCount = mcpManager.listTools().length;
771
784
  // Create subagent definitions from plugins with subagent config
772
785
  for (const plugin of options.plugins ?? []) {
773
786
  if (plugin.subagent) {
@@ -791,7 +804,8 @@ export function createAgent(options) {
791
804
  const hasSubagents = allSubagents.length > 0;
792
805
  // In proxy mode, create call_tool and configure search_tools with schema disclosure
793
806
  const isProxyMode = pluginLoadingMode === "proxy" || hasProxiedTools;
794
- if (isProxyMode && !options.disabledCoreTools?.includes("call_tool")) {
807
+ const hasProxyTargets = hasProxiedTools || mcpManager.hasExternalServers();
808
+ if (isProxyMode && hasProxyTargets && !options.disabledCoreTools?.includes("call_tool")) {
795
809
  coreTools.call_tool = createCallToolTool({
796
810
  mcpManager,
797
811
  });
@@ -803,10 +817,11 @@ export function createAgent(options) {
803
817
  // - Create when proxy mode is active (for call_tool discovery)
804
818
  // - Create when external MCP servers exist (for MCP tool search)
805
819
  // - Always auto-load tools when found (no manual load step) — unless proxy mode
806
- const shouldCreateSearchToolsForAutoThreshold = toolSearchEnabled === "auto" && totalPluginToolCount > toolSearchThreshold;
820
+ const shouldCreateSearchToolsForAutoThreshold = toolSearchEnabled === "auto" && discoverablePluginToolCount > toolSearchThreshold;
807
821
  const shouldCreateSearchTools = !options.disabledCoreTools?.includes("search_tools") &&
808
- (deferredLoadingActive ||
809
- isProxyMode ||
822
+ ((deferredLoadingActive &&
823
+ (discoverablePluginToolCount > 0 || mcpManager.hasExternalServers())) ||
824
+ (isProxyMode && hasProxyTargets) ||
810
825
  shouldCreateSearchToolsForAutoThreshold ||
811
826
  (mcpManager.hasExternalServers() && toolSearchEnabled !== "never"));
812
827
  if (shouldCreateSearchTools) {
@@ -919,6 +934,28 @@ export function createAgent(options) {
919
934
  // Build prompt using prompt builder
920
935
  return promptBuilder.build(context);
921
936
  };
937
+ const responseMessagesToModelMessages = (messages) => Array.isArray(messages) ? messages : [];
938
+ const appendResponseMessages = (baseMessages, responseMessages, fallbackAssistantText) => {
939
+ const normalized = responseMessagesToModelMessages(responseMessages);
940
+ if (normalized.length > 0) {
941
+ return [...baseMessages, ...normalized];
942
+ }
943
+ return fallbackAssistantText
944
+ ? [...baseMessages, { role: "assistant", content: fallbackAssistantText }]
945
+ : baseMessages;
946
+ };
947
+ const buildMessagesFromStepResponses = (baseMessages, steps, fallbackAssistantText) => {
948
+ const stepResponseMessages = steps.flatMap((step) => responseMessagesToModelMessages(step.response?.messages));
949
+ if (stepResponseMessages.length > 0) {
950
+ return [...baseMessages, ...stepResponseMessages];
951
+ }
952
+ // Fallback for providers that do not populate response.messages.
953
+ const lastText = steps.at(-1)?.text ?? fallbackAssistantText;
954
+ return [
955
+ ...baseMessages,
956
+ ...(lastText ? [{ role: "assistant", content: lastText }] : []),
957
+ ];
958
+ };
922
959
  // Runtime tools added/removed dynamically by plugins at runtime
923
960
  const runtimeTools = {};
924
961
  // Helper to get current active tools (core + runtime + MCP + dynamically loaded from registry)
@@ -953,11 +990,16 @@ export function createAgent(options) {
953
990
  const allTools = { ...coreTools };
954
991
  // Add runtime tools (added by plugins at runtime)
955
992
  Object.assign(allTools, runtimeTools);
956
- // Process plugins - invoke function-based tools with streaming context
993
+ // Process plugins - invoke non-deferred function-based tools with streaming context.
994
+ // Deferred function-based plugins are registered in MCP and come via mcpManager.getToolSet().
957
995
  for (const plugin of options.plugins ?? []) {
958
996
  if (plugin.tools) {
959
997
  if (typeof plugin.tools === "function") {
960
- // Streaming-aware tools: invoke factory with context
998
+ // Skip if registered in MCP (deferred or deferredLoadingActive plugins)
999
+ if (mcpManager.hasStreamingFactory(plugin.name)) {
1000
+ continue;
1001
+ }
1002
+ // Non-deferred: eagerly invoke factory with streaming context (original names)
961
1003
  const streamingTools = plugin.tools(streamingContext);
962
1004
  Object.assign(allTools, streamingTools);
963
1005
  }
@@ -967,8 +1009,8 @@ export function createAgent(options) {
967
1009
  }
968
1010
  }
969
1011
  }
970
- // Add MCP tools from plugin registrations
971
- const mcpTools = mcpManager.getToolSet();
1012
+ // Add MCP tools from plugin registrations (includes streaming factory tools with live context)
1013
+ const mcpTools = mcpManager.getToolSet(undefined, streamingContext);
972
1014
  Object.assign(allTools, mcpTools);
973
1015
  // Apply allowedTools filtering
974
1016
  const filtered = filterToolsByAllowed(allTools);
@@ -999,6 +1041,33 @@ export function createAgent(options) {
999
1041
  // Then apply hooks for observability
1000
1042
  return wrapToolsWithHooks(withTaskManager, effectiveHooks, agent, threadId ?? "default");
1001
1043
  };
1044
+ /**
1045
+ * Injects request-local streaming context into tool execution options.
1046
+ *
1047
+ * This keeps deferred streaming tools isolated per request instead of
1048
+ * relying on mutable MCPManager state shared across concurrent generations.
1049
+ */
1050
+ const wrapToolsWithStreamingContext = (tools, streamingContext) => {
1051
+ const wrapped = {};
1052
+ for (const [name, tool] of Object.entries(tools)) {
1053
+ if (!tool.execute) {
1054
+ wrapped[name] = tool;
1055
+ continue;
1056
+ }
1057
+ const originalExecute = tool.execute;
1058
+ wrapped[name] = {
1059
+ ...tool,
1060
+ execute: async (input, toolOptions) => {
1061
+ const extendedOptions = {
1062
+ ...toolOptions,
1063
+ streamingContext,
1064
+ };
1065
+ return originalExecute(input, extendedOptions);
1066
+ },
1067
+ };
1068
+ }
1069
+ return wrapped;
1070
+ };
1002
1071
  /**
1003
1072
  * Adds the task and task_output tools to a toolset.
1004
1073
  *
@@ -1616,10 +1685,7 @@ export function createAgent(options) {
1616
1685
  // cannot find the checkpoint (or finds one without messages/interrupt).
1617
1686
  if (effectiveGenOptions.threadId && options.checkpointer) {
1618
1687
  const checkpointThreadId = forkedSessionId ?? effectiveGenOptions.threadId;
1619
- const finalMessages = [
1620
- ...messages,
1621
- ...(response.text ? [{ role: "assistant", content: response.text }] : []),
1622
- ];
1688
+ const finalMessages = appendResponseMessages(messages, response.response?.messages, response.text);
1623
1689
  const savedCheckpoint = await saveCheckpoint(checkpointThreadId, finalMessages, startStep + response.steps.length);
1624
1690
  if (savedCheckpoint) {
1625
1691
  const withInterrupt = updateCheckpoint(savedCheckpoint, {
@@ -1687,13 +1753,7 @@ export function createAgent(options) {
1687
1753
  // Save checkpoint - use forked session ID if forking, otherwise use original threadId
1688
1754
  const checkpointThreadId = forkedSessionId ?? effectiveGenOptions.threadId;
1689
1755
  if (checkpointThreadId && options.checkpointer) {
1690
- // Build final messages including the assistant response.
1691
- // Skip empty text to avoid invalid content blocks (e.g. when
1692
- // the model finishes with only tool calls).
1693
- const finalMessages = [
1694
- ...messages,
1695
- ...(response.text ? [{ role: "assistant", content: response.text }] : []),
1696
- ];
1756
+ const finalMessages = appendResponseMessages(messages, response.response?.messages, response.text);
1697
1757
  await saveCheckpoint(checkpointThreadId, finalMessages, startStep + response.steps.length);
1698
1758
  }
1699
1759
  // Invoke unified PostGenerate hooks
@@ -1727,12 +1787,7 @@ export function createAgent(options) {
1727
1787
  let lastResult = finalResult;
1728
1788
  let runningMessages = hasCheckpointing
1729
1789
  ? []
1730
- : [
1731
- ...messages,
1732
- ...(finalResult.text
1733
- ? [{ role: "assistant", content: finalResult.text }]
1734
- : []),
1735
- ];
1790
+ : appendResponseMessages(messages, response.response?.messages, finalResult.text);
1736
1791
  let followUpPrompt = await getNextTaskPrompt();
1737
1792
  while (followUpPrompt !== null) {
1738
1793
  lastResult = await agent.generate({
@@ -2034,11 +2089,12 @@ export function createAgent(options) {
2034
2089
  }
2035
2090
  }
2036
2091
  // Get final result for hooks - need to await all properties
2037
- const [text, usage, finishReason, steps] = await Promise.all([
2092
+ const [text, usage, finishReason, steps, responseMeta] = await Promise.all([
2038
2093
  response.text,
2039
2094
  response.usage,
2040
2095
  response.finishReason,
2041
2096
  response.steps,
2097
+ response.response ?? Promise.resolve(undefined),
2042
2098
  ]);
2043
2099
  // Only access output if an output schema was provided
2044
2100
  let output;
@@ -2060,10 +2116,10 @@ export function createAgent(options) {
2060
2116
  };
2061
2117
  // Save checkpoint if threadId is provided
2062
2118
  if (checkpointThreadId && options.checkpointer) {
2063
- const finalMessages = [
2064
- ...messages,
2065
- ...(text ? [{ role: "assistant", content: text }] : []),
2066
- ];
2119
+ const responseMessages = responseMessagesToModelMessages(responseMeta?.messages);
2120
+ const finalMessages = responseMessages.length > 0
2121
+ ? [...messages, ...responseMessages]
2122
+ : buildMessagesFromStepResponses(messages, steps, text);
2067
2123
  await saveCheckpoint(checkpointThreadId, finalMessages, startStep + steps.length);
2068
2124
  }
2069
2125
  // Save pending interrupt to checkpoint (mirrors generate() pattern)
@@ -2111,9 +2167,12 @@ export function createAgent(options) {
2111
2167
  return;
2112
2168
  }
2113
2169
  const hasCheckpointing = !!(effectiveGenOptions.threadId && options.checkpointer);
2170
+ const initialResponseMessages = responseMessagesToModelMessages(responseMeta?.messages);
2114
2171
  let currentMessages = hasCheckpointing
2115
2172
  ? []
2116
- : [...messages, ...(text ? [{ role: "assistant", content: text }] : [])];
2173
+ : initialResponseMessages.length > 0
2174
+ ? [...messages, ...initialResponseMessages]
2175
+ : buildMessagesFromStepResponses(messages, steps, text);
2117
2176
  let followUpPrompt = await getNextTaskPrompt();
2118
2177
  while (followUpPrompt !== null) {
2119
2178
  const followUpGen = agent.stream({
@@ -2213,6 +2272,7 @@ export function createAgent(options) {
2213
2272
  const signalStopCondition = () => signalState.interrupt != null;
2214
2273
  // Track step count for incremental checkpointing
2215
2274
  let currentStepCount = 0;
2275
+ let streamedMessages = [...initialParams.messages];
2216
2276
  // Execute streamText OUTSIDE createUIMessageStream so errors propagate
2217
2277
  // to the retry loop (if streamText throws synchronously on creation,
2218
2278
  // e.g. rate limit, the catch block handles retry/fallback).
@@ -2236,14 +2296,17 @@ export function createAgent(options) {
2236
2296
  ? async (stepResult) => {
2237
2297
  if (effectiveGenOptions.threadId && options.checkpointer) {
2238
2298
  currentStepCount++;
2239
- // Build messages including this step's results
2240
- const stepMessages = [
2241
- ...initialParams.messages,
2242
- ...(stepResult.text
2243
- ? [{ role: "assistant", content: stepResult.text }]
2244
- : []),
2245
- ];
2246
- await saveCheckpoint(effectiveGenOptions.threadId, stepMessages, startStep + currentStepCount);
2299
+ const stepResponseMessages = responseMessagesToModelMessages(stepResult.response?.messages);
2300
+ streamedMessages =
2301
+ stepResponseMessages.length > 0
2302
+ ? [...streamedMessages, ...stepResponseMessages]
2303
+ : stepResult.text
2304
+ ? [
2305
+ ...streamedMessages,
2306
+ { role: "assistant", content: stepResult.text },
2307
+ ]
2308
+ : streamedMessages;
2309
+ await saveCheckpoint(effectiveGenOptions.threadId, streamedMessages, startStep + currentStepCount);
2247
2310
  }
2248
2311
  }
2249
2312
  : undefined,
@@ -2258,12 +2321,7 @@ export function createAgent(options) {
2258
2321
  });
2259
2322
  }
2260
2323
  if (effectiveGenOptions.threadId && options.checkpointer) {
2261
- const finalMessages = [
2262
- ...initialParams.messages,
2263
- ...(finishResult.text
2264
- ? [{ role: "assistant", content: finishResult.text }]
2265
- : []),
2266
- ];
2324
+ const finalMessages = buildMessagesFromStepResponses(initialParams.messages, finishResult.steps, finishResult.text);
2267
2325
  await saveCheckpoint(effectiveGenOptions.threadId, finalMessages, startStep + finishResult.steps.length);
2268
2326
  }
2269
2327
  // Invoke unified PostGenerate hooks
@@ -2294,8 +2352,9 @@ export function createAgent(options) {
2294
2352
  execute: async ({ writer }) => {
2295
2353
  // Merge initial generation into the stream
2296
2354
  writer.merge(result.toUIMessageStream());
2297
- // Wait for initial generation to complete to get final text
2298
- const text = await result.text;
2355
+ // Wait for initial generation to complete
2356
+ await result.text;
2357
+ const resultResponse = await (result.response ?? Promise.resolve(undefined));
2299
2358
  // --- Background task completion loop ---
2300
2359
  if (waitForBackgroundTasks) {
2301
2360
  // Track accumulated steps for checkpoint saves
@@ -2303,7 +2362,7 @@ export function createAgent(options) {
2303
2362
  let accumulatedStepCount = initialSteps.length;
2304
2363
  let currentMessages = [
2305
2364
  ...messages,
2306
- ...(text ? [{ role: "assistant", content: text }] : []),
2365
+ ...responseMessagesToModelMessages(resultResponse?.messages),
2307
2366
  ];
2308
2367
  let followUpPrompt = await getNextTaskPrompt();
2309
2368
  while (followUpPrompt !== null) {
@@ -2327,12 +2386,16 @@ export function createAgent(options) {
2327
2386
  });
2328
2387
  writer.merge(followUpResult.toUIMessageStream());
2329
2388
  const followUpText = await followUpResult.text;
2389
+ const followUpResponse = await (followUpResult.response ??
2390
+ Promise.resolve(undefined));
2330
2391
  currentMessages = [
2331
2392
  ...currentMessages,
2332
2393
  { role: "user", content: followUpPrompt },
2333
- ...(followUpText
2334
- ? [{ role: "assistant", content: followUpText }]
2335
- : []),
2394
+ ...(responseMessagesToModelMessages(followUpResponse?.messages).length > 0
2395
+ ? responseMessagesToModelMessages(followUpResponse?.messages)
2396
+ : followUpText
2397
+ ? [{ role: "assistant", content: followUpText }]
2398
+ : []),
2336
2399
  ];
2337
2400
  // --- Post-completion bookkeeping for follow-ups ---
2338
2401
  const followUpSteps = await followUpResult.steps;
@@ -2439,6 +2502,7 @@ export function createAgent(options) {
2439
2502
  };
2440
2503
  // Track step count for incremental checkpointing
2441
2504
  let currentStepCount = 0;
2505
+ let streamedMessages = [...initialParams.messages];
2442
2506
  // Stop condition: stop when an interrupt signal was caught, OR when
2443
2507
  // the step count reaches maxSteps.
2444
2508
  const signalStopCondition = () => signalState.interrupt != null;
@@ -2463,14 +2527,17 @@ export function createAgent(options) {
2463
2527
  ? async (stepResult) => {
2464
2528
  if (effectiveGenOptions.threadId && options.checkpointer) {
2465
2529
  currentStepCount++;
2466
- // Build messages including this step's results
2467
- const stepMessages = [
2468
- ...initialParams.messages,
2469
- ...(stepResult.text
2470
- ? [{ role: "assistant", content: stepResult.text }]
2471
- : []),
2472
- ];
2473
- await saveCheckpoint(effectiveGenOptions.threadId, stepMessages, startStep + currentStepCount);
2530
+ const stepResponseMessages = responseMessagesToModelMessages(stepResult.response?.messages);
2531
+ streamedMessages =
2532
+ stepResponseMessages.length > 0
2533
+ ? [...streamedMessages, ...stepResponseMessages]
2534
+ : stepResult.text
2535
+ ? [
2536
+ ...streamedMessages,
2537
+ { role: "assistant", content: stepResult.text },
2538
+ ]
2539
+ : streamedMessages;
2540
+ await saveCheckpoint(effectiveGenOptions.threadId, streamedMessages, startStep + currentStepCount);
2474
2541
  }
2475
2542
  }
2476
2543
  : undefined,
@@ -2485,12 +2552,7 @@ export function createAgent(options) {
2485
2552
  });
2486
2553
  }
2487
2554
  if (effectiveGenOptions.threadId && options.checkpointer) {
2488
- const finalMessages = [
2489
- ...initialParams.messages,
2490
- ...(finishResult.text
2491
- ? [{ role: "assistant", content: finishResult.text }]
2492
- : []),
2493
- ];
2555
+ const finalMessages = buildMessagesFromStepResponses(initialParams.messages, finishResult.steps, finishResult.text);
2494
2556
  await saveCheckpoint(effectiveGenOptions.threadId, finalMessages, startStep + finishResult.steps.length);
2495
2557
  }
2496
2558
  // Invoke unified PostGenerate hooks
@@ -2582,7 +2644,8 @@ export function createAgent(options) {
2582
2644
  // Apply hooks AFTER adding task tool so task tool is also wrapped.
2583
2645
  // Then wrap with signal catching as the outermost layer.
2584
2646
  const hookedStreamingTools = applyToolHooks(addTaskToolIfConfigured(getActiveToolSetWithStreaming(streamingContext, effectiveGenOptions.threadId), streamingContext), effectiveGenOptions.threadId);
2585
- const streamingTools = wrapToolsWithSignalCatching(hookedStreamingTools, signalState);
2647
+ const requestScopedStreamingTools = wrapToolsWithStreamingContext(hookedStreamingTools, streamingContext);
2648
+ const streamingTools = wrapToolsWithSignalCatching(requestScopedStreamingTools, signalState);
2586
2649
  // Build prompt context and generate system prompt
2587
2650
  const promptContext = buildPromptContext(messages, effectiveGenOptions.threadId);
2588
2651
  const systemPrompt = getSystemPrompt(promptContext);
@@ -2600,6 +2663,7 @@ export function createAgent(options) {
2600
2663
  };
2601
2664
  // Track step count for incremental checkpointing
2602
2665
  let currentStepCount = 0;
2666
+ let streamedMessages = [...initialParams.messages];
2603
2667
  // Stop condition: stop when a flow-control signal was caught, OR when
2604
2668
  // the step count reaches maxSteps.
2605
2669
  const signalStopCondition = () => signalState.interrupt != null;
@@ -2624,15 +2688,17 @@ export function createAgent(options) {
2624
2688
  ? async (stepResult) => {
2625
2689
  if (effectiveGenOptions.threadId && options.checkpointer) {
2626
2690
  currentStepCount++;
2627
- // Build messages including this step's results
2628
- const stepMessages = [
2629
- ...initialParams.messages,
2630
- {
2631
- role: "assistant",
2632
- content: stepResult.text,
2633
- },
2634
- ];
2635
- await saveCheckpoint(effectiveGenOptions.threadId, stepMessages, startStep + currentStepCount);
2691
+ const stepResponseMessages = responseMessagesToModelMessages(stepResult.response?.messages);
2692
+ streamedMessages =
2693
+ stepResponseMessages.length > 0
2694
+ ? [...streamedMessages, ...stepResponseMessages]
2695
+ : stepResult.text
2696
+ ? [
2697
+ ...streamedMessages,
2698
+ { role: "assistant", content: stepResult.text },
2699
+ ]
2700
+ : streamedMessages;
2701
+ await saveCheckpoint(effectiveGenOptions.threadId, streamedMessages, startStep + currentStepCount);
2636
2702
  }
2637
2703
  }
2638
2704
  : undefined,
@@ -2647,13 +2713,7 @@ export function createAgent(options) {
2647
2713
  });
2648
2714
  }
2649
2715
  if (effectiveGenOptions.threadId && options.checkpointer) {
2650
- const finalMessages = [
2651
- ...initialParams.messages,
2652
- {
2653
- role: "assistant",
2654
- content: finishResult.text,
2655
- },
2656
- ];
2716
+ const finalMessages = buildMessagesFromStepResponses(initialParams.messages, finishResult.steps, finishResult.text);
2657
2717
  await saveCheckpoint(effectiveGenOptions.threadId, finalMessages, startStep + finishResult.steps.length);
2658
2718
  }
2659
2719
  // Invoke unified PostGenerate hooks
@@ -2681,8 +2741,9 @@ export function createAgent(options) {
2681
2741
  });
2682
2742
  // Merge the streamText output into the UI message stream
2683
2743
  writer.merge(result.toUIMessageStream());
2684
- // Wait for initial generation to complete to get final text
2685
- const text = await result.text;
2744
+ // Wait for initial generation to complete
2745
+ await result.text;
2746
+ const resultResponse = await (result.response ?? Promise.resolve(undefined));
2686
2747
  // Save pending interrupt to checkpoint (mirrors stream() pattern)
2687
2748
  if (signalState.interrupt && effectiveGenOptions.threadId && options.checkpointer) {
2688
2749
  const interrupt = signalState.interrupt.interrupt;
@@ -2710,14 +2771,14 @@ export function createAgent(options) {
2710
2771
  await invokeHooksWithTimeout(interruptRequestedHooks, hookInput, null, agent);
2711
2772
  }
2712
2773
  }
2713
- // --- Background task completion loop ---
2774
+ // --- Background task completion loop (streamDataResponse) ---
2714
2775
  if (waitForBackgroundTasks && !signalState.interrupt) {
2715
2776
  // Track accumulated steps for checkpoint saves
2716
2777
  const initialSteps = await result.steps;
2717
2778
  let accumulatedStepCount = initialSteps.length;
2718
2779
  let currentMessages = [
2719
2780
  ...messages,
2720
- ...(text ? [{ role: "assistant", content: text }] : []),
2781
+ ...responseMessagesToModelMessages(resultResponse?.messages),
2721
2782
  ];
2722
2783
  let followUpPrompt = await getNextTaskPrompt();
2723
2784
  while (followUpPrompt !== null) {
@@ -2741,12 +2802,16 @@ export function createAgent(options) {
2741
2802
  });
2742
2803
  writer.merge(followUpResult.toUIMessageStream());
2743
2804
  const followUpText = await followUpResult.text;
2805
+ const followUpResponse = await (followUpResult.response ??
2806
+ Promise.resolve(undefined));
2744
2807
  currentMessages = [
2745
2808
  ...currentMessages,
2746
2809
  { role: "user", content: followUpPrompt },
2747
- ...(followUpText
2748
- ? [{ role: "assistant", content: followUpText }]
2749
- : []),
2810
+ ...(responseMessagesToModelMessages(followUpResponse?.messages).length > 0
2811
+ ? responseMessagesToModelMessages(followUpResponse?.messages)
2812
+ : followUpText
2813
+ ? [{ role: "assistant", content: followUpText }]
2814
+ : []),
2750
2815
  ];
2751
2816
  // --- Post-completion bookkeeping for follow-ups ---
2752
2817
  const followUpSteps = await followUpResult.steps;