server: allow stream for exec_shell_command (#25526)

* init stream

* add stream for shell tool

* add test

* nits

* update docs

tools/server/README-dev.md

@@ -235,6 +235,29 @@ That requires `JSON.stringify` when formatted to message content:
 }
 ```
 
+Set `stream: true` in the request body to stream a tool's output as it runs, instead of waiting for it to finish. Only certain tools accept this (for ex. `exec_shell_command`);
+returns 404 if tool doesn't support it.
+
+Response is SSE stream, one `data: <json>` line per chunk:
+
+```json
+{"chunk": "hello\n"}
+```
+
+followed by a final event once the tool returns:
+
+```json
+{"done": true}
+```
+
+or, if `invoke()` threw:
+
+```json
+{"done": true, "error": "..."}
+```
+
+There is no `[DONE]` sentinel (unlike `/chat/completions`), the stream ends after the `done`
+
 ### Router mode: how child <--> router communicates
 
 Upon spawning a new child process using `subprocess`, both child and router listen to the stdout/stderr (combined)

tools/server/server-tools.cpp

@@ -12,6 +12,7 @@
 #include <climits>
 #include <algorithm>
 #include <unordered_set>
+#include <functional>
 
 namespace fs = std::filesystem;
 
@@ -51,7 +52,13 @@ public:
     virtual bool write_file(const std::string & path, const std::string & content) const = 0;
     // paths relative to `base`, '/'-separated; sets `err` if `base` isn't a directory
     virtual std::vector<std::string> list_files(const std::string & base, std::string & err) const = 0;
-    virtual exec_result run(const std::vector<std::string> & args, size_t max_output, int timeout_secs) const = 0;
+    // on_chunk, if set, is called with each chunk of output as it is read (before truncation cuts in);
+    // returning false terminates the process early (e.g. the client disconnected)
+    virtual exec_result run(
+            const std::vector<std::string> & args,
+            size_t max_output,
+            int timeout_secs,
+            const std::function<bool(const std::string &)> & on_chunk = nullptr) const = 0;
 };
 
 class tools_io_basic : public tools_io {
@@ -123,7 +130,11 @@ public:
         return list_files_fallback(base);
     }
 
-    exec_result run(const std::vector<std::string> & args, size_t max_output, int timeout_secs) const override {
+    exec_result run(
+            const std::vector<std::string> & args,
+            size_t max_output,
+            int timeout_secs,
+            const std::function<bool(const std::string &)> & on_chunk = nullptr) const override {
         exec_result res;
 
         subprocess_s proc;
@@ -164,8 +175,14 @@ public:
                     size_t len = strlen(buf);
                     if (output.size() + len <= max_output) {
                         output.append(buf, len);
+                        if (on_chunk && !on_chunk(std::string(buf, len))) {
+                            subprocess_terminate(&proc);
+                            break;
+                        }
                     } else {
-                        output.append(buf, max_output - output.size());
+                        size_t remaining = max_output - output.size();
+                        output.append(buf, remaining);
+                        if (on_chunk && remaining > 0) on_chunk(std::string(buf, remaining));
                         truncated = true;
                     }
                 }
@@ -287,7 +304,7 @@ struct server_tool_read_file : server_tool {
         };
     }
 
-    json invoke(json params) const override {
+    json invoke(json params, server_tool::stream *) const override {
         std::string path  = params.at("path").get<std::string>();
         int  start_line   = json_value(params, "start_line", 1);
         int  end_line     = json_value(params, "end_line",  -1); // -1 = no limit
@@ -376,7 +393,7 @@ struct server_tool_file_glob_search : server_tool {
         };
     }
 
-    json invoke(json params) const override {
+    json invoke(json params, server_tool::stream *) const override {
         std::string base    = params.at("path").get<std::string>();
         std::string include = json_value(params, "include", std::string("**"));
         std::string exclude = json_value(params, "exclude", std::string(""));
@@ -457,7 +474,7 @@ struct server_tool_grep_search : server_tool {
         };
     }
 
-    json invoke(json params) const override {
+    json invoke(json params, server_tool::stream *) const override {
         std::string path        = params.at("path").get<std::string>();
         std::string pat_str     = params.at("pattern").get<std::string>();
         std::string include     = json_value(params, "include", std::string("**"));
@@ -577,6 +594,7 @@ struct server_tool_exec_shell_command : server_tool {
         name = "exec_shell_command";
         display_name = "Execute shell command";
         permission_write = true;
+        support_stream = true;
     }
 
     json get_definition() const override {
@@ -598,7 +616,7 @@ struct server_tool_exec_shell_command : server_tool {
         };
     }
 
-    json invoke(json params) const override {
+    json invoke(json params, server_tool::stream * st) const override {
         std::string command   = params.at("command").get<std::string>();
         int    timeout        = json_value(params, "timeout",         10);
         size_t max_output     = (size_t) json_value(params, "max_output_size", (int) SERVER_TOOL_EXEC_SHELL_COMMAND_MAX_OUTPUT_SIZE);
@@ -612,7 +630,24 @@ struct server_tool_exec_shell_command : server_tool {
         std::vector<std::string> args = {"sh", "-c", command};
 #endif
 
-        auto io  = make_tools_io(params);
+        auto io = make_tools_io(params);
+
+        if (st) {
+            auto res = io->run(args, max_output, timeout, [st](const std::string & chunk) {
+                st->push(chunk);
+                return !st->alive || st->alive();
+            });
+            if (st->alive && !st->alive()) {
+                return json();
+            }
+            std::string tail = string_format("\n[exit code: %d]", res.exit_code);
+            if (res.timed_out) {
+                tail += " [exit due to timed out]";
+            }
+            st->push(tail);
+            return json();
+        }
+
         auto res = io->run(args, max_output, timeout);
 
         std::string text_output = res.output;
@@ -654,7 +689,7 @@ struct server_tool_write_file : server_tool {
         };
     }
 
-    json invoke(json params) const override {
+    json invoke(json params, server_tool::stream *) const override {
         std::string path    = params.at("path").get<std::string>();
         std::string content = params.at("content").get<std::string>();
 
@@ -710,7 +745,7 @@ struct server_tool_edit_file : server_tool {
         };
     }
 
-    json invoke(json params) const override {
+    json invoke(json params, server_tool::stream *) const override {
         std::string path = params.at("path").get<std::string>();
         const json & edits_json = params.at("edits");
 
@@ -1018,7 +1053,7 @@ struct server_tool_get_datetime : server_tool {
         };
     }
 
-    json invoke(json) const override {
+    json invoke(json, server_tool::stream *) const override {
         auto now = std::chrono::system_clock::now();
         auto time = std::chrono::system_clock::to_time_t(now);
 
@@ -1026,6 +1061,59 @@ struct server_tool_get_datetime : server_tool {
     }
 };
 
+struct server_tool_stream_result : server_task_result {
+    std::string chunk;
+    bool done = false;
+    std::string error_msg;
+
+    json to_json() override {
+        if (!done) {
+            return {{"chunk", chunk}};
+        } else {
+            json result = {{"done", true}};
+            if (!error_msg.empty()) {
+                result["error"] = error_msg;
+            }
+            return result;
+        }
+    }
+};
+
+void server_tool::stream::push(const std::string & chunk) {
+    if (chunk.empty()) return;
+    auto r = std::make_unique<server_tool_stream_result>();
+    r->id    = id;
+    r->chunk = chunk;
+    qr.send(std::move(r));
+}
+
+struct server_tools_res : server_http_res {
+    std::thread worker;
+    server_response * qr = nullptr; // set only for streaming responses
+    int id = -1;
+
+    ~server_tools_res() override {
+        if (worker.joinable()) {
+            worker.join();
+        }
+        if (qr) {
+            qr->remove_waiting_task_id(id);
+        }
+    }
+};
+
+static server_tool & find_tool(std::vector<std::unique_ptr<server_tool>> & tools, const std::string & name, bool require_stream) {
+    for (auto & t : tools) {
+        if (t->name == name) {
+            if (require_stream && !t->support_stream) {
+                throw std::invalid_argument(string_format("tool \"%s\" does not support stream = true", name.c_str()));
+            }
+            return *t;
+        }
+    }
+    throw std::invalid_argument(string_format("unknown tool \"%s\"", name.c_str()));
+}
+
 //
 // public API
 //
@@ -1090,16 +1178,63 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools) {
     };
 
     handle_post = [this](const server_http_req & req) -> server_http_res_ptr {
-        auto res = std::make_unique<server_http_res>();
+        auto res = std::make_unique<server_tools_res>();
         try {
             json body = json::parse(req.body);
             std::string tool_name = body.at("tool").get<std::string>();
             json params = body.value("params", json::object());
-            json result = invoke(tool_name, params);
-            res->data   = safe_json_to_str(result);
+            bool stream = body.value("stream", false);
+
+            server_tool & tool = find_tool(tools, tool_name, stream);
+
+            if (stream) {
+                int id = res_id.fetch_add(1);
+                queue_res.add_waiting_task_id(id);
+                res->qr = &queue_res;
+                res->id = id;
+
+                res->worker = std::thread([this, id, &req, &tool, params]() mutable {
+                    server_tool::stream st{queue_res, id, [&req]() {
+                        return !req.should_stop();
+                    }};
+
+                    auto done = std::make_unique<server_tool_stream_result>();
+                    try {
+                        tool.invoke(params, &st);
+                    } catch (const std::exception & e) {
+                        done->error_msg = e.what();
+                    } catch (...) {
+                        done->error_msg = "An unknown error occurred";
+                    }
+                    done->id    = st.id;
+                    done->done  = true;
+                    st.qr.send(std::move(done));
+                });
+
+                res->content_type = "text/event-stream";
+                res->status = 200;
+                res->next   = [this, id](std::string & output) -> bool {
+                    auto result = queue_res.recv(id);
+                    auto * r = dynamic_cast<server_tool_stream_result *>(result.get());
+                    GGML_ASSERT(r != nullptr);
+                    output = "data: " + safe_json_to_str(r->to_json()) + "\n\n";
+                    if (r->done) {
+                        queue_res.remove_waiting_task_id(id);
+                        return false;
+                    }
+                    return true;
+                };
+            } else {
+                json result = tool.invoke(params, nullptr);
+                res->status = 200;
+                res->data   = safe_json_to_str(result);
+            }
         } catch (const json::exception & e) {
             res->status = 400;
             res->data   = safe_json_to_str(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST));
+        } catch (const std::invalid_argument & e) {
+            res->status = 404;
+            res->data   = safe_json_to_str(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST));
         } catch (const std::exception & e) {
             SRV_ERR("got exception: %s\n", e.what());
             res->status = 500;
@@ -1108,12 +1243,3 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools) {
         return res;
     };
 }
-
-json server_tools::invoke(const std::string & name, const json & params) {
-    for (auto & t : tools) {
-        if (t->name == name) {
-            return t->invoke(params);
-        }
-    }
-    return {{"error", "unknown tool: " + name}};
-}

tools/server/server-tools.h

@@ -2,15 +2,27 @@
 
 #include "server-common.h"
 #include "server-http.h"
+#include "server-queue.h"
+
+#include <atomic>
+#include <functional>
 
 struct server_tool {
     std::string name;
     std::string display_name;
     bool permission_write = false;
+    bool support_stream = false; // if true, output can be streamed
 
     virtual ~server_tool() = default;
     virtual json get_definition() const = 0;
-    virtual json invoke(json params) const = 0;
+
+    struct stream {
+        server_response & qr;
+        int id;
+        std::function<bool()> alive;
+        void push(const std::string & chunk);
+    };
+    virtual json invoke(json params, stream * st = nullptr) const = 0;
 
     json to_json() const;
 };
@@ -18,8 +30,11 @@ struct server_tool {
 struct server_tools {
     std::vector<std::unique_ptr<server_tool>> tools;
 
+    // for streaming
+    server_response queue_res;
+    std::atomic<int> res_id{0};
+
     void setup(const std::vector<std::string> & enabled_tools);
-    json invoke(const std::string & name, const json & params);
 
     server_http_context::handler_t handle_get;
     server_http_context::handler_t handle_post;

tools/server/tests/unit/test_tools_builtin.py

@@ -105,6 +105,24 @@ def test_tools_builtin_edit_file_rejects_non_unique_old_text():
             os.remove(log_path)
 
 
+def test_tools_builtin_exec_shell_command_stream():
+    global server
+    server.start()
+
+    events = list(server.make_stream_request("POST", "/tools", data={
+        "tool": "exec_shell_command",
+        "params": {"command": "echo hello"},
+        "stream": True,
+    }))
+
+    assert len(events) >= 2
+    assert events[-1]["done"] is True
+    assert not events[-1].get("error")
+    chunks = "".join(e["chunk"] for e in events[:-1])
+    assert "hello" in chunks
+    assert "[exit code: 0]" in chunks
+
+
 def test_tools_builtin_edit_file_rejects_overlapping_edits():
     global server
     server.start()