@@ -42,45 +42,6 @@ func ensureAltSse(endpoint string) (string, error) {
4242 return endpoint , nil
4343}
4444
45- // UpdateToolUseData updates the tool use data for a specific tool call in the chat
46- func UpdateToolUseData (chatId string , toolCallId string , toolUseData uctypes.UIMessageDataToolUse ) error {
47- chat := chatstore .DefaultChatStore .Get (chatId )
48- if chat == nil {
49- return fmt .Errorf ("chat not found: %s" , chatId )
50- }
51-
52- for _ , genMsg := range chat .NativeMessages {
53- chatMsg , ok := genMsg .(* GeminiChatMessage )
54- if ! ok {
55- continue
56- }
57-
58- for i , part := range chatMsg .Parts {
59- if part .FunctionCall != nil && part .ToolUseData != nil && part .ToolUseData .ToolCallId == toolCallId {
60- // Update the message with new tool use data
61- updatedMsg := & GeminiChatMessage {
62- MessageId : chatMsg .MessageId ,
63- Role : chatMsg .Role ,
64- Parts : make ([]GeminiMessagePart , len (chatMsg .Parts )),
65- Usage : chatMsg .Usage ,
66- }
67- copy (updatedMsg .Parts , chatMsg .Parts )
68- updatedMsg .Parts [i ].ToolUseData = & toolUseData
69-
70- aiOpts := & uctypes.AIOptsType {
71- APIType : chat .APIType ,
72- Model : chat .Model ,
73- APIVersion : chat .APIVersion ,
74- }
75-
76- return chatstore .DefaultChatStore .PostMessage (chatId , aiOpts , updatedMsg )
77- }
78- }
79- }
80-
81- return fmt .Errorf ("tool call with ID %s not found in chat %s" , toolCallId , chatId )
82- }
83-
8445// appendPartToLastUserMessage appends a text part to the last user message in the contents slice
8546func appendPartToLastUserMessage (contents []GeminiContent , text string ) {
8647 for i := len (contents ) - 1 ; i >= 0 ; i -- {
@@ -347,6 +308,14 @@ func processGeminiStream(
347308 if errors .Is (err , io .EOF ) {
348309 break
349310 }
311+ if sseHandler .Err () != nil {
312+ partialMsg := extractPartialGeminiMessage (msgID , textBuilder .String ())
313+ return & uctypes.WaveStopReason {
314+ Kind : uctypes .StopKindCanceled ,
315+ ErrorType : "client_disconnect" ,
316+ ErrorText : "client disconnected" ,
317+ }, partialMsg , nil
318+ }
350319 _ = sseHandler .AiMsgError (fmt .Sprintf ("stream decode error: %v" , err ))
351320 return & uctypes.WaveStopReason {
352321 Kind : uctypes .StopKindError ,
@@ -512,3 +481,19 @@ func processGeminiStream(
512481
513482 return stopReason , assistantMsg , nil
514483}
484+
485+ func extractPartialGeminiMessage (msgID string , text string ) * GeminiChatMessage {
486+ if text == "" {
487+ return nil
488+ }
489+
490+ return & GeminiChatMessage {
491+ MessageId : msgID ,
492+ Role : "model" ,
493+ Parts : []GeminiMessagePart {
494+ {
495+ Text : text ,
496+ },
497+ },
498+ }
499+ }
0 commit comments