extra metadata tests

run.sh

@@ -136,6 +136,9 @@ integration_tests() {
     sleep 2
   done
 
+  # Delete existing dist files
+  rm -rf dist
+
   # Generate latest whl file for the invariant-gateway package. 
   # This is required to run the integration tests.
   pip install build

tests/integration/mcp/test_mcp_stdio.py

@@ -28,6 +28,7 @@ async def test_mcp_stdio_with_gateway(
         push_to_explorer=push_to_explorer,
         tool_name="get_last_message_from_user",
         tool_args={"username": "Alice"},
+        metadata_keys={"my-custom-key": "value1", "my-custom-key-2": "value2"},
     )
 
     assert result.isError is False
@@ -58,6 +59,10 @@ async def test_mcp_stdio_with_gateway(
             and metadata["mcp_client"] == "mcp"
             and metadata["mcp_server"] == "messenger_server"
         )
+        # ensure custom keys are present
+        assert metadata["my-custom-key"] == "value1"
+        assert metadata["my-custom-key-2"] == "value2"
+        
         assert trace["messages"][0]["role"] == "assistant"
         assert trace["messages"][0]["tool_calls"][0]["function"] == {
             "name": "get_last_message_from_user",

tests/integration/resources/mcp/stdio/client/main.py

@@ -23,6 +23,7 @@ class MCPClient:
         project_name: str,
         server_script_path: str,
         push_to_explorer: bool,
+        metadata_keys: Optional[dict[str, str]] = None,
     ):
         """
         Connect to an MCP server.
@@ -42,6 +43,12 @@ class MCPClient:
             "--project-name",
             project_name,
         ]
+        # add metadata cli args
+        if metadata_keys is not None:
+            for key, value in metadata_keys.items():
+                args.append("--metadata-" + key + "=" + value)
+
+
         if push_to_explorer:
             args.append("--push-explorer")
         args.extend(
@@ -54,6 +61,7 @@ class MCPClient:
                 os.path.basename(server_script_path),
             ]
         )
+
         server_params = StdioServerParameters(
             command="uvx",
             args=args,
@@ -103,6 +111,7 @@ async def run(
     push_to_explorer: bool,
     tool_name: str,
     tool_args: dict[str, Any],
+    metadata_keys: Optional[dict[str, str]] = None,
 ) -> types.CallToolResult:
     """
     Main function to setup the MCP client and server.
@@ -125,6 +134,7 @@ async def run(
             project_name,
             server_script_path,
             push_to_explorer,
+            metadata_keys=metadata_keys
         )
         return await client.call_tool(tool_name, tool_args)
     finally: