From e9c6de110553795eb11e09e90607ffd2c1d1fa84 Mon Sep 17 00:00:00 2001
From: Hongwei <hongwei@tesobe.com>
Date: Mon, 11 Aug 2025 17:23:50 +0200
Subject: [PATCH] Enhance chat store to support new API message types and
 improve legacy handling. Added support for assistant token streaming, tool
 lifecycle events, and approval requests.

---
 src/stores/chat.ts | 45 ++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/stores/chat.ts b/src/stores/chat.ts
index 44a62af..5f03ebe 100644
--- a/src/stores/chat.ts
+++ b/src/stores/chat.ts
@@ -284,7 +284,8 @@ export const useChat = defineStore('chat', {
                                 
                                 
                                 // This is where we process different types of messages from Opey by their 'type' field
-                                // Process pending tool calls
+                                // Support both legacy ('message','tool','token') and new streaming types
+                                // Pending tool calls (legacy)
                                 if (data.type === 'message') {
 
                                     if (content.tool_approval_request) {
@@ -310,7 +311,7 @@ export const useChat = defineStore('chat', {
                                     }
                                 }
 
-                                // Now we handle the actual messages from the completed/ failed tool calls
+                                // Now we handle the actual messages from the completed/ failed tool calls (legacy)
                                 if (data.type === 'tool') {
                                     const toolCallId = content.tool_call_id;
                                     if (!toolCallId) {
@@ -336,13 +337,51 @@ export const useChat = defineStore('chat', {
                                     
                                 }
 
-                                if (data.type === 'token' && data.content) {
+                                // Assistant token streaming (legacy and new)
+                                if ((data.type === 'token' || data.type === 'assistant_token') && data.content) {
                                     this.currentAssistantMessage.loading = false;
                                     // Append content to the current assistant message
                                     this.currentAssistantMessage.content += data.content;
                                     // Force Vue to detect the change
                                     this.messages = [...this.messages];
                                 }
+
+                                // Assistant final content (new API emits assistant_complete)
+                                if (data.type === 'assistant_complete' && data.content) {
+                                    this.currentAssistantMessage.loading = false;
+                                    this.currentAssistantMessage.content += data.content;
+                                    this.messages = [...this.messages];
+                                }
+
+                                // New API: Tool lifecycle events mapping to UI model
+                                if (data.type === 'tool_start') {
+                                    const toolMessage: OpeyToolCall = {
+                                        status: 'pending',
+                                        toolCall: {
+                                            id: data.tool_call_id,
+                                            name: data.tool_name,
+                                            args: data.tool_input,
+                                        }
+                                    } as unknown as OpeyToolCall;
+                                    this.currentAssistantMessage.toolCalls.push(toolMessage)
+                                    this.messages = [...this.messages];
+                                }
+                                if (data.type === 'tool_end') {
+                                    const toolMessage = this.getToolCallById(data.tool_call_id)
+                                    if (toolMessage) {
+                                        toolMessage.status = data.status === 'error' ? 'error' : 'success'
+                                        toolMessage.output = data.tool_output
+                                        this.messages = [...this.messages];
+                                    }
+                                }
+                                if (data.type === 'approval_request') {
+                                    // Mark awaiting approval on matched tool call if present
+                                    const toolMessage = this.getToolCallById(data.tool_call_id)
+                                    if (toolMessage) {
+                                        toolMessage.status = 'awaiting_approval'
+                                        this.messages = [...this.messages];
+                                    }
+                                }
                             } catch (e) {
                                 throw new Error(`${e}`);
                             }