From 2c113018d25ac3b350d2673923a206c766f99b60 Mon Sep 17 00:00:00 2001
From: end-4 <97237370+end-4@users.noreply.github.com>
Date: Mon, 2 Jun 2025 18:34:44 +0200
Subject: [PATCH] ai: gemini: fix message marked done too early

---
 .../quickshell/modules/sidebarLeft/AiChat.qml |  7 +++--
 .config/quickshell/services/Ai.qml            | 29 ++++++++++++-------
 .config/quickshell/services/AiMessageData.qml |  1 +
 3 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/.config/quickshell/modules/sidebarLeft/AiChat.qml b/.config/quickshell/modules/sidebarLeft/AiChat.qml
index 66d994a1d..e110a596f 100644
--- a/.config/quickshell/modules/sidebarLeft/AiChat.qml
+++ b/.config/quickshell/modules/sidebarLeft/AiChat.qml
@@ -145,7 +145,7 @@ int main(int argc, char* argv[]) {
         Item { // Messages
             Layout.fillWidth: true
             Layout.fillHeight: true
-            ListView { // Message list
+            StyledListView { // Message list
                 id: messageListView
                 anchors.fill: parent
                 spacing: 10
@@ -187,7 +187,10 @@ int main(int argc, char* argv[]) {
                 }
 
                 model: ScriptModel {
-                    values: Ai.messageIDs
+                    values: Ai.messageIDs.filter(id => {
+                        const message = Ai.messageByID[id];
+                        return message?.visibleToUser ?? true;
+                    })
                 }
                 delegate: AiMessage {
                     required property var modelData
diff --git a/.config/quickshell/services/Ai.qml b/.config/quickshell/services/Ai.qml
index cbcb020ae..3bc027f4b 100644
--- a/.config/quickshell/services/Ai.qml
+++ b/.config/quickshell/services/Ai.qml
@@ -126,7 +126,7 @@ Singleton {
             "api_format": "gemini",
             "tools": [
                 {
-                    "google_search": {}
+                    "google_search": ({})
                 },
             ]
         },
@@ -374,9 +374,11 @@ Singleton {
         function buildGeminiRequestData(model, messages) {
             let baseData = {
                 "contents": messages.filter(message => (message.role != Ai.interfaceRole)).map(message => {
-                    if (message.functionCall != undefined && message.functionCall.length > 0) {
+                    const geminiApiRoleName = (message.role === "assistant") ? "model" : message.role;
+                    const usingSearch = model.tools[0].google_search != undefined                
+                    if (!usingSearch && message.functionCall != undefined && message.functionCall.length > 0) {
                         return {
-                            "role": message.role,
+                            "role": geminiApiRoleName,
                             "parts": [{ 
                                 functionCall: {
                                     "name": message.functionName,
@@ -384,9 +386,9 @@ Singleton {
                             }]
                         }
                     }
-                    if (message.functionResponse != undefined && message.functionResponse.length > 0) {
+                    if (!usingSearch && message.functionResponse != undefined && message.functionResponse.length > 0) {
                         return {
-                            "role": message.role,
+                            "role": geminiApiRoleName,
                             "parts": [{ 
                                 functionResponse: {
                                     "name": message.functionName,
@@ -396,7 +398,7 @@ Singleton {
                         }
                     }
                     return {
-                        "role": message.role,
+                        "role": geminiApiRoleName,
                         "parts": [{ 
                             text: message.content,
                         }]
@@ -481,10 +483,13 @@ Singleton {
         }
 
         function parseGeminiBuffer() {
-            console.log("BUFFER DATA: ", requester.geminiBuffer);
+            // console.log("BUFFER DATA: ", requester.geminiBuffer);
             try {
                 if (requester.geminiBuffer.length === 0) return;
                 const dataJson = JSON.parse(requester.geminiBuffer);
+                if (dataJson.candidates[0]?.finishReason) {
+                    requester.markDone();
+                }
                 // Function call handling
                 if (dataJson.candidates[0]?.content?.parts[0]?.functionCall) {
                     const functionCall = dataJson.candidates[0]?.content?.parts[0]?.functionCall;
@@ -531,7 +536,6 @@ Singleton {
             } else if (line == "]") {
                 requester.geminiBuffer += line.slice(0, -1).trim();
                 parseGeminiBuffer();
-                requester.markDone();
             } else if (line.startsWith(",")) { // end of one entry 
                 parseGeminiBuffer();
             } else {
@@ -575,7 +579,9 @@ Singleton {
 
             requester.message.content += newContent;
 
-            if (dataJson.done) requester.markDone();
+            if (dataJson.done) {
+                requester.markDone();
+            }
         }
 
         stdout: SplitParser {
@@ -604,8 +610,8 @@ Singleton {
         }
 
         onExited: (exitCode, exitStatus) => {
-            requester.markDone();
             if (requester.apiFormat == "gemini") requester.parseGeminiBuffer();
+            else requester.markDone();
 
             try { // to parse full response into json for error handling
                 // console.log("Full response: ", requester.message.content + "]"); 
@@ -635,8 +641,9 @@ Singleton {
             "functionResponse": output,
             "thinking": false,
             "done": true,
+            "visibleToUser": false,
         });
-        console.log("Adding function output message: ", JSON.stringify(aiMessage));
+        // console.log("Adding function output message: ", JSON.stringify(aiMessage));
         const id = idForMessage(aiMessage);
         root.messageIDs = [...root.messageIDs, id];
         root.messageByID[id] = aiMessage;
diff --git a/.config/quickshell/services/AiMessageData.qml b/.config/quickshell/services/AiMessageData.qml
index a81566238..b5f208548 100644
--- a/.config/quickshell/services/AiMessageData.qml
+++ b/.config/quickshell/services/AiMessageData.qml
@@ -15,4 +15,5 @@ QtObject {
     property string functionName
     property string functionCall
     property string functionResponse
+    property bool visibleToUser: true
 }