Make reply use the API (#5389)

Co-authored-by: Douwe Osinga <douwe@squareup.com>

crates/goose-cli/src/cli.rs

@@ -27,6 +27,7 @@ use goose_bench::runners::metric_aggregator::MetricAggregator;
 use goose_bench::runners::model_runner::ModelRunner;
 use std::io::Read;
 use std::path::PathBuf;
+use tracing::warn;
 
 #[derive(Parser)]
 #[command(author, version, display_name = "", about, long_about = None)]
@@ -820,9 +821,8 @@ pub struct RecipeInfo {
 pub async fn cli() -> Result<()> {
     let cli = Cli::parse();
 
-    // Track the current directory in projects.json
     if let Err(e) = crate::project_tracker::update_project_tracker(None, None) {
-        eprintln!("Warning: Failed to update project tracker: {}", e);
+        warn!("Warning: Failed to update project tracker: {}", e);
     }
 
     let command_name = match &cli.command {

crates/goose-server/src/openapi.rs

@@ -23,6 +23,7 @@ use goose::conversation::message::{
     ToolRequest, ToolResponse,
 };
 
+use crate::routes::reply::MessageEvent;
 use utoipa::openapi::schema::{
     AdditionalProperties, AnyOfBuilder, ArrayBuilder, ObjectBuilder, OneOfBuilder, Schema,
     SchemaFormat, SchemaType,
@@ -420,6 +421,7 @@ derive_utoipa!(Icon as IconSchema);
         ResourceContentsSchema,
         SystemNotificationType,
         SystemNotificationContent,
+        MessageEvent,
         JsonObjectSchema,
         RoleSchema,
         ProviderMetadata,

crates/goose-server/src/routes/reply.rs

@@ -139,6 +139,7 @@ pub enum MessageEvent {
     },
     Notification {
         request_id: String,
+        #[schema(value_type = Object)]
         message: ServerNotification,
     },
     UpdateConversation {
@@ -170,7 +171,9 @@ async fn stream_event(
     path = "/reply",
     request_body = ChatRequest,
     responses(
-        (status = 200, description = "Streaming response initiated", content_type = "text/event-stream"),
+        (status = 200, description = "Streaming response initiated",
+         body = MessageEvent,
+         content_type = "text/event-stream"),
         (status = 424, description = "Agent not initialized"),
         (status = 500, description = "Internal server error")
     )

ui/desktop/openapi.json

@@ -1247,7 +1247,14 @@
         },
         "responses": {
           "200": {
-            "description": "Streaming response initiated"
+            "description": "Streaming response initiated",
+            "content": {
+              "text/event-stream": {
+                "schema": {
+                  "$ref": "#/components/schemas/MessageEvent"
+                }
+              }
+            }
           },
           "424": {
             "description": "Agent not initialized"
@@ -3172,6 +3179,143 @@
           "propertyName": "type"
         }
       },
+      "MessageEvent": {
+        "oneOf": [
+          {
+            "type": "object",
+            "required": [
+              "message",
+              "type"
+            ],
+            "properties": {
+              "message": {
+                "$ref": "#/components/schemas/Message"
+              },
+              "type": {
+                "type": "string",
+                "enum": [
+                  "Message"
+                ]
+              }
+            }
+          },
+          {
+            "type": "object",
+            "required": [
+              "error",
+              "type"
+            ],
+            "properties": {
+              "error": {
+                "type": "string"
+              },
+              "type": {
+                "type": "string",
+                "enum": [
+                  "Error"
+                ]
+              }
+            }
+          },
+          {
+            "type": "object",
+            "required": [
+              "reason",
+              "type"
+            ],
+            "properties": {
+              "reason": {
+                "type": "string"
+              },
+              "type": {
+                "type": "string",
+                "enum": [
+                  "Finish"
+                ]
+              }
+            }
+          },
+          {
+            "type": "object",
+            "required": [
+              "model",
+              "mode",
+              "type"
+            ],
+            "properties": {
+              "mode": {
+                "type": "string"
+              },
+              "model": {
+                "type": "string"
+              },
+              "type": {
+                "type": "string",
+                "enum": [
+                  "ModelChange"
+                ]
+              }
+            }
+          },
+          {
+            "type": "object",
+            "required": [
+              "request_id",
+              "message",
+              "type"
+            ],
+            "properties": {
+              "message": {
+                "type": "object"
+              },
+              "request_id": {
+                "type": "string"
+              },
+              "type": {
+                "type": "string",
+                "enum": [
+                  "Notification"
+                ]
+              }
+            }
+          },
+          {
+            "type": "object",
+            "required": [
+              "conversation",
+              "type"
+            ],
+            "properties": {
+              "conversation": {
+                "$ref": "#/components/schemas/Conversation"
+              },
+              "type": {
+                "type": "string",
+                "enum": [
+                  "UpdateConversation"
+                ]
+              }
+            }
+          },
+          {
+            "type": "object",
+            "required": [
+              "type"
+            ],
+            "properties": {
+              "type": {
+                "type": "string",
+                "enum": [
+                  "Ping"
+                ]
+              }
+            }
+          }
+        ],
+        "discriminator": {
+          "propertyName": "type"
+        }
+      },
       "MessageMetadata": {
         "type": "object",
         "description": "Metadata for message visibility",

ui/desktop/src/api/sdk.gen.ts

@@ -351,7 +351,7 @@ export const scanRecipe = <ThrowOnError extends boolean = false>(options: Option
 };
 
 export const reply = <ThrowOnError extends boolean = false>(options: Options<ReplyData, ThrowOnError>) => {
-    return (options.client ?? client).post<ReplyResponses, ReplyErrors, ThrowOnError>({
+    return (options.client ?? client).sse.post<ReplyResponses, ReplyErrors, ThrowOnError>({
         url: '/reply',
         ...options,
         headers: {

ui/desktop/src/api/types.gen.ts

@@ -360,6 +360,32 @@ export type MessageContent = (TextContent & {
     type: 'systemNotification';
 });
 
+export type MessageEvent = {
+    message: Message;
+    type: 'Message';
+} | {
+    error: string;
+    type: 'Error';
+} | {
+    reason: string;
+    type: 'Finish';
+} | {
+    mode: string;
+    model: string;
+    type: 'ModelChange';
+} | {
+    message: {
+        [key: string]: unknown;
+    };
+    request_id: string;
+    type: 'Notification';
+} | {
+    conversation: Conversation;
+    type: 'UpdateConversation';
+} | {
+    type: 'Ping';
+};
+
 /**
  * Metadata for message visibility
  */
@@ -1825,9 +1851,11 @@ export type ReplyResponses = {
     /**
      * Streaming response initiated
      */
-    200: unknown;
+    200: MessageEvent;
 };
 
+export type ReplyResponse = ReplyResponses[keyof ReplyResponses];
+
 export type CreateScheduleData = {
     body: CreateScheduleRequest;
     path?: never;

ui/desktop/src/hooks/useChatStream.ts

@@ -1,17 +1,18 @@
 import { useCallback, useEffect, useRef, useState } from 'react';
 import { ChatState } from '../types/chatState';
+
 import {
-  Conversation,
   Message,
+  MessageEvent,
+  reply,
   resumeAgent,
   Session,
   updateFromSession,
   updateSessionUserRecipeValues,
 } from '../api';
-import { getApiUrl } from '../config';
+
 import { createUserMessage, getCompactingMessage, getThinkingMessage } from '../types/message';
 
-const TextDecoder = globalThis.TextDecoder;
 const resultsCache = new Map<string, { messages: Message[]; session: Session }>();
 
 // Debug logging - set to false in production
@@ -45,28 +46,6 @@ const log = {
   },
 };
 
-type JsonValue = string | number | boolean | null | JsonValue[] | { [key: string]: JsonValue };
-
-interface NotificationEvent {
-  type: 'Notification';
-  request_id: string;
-  message: {
-    method: string;
-    params: {
-      [key: string]: JsonValue;
-    };
-  };
-}
-
-type MessageEvent =
-  | { type: 'Message'; message: Message }
-  | { type: 'Error'; error: string }
-  | { type: 'Ping' }
-  | { type: 'Finish'; reason: string }
-  | { type: 'ModelChange'; model: string; mode: string }
-  | { type: 'UpdateConversation'; conversation: Conversation }
-  | NotificationEvent;
-
 interface UseChatStreamProps {
   sessionId: string;
   onStreamFinish: () => void;
@@ -106,121 +85,78 @@ function pushMessage(currentMessages: Message[], incomingMsg: Message): Message[
 }
 
 async function streamFromResponse(
-  response: Response,
+  stream: AsyncIterable<MessageEvent>,
   initialMessages: Message[],
   updateMessages: (messages: Message[]) => void,
   updateChatState: (state: ChatState) => void,
   onFinish: (error?: string) => void
 ): Promise<void> {
-  let chunkCount = 0;
   let messageEventCount = 0;
+  let currentMessages = initialMessages;
 
   try {
-    if (!response.ok) throw new Error(`HTTP ${response.status}`);
+    log.stream('reading-events');
-    if (!response.body) throw new Error('No response body');
 
-    const reader = response.body.getReader();
+    for await (const event of stream) {
-    const decoder = new TextDecoder();
+      switch (event.type) {
-    let currentMessages = initialMessages;
+        case 'Message': {
+          messageEventCount++;
+          const msg = event.message;
+          currentMessages = pushMessage(currentMessages, msg);
 
-    log.stream('reading-chunks');
+          if (getCompactingMessage(msg)) {
-
+            log.state(ChatState.Compacting, { reason: 'compacting notification' });
-    while (true) {
+            updateChatState(ChatState.Compacting);
-      const { done, value } = await reader.read();
+          } else if (getThinkingMessage(msg)) {
-      if (done) {
+            log.state(ChatState.Thinking, { reason: 'thinking notification' });
-        log.stream('chunks-complete', {
+            updateChatState(ChatState.Thinking);
-          totalChunks: chunkCount,
-          messageEvents: messageEventCount,
-        });
-        break;
-      }
-
-      chunkCount++;
-      const chunk = decoder.decode(value);
-      const lines = chunk.split('\n');
-
-      for (const line of lines) {
-        if (!line.startsWith('data: ')) continue;
-
-        const data = line.slice(6);
-        if (data === '[DONE]') continue;
-
-        try {
-          const event = JSON.parse(data) as MessageEvent;
-
-          switch (event.type) {
-            case 'Message': {
-              messageEventCount++;
-              const msg = event.message;
-              currentMessages = pushMessage(currentMessages, msg);
-
-              if (getCompactingMessage(msg)) {
-                log.state(ChatState.Compacting, { reason: 'compacting notification' });
-                updateChatState(ChatState.Compacting);
-              } else if (getThinkingMessage(msg)) {
-                log.state(ChatState.Thinking, { reason: 'thinking notification' });
-                updateChatState(ChatState.Thinking);
-              }
-
-              // Only log every 10th message event to avoid spam
-              if (messageEventCount % 10 === 0) {
-                log.stream('message-chunk', {
-                  eventCount: messageEventCount,
-                  messageCount: currentMessages.length,
-                });
-              }
-
-              // This calls the wrapped setMessagesAndLog with 'streaming' context
-              updateMessages(currentMessages);
-              break;
-            }
-            case 'Error': {
-              log.error('stream event error', event.error);
-              onFinish('Stream error: ' + event.error);
-              return;
-            }
-            case 'Finish': {
-              log.stream('finish-event', { reason: event.reason });
-              onFinish();
-              return;
-            }
-            case 'ModelChange': {
-              log.stream('model-change', {
-                model: event.model,
-                mode: event.mode,
-              });
-              break;
-            }
-            case 'UpdateConversation': {
-              log.messages('conversation-update', event.conversation.length);
-              currentMessages = event.conversation;
-              // This calls the wrapped setMessagesAndLog with 'streaming' context
-              updateMessages(event.conversation);
-              break;
-            }
-            case 'Notification': {
-              // Don't log notifications, too noisy
-              break;
-            }
-            case 'Ping': {
-              // Don't log pings
-              break;
-            }
-            default: {
-              console.warn('Unhandled event type:', event['type']);
-              break;
-            }
           }
-        } catch (e) {
+
-          log.error('SSE parse failed', e);
+          if (messageEventCount % 10 === 0) {
-          onFinish('Failed to parse SSE:' + e);
+            log.stream('message-chunk', {
+              eventCount: messageEventCount,
+              messageCount: currentMessages.length,
+            });
+          }
+
+          updateMessages(currentMessages);
+          break;
         }
+        case 'Error': {
+          log.error('stream event error', event.error);
+          onFinish('Stream error: ' + event.error);
+          return;
+        }
+        case 'Finish': {
+          log.stream('finish-event', { reason: event.reason });
+          onFinish();
+          return;
+        }
+        case 'ModelChange': {
+          log.stream('model-change', {
+            model: event.model,
+            mode: event.mode,
+          });
+          break;
+        }
+        case 'UpdateConversation': {
+          log.messages('conversation-update', event.conversation.length);
+          currentMessages = event.conversation;
+          updateMessages(event.conversation);
+          break;
+        }
+        case 'Notification':
+        case 'Ping':
+          break;
       }
     }
+
+    log.stream('events-complete', { messageEvents: messageEventCount });
+    onFinish();
   } catch (error) {
     if (error instanceof Error && error.name !== 'AbortError') {
       log.error('stream read error', error);
-      onFinish('Stream error:' + error);
+      onFinish('Stream error: ' + error);
     }
   }
 }
@@ -337,26 +273,19 @@ export function useChatStream({
       try {
         log.stream('request-start', { sessionId: sessionId.slice(0, 8) });
 
-        const response = await fetch(getApiUrl('/reply'), {
+        const { stream } = await reply({
-          method: 'POST',
+          body: {
-          headers: {
-            'Content-Type': 'application/json',
-            'X-Secret-Key': await window.electron.getSecretKey(),
-          },
-          body: JSON.stringify({
             session_id: sessionId,
             messages: currentMessages,
-          }),
+          },
+          throwOnError: true,
           signal: abortControllerRef.current.signal,
         });
 
-        log.stream('response-received', {
+        log.stream('stream-started');
-          status: response.status,
-          ok: response.ok,
-        });
 
         await streamFromResponse(
-          response,
+          stream,
           currentMessages,
           (messages: Message[]) => setMessagesAndLog(messages, 'streaming'),
           setChatState,