From d9a787f8b7cb60cf51da0347b881df13b8ea7c69 Mon Sep 17 00:00:00 2001
From: Richard Gebhardt <richard.gebhardt@oracle.com>
Date: Fri, 31 Oct 2025 16:01:32 -0400
Subject: [PATCH] fix: properly handle malformed command output

Signed-off-by: Richard Gebhardt <richard.gebhardt@oracle.com>
---
 .../oracle/oci_api_mcp_server/server.py       | 29 ++++++---
 .../tests/test_oci_api_tools.py               | 61 ++++++++++++++-----
 2 files changed, 67 insertions(+), 23 deletions(-)

diff --git a/src/oci-api-mcp-server/oracle/oci_api_mcp_server/server.py b/src/oci-api-mcp-server/oracle/oci_api_mcp_server/server.py
index 039c25f..2368378 100644
--- a/src/oci-api-mcp-server/oracle/oci_api_mcp_server/server.py
+++ b/src/oci-api-mcp-server/oracle/oci_api_mcp_server/server.py
@@ -161,17 +161,30 @@ def run_oci_command(
             check=True,
             shell=False,
         )
-        return (
-            json.loads(result.stdout)
-            if result.stdout
-            else {
-                "error": result.stderr,
-            }
-        )
+
+        result.check_returncode()
+
+        response = {
+            "command": command,
+            "output": result.stdout,
+            "error": result.stderr,
+            "returncode": result.returncode,
+        }
+
+        try:
+            response["output"] = json.loads(result.stdout)
+        except TypeError:
+            pass
+        except json.JSONDecodeError:
+            pass
+
+        return response
     except subprocess.CalledProcessError as e:
         return {
-            "error": e.stderr,
+            "command": command,
             "output": e.stdout,
+            "error": e.stderr,
+            "returncode": e.returncode,
         }
 
 
diff --git a/src/oci-api-mcp-server/oracle/oci_api_mcp_server/tests/test_oci_api_tools.py b/src/oci-api-mcp-server/oracle/oci_api_mcp_server/tests/test_oci_api_tools.py
index 28024ef..c563c4f 100644
--- a/src/oci-api-mcp-server/oracle/oci_api_mcp_server/tests/test_oci_api_tools.py
+++ b/src/oci-api-mcp-server/oracle/oci_api_mcp_server/tests/test_oci_api_tools.py
@@ -5,6 +5,7 @@
 """
 
 import importlib.metadata
+import json
 import subprocess
 from unittest.mock import ANY, MagicMock, patch
 
@@ -72,47 +73,77 @@ async def test_get_oci_command_help_failure(self, mock_run):
 
     @pytest.mark.asyncio
     @patch("oracle.oci_api_mcp_server.server.subprocess.run")
-    @patch("oracle.oci_api_mcp_server.server.json.loads")
-    async def test_run_oci_command_success(self, mock_json_loads, mock_run):
+    async def test_run_oci_command_success(self, mock_run):
+        command = "compute instance list"
+
         mock_result = MagicMock()
         mock_result.stdout = '{"key": "value"}'
         mock_result.stderr = ""
+        mock_result.returncode = 0
         mock_run.return_value = mock_result
-        mock_json_loads.return_value = {"key": "value"}
 
         async with Client(mcp) as client:
             result = (
-                await client.call_tool(
-                    "run_oci_command", {"command": "compute instance list"}
-                )
+                await client.call_tool("run_oci_command", {"command": command})
+            ).data
+
+            assert result == {
+                "command": command,
+                "output": json.loads(mock_result.stdout),
+                "error": mock_result.stderr,
+                "returncode": mock_result.returncode,
+            }
+
+    @pytest.mark.asyncio
+    @patch("oracle.oci_api_mcp_server.server.subprocess.run")
+    async def test_run_oci_command_string_success(self, mock_run):
+        command = "compute instance list"
+
+        mock_result = MagicMock()
+        mock_result.stdout = "This is not JSON"
+        mock_result.stderr = ""
+        mock_result.returncode = 0
+        mock_run.return_value = mock_result
+
+        async with Client(mcp) as client:
+            result = (
+                await client.call_tool("run_oci_command", {"command": command})
             ).data
 
-            assert result == {"key": "value"}
-            mock_json_loads.assert_called_once_with('{"key": "value"}')
+            assert result == {
+                "command": command,
+                "output": mock_result.stdout,
+                "error": mock_result.stderr,
+                "returncode": mock_result.returncode,
+            }
 
     @pytest.mark.asyncio
     @patch("oracle.oci_api_mcp_server.server.subprocess.run")
     async def test_run_oci_command_failure(self, mock_run):
+        command = "compute instance list"
+
         mock_result = MagicMock()
         mock_result.stdout = "Some output"
         mock_result.stderr = "Some error"
+        mock_result.returncode = 1
+
         mock_run.side_effect = subprocess.CalledProcessError(
-            returncode=1,
-            cmd=["oci", "compute", "instance", "list"],
+            returncode=mock_result.returncode,
+            cmd=["oci"] + command.split(),
             output=mock_result.stdout,
             stderr=mock_result.stderr,
         )
 
         async with Client(mcp) as client:
             result = (
-                await client.call_tool(
-                    "run_oci_command", {"command": "compute instance list"}
-                )
+                await client.call_tool("run_oci_command", {"command": command})
             ).data
 
             assert result == {
-                "error": "Some error",
-                "output": "Some output",
+                "command": command,
+                "output": mock_result.stdout,
+                "error": mock_result.stderr,
+                "returncode": mock_result.returncode,
             }
 
     @pytest.mark.asyncio