From 3a2e2796c8f815dd4cfd1d86253745c8bb601be8 Mon Sep 17 00:00:00 2001
From: mikewong23571 <mikewong23571@gmail.com>
Date: Fri, 13 Mar 2026 22:45:00 +0800
Subject: [PATCH] fix: correct storage local key lookup parsing and text output
 (#761)

---
 cli/src/commands.rs | 28 ++++++++++++++---
 cli/src/output.rs   | 74 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+), 4 deletions(-)

diff --git a/cli/src/commands.rs b/cli/src/commands.rs
index f4b57ea..a898339 100644
--- a/cli/src/commands.rs
+++ b/cli/src/commands.rs
@@ -2084,10 +2084,14 @@ fn parse_storage(rest: &[&str], id: &str) -> Result<Value, ParseError> {
     match rest.first().copied() {
         Some("local") | Some("session") => {
             let storage_type = rest.first().unwrap();
-            let op = rest.get(1).unwrap_or(&"get");
-            let key = rest.get(2);
-            let value = rest.get(3);
-            match *op {
+            let (op, key, value) = match rest.get(1) {
+                Some(&"get") => ("get", rest.get(2), rest.get(3)),
+                Some(&"set") => ("set", rest.get(2), rest.get(3)),
+                Some(&"clear") => ("clear", rest.get(2), rest.get(3)),
+                Some(_) => ("get", rest.get(1), rest.get(2)),
+                None => ("get", None, None),
+            };
+            match op {
                 "set" => {
                     let k = key.ok_or_else(|| ParseError::MissingArguments {
                         context: format!("storage {} set", storage_type),
@@ -2365,6 +2369,14 @@ mod tests {
         assert_eq!(cmd["key"], "mykey");
     }
 
+    #[test]
+    fn test_storage_local_get_implicit_key() {
+        let cmd = parse_command(&args("storage local mykey"), &default_flags()).unwrap();
+        assert_eq!(cmd["action"], "storage_get");
+        assert_eq!(cmd["type"], "local");
+        assert_eq!(cmd["key"], "mykey");
+    }
+
     #[test]
     fn test_storage_session_get() {
         let cmd = parse_command(&args("storage session"), &default_flags()).unwrap();
@@ -2372,6 +2384,14 @@ mod tests {
         assert_eq!(cmd["type"], "session");
     }
 
+    #[test]
+    fn test_storage_session_get_implicit_key() {
+        let cmd = parse_command(&args("storage session mykey"), &default_flags()).unwrap();
+        assert_eq!(cmd["action"], "storage_get");
+        assert_eq!(cmd["type"], "session");
+        assert_eq!(cmd["key"], "mykey");
+    }
+
     #[test]
     fn test_storage_local_set() {
         let cmd =
diff --git a/cli/src/output.rs b/cli/src/output.rs
index 9b47f54..346db8d 100644
--- a/cli/src/output.rs
+++ b/cli/src/output.rs
@@ -64,6 +64,31 @@ fn print_with_boundaries(content: &str, origin: Option<&str>, opts: &OutputOptio
     }
 }
 
+fn format_storage_value(value: &serde_json::Value) -> String {
+    value
+        .as_str()
+        .map(ToString::to_string)
+        .unwrap_or_else(|| serde_json::to_string(value).unwrap_or_default())
+}
+
+fn format_storage_text(data: &serde_json::Value) -> Option<String> {
+    if let Some(entries) = data.get("data").and_then(|v| v.as_object()) {
+        if entries.is_empty() {
+            return Some("No storage entries".to_string());
+        }
+
+        let lines = entries
+            .iter()
+            .map(|(key, value)| format!("{}: {}", key, format_storage_value(value)))
+            .collect::<Vec<_>>();
+        return Some(lines.join("\n"));
+    }
+
+    let key = data.get("key").and_then(|v| v.as_str())?;
+    let value = data.get("value")?;
+    Some(format!("{}: {}", key, format_storage_value(value)))
+}
+
 pub fn print_response_with_opts(resp: &Response, action: Option<&str>, opts: &OutputOptions) {
     if opts.json {
         if opts.content_boundaries {
@@ -100,6 +125,12 @@ pub fn print_response_with_opts(resp: &Response, action: Option<&str>, opts: &Ou
     }
 
     if let Some(data) = &resp.data {
+        if action == Some("storage_get") {
+            if let Some(output) = format_storage_text(data) {
+                println!("{}", output);
+                return;
+            }
+        }
         // Inspect response (check before generic URL handler since it also has a "url" field)
         if action == Some("inspect") {
             let opened = data
@@ -2727,3 +2758,46 @@ fn print_screenshot_diff(data: &serde_json::Map<String, serde_json::Value>) {
 pub fn print_version() {
     println!("agent-browser {}", env!("CARGO_PKG_VERSION"));
 }
+
+#[cfg(test)]
+mod tests {
+    use super::format_storage_text;
+    use serde_json::json;
+
+    #[test]
+    fn test_format_storage_text_for_all_entries() {
+        let data = json!({
+            "data": {
+                "token": "abc123",
+                "user": "alice"
+            }
+        });
+
+        let rendered = format_storage_text(&data).unwrap();
+
+        assert_eq!(rendered, "token: abc123\nuser: alice");
+    }
+
+    #[test]
+    fn test_format_storage_text_for_key_lookup() {
+        let data = json!({
+            "key": "token",
+            "value": "abc123"
+        });
+
+        let rendered = format_storage_text(&data).unwrap();
+
+        assert_eq!(rendered, "token: abc123");
+    }
+
+    #[test]
+    fn test_format_storage_text_for_empty_store() {
+        let data = json!({
+            "data": {}
+        });
+
+        let rendered = format_storage_text(&data).unwrap();
+
+        assert_eq!(rendered, "No storage entries");
+    }
+}