Skip to content

Commit 5429811

Browse files
vorburgercopybara-github
authored andcommitted
fix: Remove copy/pasta 🍝 in Mcp[Aync]Tool
PiperOrigin-RevId: 795476458
1 parent d094cc2 commit 5429811

File tree

3 files changed

+173
-169
lines changed

3 files changed

+173
-169
lines changed
Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.tools.mcp;
18+
19+
import com.fasterxml.jackson.core.JsonProcessingException;
20+
import com.fasterxml.jackson.core.type.TypeReference;
21+
import com.fasterxml.jackson.databind.ObjectMapper;
22+
import com.google.adk.tools.BaseTool;
23+
import com.google.common.collect.ImmutableMap;
24+
import com.google.genai.types.FunctionDeclaration;
25+
import com.google.genai.types.Schema;
26+
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
27+
import io.modelcontextprotocol.spec.McpSchema.Content;
28+
import io.modelcontextprotocol.spec.McpSchema.JsonSchema;
29+
import io.modelcontextprotocol.spec.McpSchema.TextContent;
30+
import io.modelcontextprotocol.spec.McpSchema.Tool;
31+
import java.io.IOException;
32+
import java.util.ArrayList;
33+
import java.util.List;
34+
import java.util.Map;
35+
import java.util.Optional;
36+
37+
/**
38+
* Base class for MCP tools.
39+
*
40+
* @param <T> The type of the MCP session client.
41+
*/
42+
public abstract class AbstractMcpTool<T> extends BaseTool {
43+
44+
protected final Tool mcpTool;
45+
protected final McpSessionManager mcpSessionManager;
46+
protected final ObjectMapper objectMapper;
47+
48+
// Volatile ensures write visibility in the asynchronous chain for McpAsyncTool.
49+
protected volatile T mcpSession;
50+
51+
protected AbstractMcpTool(
52+
Tool mcpTool, T mcpSession, McpSessionManager mcpSessionManager, ObjectMapper objectMapper) {
53+
super(
54+
mcpTool == null ? "" : mcpTool.name(),
55+
mcpTool == null ? "" : (mcpTool.description().isEmpty() ? "" : mcpTool.description()));
56+
57+
if (mcpTool == null) {
58+
throw new IllegalArgumentException("mcpTool cannot be null");
59+
}
60+
if (mcpSession == null) {
61+
throw new IllegalArgumentException("mcpSession cannot be null");
62+
}
63+
if (mcpSessionManager == null) {
64+
throw new IllegalArgumentException("mcpSessionManager cannot be null");
65+
}
66+
if (objectMapper == null) {
67+
throw new IllegalArgumentException("objectMapper cannot be null");
68+
}
69+
this.mcpTool = mcpTool;
70+
this.mcpSession = mcpSession;
71+
this.mcpSessionManager = mcpSessionManager;
72+
this.objectMapper = objectMapper;
73+
}
74+
75+
public T getMcpSession() {
76+
return this.mcpSession;
77+
}
78+
79+
protected Schema toGeminiSchema(JsonSchema openApiSchema) {
80+
try {
81+
return GeminiSchemaUtil.toGeminiSchema(openApiSchema, this.objectMapper);
82+
} catch (IOException | IllegalArgumentException e) {
83+
throw new IllegalArgumentException(
84+
"Error generating function declaration for tool '" + this.name() + "': " + e.getMessage(),
85+
e);
86+
}
87+
}
88+
89+
@Override
90+
public Optional<FunctionDeclaration> declaration() {
91+
try {
92+
Schema schema = toGeminiSchema(this.mcpTool.inputSchema());
93+
return Optional.ofNullable(schema)
94+
.map(
95+
value ->
96+
FunctionDeclaration.builder()
97+
.name(this.name())
98+
.description(this.description())
99+
.parameters(value)
100+
.build());
101+
} catch (IllegalArgumentException e) {
102+
System.err.println(e.getMessage());
103+
return Optional.empty();
104+
}
105+
}
106+
107+
@SuppressWarnings("PreferredInterfaceType") // BaseTool.runAsync() returns Map<String, Object>
108+
protected static Map<String, Object> wrapCallResult(
109+
ObjectMapper objectMapper, String mcpToolName, CallToolResult callResult) {
110+
if (callResult == null) {
111+
return ImmutableMap.of("error", "MCP framework error: CallToolResult was null");
112+
}
113+
114+
List<Content> contents = callResult.content();
115+
Boolean isToolError = callResult.isError();
116+
117+
if (isToolError != null && isToolError) {
118+
String errorMessage = "Tool execution failed.";
119+
if (contents != null
120+
&& !contents.isEmpty()
121+
&& contents.get(0) instanceof TextContent textContent) {
122+
if (textContent.text() != null && !textContent.text().isEmpty()) {
123+
errorMessage += " Details: " + textContent.text();
124+
}
125+
}
126+
return ImmutableMap.of("error", errorMessage);
127+
}
128+
129+
if (contents == null || contents.isEmpty()) {
130+
return ImmutableMap.of();
131+
}
132+
133+
List<String> textOutputs = new ArrayList<>();
134+
for (Content content : contents) {
135+
if (content instanceof TextContent textContent) {
136+
if (textContent.text() != null) {
137+
textOutputs.add(textContent.text());
138+
}
139+
}
140+
}
141+
142+
if (textOutputs.isEmpty()) {
143+
return ImmutableMap.of(
144+
"error",
145+
"Tool '" + mcpToolName + "' returned content that is not TextContent.",
146+
"content_details",
147+
contents.toString());
148+
}
149+
150+
List<Map<String, Object>> resultMaps = new ArrayList<>();
151+
for (String textOutput : textOutputs) {
152+
try {
153+
resultMaps.add(
154+
objectMapper.readValue(textOutput, new TypeReference<Map<String, Object>>() {}));
155+
} catch (JsonProcessingException e) {
156+
resultMaps.add(ImmutableMap.of("text", textOutput));
157+
}
158+
}
159+
return ImmutableMap.of("text_output", resultMaps);
160+
}
161+
}

core/src/main/java/com/google/adk/tools/mcp/McpAsyncTool.java

Lines changed: 5 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -20,20 +20,15 @@
2020

2121
import com.fasterxml.jackson.databind.ObjectMapper;
2222
import com.google.adk.JsonBaseModel;
23-
import com.google.adk.tools.BaseTool;
2423
import com.google.adk.tools.ToolContext;
2524
import com.google.common.collect.ImmutableMap;
26-
import com.google.genai.types.FunctionDeclaration;
27-
import com.google.genai.types.Schema;
2825
import io.modelcontextprotocol.client.McpAsyncClient;
2926
import io.modelcontextprotocol.spec.McpSchema;
3027
import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
31-
import io.modelcontextprotocol.spec.McpSchema.JsonSchema;
3228
import io.modelcontextprotocol.spec.McpSchema.Tool;
3329
import io.reactivex.rxjava3.core.Maybe;
3430
import io.reactivex.rxjava3.core.Single;
3531
import java.util.Map;
36-
import java.util.Optional;
3732
import org.slf4j.Logger;
3833
import org.slf4j.LoggerFactory;
3934

@@ -45,16 +40,10 @@
4540
* <p>This wraps a MCP Tool interface and an active MCP Session. It invokes the MCP Tool through
4641
* executing the tool from remote MCP Session.
4742
*/
48-
public final class McpAsyncTool extends BaseTool {
43+
public final class McpAsyncTool extends AbstractMcpTool<McpAsyncClient> {
4944

5045
private static final Logger logger = LoggerFactory.getLogger(McpAsyncTool.class);
5146

52-
Tool mcpTool;
53-
// Volatile ensures write visibility in the asynchronous chain.
54-
volatile McpAsyncClient mcpSession;
55-
McpSessionManager mcpSessionManager;
56-
ObjectMapper objectMapper;
57-
5847
/**
5948
* Creates a new McpAsyncTool with the default ObjectMapper.
6049
*
@@ -65,7 +54,7 @@ public final class McpAsyncTool extends BaseTool {
6554
*/
6655
public McpAsyncTool(
6756
Tool mcpTool, McpAsyncClient mcpSession, McpSessionManager mcpSessionManager) {
68-
this(mcpTool, mcpSession, mcpSessionManager, JsonBaseModel.getMapper());
57+
super(mcpTool, mcpSession, mcpSessionManager, JsonBaseModel.getMapper());
6958
}
7059

7160
/**
@@ -82,31 +71,7 @@ public McpAsyncTool(
8271
McpAsyncClient mcpSession,
8372
McpSessionManager mcpSessionManager,
8473
ObjectMapper objectMapper) {
85-
super(
86-
mcpTool == null ? "" : mcpTool.name(),
87-
mcpTool == null ? "" : (mcpTool.description().isEmpty() ? "" : mcpTool.description()));
88-
89-
if (mcpTool == null) {
90-
throw new IllegalArgumentException("mcpTool cannot be null");
91-
}
92-
if (mcpSession == null) {
93-
throw new IllegalArgumentException("mcpSession cannot be null");
94-
}
95-
if (objectMapper == null) {
96-
throw new IllegalArgumentException("objectMapper cannot be null");
97-
}
98-
this.mcpTool = mcpTool;
99-
this.mcpSession = mcpSession;
100-
this.mcpSessionManager = mcpSessionManager;
101-
this.objectMapper = objectMapper;
102-
}
103-
104-
public McpAsyncClient getMcpSession() {
105-
return this.mcpSession;
106-
}
107-
108-
public Schema toGeminiSchema(JsonSchema openApiSchema) {
109-
return Schema.fromJson(objectMapper.valueToTree(openApiSchema).toString());
74+
super(mcpTool, mcpSession, mcpSessionManager, objectMapper);
11075
}
11176

11277
private Single<McpSchema.InitializeResult> reintializeSession() {
@@ -129,16 +94,6 @@ private Single<McpSchema.InitializeResult> reintializeSession() {
12994
.toFuture());
13095
}
13196

132-
@Override
133-
public Optional<FunctionDeclaration> declaration() {
134-
return Optional.of(
135-
FunctionDeclaration.builder()
136-
.name(this.name())
137-
.description(this.description())
138-
.parameters(toGeminiSchema(this.mcpTool.inputSchema()))
139-
.build());
140-
}
141-
14297
@Override
14398
public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
14499
return Single.defer(
@@ -147,12 +102,10 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
147102
this.mcpSession
148103
.callTool(new CallToolRequest(this.name(), ImmutableMap.copyOf(args)))
149104
.toFuture())
150-
.map(
151-
callResult ->
152-
McpTool.wrapCallResult(this.objectMapper, this.name(), callResult))
105+
.map(callResult -> wrapCallResult(this.objectMapper, this.name(), callResult))
153106
.switchIfEmpty(
154107
Single.fromCallable(
155-
() -> McpTool.wrapCallResult(this.objectMapper, this.name(), null))))
108+
() -> wrapCallResult(this.objectMapper, this.name(), null))))
156109
.retryWhen(
157110
errors ->
158111
errors

0 commit comments

Comments
 (0)