diff --git a/api/src/main/java/org/apache/flink/agents/api/resource/python/PythonResourceAdapter.java b/api/src/main/java/org/apache/flink/agents/api/resource/python/PythonResourceAdapter.java index 9c4ad0bb6..272e1eeb7 100644 --- a/api/src/main/java/org/apache/flink/agents/api/resource/python/PythonResourceAdapter.java +++ b/api/src/main/java/org/apache/flink/agents/api/resource/python/PythonResourceAdapter.java @@ -175,7 +175,8 @@ Map getPythonToolMetadata( * @param qualName the qualified name of the callable inside the module * @param kwargs keyword arguments to pass to the callable; LLM tool calls always arrive as * keyword arguments - * @return the raw return value from the Python callable + * @return the bridge-encoded tool result; implementations may return a raw value for backward + * compatibility */ Object invokePythonTool(String module, String qualName, Map kwargs); } diff --git a/docs/content/docs/development/tool_use.md b/docs/content/docs/development/tool_use.md index 12a35810c..c468ba833 100644 --- a/docs/content/docs/development/tool_use.md +++ b/docs/content/docs/development/tool_use.md @@ -336,6 +336,8 @@ See [MCP]({{< ref "docs/development/mcp" >}}) for details. The built-in `tool_call_action` listens to `ToolRequestEvent`. For each tool call, it looks up the tool resource by function name, executes it through durable execution, and records whether it succeeded. After all tool calls in the batch have been processed, it sends a `ToolResponseEvent`. +Python tools can continue returning raw values, which are treated as successful results. A tool that completes normally but cannot perform the requested operation can return `ToolResponse.error("reason")`; an exception still represents an invocation failure. Both failure forms are recorded as failed tool calls in `ToolResponseEvent`. + When the tool request comes from `chat_model_action`, the emitted `ToolResponseEvent` is automatically consumed by `chat_model_action` to continue the chat. See [Built-in Events and Actions in Chat Models]({{< ref "docs/development/chat_models#built-in-events-and-actions" >}}) for details on how `chat_model_action` handles tool responses. Users can also send `ToolRequestEvent` directly when they want to invoke tools programmatically. diff --git a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/MCPCrossLanguageAgent.java b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/MCPCrossLanguageAgent.java index 30b921e8c..74ecfecbc 100644 --- a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/MCPCrossLanguageAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/MCPCrossLanguageAgent.java @@ -59,6 +59,13 @@ public static void process(Event event, RunnerContext ctx) throws Exception { ToolResponse response = add.call(new ToolParameters(Map.of("a", 1, "b", 2))); Assertions.assertTrue(response.getResult().toString().contains("3")); + + Tool failingTool = (Tool) ctx.getResource("fail_with_recovery_hint", ResourceType.TOOL); + ToolResponse failedResponse = + failingTool.call(new ToolParameters(Map.of("query", "all records"))); + Assertions.assertTrue(failedResponse.isError()); + Assertions.assertTrue( + failedResponse.getError().contains("retry with a narrower query")); System.out.println("[TEST] MCP Tools PASSED"); Prompt askSum = (Prompt) ctx.getResource("ask_sum", ResourceType.PROMPT); diff --git a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/resources/mcp_server.py b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/resources/mcp_server.py index 5c1d90a81..1f8bb7c62 100644 --- a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/resources/mcp_server.py +++ b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/resources/mcp_server.py @@ -50,4 +50,10 @@ async def add(a: int, b: int) -> int: return a + b +@mcp.tool() +async def fail_with_recovery_hint(query: str) -> str: + """Return a protocol-level tool error with a model-facing recovery hint.""" + raise ValueError(f"retry with a narrower query than '{query}'") + + mcp.run("streamable-http") diff --git a/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPServer.java b/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPServer.java index 5692de7a0..217bdecfc 100644 --- a/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPServer.java +++ b/integrations/mcp/src/main/java/org/apache/flink/agents/integrations/mcp/MCPServer.java @@ -443,26 +443,37 @@ public ToolMetadata getToolMetadata(String name) { * @param toolName The name of the tool to call * @param arguments The arguments to pass to the tool * @return The result as a list of content items + * @throws IllegalStateException if the MCP response reports a protocol-level tool error */ public List callTool(String toolName, Map arguments) { - return getRetryExecutor() - .execute( - () -> { - McpSyncClient mcpClient = getClient(); - McpSchema.CallToolRequest request = - new McpSchema.CallToolRequest( - toolName, - arguments != null ? arguments : new HashMap<>()); - McpSchema.CallToolResult result = mcpClient.callTool(request); - - List content = new ArrayList<>(); - for (var item : result.content()) { - content.add(MCPContentExtractor.extractContentItem(item)); - } + McpSchema.CallToolResult result = + getRetryExecutor() + .execute( + () -> { + McpSyncClient mcpClient = getClient(); + McpSchema.CallToolRequest request = + new McpSchema.CallToolRequest( + toolName, + arguments != null + ? arguments + : new HashMap<>()); + return mcpClient.callTool(request); + }, + "callTool:" + toolName); + return extractToolContent(toolName, result); + } - return content; - }, - "callTool:" + toolName); + static List extractToolContent( + String toolName, McpSchema.CallToolResult callToolResult) { + List content = new ArrayList<>(); + for (var item : callToolResult.content()) { + content.add(MCPContentExtractor.extractContentItem(item)); + } + if (Boolean.TRUE.equals(callToolResult.isError())) { + throw new IllegalStateException( + "MCP tool '" + toolName + "' returned an error: " + content); + } + return content; } /** diff --git a/integrations/mcp/src/test/java/org/apache/flink/agents/integrations/mcp/MCPServerTest.java b/integrations/mcp/src/test/java/org/apache/flink/agents/integrations/mcp/MCPServerTest.java index c63bbd187..847fd2248 100644 --- a/integrations/mcp/src/test/java/org/apache/flink/agents/integrations/mcp/MCPServerTest.java +++ b/integrations/mcp/src/test/java/org/apache/flink/agents/integrations/mcp/MCPServerTest.java @@ -19,6 +19,7 @@ package org.apache.flink.agents.integrations.mcp; import com.fasterxml.jackson.databind.ObjectMapper; +import io.modelcontextprotocol.spec.McpSchema; import org.apache.flink.agents.api.resource.ResourceContext; import org.apache.flink.agents.api.resource.ResourceDescriptor; import org.apache.flink.agents.api.resource.ResourceName; @@ -306,4 +307,28 @@ void testListPromptsReturnsEmptyWhenNotSupported() { List prompts = server.listPrompts(); assertThat(prompts).isEmpty(); } + + @Test + @DisabledOnJre(JRE.JAVA_11) + void protocolToolErrorIsNotReturnedAsSuccessfulContent() { + McpSchema.CallToolResult result = + McpSchema.CallToolResult.builder() + .addTextContent("business failure") + .isError(true) + .build(); + + assertThatThrownBy(() -> MCPServer.extractToolContent("lookup", result)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("lookup") + .hasMessageContaining("business failure"); + } + + @Test + @DisabledOnJre(JRE.JAVA_11) + void successfulProtocolToolResultReturnsContent() { + McpSchema.CallToolResult result = + McpSchema.CallToolResult.builder().addTextContent("result").build(); + + assertThat(MCPServer.extractToolContent("lookup", result)).containsExactly("result"); + } } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java index 8a831db5a..ecd16286a 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java @@ -39,6 +39,8 @@ public class PythonMCPTool extends Tool implements PythonResourceWrapper, ToolExecutionMetadataProvider { private static final String GET_JAVA_TOOL_META = "python_java_utils.get_java_tool_metadata_from_tool"; + private static final String INVOKE_PYTHON_TOOL = + "python_java_utils.invoke_python_tool_instance"; private final PyObject tool; private final PythonResourceAdapter adapter; @Nullable private final String mcpServerName; @@ -84,8 +86,8 @@ public ToolResponse call(ToolParameters parameters) { kwargs.put(paramName, parameters.getParameter(paramName)); } try { - Object result = adapter.callMethod(tool, "call", kwargs); - return ToolResponse.success(result); + Object result = adapter.invoke(INVOKE_PYTHON_TOOL, tool, kwargs); + return PythonToolResultConverter.fromBridgeResult(result); } catch (Exception e) { return ToolResponse.error(e); } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonToolResultConverter.java b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonToolResultConverter.java new file mode 100644 index 000000000..0e1218782 --- /dev/null +++ b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonToolResultConverter.java @@ -0,0 +1,63 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.agents.plan.resource.python; + +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.annotation.Internal; + +import java.util.Map; + +/** Converts the internal Python bridge representation into a Java {@link ToolResponse}. */ +@Internal +public final class PythonToolResultConverter { + + private static final String RESULT_MARKER = "__flink_agents_tool_result__"; + + public static ToolResponse fromBridgeResult(Object result) { + if (!(result instanceof Map)) { + return ToolResponse.success(result); + } + + Map response = (Map) result; + Object resultKind = response.get(RESULT_MARKER); + if ("raw".equals(resultKind)) { + return ToolResponse.success(response.get("result")); + } + if (!"response".equals(resultKind)) { + return ToolResponse.success(result); + } + + long executionTimeMs = numberValue(response.get("execution_time_ms")); + String toolName = stringValue(response.get("tool_name")); + if (Boolean.TRUE.equals(response.get("success"))) { + return ToolResponse.success(response.get("result"), executionTimeMs, toolName); + } + return ToolResponse.error(stringValue(response.get("error")), executionTimeMs, toolName); + } + + private static long numberValue(Object value) { + return value instanceof Number ? ((Number) value).longValue() : 0L; + } + + private static String stringValue(Object value) { + return value == null ? null : String.valueOf(value); + } + + private PythonToolResultConverter() {} +} diff --git a/plan/src/main/java/org/apache/flink/agents/plan/tools/FunctionTool.java b/plan/src/main/java/org/apache/flink/agents/plan/tools/FunctionTool.java index 83534d31b..4c0178af2 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/tools/FunctionTool.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/tools/FunctionTool.java @@ -36,6 +36,7 @@ import org.apache.flink.agents.plan.Function; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.PythonFunction; +import org.apache.flink.agents.plan.resource.python.PythonToolResultConverter; import org.apache.flink.agents.plan.tools.serializer.FunctionToolJsonDeserializer; import org.apache.flink.agents.plan.tools.serializer.FunctionToolJsonSerializer; @@ -57,7 +58,6 @@ public class FunctionTool extends Tool { private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); - private final Function function; private Map injectedArgs; @@ -185,7 +185,7 @@ private ToolResponse callPython(PythonFunction pf, ToolParameters parameters) { } Object result = pythonResourceAdapter.invokePythonTool(pf.getModule(), pf.getQualName(), kwargs); - return ToolResponse.success(result); + return PythonToolResultConverter.fromBridgeResult(result); } public Function getFunction() { diff --git a/plan/src/test/java/org/apache/flink/agents/plan/resource/python/PythonMCPToolTest.java b/plan/src/test/java/org/apache/flink/agents/plan/resource/python/PythonMCPToolTest.java new file mode 100644 index 000000000..c6ce28b84 --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/resource/python/PythonMCPToolTest.java @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.agents.plan.resource.python; + +import org.apache.flink.agents.api.resource.python.PythonResourceAdapter; +import org.apache.flink.agents.api.tools.ToolParameters; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.junit.jupiter.api.Test; +import pemja.core.object.PyObject; + +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class PythonMCPToolTest { + + @Test + void preservesExplicitPythonMcpToolFailure() { + PythonResourceAdapter adapter = mock(PythonResourceAdapter.class); + PyObject pythonTool = mock(PyObject.class); + when(adapter.invoke("python_java_utils.get_java_tool_metadata_from_tool", pythonTool)) + .thenReturn( + Map.of( + "name", "lookup", + "description", "Lookup a value.", + "inputSchema", "{\"type\":\"object\"}")); + when(adapter.invoke( + "python_java_utils.invoke_python_tool_instance", + pythonTool, + Map.of("query", "flink"))) + .thenReturn( + Map.of( + "__flink_agents_tool_result__", "response", + "success", false, + "error", "retry with a narrower query", + "execution_time_ms", 7L, + "tool_name", "lookup")); + PythonMCPTool tool = new PythonMCPTool(adapter, pythonTool, "search-server"); + + ToolResponse response = tool.call(new ToolParameters(Map.of("query", "flink"))); + + assertThat(response.isError()).isTrue(); + assertThat(response.getError()).isEqualTo("retry with a narrower query"); + assertThat(response.getExecutionTimeMs()).isEqualTo(7L); + assertThat(response.getToolName()).isEqualTo("lookup"); + verify(adapter) + .invoke( + "python_java_utils.invoke_python_tool_instance", + pythonTool, + Map.of("query", "flink")); + } +} diff --git a/plan/src/test/java/org/apache/flink/agents/plan/tools/FunctionToolSetPythonAdapterTest.java b/plan/src/test/java/org/apache/flink/agents/plan/tools/FunctionToolSetPythonAdapterTest.java index 88ab75375..876095c5f 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/tools/FunctionToolSetPythonAdapterTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/tools/FunctionToolSetPythonAdapterTest.java @@ -20,6 +20,8 @@ import org.apache.flink.agents.api.resource.python.PythonResourceAdapter; import org.apache.flink.agents.api.tools.ToolMetadata; import org.apache.flink.agents.api.tools.ToolParameterInjection; +import org.apache.flink.agents.api.tools.ToolParameters; +import org.apache.flink.agents.api.tools.ToolResponse; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.PythonFunction; import org.junit.jupiter.api.Test; @@ -37,6 +39,9 @@ class FunctionToolSetPythonAdapterTest { + private static final ToolMetadata PYTHON_TOOL_METADATA = + new ToolMetadata("notify", "Send a notification.", "{\"properties\":{}}"); + @Test void replacesPlaceholderMetadataForPythonFunction() { ToolMetadata placeholder = new ToolMetadata("notify", "", "{}"); @@ -111,6 +116,80 @@ void noOpForJavaFunction() throws Exception { .getPythonToolMetadata(Mockito.anyString(), Mockito.anyString(), anyList()); } + @Test + void preservesExplicitPythonToolFailure() { + PythonFunction function = new PythonFunction("pkg.mod", "notify"); + FunctionTool tool = new FunctionTool(PYTHON_TOOL_METADATA, function); + PythonResourceAdapter adapter = pythonAdapter(); + when(adapter.invokePythonTool(eq("pkg.mod"), eq("notify"), eq(Map.of("id", "1")))) + .thenReturn( + Map.of( + "__flink_agents_tool_result__", "response", + "success", false, + "error", "recipient not found", + "execution_time_ms", 7L, + "tool_name", "notify")); + tool.setPythonResourceAdapter(adapter); + + ToolResponse response = tool.call(new ToolParameters(Map.of("id", "1"))); + + assertThat(response.isError()).isTrue(); + assertThat(response.getError()).isEqualTo("recipient not found"); + assertThat(response.getExecutionTimeMs()).isEqualTo(7L); + assertThat(response.getToolName()).isEqualTo("notify"); + } + + @Test + void preservesExplicitPythonToolSuccess() { + PythonFunction function = new PythonFunction("pkg.mod", "notify"); + FunctionTool tool = new FunctionTool(PYTHON_TOOL_METADATA, function); + PythonResourceAdapter adapter = pythonAdapter(); + when(adapter.invokePythonTool(eq("pkg.mod"), eq("notify"), eq(Map.of("id", "1")))) + .thenReturn( + Map.of( + "__flink_agents_tool_result__", "response", + "result", "sent", + "success", true, + "execution_time_ms", 5L, + "tool_name", "notify")); + tool.setPythonResourceAdapter(adapter); + + ToolResponse response = tool.call(new ToolParameters(Map.of("id", "1"))); + + assertThat(response.isSuccess()).isTrue(); + assertThat(response.getResult()).isEqualTo("sent"); + assertThat(response.getExecutionTimeMs()).isEqualTo(5L); + assertThat(response.getToolName()).isEqualTo("notify"); + } + + @Test + void unwrapsRawPythonToolResultEnvelope() { + PythonFunction function = new PythonFunction("pkg.mod", "notify"); + FunctionTool tool = new FunctionTool(PYTHON_TOOL_METADATA, function); + PythonResourceAdapter adapter = pythonAdapter(); + Map rawResult = + Map.of("__flink_agents_tool_result__", "response", "value", "raw"); + when(adapter.invokePythonTool(eq("pkg.mod"), eq("notify"), eq(Map.of("id", "1")))) + .thenReturn(Map.of("__flink_agents_tool_result__", "raw", "result", rawResult)); + tool.setPythonResourceAdapter(adapter); + + ToolResponse response = tool.call(new ToolParameters(Map.of("id", "1"))); + + assertThat(response.isSuccess()).isTrue(); + assertThat(response.getResult()).isEqualTo(rawResult); + } + + private static PythonResourceAdapter pythonAdapter() { + PythonResourceAdapter adapter = Mockito.mock(PythonResourceAdapter.class); + when(adapter.getPythonToolMetadata(eq("pkg.mod"), eq("notify"), anyList())) + .thenReturn( + Map.of( + "name", "notify", + "description", "Send a notification.", + "inputSchema", "{\"properties\":{}}")); + return adapter; + } + /** Helper static method to back JavaFunction in the no-op test. */ public static int stubMethod(int x) { return x; diff --git a/python/flink_agents/api/tests/test_tool_response.py b/python/flink_agents/api/tests/test_tool_response.py new file mode 100644 index 000000000..73b9228fa --- /dev/null +++ b/python/flink_agents/api/tests/test_tool_response.py @@ -0,0 +1,36 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################# +from flink_agents.api.tools import ToolResponse + + +def test_tool_response_represents_success() -> None: + response = ToolResponse.success({"answer": 42}, tool_name="calculator") + + assert response.is_success() + assert not response.is_error() + assert response.result == {"answer": 42} + assert str(response) == "{'answer': 42}" + + +def test_tool_response_represents_failure() -> None: + response = ToolResponse.error("not found", tool_name="lookup") + + assert response.is_error() + assert not response.is_success() + assert response.error_message == "not found" + assert str(response) == "not found" diff --git a/python/flink_agents/api/tools/__init__.py b/python/flink_agents/api/tools/__init__.py index 78d07b8ba..90f8b9c8f 100644 --- a/python/flink_agents/api/tools/__init__.py +++ b/python/flink_agents/api/tools/__init__.py @@ -22,9 +22,11 @@ InjectedArg, ToolParameterSource, ) +from flink_agents.api.tools.tool_response import ToolResponse __all__ = [ "InjectedArg", "ToolExecutionMetadataProvider", "ToolParameterSource", + "ToolResponse", ] diff --git a/python/flink_agents/api/tools/tool.py b/python/flink_agents/api/tools/tool.py index 5e399dd13..9fcd28de4 100644 --- a/python/flink_agents/api/tools/tool.py +++ b/python/flink_agents/api/tools/tool.py @@ -157,4 +157,7 @@ def call( """Call the tools with arguments. This is the method that should be implemented by the tools' developer. + A raw return value represents success. Return + :class:`flink_agents.api.tools.ToolResponse` to report an explicit + tool-level failure without raising an exception. """ diff --git a/python/flink_agents/api/tools/tool_response.py b/python/flink_agents/api/tools/tool_response.py new file mode 100644 index 000000000..9d403d573 --- /dev/null +++ b/python/flink_agents/api/tools/tool_response.py @@ -0,0 +1,76 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################# +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class ToolResponse: + """Represents the result and status of one Python tool execution. + + Python tools may continue returning raw values, which the runtime treats as + successful results. Return ``ToolResponse.error(...)`` when a tool call + completed normally but the tool operation itself failed. + """ + + result: Any = None + error_message: str | None = None + execution_time_ms: int = 0 + tool_name: str | None = None + + @classmethod + def success( + cls, + result: Any, + execution_time_ms: int = 0, + tool_name: str | None = None, + ) -> "ToolResponse": + """Create a successful tool response.""" + return cls( + result=result, + execution_time_ms=execution_time_ms, + tool_name=tool_name, + ) + + @classmethod + def error( + cls, + error: str, + execution_time_ms: int = 0, + tool_name: str | None = None, + ) -> "ToolResponse": + """Create a failed tool response.""" + if error is None: + msg = "error cannot be None" + raise ValueError(msg) + return cls( + error_message=error, + execution_time_ms=execution_time_ms, + tool_name=tool_name, + ) + + def is_success(self) -> bool: + """Return whether the tool operation succeeded.""" + return self.error_message is None + + def is_error(self) -> bool: + """Return whether the tool operation failed.""" + return not self.is_success() + + def __str__(self) -> str: + return str(self.result) if self.is_success() else str(self.error_message) diff --git a/python/flink_agents/integrations/mcp/mcp.py b/python/flink_agents/integrations/mcp/mcp.py index aa80073f2..f6753760d 100644 --- a/python/flink_agents/integrations/mcp/mcp.py +++ b/python/flink_agents/integrations/mcp/mcp.py @@ -39,7 +39,7 @@ from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.prompts.prompt import Prompt from flink_agents.api.resource import Resource, ResourceType -from flink_agents.api.tools import ToolExecutionMetadataProvider +from flink_agents.api.tools import ToolExecutionMetadataProvider, ToolResponse from flink_agents.api.tools.tool import Tool, ToolMetadata, ToolType from flink_agents.api.trace import ( ToolExecutionMetadataKeys, @@ -65,12 +65,20 @@ def tool_type(cls) -> ToolType: def call(self, *args: Any, **kwargs: Any) -> Any: """Call the MCP tool with the given arguments.""" if self.mcp_server is None: - msg = "MCP tool call requires a reference to the MCP server" - raise ValueError(msg) + return ToolResponse.error( + "MCP tool call requires a reference to the MCP server", + tool_name=self.metadata.name, + ) - return asyncio.run( - self.mcp_server.call_tool_async(self.metadata.name, *args, **kwargs) - ) + try: + return asyncio.run( + self.mcp_server.call_tool_async(self.metadata.name, *args, **kwargs) + ) + except Exception as e: + return ToolResponse.error( + f"Error calling MCP tool '{self.metadata.name}': {e}", + tool_name=self.metadata.name, + ) @override def get_tool_execution_metadata( @@ -225,7 +233,11 @@ async def _cleanup_connection(self) -> None: pass async def call_tool_async(self, tool_name: str, *args: Any, **kwargs: Any) -> Any: - """Call a tool on the MCP server asynchronously.""" + """Call a tool on the MCP server asynchronously. + + Raises: + RuntimeError: If the MCP response marks the tool result as an error. + """ async with self._get_session() as session: arguments = kwargs if kwargs else (args[0] if args else {}) @@ -235,7 +247,10 @@ async def call_tool_async(self, tool_name: str, *args: Any, **kwargs: Any) -> An read_timeout_seconds=timedelta(seconds=self.timeout), ) - content = [extract_mcp_content_item(item) for item in result.content] + content = [extract_mcp_content_item(item) for item in result.content] + if result.isError: + msg = f"MCP tool '{tool_name}' returned an error: {content}" + raise RuntimeError(msg) return content diff --git a/python/flink_agents/integrations/mcp/tests/test_mcp.py b/python/flink_agents/integrations/mcp/tests/test_mcp.py index 1b013ebb3..39dd35b08 100644 --- a/python/flink_agents/integrations/mcp/tests/test_mcp.py +++ b/python/flink_agents/integrations/mcp/tests/test_mcp.py @@ -15,17 +15,25 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################# +import asyncio import multiprocessing import runpy import time +from contextlib import asynccontextmanager from pathlib import Path +from typing import AsyncIterator from urllib.parse import parse_qs, urlparse +import anyio +import pytest from mcp.client.auth import OAuthClientProvider, TokenStorage +from mcp.client.session import ClientSession from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken +from mcp.types import CallToolResult, TextContent from pydantic import AnyUrl from flink_agents.api.chat_message import ChatMessage, MessageRole +from flink_agents.api.tools import ToolResponse from flink_agents.api.tools.tool import ToolMetadata from flink_agents.api.trace import ToolExecutionMetadataKeys from flink_agents.integrations.mcp.mcp import MCPServer, MCPTool @@ -158,3 +166,74 @@ def test_mcp_tool_roundtrip_preserves_metadata() -> None: assert restored.get_tool_execution_metadata({}) == { ToolExecutionMetadataKeys.MCP_SERVER: "calculator_server" } + + +class _ProtocolErrorClientSession(ClientSession): + async def call_tool(self, *args: object, **kwargs: object) -> CallToolResult: + return CallToolResult( + content=[TextContent(type="text", text="business failure")], + isError=True, + ) + + +class _ProtocolErrorServer(MCPServer): + @asynccontextmanager + async def _get_session(self) -> AsyncIterator[ClientSession]: + server_send, client_receive = anyio.create_memory_object_stream(1) + client_send, server_receive = anyio.create_memory_object_stream(1) + async with ( + server_send, + client_receive, + client_send, + server_receive, + _ProtocolErrorClientSession(client_receive, client_send) as session, + ): + yield session + + +class _ProtocolSuccessSession: + async def call_tool(self, *args: object, **kwargs: object) -> CallToolResult: + return CallToolResult( + content=[TextContent(type="text", text="result")], + isError=False, + ) + + +class _ProtocolSuccessServer(MCPServer): + @asynccontextmanager + async def _get_session(self) -> AsyncIterator[_ProtocolSuccessSession]: + yield _ProtocolSuccessSession() + + +def test_mcp_protocol_error_maps_to_tool_failure() -> None: + server = _ProtocolErrorServer(endpoint="http://localhost/mcp") + + with pytest.raises(RuntimeError, match="business failure"): + asyncio.run(server.call_tool_async("lookup", query="flink")) + + tool = MCPTool( + metadata=ToolMetadata( + name="lookup", + description="Lookup a value.", + args_schema={"type": "object", "properties": {}}, + ), + mcp_server=server, + ) + response = tool.call(query="flink") + + assert isinstance(response, ToolResponse) + assert response.is_error() + assert "business failure" in response.error_message + + +def test_mcp_tool_success_preserves_raw_result() -> None: + tool = MCPTool( + metadata=ToolMetadata( + name="lookup", + description="Lookup a value.", + args_schema={"type": "object", "properties": {}}, + ), + mcp_server=_ProtocolSuccessServer(endpoint="http://localhost/mcp"), + ) + + assert tool.call(query="flink") == ["result"] diff --git a/python/flink_agents/plan/actions/tool_call_action.py b/python/flink_agents/plan/actions/tool_call_action.py index c6b68b56c..4952bfd8d 100644 --- a/python/flink_agents/plan/actions/tool_call_action.py +++ b/python/flink_agents/plan/actions/tool_call_action.py @@ -25,7 +25,7 @@ from flink_agents.api.memory_object import MemoryObject from flink_agents.api.resource import ResourceType from flink_agents.api.runner_context import DurableCall, Outcome, RunnerContext -from flink_agents.api.tools import ToolExecutionMetadataProvider +from flink_agents.api.tools import ToolExecutionMetadataProvider, ToolResponse from flink_agents.api.tools.tool_parameter_injection import ( InjectedArg, ToolParameterSource, @@ -247,14 +247,7 @@ async def _execute_sequentially( *call.args, **(call.kwargs or {}), ) - responses[execution.id] = response - success[execution.id] = True - ExecutionReporters.succeeded( - ctx, - ExecutionEntityTypes.TOOL, - execution.name, - execution.entity_metadata, - ) + _record_tool_response(execution, response, ctx, responses, success, error) except Exception as e: # noqa: PERF203 _record_execution_exception(execution, e, ctx, responses, success, error) @@ -272,14 +265,41 @@ def _record_outcome( execution, outcome.error, ctx, responses, success, error ) else: - responses[execution.id] = outcome.value - success[execution.id] = True - ExecutionReporters.succeeded( + _record_tool_response(execution, outcome.value, ctx, responses, success, error) + + +def _record_tool_response( + execution: _ToolCallExecution, + value: Any, + ctx: RunnerContext, + responses: dict, + success: dict, + error: dict, +) -> None: + response = value if isinstance(value, ToolResponse) else ToolResponse.success(value) + if response.is_error(): + message = response.error_message + responses[execution.id] = message + success[execution.id] = False + error[execution.id] = message + ExecutionReporters.failed( ctx, ExecutionEntityTypes.TOOL, execution.name, execution.entity_metadata, + RuntimeError(message), + ExecutionProblemCategories.TOOL_CALL_FAILED, ) + return + + responses[execution.id] = response.result + success[execution.id] = True + ExecutionReporters.succeeded( + ctx, + ExecutionEntityTypes.TOOL, + execution.name, + execution.entity_metadata, + ) def _record_execution_exception( diff --git a/python/flink_agents/plan/function.py b/python/flink_agents/plan/function.py index 7988084ae..2d82e7f4f 100644 --- a/python/flink_agents/plan/function.py +++ b/python/flink_agents/plan/function.py @@ -24,10 +24,12 @@ from pydantic import BaseModel, PrivateAttr, model_serializer +from flink_agents.api.tools import ToolResponse from flink_agents.plan.utils import check_type_match # Global cache for PythonFunction instances to avoid repeated creation _PYTHON_FUNCTION_CACHE: Dict[Tuple[str, str], "PythonFunction"] = {} +_TOOL_RESULT_MARKER = "__flink_agents_tool_result__" logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" @@ -321,11 +323,13 @@ def __call__(self, *args: Tuple[Any, ...], **kwargs: Dict[str, Any]) -> Any: self.parameter_types, list(args), ) - return self._j_resource_adapter.invokeJavaTool( - self.qualname, - self.method_name, - self.parameter_types, - kwargs, + return _decode_java_tool_result( + self._j_resource_adapter.invokeJavaTool( + self.qualname, + self.method_name, + self.parameter_types, + kwargs, + ) ) def check_signature(self, *args: Tuple[Any, ...]) -> None: @@ -340,6 +344,22 @@ def check_signature(self, *args: Tuple[Any, ...]) -> None: raise TypeError(msg) +def _decode_java_tool_result(value: Any) -> Any: + if not isinstance(value, dict) or value.get(_TOOL_RESULT_MARKER) != "response": + return value + if value.get("success") is True: + return ToolResponse.success( + value.get("result"), + execution_time_ms=value.get("execution_time_ms", 0), + tool_name=value.get("tool_name"), + ) + return ToolResponse.error( + value.get("error"), + execution_time_ms=value.get("execution_time_ms", 0), + tool_name=value.get("tool_name"), + ) + + def call_python_function(module: str, qualname: str, func_args: Tuple[Any, ...]) -> Any: """Used to call a Python function in the Pemja environment. diff --git a/python/flink_agents/plan/tests/actions/test_tool_call_action.py b/python/flink_agents/plan/tests/actions/test_tool_call_action.py index cc25d2438..f119e31e7 100644 --- a/python/flink_agents/plan/tests/actions/test_tool_call_action.py +++ b/python/flink_agents/plan/tests/actions/test_tool_call_action.py @@ -24,7 +24,11 @@ from flink_agents.api.memory_object import MemoryObject from flink_agents.api.resource import ResourceType from flink_agents.api.runner_context import Outcome -from flink_agents.api.tools import InjectedArg, ToolExecutionMetadataProvider +from flink_agents.api.tools import ( + InjectedArg, + ToolExecutionMetadataProvider, + ToolResponse, +) from flink_agents.api.tools.tool import ToolType from flink_agents.api.trace import ( ExecutionEntityTypes, @@ -424,6 +428,24 @@ def test_tool_call_action_records_parallel_outcome_failure() -> None: assert response.error["call-2"] == "boom" +def test_tool_call_action_records_parallel_tool_response_failure() -> None: + config = AgentConfiguration({"tenant_id": "tenant-1"}) + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, True) + config.set(AgentExecutionOptions.TOOL_CALL_PARALLELISM, 4) + ctx = _Context(config=config) + ctx.durable_execute_all_async_outcomes = [ + Outcome.success("ok"), + Outcome.success(ToolResponse.error("business failure")), + ] + + asyncio.run(process_tool_request(tool_request("call-1", "call-2"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.success == {"call-1": True, "call-2": False} + assert response.responses["call-2"] == "business failure" + assert response.error["call-2"] == "business failure" + + def test_tool_call_action_uses_sync_when_async_disabled_multi_tool() -> None: config = AgentConfiguration({"tenant_id": "tenant-1"}) config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, False) @@ -544,6 +566,38 @@ def test_tool_call_reports_failed() -> None: assert args[4] == ExecutionProblemCategories.TOOL_CALL_FAILED +def test_tool_call_reports_explicit_tool_response_failure() -> None: + tool = MagicMock() + tool.tool_type.return_value = ToolType.FUNCTION + tool.call = MagicMock(return_value=ToolResponse.error("business failure")) + ctx, sent_events = trace_context(tool) + request = ToolRequestEvent(model="model-a", tool_calls=[trace_tool_call()]) + + asyncio.run(process_tool_request(request, ctx)) + + response = ToolResponseEvent.from_event(sent_events[0]) + assert response.responses["call-1"] == "business failure" + assert response.success["call-1"] is False + assert response.error["call-1"] == "business failure" + ctx.report_execution_failed.assert_called_once() + ctx.report_execution_succeeded.assert_not_called() + + +def test_tool_call_preserves_empty_tool_response_error() -> None: + tool = MagicMock() + tool.tool_type.return_value = ToolType.FUNCTION + tool.call = MagicMock(return_value=ToolResponse.error("")) + ctx, sent_events = trace_context(tool) + request = ToolRequestEvent(model="model-a", tool_calls=[trace_tool_call()]) + + asyncio.run(process_tool_request(request, ctx)) + + response = ToolResponseEvent.from_event(sent_events[0]) + assert response.responses["call-1"] == "" + assert response.success["call-1"] is False + assert response.error["call-1"] == "" + + def test_tool_call_includes_provider_metadata() -> None: class MetadataTool(ToolExecutionMetadataProvider): @staticmethod diff --git a/python/flink_agents/plan/tests/tools/test_function_tool.py b/python/flink_agents/plan/tests/tools/test_function_tool.py index 291de408e..dce0a841a 100644 --- a/python/flink_agents/plan/tests/tools/test_function_tool.py +++ b/python/flink_agents/plan/tests/tools/test_function_tool.py @@ -21,7 +21,7 @@ import pytest -from flink_agents.api.tools import InjectedArg +from flink_agents.api.tools import InjectedArg, ToolResponse from flink_agents.plan.function import JavaFunction, PythonFunction from flink_agents.plan.tools.function_tool import FunctionTool @@ -120,6 +120,7 @@ def _java_func() -> JavaFunction: parameter_types=["int", "int"], ) + _FAKE_JAVA_SCHEMA = json.dumps( { "type": "object", @@ -218,11 +219,19 @@ def test_java_function_tool_metadata_is_none_without_adapter() -> None: def test_java_function_tool_call_dispatches_through_adapter() -> None: tool = FunctionTool(func=_java_func()) adapter = _fake_adapter() + adapter.invokeJavaTool.return_value = { + "__flink_agents_tool_result__": "response", + "result": 1065, + "success": True, + "error": None, + "execution_time_ms": 7, + "tool_name": "add", + } tool.set_java_resource_adapter(adapter) result = tool.call(a=377, b=688) - assert result == 1065 + assert result == ToolResponse.success(1065, execution_time_ms=7, tool_name="add") adapter.invokeJavaTool.assert_called_once_with( "com.example.Tools", "add", @@ -231,6 +240,35 @@ def test_java_function_tool_call_dispatches_through_adapter() -> None: ) +def test_java_function_tool_preserves_error_response() -> None: + tool = FunctionTool(func=_java_func()) + adapter = _fake_adapter() + adapter.invokeJavaTool.return_value = { + "__flink_agents_tool_result__": "response", + "result": None, + "success": False, + "error": "calculation rejected", + "execution_time_ms": 9, + "tool_name": "add", + } + tool.set_java_resource_adapter(adapter) + + result = tool.call(a=377, b=688) + + assert result == ToolResponse.error( + "calculation rejected", execution_time_ms=9, tool_name="add" + ) + + +def test_java_function_tool_keeps_legacy_raw_adapter_result() -> None: + tool = FunctionTool(func=_java_func()) + adapter = _fake_adapter() + adapter.invokeJavaTool.return_value = 1065 + tool.set_java_resource_adapter(adapter) + + assert tool.call(a=377, b=688) == 1065 + + def test_java_function_tool_call_without_adapter_raises() -> None: tool = FunctionTool(func=_java_func()) with pytest.raises(RuntimeError, match="JVM resource adapter"): diff --git a/python/flink_agents/runtime/python_java_utils.py b/python/flink_agents/runtime/python_java_utils.py index e1ecac32b..d13db0b79 100644 --- a/python/flink_agents/runtime/python_java_utils.py +++ b/python/flink_agents/runtime/python_java_utils.py @@ -191,12 +191,39 @@ def invoke_python_tool(module: str, qual_name: str, kwargs: Dict[str, Any]) -> A Used by the Java-side ``PythonResourceAdapter.invokePythonTool`` so a Java host can dispatch a Python function tool from a Java chat model without the Python side - needing to know about Pemja's threading model. + needing to know about Pemja's threading model. The return value is wrapped in + an internal envelope so Java can distinguish a raw result from an explicit + ``ToolResponse`` without inspecting user payloads. """ from flink_agents.api.function import PythonFunction descriptor = PythonFunction(module=module, qualname=qual_name) - return descriptor.as_callable()(**kwargs) + result = descriptor.as_callable()(**kwargs) + return _encode_python_tool_result(result) + + +def invoke_python_tool_instance(tool: Tool, kwargs: Dict[str, Any]) -> Any: + """Invoke a Python Tool instance and encode its result for the Java bridge.""" + return _encode_python_tool_result(tool.call(**kwargs)) + + +def _encode_python_tool_result(result: Any) -> Dict[str, Any]: + """Encode raw values and explicit ToolResponses without inspecting user payloads.""" + from flink_agents.api.tools import ToolResponse + + if not isinstance(result, ToolResponse): + return { + "__flink_agents_tool_result__": "raw", + "result": result, + } + return { + "__flink_agents_tool_result__": "response", + "result": result.result, + "success": result.is_success(), + "error": result.error_message, + "execution_time_ms": result.execution_time_ms, + "tool_name": result.tool_name, + } def from_java_prompt(j_prompt: Any) -> JavaPrompt: diff --git a/python/flink_agents/runtime/skill/skill_tools.py b/python/flink_agents/runtime/skill/skill_tools.py index 5178ab308..f26221246 100644 --- a/python/flink_agents/runtime/skill/skill_tools.py +++ b/python/flink_agents/runtime/skill/skill_tools.py @@ -29,7 +29,7 @@ from pydantic import BaseModel, Field -from flink_agents.api.tools import ToolExecutionMetadataProvider +from flink_agents.api.tools import ToolExecutionMetadataProvider, ToolResponse from flink_agents.api.tools.tool import Tool, ToolMetadata, ToolType from flink_agents.api.trace import ToolExecutionMetadataKeys @@ -100,7 +100,7 @@ def get_tool_execution_metadata( ) return metadata - def call(self, *args: Any, **kwargs: Any) -> str: + def call(self, *args: Any, **kwargs: Any) -> str | ToolResponse: """Call the tool to load a skill.""" if args: parsed_args = LoadSkillArgs(name=args[0], **kwargs) @@ -113,7 +113,9 @@ def call(self, *args: Any, **kwargs: Any) -> str: manager = self._get_skill_manager() if manager is None: - return "Skill manager not available. No skills have been registered." + return ToolResponse.error( + "Skill manager not available. No skills have been registered." + ) try: skill = manager.get_skill(skill_name) @@ -122,7 +124,9 @@ def call(self, *args: Any, **kwargs: Any) -> str: available_str = ( ", ".join(available) if available else "No skills available." ) - return f"Skill '{skill_name}' not found. Available skills: {available_str}" + return ToolResponse.error( + f"Skill '{skill_name}' not found. Available skills: {available_str}" + ) if resource_path is None or resource_path == "SKILL.md": skill_dir = manager.get_skill_dir(skill_name) @@ -148,7 +152,9 @@ def call(self, *args: Any, **kwargs: Any) -> str: content = skill.get_resource(resource_path) if content is None: available = sorted(skill.get_resource_paths()) - return f"Resource '{resource_path}' not found in skill '{skill_name}', Available resources: {available}" + return ToolResponse.error( + f"Resource '{resource_path}' not found in skill '{skill_name}', Available resources: {available}" + ) return content def _get_skill_manager(self) -> SkillManager | None: diff --git a/python/flink_agents/runtime/skill/tests/test_load_skill.py b/python/flink_agents/runtime/skill/tests/test_load_skill.py index be0246ce8..e428f8e70 100644 --- a/python/flink_agents/runtime/skill/tests/test_load_skill.py +++ b/python/flink_agents/runtime/skill/tests/test_load_skill.py @@ -24,6 +24,7 @@ from flink_agents.api.resource_context import ResourceContext from flink_agents.api.skills import Skills +from flink_agents.api.tools import ToolResponse from flink_agents.api.trace import ToolExecutionMetadataKeys from flink_agents.runtime.skill.skill_manager import SkillManager from flink_agents.runtime.skill.skill_tools import LoadSkillTool @@ -84,8 +85,10 @@ def test_load_resource(self, tool: LoadSkillTool) -> None: def test_load_resource_not_found(self, tool: LoadSkillTool) -> None: """Loading a nonexistent resource returns an error with available list.""" result = tool.call(name="nano-banana-pro", path="nonexistent.txt") - assert "not found" in result.lower() - assert "scripts/generate_image.py" in result + assert isinstance(result, ToolResponse) + assert result.is_error() + assert "not found" in result.error_message.lower() + assert "scripts/generate_image.py" in result.error_message def test_execution_metadata_describes_requested_resource( self, tool: LoadSkillTool @@ -97,7 +100,9 @@ def test_execution_metadata_describes_requested_resource( assert metadata[ToolExecutionMetadataKeys.SKILL_NAME] == "github" assert metadata[ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH] == "README.md" - def test_execution_metadata_normalizes_omitted_path(self, tool: LoadSkillTool) -> None: + def test_execution_metadata_normalizes_omitted_path( + self, tool: LoadSkillTool + ) -> None: metadata = tool.get_tool_execution_metadata({"name": "github"}) assert metadata[ToolExecutionMetadataKeys.SKILL_NAME] == "github" assert metadata[ToolExecutionMetadataKeys.SKILL_RESOURCE_PATH] == "SKILL.md" @@ -114,9 +119,11 @@ def test_execution_metadata_normalizes_explicit_none_path( def test_skill_not_found(self, tool: LoadSkillTool) -> None: """A nonexistent skill returns an error listing available skills.""" result = tool.call(name="nonexistent-skill") - assert "not found" in result.lower() - assert "github" in result - assert "nano-banana-pro" in result + assert isinstance(result, ToolResponse) + assert result.is_error() + assert "not found" in result.error_message.lower() + assert "github" in result.error_message + assert "nano-banana-pro" in result.error_message # -- no skill manager ---------------------------------------------------- @@ -125,7 +132,9 @@ def test_no_skill_manager(self) -> None: mock_ctx = MagicMock(spec=ResourceContext) tool = LoadSkillTool(resource_context=mock_ctx) result = tool.call(name="github") - assert "not available" in result + assert isinstance(result, ToolResponse) + assert result.is_error() + assert "not available" in result.error_message # -- positional args ----------------------------------------------------- diff --git a/python/flink_agents/runtime/tests/test_python_java_utils.py b/python/flink_agents/runtime/tests/test_python_java_utils.py index d4264d46f..2555453df 100644 --- a/python/flink_agents/runtime/tests/test_python_java_utils.py +++ b/python/flink_agents/runtime/tests/test_python_java_utils.py @@ -25,11 +25,13 @@ EmbeddingTokenUsage, ) from flink_agents.api.events.event import Event -from flink_agents.api.tools import InjectedArg +from flink_agents.api.tools import InjectedArg, ToolResponse from flink_agents.runtime.python_java_utils import ( call_embedding_with_usage, convert_to_python_key_text, get_python_tool_metadata, + invoke_python_tool, + invoke_python_tool_instance, to_python_memory_set, wrap_to_input_event, ) @@ -41,6 +43,23 @@ def decorated_python_tool(order_id: str, tenant_id: str, request_id: str) -> str return f"{tenant_id}:{request_id}:{order_id}" +def raw_python_tool(value: str) -> dict[str, object]: + return {"__flink_agents_tool_result__": "response", "value": value} + + +def failed_python_tool(value: str) -> ToolResponse: + return ToolResponse.error(value, execution_time_ms=7, tool_name="failed") + + +def successful_python_tool(value: str) -> ToolResponse: + return ToolResponse.success(value, execution_time_ms=5, tool_name="successful") + + +class _FailedTool: + def call(self, value: str) -> ToolResponse: + return failed_python_tool(value) + + def test_get_python_tool_metadata_merges_callable_injected_args() -> None: flat = get_python_tool_metadata( __name__, "decorated_python_tool", injected_args=["request_id"] @@ -52,6 +71,54 @@ def test_get_python_tool_metadata_merges_callable_injected_args() -> None: assert injected_args == {"tenant_id": {"source": "config", "key": "tenant.id"}} +def test_invoke_python_tool_wraps_raw_payload_without_inspecting_it() -> None: + result = invoke_python_tool(__name__, "raw_python_tool", {"value": "raw"}) + + assert result == { + "__flink_agents_tool_result__": "raw", + "result": {"__flink_agents_tool_result__": "response", "value": "raw"}, + } + + +def test_invoke_python_tool_preserves_explicit_failure() -> None: + result = invoke_python_tool(__name__, "failed_python_tool", {"value": "failed"}) + + assert result == { + "__flink_agents_tool_result__": "response", + "result": None, + "success": False, + "error": "failed", + "execution_time_ms": 7, + "tool_name": "failed", + } + + +def test_invoke_python_tool_preserves_explicit_success() -> None: + result = invoke_python_tool(__name__, "successful_python_tool", {"value": "ok"}) + + assert result == { + "__flink_agents_tool_result__": "response", + "result": "ok", + "success": True, + "error": None, + "execution_time_ms": 5, + "tool_name": "successful", + } + + +def test_invoke_python_tool_instance_preserves_explicit_failure() -> None: + result = invoke_python_tool_instance(_FailedTool(), {"value": "failed"}) + + assert result == { + "__flink_agents_tool_result__": "response", + "result": None, + "success": False, + "error": "failed", + "execution_time_ms": 7, + "tool_name": "failed", + } + + class _UsageAwareEmbeddingModel: def embed_with_usage( self, text: str, **kwargs: object diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapter.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapter.java index 2585145b8..4a03a827b 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapter.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapter.java @@ -40,6 +40,8 @@ /** Adapter for managing Java resources and facilitating Python-Java interoperability. */ public class JavaResourceAdapter { + private static final String TOOL_RESULT_MARKER = "__flink_agents_tool_result__"; + private final ResourceContext resourceContext; private final transient PythonInterpreter interpreter; @@ -176,10 +178,11 @@ public Map getJavaToolMetadata( * org.apache.flink.agents.api.annotation.ToolParam} name override, {@link ToolParameters} * numeric coercion (covers the LLM-emitted JSON Number → Java box type mismatch that reflective * {@code Method.invoke} otherwise rejects), required-parameter checking, and {@link - * ToolResponse} success / error semantics. The success result is unwrapped for the Python - * caller; an unsuccessful response is re-thrown as a {@link RuntimeException}. + * ToolResponse} success / error semantics. The response is returned in an internal envelope so + * the Python caller can distinguish an explicit Tool error from an invocation exception without + * inspecting user payloads. */ - public Object invokeJavaTool( + public Map invokeJavaTool( String className, String methodName, List parameterTypes, @@ -189,10 +192,14 @@ public Object invokeJavaTool( FunctionTool tool = FunctionTool.fromStaticMethod(method); ToolResponse response = tool.call(new ToolParameters(arguments == null ? new HashMap<>() : arguments)); - if (!response.isSuccess()) { - throw new RuntimeException(response.getError()); - } - return response.getResult(); + Map result = new HashMap<>(); + result.put(TOOL_RESULT_MARKER, "response"); + result.put("result", response.getResult()); + result.put("success", response.isSuccess()); + result.put("error", response.getError()); + result.put("execution_time_ms", response.getExecutionTimeMs()); + result.put("tool_name", response.getToolName()); + return result; } /** Invoke a Java static action method with positional arguments from Python. */ diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/skill/LoadSkillTool.java b/runtime/src/main/java/org/apache/flink/agents/runtime/skill/LoadSkillTool.java index abfeb306f..7891780f6 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/skill/LoadSkillTool.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/skill/LoadSkillTool.java @@ -87,11 +87,11 @@ public ToolResponse call(ToolParameters parameters) { try { manager = resolveSkillManager(); } catch (Exception e) { - return ToolResponse.success( + return ToolResponse.error( "Skill manager not available. No skills have been registered."); } if (manager == null) { - return ToolResponse.success( + return ToolResponse.error( "Skill manager not available. No skills have been registered."); } @@ -102,7 +102,7 @@ public ToolResponse call(ToolParameters parameters) { List available = manager.getAllSkillNames(); String availableStr = available.isEmpty() ? "No skills available." : String.join(", ", available); - return ToolResponse.success( + return ToolResponse.error( "Skill '" + name + "' not found. Available skills: " + availableStr); } @@ -140,7 +140,7 @@ public ToolResponse call(ToolParameters parameters) { String content = skill.getResource(path); if (content == null) { - return ToolResponse.success( + return ToolResponse.error( "Resource '" + path + "' not found in skill '" diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapterTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapterTest.java index f1d4a82ef..dd8ddd390 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapterTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapterTest.java @@ -30,7 +30,7 @@ import static org.assertj.core.api.Assertions.assertThat; -class JavaResourceAdapterTest { +public class JavaResourceAdapterTest { @Test void getJavaToolMetadataHidesInjectedArgsAndReturnsAnnotatedDeclaration() throws Exception { @@ -58,6 +58,52 @@ void getJavaToolMetadataHidesInjectedArgsAndReturnsAnnotatedDeclaration() throws assertThat(injectedArgs.get("tenant_id").get("key").asText()).isEqualTo("tenant.id"); } + @Test + void invokeJavaToolPreservesSuccessResponseForPythonCaller() throws Exception { + JavaResourceAdapter adapter = + new JavaResourceAdapter(null, null, Thread.currentThread().getContextClassLoader()); + + Map result = + adapter.invokeJavaTool( + JavaResourceAdapterTest.class.getName(), + "queryOrder", + List.of( + String.class.getName(), + String.class.getName(), + String.class.getName()), + Map.of( + "order_id", "order-1", + "tenant_id", "tenant-1", + "request_id", "request-1")); + + assertThat(result) + .containsEntry("__flink_agents_tool_result__", "response") + .containsEntry("success", true) + .containsEntry("result", "tenant-1:request-1:order-1") + .containsEntry("execution_time_ms", 0L); + assertThat(result.get("error")).isNull(); + } + + @Test + void invokeJavaToolPreservesErrorResponseForPythonCaller() throws Exception { + JavaResourceAdapter adapter = + new JavaResourceAdapter(null, null, Thread.currentThread().getContextClassLoader()); + + Map result = + adapter.invokeJavaTool( + JavaResourceAdapterTest.class.getName(), + "failingTool", + List.of(String.class.getName()), + Map.of("value", "input")); + + assertThat(result) + .containsEntry("__flink_agents_tool_result__", "response") + .containsEntry("success", false) + .containsEntry("error", "tool rejected input") + .containsEntry("execution_time_ms", 0L); + assertThat(result.get("result")).isNull(); + } + @Tool(description = "Query order.") public static String queryOrder( @ToolParam(name = "order_id") String orderId, @@ -70,4 +116,9 @@ public static String queryOrder( @ToolParam(name = "request_id") String requestId) { return tenantId + ":" + requestId + ":" + orderId; } + + @Tool(description = "Fail a tool call.") + public static String failingTool(@ToolParam(name = "value") String value) { + throw new IllegalStateException("tool rejected " + value); + } } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/skill/LoadSkillToolTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/skill/LoadSkillToolTest.java index f131cd07e..71d6ff6e3 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/skill/LoadSkillToolTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/skill/LoadSkillToolTest.java @@ -70,7 +70,8 @@ private static ToolParameters args(String name, String path) { void unknownSkillReturnsAvailableList() { LoadSkillTool t = tool(contextWithSkills()); ToolResponse resp = t.call(args("does-not-exist", null)); - String out = (String) resp.getResult(); + assertTrue(resp.isError()); + String out = resp.getError(); assertTrue(out.contains("not found")); assertTrue(out.contains("github")); assertTrue(out.contains("nano-banana-pro")); @@ -145,7 +146,8 @@ void explicitNullPathLoadsSkillContentEnvelope() { void missingResourceReportsAvailable() { LoadSkillTool t = tool(contextWithSkills()); ToolResponse resp = t.call(args("nano-banana-pro", "no-such.txt")); - String out = (String) resp.getResult(); + assertTrue(resp.isError()); + String out = resp.getError(); assertTrue(out.contains("Resource 'no-such.txt' not found")); assertTrue(out.contains("Available resources")); } @@ -156,7 +158,8 @@ void noSkillsRegisteredReturnsFriendlyMessage() { ResourceContextImpl ctx = new ResourceContextImpl((name, type) -> null); LoadSkillTool t = tool(ctx); ToolResponse resp = t.call(args("anything", null)); + assertTrue(resp.isError()); assertEquals( - "Skill manager not available. No skills have been registered.", resp.getResult()); + "Skill manager not available. No skills have been registered.", resp.getError()); } }