merge main

Author: alkant

cli/CHANGELOG.md

@@ -98,4 +98,4 @@
 
 No content.
 
-<!-- Increment to skip CHANGELOG.md test: 6 -->
+<!-- Increment to skip CHANGELOG.md test: 7 -->

cli/Cargo.lock

@@ -33,7 +33,7 @@ reqwest = { version = "0.12.15", features = ["json"] }
 rmcp = "0.8.2"
 schemars = "1.0.4"
 serde = { version = "1.0.219", features = ["derive"] }
-serde_json = "1.0.140"
+serde_json = { version = "1.0.140", features = ["preserve_order"] }
 sha2 = "0.10.9"
 tokio = { version = "1.44.2", features = ["fs", "macros", "process", "rt-multi-thread"] }
 tokio-tungstenite = { version = "0.26.2", features = ["native-tls"] }

cli/Cargo.toml

@@ -277,33 +277,75 @@ async fn call_tool_(name: &mut String, content: &str, session: &Session) -> Resu
     let args: LocalToolRequest = serde_json::from_str(content).map_err(|e| e.to_string())?;
     *name = args.tool_name;
     let tool = session.tool(name).ok_or_else(|| "tool not found".to_string())?;
-    let command = format!(
-        "{} {}",
-        name.bold().yellow(),
-        serde_json::to_string(&args.tool_args).unwrap().yellow()
-    );
+    let args = (tool.reorder)(args.tool_args).await?.0;
+    let command = display_tool_request(name, &args);
     let (kind, do_not_ask) = match tool.effect {
         tool::Effect::ReadOnly => ("read-only".green(), config::LocalToolAsk::Mutating),
         tool::Effect::Mutating => ("mutating".yellow(), config::LocalToolAsk::Destructive),
         tool::Effect::Destructive => ("destructive".red(), config::LocalToolAsk::Never),
     };
-    println!("Sec-Gemini wants to execute a {kind} local tool on your machine:\n{command}");
+    print!("Sec-Gemini wants to execute a {kind} local tool on your machine:\n{command}");
     authorize(&config::LOCAL_TOOL_ASK_BEFORE, |x| authorize_tool(x, tool.effect), do_not_ask)
         .await?;
-    let result = (tool.call)(args.tool_args).await.and_then(convert_tool_response);
+    let result = (tool.call)(args).await;
     let (success, output) = match &result {
-        Ok(x) => ("successful".green(), x.blue()),
-        Err(e) => ("failed".bold().red(), e.blue()),
+        Ok(x) => ("successful".green(), display_tool_response(&x.0)),
+        Err(e) => ("failed".bold().red(), format!("{}\n", e.red())),
     };
-    println!("Sec-Gemini wants to access the result of the {success} execution:\n{output}");
+    print!("Sec-Gemini wants to access the result of the {success} execution:\n{output}");
     authorize(&config::LOCAL_TOOL_ASK_AFTER, |x| !x, false).await?;
+    result.map(convert_tool_response)
+}
+
+fn display_tool_request(name: &str, args: &serde_json::Value) -> String {
+    let mut result = format!("{}()\n", name.bold().yellow());
+    if let Some(obj) = args.as_object() {
+        display_json_object(&mut result, colored::Color::Yellow, obj);
+    } else {
+        writeln!(result, " {}", serde_json::to_string(args).unwrap().yellow()).unwrap();
+    }
+    result
+}
+
+fn display_tool_response(resp: &serde_json::Value) -> String {
+    let mut result = String::new();
+    match resp {
+        serde_json::Value::String(x) => {
+            write!(result, "{}", x.blue()).unwrap();
+            ensure_newline(x, &mut result);
+        }
+        serde_json::Value::Object(x) => display_json_object(&mut result, colored::Color::Blue, x),
+        x => fail!("unexpected response {x:?}"),
+    }
     result
 }
 
-fn convert_tool_response(response: rmcp::Json<serde_json::Value>) -> Result<String, String> {
+fn display_json_object(
+    out: &mut String, color: colored::Color, map: &serde_json::Map<String, serde_json::Value>,
+) {
+    for (key, val) in map {
+        write!(out, "{}:", key.bold().color(color)).unwrap();
+        if let Some(val) = val.as_str() {
+            if val.find('\n').is_some_and(|i| i + 1 < val.len()) {
+                write!(out, "\n{}", val.color(color)).unwrap();
+                ensure_newline(val, out);
+                continue;
+            }
+        }
+        writeln!(out, " {}", serde_json::to_string(val).unwrap().color(color)).unwrap();
+    }
+}
+
+fn ensure_newline(val: &str, out: &mut String) {
+    if !val.ends_with('\n') {
+        or_fail(writeln!(out, "%"));
+    }
+}
+
+fn convert_tool_response(response: rmcp::Json<serde_json::Value>) -> String {
     match response.0 {
-        serde_json::Value::String(x) => Ok(x),
-        x => serde_json::to_string(&x).map_err(|e| e.to_string()),
+        serde_json::Value::String(x) => x,
+        x => serde_json::to_string(&x).unwrap(),
     }
 }
 

cli/src/cli/interact.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 use std::pin::Pin;
 
 use rmcp::Json;
@@ -30,12 +30,13 @@ pub type CallOutput<R> = Result<Json<R>, String>;
 pub type CallResult<R> = Pin<Box<dyn Future<Output = CallOutput<R>>>>;
 
 pub struct Tools {
-    tools: HashMap<String, Tool>,
+    tools: BTreeMap<String, Tool>,
 }
 
 pub struct Tool {
     pub desc: String,
     pub input: JsonObject,
+    pub reorder: Box<dyn Fn(Value) -> CallResult<Value> + Send + Sync>,
     pub call: Box<dyn Fn(Value) -> CallResult<Value> + Send + Sync>,
     pub effect: Effect,
 }
@@ -75,7 +76,7 @@ impl Enable {
 
 impl Tools {
     pub fn list() -> Tools {
-        let mut tools = Tools { tools: HashMap::new() };
+        let mut tools = Tools { tools: BTreeMap::new() };
         exec::list(&mut tools);
         file::list(&mut tools);
         net::list(&mut tools);
@@ -96,13 +97,19 @@ impl Tools {
 
     fn push<P, R, F>(&mut self, tool: rmcp::model::Tool, call: fn(Parameters<P>) -> F)
     where
-        P: DeserializeOwned + 'static,
+        P: Serialize + DeserializeOwned + 'static,
         R: Serialize,
         F: Future<Output = CallOutput<R>> + 'static,
     {
         let name = tool.name.to_string();
         let desc = tool.description.map_or(String::new(), |x| x.to_string());
         let input = (*tool.input_schema).clone();
+        let reorder = Box::new(move |x| {
+            Box::pin(async move {
+                let x: P = serde_json::from_value(x).map_err(|e| e.to_string())?;
+                Ok(Json(serde_json::to_value(x).map_err(|e| e.to_string())?))
+            }) as CallResult<Value>
+        });
         let call = Box::new(move |x| {
             Box::pin(async move {
                 let x = serde_json::from_value(x).map_err(|e| e.to_string())?;
@@ -117,7 +124,7 @@ impl Tools {
                 (false, false) => Effect::Mutating,
                 (false, true) => Effect::Destructive,
             };
-        let tool = Tool { desc, input, call, effect };
+        let tool = Tool { desc, input, reorder, call, effect };
         assert!(self.tools.insert(name, tool).is_none());
     }
 }

cli/src/tool.rs

@@ -35,14 +35,15 @@ pub fn list(tools: &mut Tools) {
 pub struct SpawnRequest {
     program: String,
     arguments: Vec<String>,
-    stdin: String,
+    #[serde(flatten)]
+    interact: InteractRequest,
 }
 
 /// Spawns a process with input and output interaction.
 ///
 /// To interact with the process, use the `exec_interact` tool. To kill the program, use the
-/// `exec_kill` tool. The initial interaction is built-in by taking the same `stdin` field as
-/// `exec_interact` and returning the same fields.
+/// `exec_kill` tool. The initial interaction is built-in by taking the same `stdin` and
+/// `close_stdin` fields as `exec_interact` and returning the same fields.
 ///
 /// This tool can be used together with `file_write` (setting `executable` to create an executable
 /// file) to execute a Python or shell script.
@@ -53,7 +54,7 @@ async fn _spawn(_: Parameters<SpawnRequest>) -> Result<Json<InteractResponse>, S
 }
 
 async fn spawn(params: Parameters<SpawnRequest>) -> Result<Json<InteractResponse>, String> {
-    let SpawnRequest { program, arguments, stdin } = params.0;
+    let SpawnRequest { program, arguments, interact: request } = params.0;
     let mut child = Command::new(program)
         .args(arguments)
         .stdin(Stdio::piped())
@@ -62,18 +63,21 @@ async fn spawn(params: Parameters<SpawnRequest>) -> Result<Json<InteractResponse
         .spawn()
         .map_err(|e| e.to_string())?;
     let running = Running {
-        stdin: child.stdin.take().unwrap(),
+        stdin: Some(child.stdin.take().unwrap()),
         stdout: child.stdout.take().unwrap(),
         stderr: child.stderr.take().unwrap(),
         child,
     };
     *STATE.lock().unwrap() = Some(running);
-    interact(Parameters(InteractRequest { stdin })).await
+    interact(Parameters(request)).await
 }
 
 #[derive(Serialize, Deserialize, JsonSchema)]
 pub struct InteractRequest {
+    #[serde(default)]
     stdin: String,
+    #[serde(default)]
+    close_stdin: bool,
 }
 
 #[derive(Serialize, Deserialize, JsonSchema)]
@@ -87,12 +91,18 @@ pub struct InteractResponse {
 /// Interacts with a running process in textual form.
 ///
 /// The `stdin` parameter will be written to the standard input of the process. It can be empty if
-/// interacting for output only. In the response, the `running` field indicates whether the process
-/// is still running. When it's not running anymore, the `success` field indicates if it terminated
-/// successfully or not. The `stdout` and `stderr` fields contain the output read from the standard
-/// output and error since the last interaction. The tool will listen for output until some amount
-/// of inactivity or some fixed deadline, whichever happens first. The tool will fail if the process
-/// outputs binary data that is not UTF-8.
+/// interacting for output only. When the `close_stdin` parameter is set, the standard input will be
+/// closed. This is necessary for some programs to start processing. Note that once the standard
+/// input is closed, future interactions must not set `stdin` to non-empty content.
+///
+/// In the response, the `running` field indicates whether the process is still running. When it's
+/// not running anymore, the `success` field indicates if it terminated successfully or not. The
+/// `stdout` and `stderr` fields contain the output read from the standard output and error since
+/// the last interaction.
+///
+/// The tool will listen for output until some amount of inactivity or some fixed deadline,
+/// whichever happens first. The tool will fail if the process outputs binary data that is not
+/// UTF-8.
 #[rmcp::tool(name = "exec_interact")]
 async fn _interact(_: Parameters<InteractRequest>) -> Result<Json<InteractResponse>, String> {
     // TODO(https://github.com/modelcontextprotocol/rust-sdk/issues/495): Remove when fixed.
@@ -101,13 +111,20 @@ async fn _interact(_: Parameters<InteractRequest>) -> Result<Json<InteractRespon
 
 #[allow(clippy::await_holding_lock)]
 async fn interact(params: Parameters<InteractRequest>) -> Result<Json<InteractResponse>, String> {
-    let InteractRequest { stdin } = params.0;
+    let InteractRequest { stdin, close_stdin } = params.0;
     let mut state = STATE.lock().unwrap();
     let Some(running) = &mut *state else {
         return Err("no running process".to_string());
     };
     if !stdin.is_empty() {
-        running.stdin.write_all(stdin.as_bytes()).await.map_err(|e| e.to_string())?;
+        let Some(pipe) = running.stdin.as_mut() else { return Err("stdin is closed".to_string()) };
+        pipe.write_all(stdin.as_bytes()).await.map_err(|e| e.to_string())?;
+    }
+    if close_stdin {
+        let Some(mut pipe) = running.stdin.take() else {
+            return Err("stdin is already closed".to_string());
+        };
+        pipe.shutdown().await.map_err(|e| e.to_string())?;
     }
     const SIZE: usize = 1024;
     let mut stdout = vec![0; SIZE];
@@ -179,7 +196,7 @@ async fn kill(params: Parameters<KillRequest>) -> Result<Json<String>, String> {
 static STATE: Mutex<Option<Running>> = Mutex::new(None);
 
 struct Running {
-    stdin: ChildStdin,
+    stdin: Option<ChildStdin>,
     stdout: ChildStdout,
     stderr: ChildStderr,
     child: Child,

cli/src/tool/exec.rs

@@ -80,6 +80,7 @@ async fn read(params: Parameters<ReadRequest>) -> Result<Json<String>, String> {
 pub struct WriteRequest {
     path: String,
     content: String,
+    #[serde(default)]
     executable: bool,
 }
 

cli/src/tool/file.rs

@@ -31,7 +31,7 @@ class NetResponse(BaseModel):
   url: str = Field(title="Request URL")
   ok: bool
   error_message: str = Field("", title="Error Message")
-  data: dict = Field({}, title="Response Data")
+  data: dict = Field(default_factory=dict, title="Response Data")
   latency: float = Field(
     0.0,
     title="Latency",
@@ -45,8 +45,10 @@ def __init__(self, base_url: str, api_key: str):
     self.api_key = api_key
     self.client = httpx.Client(timeout=90)
 
-  def post(self, endpoint: str, model: T, headers: dict = {}) -> NetResponse:
-    """Post Request to the API
+  def post(
+    self, endpoint: str, model: T, headers: dict | None = None
+  ) -> NetResponse:
+    """Post Request to the API.
 
     Args:
         endpoint: The API endpoint to post to.
@@ -56,6 +58,9 @@ def post(self, endpoint: str, model: T, headers: dict = {}) -> NetResponse:
     Returns:
         HTTPResponse: The response from the API.
     """
+    if headers is None:
+      headers = {}
+
     data = model.model_dump()
     url = self._make_url(endpoint)
     headers = self._make_headers(headers)

sec-gemini-python/sec_gemini/http.py

@@ -1,5 +1,21 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import logging
 
+from rich.logging import RichHandler
+
 _LOGGER = None
 
 
@@ -8,4 +24,14 @@ def get_logger():
   if _LOGGER is None:
     _LOGGER = logging.getLogger("secgemini")
     _LOGGER.setLevel(level=logging.WARNING)
+
+    rich_handler = RichHandler(
+      show_time=True,  # Ensure timestamp is shown
+      show_path=False,  # Hides file path for a cleaner log
+    )
+
+    _LOGGER.addHandler(rich_handler)
+
+    _LOGGER.propagate = False
+
   return _LOGGER

sec-gemini-python/sec_gemini/logger.py

@@ -120,7 +120,9 @@ class ModelInfo(BaseModel):
 
   @staticmethod
   def parse_model_string(model_string: str) -> tuple[str, str, bool]:
-    """Parse a given model string and splits it in:
+    """Parse a given model string.
+
+    Parsing the model string means splitting in three parts:
     - model_name
     - version
     - use_experimental