From f5fb3275c0916e8ba492f7bb0a1c6273d6e3e40b Mon Sep 17 00:00:00 2001
From: janole <ole@mobileways.de>
Date: Wed, 24 Sep 2025 12:49:30 +0200
Subject: [PATCH 1/2] refactor: change work function to loop and streamline
 message streaming and handling logic

---
 src/ai/work.ts | 97 +++++++++++++++++++++++++-------------------------
 1 file changed, 49 insertions(+), 48 deletions(-)

diff --git a/src/ai/work.ts b/src/ai/work.ts
index 4664e9e..f13b7e7 100644
--- a/src/ai/work.ts
+++ b/src/ai/work.ts
@@ -17,77 +17,78 @@ async function work(props: WorkProps)
 {
     const { session, chatService, signal } = props;
 
-    const messages = await workTools({ ...props, session });
-
-    if (needsToolConfirmation(messages))
-    {
-        return messages;
-    }
-
-    let aiMessage: AIMessageChunk | undefined = undefined;
-    let toolProgressMessages: ToolProgressMessage[] = [];
-
-    if (session.streaming)
+    while (true)
     {
-        const { result: stream, error } = await tryCatch(chatService.stream(session, signal));
+        const messages = await workTools({ ...props, session });
 
-        if (!stream)
+        if (needsToolConfirmation(messages))
         {
-            messages.push(new ErrorMessage(`ERROR: ${error?.message || error?.toString() || "llm.stream(...) failed."}`, "error", error));
             return messages;
         }
 
-        for await (const chunk of stream)
+        let aiMessage: AIMessageChunk | undefined = undefined;
+        let toolProgressMessages: ToolProgressMessage[] = [];
+
+        if (session.streaming)
         {
-            aiMessage = aiMessage !== undefined ? concat(aiMessage, chunk) : chunk;
-            setMessageIsStreaming(aiMessage, true);
+            const { result: stream, error } = await tryCatch(chatService.stream(session, signal));
 
-            if (!aiMessage?.tool_calls?.length && !aiMessage?.tool_call_chunks?.length)
+            if (!stream)
             {
-                session.setMessages([...messages, aiMessage]);
+                messages.push(new ErrorMessage(`ERROR: ${error?.message || error?.toString() || "llm.stream(...) failed."}`, "error", error));
+                return messages;
             }
-            else
+
+            for await (const chunk of stream)
             {
-                toolProgressMessages = ToolProgressMessage.createFromChunks(aiMessage.tool_call_chunks);
+                aiMessage = aiMessage !== undefined ? concat(aiMessage, chunk) : chunk;
+                setMessageIsStreaming(aiMessage, true);
 
-                if (toolProgressMessages.length)
+                if (!aiMessage?.tool_calls?.length && !aiMessage?.tool_call_chunks?.length)
                 {
-                    session.setMessages([...messages, aiMessage, ...toolProgressMessages]);
+                    session.setMessages([...messages, aiMessage]);
+                }
+                else
+                {
+                    toolProgressMessages = ToolProgressMessage.createFromChunks(aiMessage.tool_call_chunks);
+
+                    if (toolProgressMessages.length)
+                    {
+                        session.setMessages([...messages, aiMessage, ...toolProgressMessages]);
+                    }
                 }
             }
         }
-    }
-    else
-    {
-        const { result, error } = await tryCatch(chatService.generate(session, signal));
-
-        if (!result || error)
+        else
         {
-            messages.push(new ErrorMessage(`ERROR: ${error?.message || error?.toString() || "llm.generate(...) failed."}`, "error", error));
-            return messages;
-        }
+            const { result, error } = await tryCatch(chatService.generate(session, signal));
 
-        aiMessage = result;
-    }
+            if (!result || error)
+            {
+                messages.push(new ErrorMessage(`ERROR: ${error?.message || error?.toString() || "llm.generate(...) failed."}`, "error", error));
+                return messages;
+            }
 
-    if (aiMessage)
-    {
-        setMessageIsStreaming(aiMessage, false);
-        messages.push(aiMessage);
-    }
+            aiMessage = result;
+        }
 
-    if (!aiMessage?.tool_calls?.length)
-    {
-        return messages;
-    }
+        if (aiMessage)
+        {
+            setMessageIsStreaming(aiMessage, false);
+            messages.push(aiMessage);
+        }
 
-    // add pending toolCall(s)
-    toolProgressMessages = aiMessage.tool_calls.map(toolCall => new ToolProgressMessage(toolCall)) || [];
-    session.setMessages([...messages, ...toolProgressMessages]);
+        if (!aiMessage?.tool_calls?.length)
+        {
+            return messages;
+        }
 
-    await session.flush();
+        // add pending toolCall(s)
+        toolProgressMessages = aiMessage.tool_calls.map(toolCall => new ToolProgressMessage(toolCall)) || [];
+        session.setMessages([...messages, ...toolProgressMessages]);
 
-    return work(props);
+        // await session.flush();
+    }
 }
 
 async function workTools(props: Pick<WorkProps, "session" | "chatService" | "signal">)

From a5969d0b493aab32c0f47cae59837bd24b5ab162 Mon Sep 17 00:00:00 2001
From: janole <ole@mobileways.de>
Date: Wed, 24 Sep 2025 13:47:46 +0200
Subject: [PATCH 2/2] refactor: remove unnecessary session.flush call in work
 function

---
 src/ai/work.ts | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/ai/work.ts b/src/ai/work.ts
index f13b7e7..0f034fb 100644
--- a/src/ai/work.ts
+++ b/src/ai/work.ts
@@ -86,8 +86,6 @@ async function work(props: WorkProps)
         // add pending toolCall(s)
         toolProgressMessages = aiMessage.tool_calls.map(toolCall => new ToolProgressMessage(toolCall)) || [];
         session.setMessages([...messages, ...toolProgressMessages]);
-
-        // await session.flush();
     }
 }