yunfengzhou-hub commented on code in PR #1114: URL: https://github.com/apache/flink-agents/pull/1114#discussion_r4036965251
########## plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionSubagentTest.java: ########## @@ -0,0 +1,673 @@ +/* + * 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.actions; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.configuration.ReadableConfiguration; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.MemoryObject; +import org.apache.flink.agents.api.context.Outcome; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.ToolRequestEvent; +import org.apache.flink.agents.api.event.ToolResponseEvent; +import org.apache.flink.agents.api.memory.BaseLongTermMemory; +import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; +import org.apache.flink.agents.api.resource.Resource; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.subagent.SubagentFuture; +import org.apache.flink.agents.api.subagent.SubagentFutures; +import org.apache.flink.agents.api.subagent.SubagentResult; +import org.apache.flink.agents.api.subagent.SubagentSetup; +import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolMetadata; +import org.apache.flink.agents.api.tools.ToolParameters; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.tools.ToolType; +import org.apache.flink.agents.plan.AgentConfiguration; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatExceptionOfType; + +/** Tests for dispatching a tool call to an {@code AGENT} resource. */ +class ToolCallActionSubagentTest { + + @Test + void delegatesToTheSubagentAndReportsItsNormalizedResult() throws Exception { + Map<String, Object> payload = new LinkedHashMap<>(); + payload.put("verdict", "approved"); + payload.put("findings", List.of("style")); + RecordingSubagentSetup agent = new RecordingSubagentSetup(SubagentResult.ok(payload)); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("_subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", true); + assertThat(response.getResponses().get("call-1").getResult()) + .isEqualTo("{\"verdict\":\"approved\",\"findings\":[\"style\"]}"); + assertThat(response.getError()).doesNotContainKey("call-1"); + } + + @Test + void handsTheModelArgumentsToTheSubagentAsThePrompt() throws Exception { + RecordingSubagentSetup agent = new RecordingSubagentSetup(SubagentResult.ok("done")); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("_subagent_reviewer"), ctx); + + assertThat(agent.prompts).containsExactly(Map.of("prompt", "review the diff")); + // A sub-agent call resolves through the setup, which owns its own durable execution. + assertThat(ctx.durableExecutions).isZero(); + } + + @Test + void reportsAFailedSubagentResultWithTheDetailExposedToTheModel() throws Exception { + RecordingSubagentSetup agent = + new RecordingSubagentSetup(SubagentResult.error("upstream refused")); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("_subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", false); + assertThat(response.getResponses().get("call-1").getError()) + .isEqualTo("Sub-agent _subagent_reviewer execute failed: upstream refused"); + assertThat(response.getError()).containsEntry("call-1", "upstream refused"); + } + + @Test + void reportsAFailureRaisedWhileSubmitting() throws Exception { + RecordingSubagentSetup agent = new RecordingSubagentSetup(SubagentResult.ok("unreachable")); + agent.submitFailure = new IllegalStateException("mailbox is full"); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("_subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", false); + assertThat(response.getResponses().get("call-1").getError()) + .isEqualTo("Sub-agent _subagent_reviewer execute failed: mailbox is full"); + assertThat(response.getError()).containsEntry("call-1", "mailbox is full"); + } + + @Test + void rejectsAResultJsonCannotExpress() throws Exception { + RecordingSubagentSetup agent = + new RecordingSubagentSetup(SubagentResult.ok(Map.of("handle", new Object()))); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("_subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", false); + assertThat(response.getResponses().get("call-1").getError()) + .startsWith("Sub-agent _subagent_reviewer execute failed") + .contains("result.handle"); + assertThat(response.getError().get("call-1")).contains("result.handle"); + } + + /** A declared result type is what admits a result JSON cannot express on its own. */ + @Test + void readsAResultThroughTheTypeTheSubagentDeclares() throws Exception { + RecordingSubagentSetup agent = + new TypedRecordingSubagentSetup(SubagentResult.ok(new Verdict(true, "clean"))); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + ToolCallAction.processToolRequest(toolRequest("_subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", true); + assertThat(response.getResponses().get("call-1").getResult()) + .isEqualTo("{\"approved\":true,\"note\":\"clean\"}"); + } + + /** The reserved prefix routes each namespace on its own, even under one shared name. */ + @Test + void routesAToolAndASubagentSharingANameToTheirOwnNamespace() throws Exception { + FakeRunnerContext ctx = + new FakeRunnerContext() + .withAgent( + "reviewer", new RecordingSubagentSetup(SubagentResult.ok("done"))) + .withTool("reviewer", new StubTool("reviewer")); + + ToolCallAction.processToolRequest(toolRequest("_subagent_reviewer"), ctx); + + ToolResponseEvent delegated = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(delegated.getSuccess()).containsEntry("call-1", true); + assertThat(delegated.getResponses().get("call-1").getResult()).isEqualTo("done"); + // A sub-agent call resolves through the setup, which owns its own durable execution. + assertThat(ctx.durableExecutions).isZero(); + + ToolCallAction.processToolRequest(toolRequest("reviewer"), ctx); + + ToolResponseEvent direct = ToolResponseEvent.fromEvent(ctx.sentEvents.get(1)); + assertThat(direct.getSuccess()).containsEntry("call-1", true); + assertThat(direct.getResponses().get("call-1").getResult()).isEqualTo("reviewer called"); + assertThat(ctx.durableExecutions).isOne(); + } + + @Test + void refusesAnAgentResourceThatCarriesNoCallableSetup() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext(); + ctx.agents.put("reviewer", new StubTool("reviewer")); + + ToolCallAction.processToolRequest(toolRequest("_subagent_reviewer"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", false); + assertThat(response.getResponses().get("call-1").getError()) + .isEqualTo( + "Sub-agent _subagent_reviewer execute failed: Sub-agent reviewer must" + + " resolve to a SubagentSetup, but was " + + StubTool.class.getName() + + "."); + assertThat(response.getError().get("call-1")) + .isEqualTo( + "Sub-agent reviewer must resolve to a SubagentSetup, but was " + + StubTool.class.getName() + + "."); + } + + @Test + void stillDispatchesAToolWhenBothKindsAreRegisteredUnderDifferentNames() throws Exception { + FakeRunnerContext ctx = + new FakeRunnerContext() + .withAgent( + "reviewer", new RecordingSubagentSetup(SubagentResult.ok("done"))) + .withTool("queryOrder", new StubTool("queryOrder")); + + ToolCallAction.processToolRequest(toolRequest("queryOrder"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", true); + assertThat(response.getResponses().get("call-1").getResult()) + .isEqualTo("queryOrder called"); + assertThat(ctx.durableExecutions).isOne(); + } + + /** + * The batched path runs sub-agent calls concurrently: every call is submitted before any is + * awaited, so the async setups' remote runs overlap instead of blocking one behind the next. + * The serial path interleaves submit and await per call, which this order assertion rejects. + */ + @Test + void submitsEverySubagentCallBeforeAwaitingAnyUnderParallelDispatch() throws Exception { + List<String> ops = new ArrayList<>(); + FakeRunnerContext ctx = + new FakeRunnerContext() + .withParallelToolCalls() + .withAgent("a", new OrderRecordingSubagentSetup("a", ops)) + .withAgent("b", new OrderRecordingSubagentSetup("b", ops)); + + ToolCallAction.processToolRequest(twoSubagentRequest("a", "b"), ctx); + + assertThat(ops).containsExactly("submit:a", "submit:b", "await:a", "await:b"); + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()) + .containsEntry("call-1", true) + .containsEntry("call-2", true); + assertThat(response.getResponses().get("call-1").getResult()).isEqualTo("a done"); + assertThat(response.getResponses().get("call-2").getResult()).isEqualTo("b done"); + } + + /** + * A cancelled sub-agent call must propagate like a cancelled tool call (#1111), not be folded + * into a tool-error response: no ToolResponseEvent goes out, so no further chat call is driven + * off a cancelled delegation and the action is not persisted as completed on the back of it. + */ + @Test + void propagatesInterruptionFromASubagentCallInsteadOfRecordingAFailure() throws Exception { + RecordingSubagentSetup agent = new RecordingSubagentSetup(SubagentResult.ok("unreachable")); + agent.submitFailure = new InterruptedException("cancelled"); + FakeRunnerContext ctx = new FakeRunnerContext().withAgent("reviewer", agent); + + Thread.interrupted(); + + assertThatExceptionOfType(InterruptedException.class) + .isThrownBy( + () -> + ToolCallAction.processToolRequest( + toolRequest("_subagent_reviewer"), ctx)); + + assertThat(Thread.interrupted()).as("interrupt status should be restored").isTrue(); + assertThat(ctx.sentEvents).isEmpty(); + } + + /** + * Under the batched path, a cancellation while awaiting one sub-agent must propagate (#1111) + * and must not leave the other already-submitted handles dangling: the interrupted handle and + * every later one, submitted but now never awaited, are cancelled on the way out. + */ + @Test + void propagatesInterruptionUnderParallelDispatchAndCancelsSubmittedHandles() throws Exception { + List<String> ops = new ArrayList<>(); + FakeRunnerContext ctx = + new FakeRunnerContext() + .withParallelToolCalls() + .withAgent("a", new InterruptingSubagentSetup("a", ops)) + .withAgent("b", new OrderRecordingSubagentSetup("b", ops)); + + Thread.interrupted(); + + assertThatExceptionOfType(InterruptedException.class) + .isThrownBy( + () -> ToolCallAction.processToolRequest(twoSubagentRequest("a", "b"), ctx)); + + assertThat(Thread.interrupted()).as("interrupt status should be restored").isTrue(); + assertThat(ctx.sentEvents).isEmpty(); + assertThat(ops).containsExactly("submit:a", "submit:b", "await:a", "cancel:a", "cancel:b"); + } + + private static ToolRequestEvent toolRequest(String callableName) { + return new ToolRequestEvent( + "model", + List.of( + Map.of( + "id", + "call-1", + "type", + "function", + "function", + Map.of( + "name", + callableName, + "arguments", + Map.of("prompt", "review the diff"))))); + } + + /** One request carrying two sub-agent calls, so the batched path has more than one to run. */ + private static ToolRequestEvent twoSubagentRequest(String first, String second) { Review Comment: Added — `parallelDispatchKeepsToolAndSubagentResultsOnTheirOwnIds` runs one `_subagent_` call alongside two tools under withParallelToolCalls(), and asserts call-1/2/3 each land on their own result ("agent-result", "alpha called", "beta called"). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
