From c92e806d1c81091c9035edce99c35374da1b465e Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Sat, 11 Jul 2026 12:42:55 +0200 Subject: [PATCH] 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 | 23 +++ tools/server/server-tools.cpp | 172 +++++++++++++++--- tools/server/server-tools.h | 19 +- tools/server/tests/unit/test_tools_builtin.py | 18 ++ 4 files changed, 207 insertions(+), 25 deletions(-) diff --git a/tools/server/README-dev.md b/tools/server/README-dev.md index a5619485f0..e81336e5e2 100644 --- a/tools/server/README-dev.md +++ b/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: ` 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) diff --git a/tools/server/server-tools.cpp b/tools/server/server-tools.cpp index 3f162a13e0..a8216d7dbc 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -12,6 +12,7 @@ #include #include #include +#include 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 list_files(const std::string & base, std::string & err) const = 0; - virtual exec_result run(const std::vector & 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 & args, + size_t max_output, + int timeout_secs, + const std::function & 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 & args, size_t max_output, int timeout_secs) const override { + exec_result run( + const std::vector & args, + size_t max_output, + int timeout_secs, + const std::function & 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(); 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 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 pat_str = params.at("pattern").get(); 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(); 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 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 content = params.at("content").get(); @@ -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(); 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(); + 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> & 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 & enabled_tools) { }; handle_post = [this](const server_http_req & req) -> server_http_res_ptr { - auto res = std::make_unique(); + auto res = std::make_unique(); try { json body = json::parse(req.body); std::string tool_name = body.at("tool").get(); 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(); + 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(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 & 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}}; -} diff --git a/tools/server/server-tools.h b/tools/server/server-tools.h index cc147379bf..6f6528f484 100644 --- a/tools/server/server-tools.h +++ b/tools/server/server-tools.h @@ -2,15 +2,27 @@ #include "server-common.h" #include "server-http.h" +#include "server-queue.h" + +#include +#include 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 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> tools; + // for streaming + server_response queue_res; + std::atomic res_id{0}; + void setup(const std::vector & 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; diff --git a/tools/server/tests/unit/test_tools_builtin.py b/tools/server/tests/unit/test_tools_builtin.py index d4fd5dc9b7..1b2d0db432 100755 --- a/tools/server/tests/unit/test_tools_builtin.py +++ b/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()