From 211c01503b9ca0addca34b8901f7ae4f803ececb Mon Sep 17 00:00:00 2001
From: standujar <s.andujar@proton.me>
Date: Mon, 10 Nov 2025 23:39:37 +0100
Subject: [PATCH 1/2] feat: add onStepFinish callback to text generation
 functions

- Introduced an optional onStepFinish parameter to generateTextWithModel, handleTextSmall, and handleTextLarge functions.
- Enhanced the text generation process by allowing users to provide a callback that executes after each step, enabling custom handling of step results.
---
 src/models/text.ts | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/models/text.ts b/src/models/text.ts
index 921b5d9..b240bf0 100644
--- a/src/models/text.ts
+++ b/src/models/text.ts
@@ -18,9 +18,10 @@ async function generateTextWithModel(
   params: GenerateTextParams & {
     tools?: Record<string, Tool>;
     toolChoice?: ToolChoice<Record<string, Tool>>;
+    onStepFinish?: (stepResult: any) => void | Promise<void>;
   },
 ): Promise<string | ToolResponse> {
-  const { prompt, stopSequences = [], tools, toolChoice } = params;
+  const { prompt, stopSequences = [], tools, toolChoice, onStepFinish } = params;
   const temperature = params.temperature ?? 0.7;
   const frequencyPenalty = params.frequencyPenalty ?? 0.7;
   const presencePenalty = params.presencePenalty ?? 0.7;
@@ -72,6 +73,7 @@ async function generateTextWithModel(
 
   if (tools) {
     (generateParams as any).onStepFinish = async (stepResult: any) => {
+      // Internal capture logic
       if (stepResult.toolCalls && stepResult.toolCalls.length > 0) {
         capturedToolCalls = [
           ...capturedToolCalls,
@@ -92,6 +94,11 @@ async function generateTextWithModel(
           capturedToolResults = [...capturedToolResults, ...toolResultsFromContent];
         }
       }
+
+      // Call user-provided callback if provided
+      if (onStepFinish) {
+        await onStepFinish(stepResult);
+      }
     };
   }
 
@@ -137,6 +144,7 @@ export async function handleTextSmall(
   params: GenerateTextParams & {
     tools?: Record<string, Tool>;
     toolChoice?: ToolChoice<Record<string, Tool>>;
+    onStepFinish?: (stepResult: any) => void | Promise<void>;
   },
 ): Promise<string | ToolResponse> {
   return generateTextWithModel(runtime, ModelType.TEXT_SMALL, params);
@@ -150,6 +158,7 @@ export async function handleTextLarge(
   params: GenerateTextParams & {
     tools?: Record<string, Tool>;
     toolChoice?: ToolChoice<Record<string, Tool>>;
+    onStepFinish?: (stepResult: any) => void | Promise<void>;
   },
 ): Promise<string | ToolResponse> {
   return generateTextWithModel(runtime, ModelType.TEXT_LARGE, params);

From 12b084d47c57d74f1c421ae761e68a3f16aefdd8 Mon Sep 17 00:00:00 2001
From: standujar <s.andujar@proton.me>
Date: Tue, 11 Nov 2025 02:02:46 +0100
Subject: [PATCH 2/2] fix: handle errors in onStepFinish callback during text
 generation

---
 src/models/text.ts | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/models/text.ts b/src/models/text.ts
index b240bf0..bda1688 100644
--- a/src/models/text.ts
+++ b/src/models/text.ts
@@ -97,7 +97,11 @@ async function generateTextWithModel(
 
       // Call user-provided callback if provided
       if (onStepFinish) {
-        await onStepFinish(stepResult);
+        try {
+          await onStepFinish(stepResult);
+        } catch (error) {
+          logger.error('[OpenRouter] Error in onStepFinish callback:', error instanceof Error ? error.message : String(error));
+        }
       }
     };
   }