Make ChatPlugin support string only response

ChatPlugins/Sources/ShortcutChatPlugin/ShortcutChatPlugin.swift

@@ -20,7 +20,36 @@ public final class ShortcutChatPlugin: ChatPlugin {
         terminal = Terminal()
     }
 
-    public func send(_ request: Request) async -> AsyncThrowingStream<Response, any Error> {
+    public func sendForTextResponse(_ request: Request) async
+        -> AsyncThrowingStream<String, any Error>
+    {
+        let stream = await sendForComplicatedResponse(request)
+        return .init { continuation in
+            let task = Task {
+                do {
+                    for try await response in stream {
+                        switch response {
+                        case let .content(.text(content)):
+                            continuation.yield(content)
+                        default:
+                            break
+                        }
+                    }
+                    continuation.finish()
+                } catch {
+                    continuation.finish(throwing: error)
+                }
+            }
+
+            continuation.onTermination = { _ in
+                task.cancel()
+            }
+        }
+    }
+
+    public func sendForComplicatedResponse(_ request: Request) async
+        -> AsyncThrowingStream<Response, any Error>
+    {
         return .init { continuation in
             let task = Task {
                 let id = "\(Self.command)-\(UUID().uuidString)"

ChatPlugins/Sources/TerminalChatPlugin/TerminalChatPlugin.swift

@@ -9,7 +9,7 @@ public final class TerminalChatPlugin: ChatPlugin {
     public static var name: String { "Terminal" }
     public static var description: String { """
     Run the command in the message from terminal.
-    
+
     You can use environment variable `$FILE_PATH` and `$PROJECT_ROOT` to access the current file path and project root.
     """ }
 
@@ -23,41 +23,11 @@ public final class TerminalChatPlugin: ChatPlugin {
         terminal = Terminal()
     }
 
-    public func formatContent(_ content: Response.Content) -> Response.Content {
-        switch content {
-        case let .text(content):
-            return .text("""
-            ```sh
-            \(content)
-            ```
-            """)
-        }
-    }
-
-    public func send(_ request: Request) async -> AsyncThrowingStream<Response, any Error> {
+    public func sendForTextResponse(_ request: Request) async
+        -> AsyncThrowingStream<String, any Error>
+    {
         return .init { continuation in
             let task = Task {
-                var updateTime = Date()
-
-                func streamOutput(_ content: String) {
-                    defer { updateTime = Date() }
-                    if Date().timeIntervalSince(updateTime) > 60 * 2 {
-                        continuation.yield(.startNewMessage)
-                        continuation.yield(.startAction(
-                            id: "run",
-                            task: "Continue `\(request.text)`"
-                        ))
-                        continuation.yield(.finishAction(
-                            id: "run",
-                            result: .success("Executed.")
-                        ))
-                        continuation.yield(.content(.text("[continue]\n")))
-                        continuation.yield(.content(.text(content)))
-                    } else {
-                        continuation.yield(.content(.text(content)))
-                    }
-                }
-
                 do {
                     let fileURL = XcodeInspector.shared.realtimeActiveDocumentURL
                     let projectURL = XcodeInspector.shared.realtimeActiveProjectURL
@@ -75,34 +45,25 @@ public final class TerminalChatPlugin: ChatPlugin {
                     let env = ProcessInfo.processInfo.environment
                     let shell = env["SHELL"] ?? "/bin/bash"
 
-                    continuation.yield(.startAction(id: "run", task: "Run `\(request.text)`"))
-
                     let output = terminal.streamCommand(
                         shell,
                         arguments: ["-i", "-l", "-c", request.text],
                         currentDirectoryURL: projectURL,
                         environment: environment
                     )
 
-                    continuation.yield(.finishAction(
-                        id: "run",
-                        result: .success("Executed.")
-                    ))
-
+                    var accumulatedOutput = ""
                     for try await content in output {
                         try Task.checkCancellation()
-                        streamOutput(content)
+                        accumulatedOutput += content
+                        continuation.yield(accumulatedOutput)
                     }
                 } catch let error as Terminal.TerminationError {
-                    continuation.yield(.content(.text("""
-
-                    [error: \(error.reason)]
-                    """)))
+                    let errorMessage = "\n\n[error: \(error.reason)]"
+                    continuation.yield(errorMessage)
                 } catch {
-                    continuation.yield(.content(.text("""
-
-                    [error: \(error.localizedDescription)]
-                    """)))
+                    let errorMessage = "\n\n[error: \(error.localizedDescription)]"
+                    continuation.yield(errorMessage)
                 }
 
                 continuation.finish()
@@ -116,5 +77,71 @@ public final class TerminalChatPlugin: ChatPlugin {
             }
         }
     }
+
+    public func formatContent(_ content: Response.Content) -> Response.Content {
+        switch content {
+        case let .text(content):
+            return .text("""
+            ```sh
+            \(content)
+            ```
+            """)
+        }
+    }
+
+    public func sendForComplicatedResponse(_ request: Request) async
+        -> AsyncThrowingStream<Response, any Error>
+    {
+        return .init { continuation in
+            let task = Task {
+                var updateTime = Date()
+
+                continuation.yield(.startAction(id: "run", task: "Run `\(request.text)`"))
+
+                let textStream = await sendForTextResponse(request)
+                var previousOutput = ""
+
+                continuation.yield(.finishAction(
+                    id: "run",
+                    result: .success("Executed.")
+                ))
+
+                do {
+                    for try await accumulatedOutput in textStream {
+                        try Task.checkCancellation()
+
+                        let newContent = accumulatedOutput.dropFirst(previousOutput.count)
+                        previousOutput = accumulatedOutput
+
+                        if !newContent.isEmpty {
+                            if Date().timeIntervalSince(updateTime) > 60 * 2 {
+                                continuation.yield(.startNewMessage)
+                                continuation.yield(.startAction(
+                                    id: "run",
+                                    task: "Continue `\(request.text)`"
+                                ))
+                                continuation.yield(.finishAction(
+                                    id: "run",
+                                    result: .success("Executed.")
+                                ))
+                                continuation.yield(.content(.text("[continue]\n")))
+                                updateTime = Date()
+                            }
+
+                            continuation.yield(.content(.text(String(newContent))))
+                        }
+                    }
+
+                    continuation.finish()
+                } catch {
+                    continuation.finish(throwing: error)
+                }
+            }
+
+            continuation.onTermination = { _ in
+                task.cancel()
+            }
+        }
+    }
 }
 

Core/Sources/ChatService/AllPlugins.swift

@@ -56,7 +56,7 @@ final class LegacyChatPluginWrapper<Plugin: ChatPlugin>: LegacyChatPlugin {
 
         let plugin = Plugin()
 
-        let stream = await plugin.send(.init(
+        let stream = await plugin.sendForComplicatedResponse(.init(
             text: content,
             arguments: [],
             history: chatGPTService.memory.history

Tool/Sources/ChatBasic/ChatPlugin.swift

@@ -19,7 +19,13 @@ public protocol ChatPlugin {
     static var command: String { get }
     static var name: String { get }
     static var description: String { get }
-    func send(_ request: Request) async -> AsyncThrowingStream<Response, any Error>
+    // In this method, the plugin is able to send more complicated response. It also enables it to
+    // perform special tasks like starting a new message or reporting progress.
+    func sendForComplicatedResponse(
+        _ request: Request
+    ) async -> AsyncThrowingStream<Response, any Error>
+    // This method allows the plugin to respond a stream of text content only.
+    func sendForTextResponse(_ request: Request) async -> AsyncThrowingStream<String, any Error>
     func formatContent(_ content: Response.Content) -> Response.Content
     init()
 }
@@ -28,5 +34,26 @@ public extension ChatPlugin {
     func formatContent(_ content: Response.Content) -> Response.Content {
         return content
     }
+
+    func sendForComplicatedResponse(
+        _ request: Request
+    ) async -> AsyncThrowingStream<Response, any Error> {
+        let textStream = await sendForTextResponse(request)
+        return AsyncThrowingStream<Response, any Error> { continuation in
+            let task = Task {
+                do {
+                    for try await text in textStream {
+                        continuation.yield(Response.content(.text(text)))
+                    }
+                    continuation.finish()
+                } catch {
+                    continuation.finish(throwing: error)
+                }
+            }
+
+            continuation.onTermination = { _ in
+                task.cancel()
+            }
+        }
+    }
 }
-