Add MCP (Model Context Protocol) client implementation (Phase 2) #161

xerial · 2025-06-30T03:55:41Z

Summary

This PR implements Phase 2 of the Model Context Protocol (MCP) support (#159) by adding the MCP client implementation with stdio transport. This enables scala-ai to communicate with MCP servers via JSON-RPC protocol.

What's Implemented

JSON-RPC Protocol Layer

JsonRpc object with support for Request/Response/Notification messages
Standard JSON-RPC 2.0 error codes and message parsing
Type-safe message construction with proper serialization

MCP Protocol Messages

MCPMessages with all MCP-specific message types
Initialize, ListTools, and CallTool protocol support
Content types for tool results (text, image, resource)

Client Implementation

MCPClient trait defining the client interface
StdioMCPClient for process-based MCP server communication
Asynchronous message handling with Rx streams
Proper process lifecycle management

Tool Integration

MCPToolExecutor bridges MCP tools with ToolExecutor interface
Automatic tool discovery from MCP servers
JSON Schema parameter extraction and type mapping
Support for multiple content types in tool results

Features

✅ Full JSON-RPC 2.0 protocol support
✅ Stdio transport for subprocess-based MCP servers
✅ Tool discovery and automatic parameter extraction
✅ Asynchronous execution using Airframe Rx
✅ Proper error handling and recovery
✅ Comprehensive test coverage

Usage Example

// Create MCP client for a stdio-based server
val client = StdioMCPClient("mcp-server", args = Seq("--port", "8080"))

// Create tool executor
val executor = MCPToolExecutor(client)
executor.initialize().map { _ =>
  // Tools are now available
  val tools = executor.availableTools
  println(s"Discovered ${tools.size} tools")
}

// Or use the factory method
MCPToolExecutor.fromCommand("mcp-server").map { executor =>
  // Ready to execute tools
  val result = executor.executeToolCall(
    ToolCallRequest("1", "get_weather", Map("location" -> "Tokyo"))
  )
}

Testing

Added comprehensive test suites:

JsonRpcTest - Tests JSON-RPC message parsing and serialization
MCPMessagesTest - Tests MCP protocol messages
MCPToolExecutorTest - Tests tool discovery and execution with mock client

All tests pass:

[info] Passed: Total 19, Failed 0, Errors 0, Passed 19

Next Steps

This completes Phase 2 of MCP support. Next phases will include:

Phase 3: Tool adaptation and parameter validation
Phase 4: Integration tests with real MCP servers

Related Issues

Support MCP (Model Context Protocol) tool calling #159 - Support MCP (Model Context Protocol) tool calling

🤖 Generated with Claude Code

…e 2) - Implement JSON-RPC 2.0 protocol handling with request/response/notification types - Create StdioMCPClient for communication with process-based MCP servers - Add MCPToolExecutor to bridge MCP tools with ToolExecutor interface - Support tool discovery and automatic parameter extraction from JSON Schema - Handle multiple content types (text, image, resource) in tool results - Add comprehensive tests for JSON-RPC parsing and MCP functionality This is Phase 2 of #159 - MCP client implementation with stdio transport. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>

gemini-code-assist

Summary of Changes

Hello @xerial, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly expands the scala-ai project's capabilities by introducing a robust client-side implementation for the Model Context Protocol (MCP). It establishes the foundational communication layer using JSON-RPC over stdio, allowing scala-ai to discover and execute tools provided by external MCP servers. This integration streamlines the process of extending scala-ai with new functionalities by connecting to a broader ecosystem of MCP-compliant tools.

Highlights

Model Context Protocol (MCP) Client Implementation: I've added the core client-side implementation for the Model Context Protocol (MCP), enabling scala-ai to communicate with MCP servers. This is Phase 2 of the MCP support, focusing on the client aspect.
JSON-RPC 2.0 Protocol Layer: A new JsonRpc object has been introduced to handle JSON-RPC 2.0 messages, including requests, responses, notifications, and error objects. It provides type-safe message construction and parsing capabilities.
Stdio Transport for MCP: The StdioMCPClient provides a concrete implementation of the MCPClient trait, allowing communication with MCP servers via standard input/output streams. This is crucial for subprocess-based MCP server interactions.
MCP Tool Integration: The MCPToolExecutor bridges MCP tools with the existing ToolExecutor interface. It handles automatic tool discovery from MCP servers, extracts JSON Schema parameters, and supports various content types (text, image, resource) in tool results.
Asynchronous Communication with Airframe Rx: All client-server interactions are designed to be asynchronous, leveraging Airframe Rx streams for reactive handling of messages and responses, ensuring non-blocking operations.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

The pull request implements the MCP client with stdio transport. The feedback focuses on improving robustness by suggesting safer alternatives to type casts and string manipulations.

gemini-code-assist · 2025-06-30T03:58:49Z

ai-agent/src/main/scala/wvlet/ai/agent/mcp/MCPToolExecutor.scala

+              contentMap.get("type") match
+                case Some("text") =>
+                  contentMap.getOrElse("text", "").toString
+                case Some("image") =>
+                  val data     = contentMap.getOrElse("data", "").toString
+                  val mimeType = contentMap.getOrElse("mimeType", "").toString
+                  s"""{"type": "image", "mimeType": "$mimeType", "data": "$data"}"""
+                case Some("resource") =>
+                  contentMap.get("resource") match
+                    case Some(res: Map[String, Any]) =>
+                      val uri = res.getOrElse("uri", "").toString
+                      s"""{"type": "resource", "uri": "$uri"}"""
+                    case _ =>
+                      """{"type": "resource", "uri": ""}"""
+                case _ =>
+                  MessageCodec.toJson(contentMap)
+            }


Instead of constructing JSON strings via string interpolation, use MessageCodec.toJson to serialize the content map into a valid JSON string. This ensures proper escaping of special characters and prevents malformed JSON.

contentMap.get("type") match case Some("text") => contentMap.getOrElse("text", "").toString case _ => MessageCodec.toJson(contentMap)

gemini-code-assist · 2025-06-30T03:58:50Z

ai-agent/src/main/scala/wvlet/ai/agent/mcp/MCPToolExecutor.scala

+          ToolResultMessage(
+            id = toolCall.id,
+            toolName = toolCall.name,
+            text = s"""{"error": "${e.getMessage}"}"""


Use MessageCodec.toJson to serialize the error message to ensure proper escaping and valid JSON, instead of string interpolation.

text = MessageCodec.toJson(Map("error" -> e.getMessage))

gemini-code-assist · 2025-06-30T03:58:50Z

ai-agent/src/main/scala/wvlet/ai/agent/mcp/MCPToolExecutor.scala

+    schema.get("properties") match
+      case Some(props: Map[String, Any]) =>
+        val required =
+          schema.get("required") match
+            case Some(req: List[String]) =>
+              req.toSet
+            case _ =>
+              Set.empty[String]
+
+        props
+          .map { case (name, propSchemaAny) =>
+            val propSchema = propSchemaAny.asInstanceOf[Map[String, Any]]
+            val description =
+              propSchema.get("description") match
+                case Some(desc: String) =>
+                  desc
+                case _ =>
+                  ""
+
+            val paramType: DataType =
+              propSchema.get("type") match
+                case Some(t: String) =>
+                  t match
+                    case "string" =>
+                      DataType.StringType
+                    case "number" =>
+                      DataType.FloatType
+                    case "integer" =>
+                      DataType.IntegerType
+                    case "boolean" =>
+                      DataType.BooleanType
+                    case "array" =>
+                      DataType.ArrayType(DataType.AnyType)
+                    case "object" =>
+                      DataType.JsonType
+                    case _ =>
+                      DataType.JsonType
+                case _ =>
+                  DataType.JsonType
+
+            val defaultValue = propSchema.get("default")
+
+            ToolParameter(
+              name = name,
+              description = description,
+              dataType = paramType,
+              defaultValue = defaultValue
+            )
+          }
+          .toList
+
+      case _ =>
+        List.empty
+


The method uses asInstanceOf calls which could lead to runtime exceptions if the MCP server provides a schema that doesn't match the expected structure. Use pattern matching with collect or match blocks to safely check and cast types.

schema.get("properties").collect { case m: Map[_, _] => m.asInstanceOf[Map[String, Any]] } match case Some(props) => val required = schema.get("required").collect { case l: List[_] => l.map(_.toString).toSet }.getOrElse(Set.empty) props.flatMap { case (name, propSchemaAny) => propSchemaAny match case propSchema: Map[_, _] => val propMap = propSchema.asInstanceOf[Map[String, Any]] val description = propMap.get("description").map(_.toString).getOrElse("") val paramType: DataType = propMap.get("type").map(_.toString) match case Some("string") => DataType.StringType case Some("number") => DataType.FloatType case Some("integer") => DataType.IntegerType case Some("boolean") => DataType.BooleanType case Some("array") => DataType.ArrayType(DataType.AnyType) case Some("object") => DataType.JsonType case _ => DataType.JsonType val defaultValue = propMap.get("default") Some(ToolParameter( name = name, description = description, dataType = paramType, defaultValue = defaultValue )) case _ => None }.toList case _ => List.empty

gemini-code-assist · 2025-06-30T03:58:50Z

ai-agent/src/main/scala/wvlet/ai/agent/mcp/StdioMCPClient.scala

+      response.result match
+        case Some(result) =>
+          val codec = MessageCodec.of[InitializeResult]
+          codec.fromMap(result.asInstanceOf[Map[String, Any]])
+        case None =>
+          throw new RuntimeException(s"Initialize failed: ${response.error}")


The code casts the result from a JSON-RPC response to Map[String, Any] without verifying its type. Use pattern matching to ensure the result is a Map before casting, to prevent ClassCastException.

response.result match case Some(result: Map[_, _]) => val codec = MessageCodec.of[InitializeResult] codec.fromMap(result.asInstanceOf[Map[String, Any]]) case Some(other) => throw new RuntimeException(s"Initialize failed: unexpected result type. Expected a JSON object, but got $other") case None => throw new RuntimeException(s"Initialize failed: ${response.error}")

gemini-code-assist · 2025-06-30T03:58:50Z

ai-agent/src/main/scala/wvlet/ai/agent/mcp/MCPMessages.scala

+    */
+  def createCallToolRequest(toolName: String, arguments: Map[String, Any]): JsonRpc.Request =
+    JsonRpc.Request(
+      id = Some(s"call-$toolName-${System.currentTimeMillis()}"),


Consider using java.util.UUID.randomUUID() for generating unique request IDs instead of System.currentTimeMillis(), to avoid potential collisions if requests are made in rapid succession.

id = Some(s"call-$toolName-${java.util.UUID.randomUUID()}"),

- Use MessageCodec.toJson for proper JSON serialization instead of string interpolation - Add type safety with pattern matching for JSON-RPC response parsing - Replace unsafe asInstanceOf calls with safe pattern matching in extractParameters - Use UUID.randomUUID() instead of timestamp for unique request IDs - Improve error messages for unexpected result types 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>

- Remove Either return type from JsonRpc.parse method - Throw AIException with StatusCode.INVALID_MESSAGE_TYPE for parsing errors - Update StdioMCPClient to use try-catch instead of pattern matching on Either - Update tests to use intercept[AIException] for error cases - Add test for invalid JSON-RPC message format This change aligns with the codebase convention of using exceptions rather than Either types for error handling. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>

- Add new JSON_RPC_ERROR status code as a UserError type - Update JsonRpc.parse to use JSON_RPC_ERROR instead of INVALID_MESSAGE_TYPE - This provides more specific error categorization for JSON-RPC protocol errors 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>

@throws

- Add import for AIException and StatusCode - Replace fully qualified StatusCode references with imported name - Update @throws documentation to use short name 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>

gemini-code-assist bot reviewed Jun 30, 2025

View reviewed changes

xerial and others added 4 commits June 30, 2025 13:03

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add MCP (Model Context Protocol) client implementation (Phase 2) #161

Add MCP (Model Context Protocol) client implementation (Phase 2) #161

Uh oh!

xerial commented Jun 30, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jun 30, 2025

Uh oh!

gemini-code-assist bot Jun 30, 2025

Uh oh!

gemini-code-assist bot Jun 30, 2025

Uh oh!

gemini-code-assist bot Jun 30, 2025

Uh oh!

gemini-code-assist bot Jun 30, 2025

Uh oh!

Uh oh!

Add MCP (Model Context Protocol) client implementation (Phase 2) #161

Are you sure you want to change the base?

Add MCP (Model Context Protocol) client implementation (Phase 2) #161

Uh oh!

Conversation

xerial commented Jun 30, 2025

Summary

What's Implemented

JSON-RPC Protocol Layer

MCP Protocol Messages

Client Implementation

Tool Integration

Features

Usage Example

Testing

Next Steps

Related Issues

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jun 30, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jun 30, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jun 30, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jun 30, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jun 30, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!