diff --git a/api/pom.xml b/api/pom.xml index f7bdbbabd..ccaebafd4 100644 --- a/api/pom.xml +++ b/api/pom.xml @@ -68,4 +68,21 @@ under the License. + + + + org.apache.maven.plugins + maven-jar-plugin + 3.4.2 + + + + test-jar + + + + + + + \ No newline at end of file diff --git a/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java b/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java index c3e5d19ba..fb52b15d5 100644 --- a/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java +++ b/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java @@ -146,6 +146,12 @@ public interface RunnerContext { */ T durableExecuteAsync(DurableCallable callable) throws Exception; + /** Creates a new session id for a sub-agent call. */ + String nextSessionId(); + + /** Creates a new call id for a sub-agent invocation under the given session. */ + String nextCallId(String sessionId); + /** Clean up the resource. */ void close() throws Exception; } diff --git a/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java b/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java index 79e4ab52e..b2cd3bd6d 100644 --- a/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java +++ b/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java @@ -32,7 +32,8 @@ public enum ResourceType { PROMPT("prompt"), TOOL("tool"), MCP_SERVER("mcp_server"), - SKILLS("skills"); + SKILLS("skills"), + AGENT("agent"); private final String value; diff --git a/api/src/main/java/org/apache/flink/agents/api/subagent/BaseSubagentCallable.java b/api/src/main/java/org/apache/flink/agents/api/subagent/BaseSubagentCallable.java new file mode 100644 index 000000000..c53dc99a9 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/subagent/BaseSubagentCallable.java @@ -0,0 +1,65 @@ +/* + * 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.api.subagent; + +import org.apache.flink.agents.api.context.DurableCallable; + +/** + * Convenience base for the {@link DurableCallable} returned by {@code asAsyncCallable}. + * + *

Keys the durable call by the framework-assigned identity as {@code sessionId#callId} (the + * {@link SubagentSetup} contract) and captures exceptions thrown by {@link #callInternal()} into + * {@link Result#error(Exception)}, so failures are reported through the result rather than thrown. + * Implementations only provide {@link #callInternal()}. + */ +public abstract class BaseSubagentCallable implements DurableCallable { + + private final String sessionId; + private final String callId; + + protected BaseSubagentCallable(String sessionId, String callId) { + this.sessionId = sessionId; + this.callId = callId; + } + + @Override + public String getId() { + return sessionId + "#" + callId; + } + + @Override + public Class getResultClass() { + return Result.class; + } + + @Override + public final Result call() { + try { + return Result.ok(callInternal()); + } catch (Exception e) { + return Result.error(e); + } + } + + /** + * Performs the invocation and returns the JSON-serializable payload. Thrown exceptions are + * captured into a failed {@link Result}. + */ + protected abstract Object callInternal() throws Exception; +} diff --git a/api/src/main/java/org/apache/flink/agents/api/subagent/Result.java b/api/src/main/java/org/apache/flink/agents/api/subagent/Result.java new file mode 100644 index 000000000..9a1593b2f --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/subagent/Result.java @@ -0,0 +1,98 @@ +/* + * 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.api.subagent; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.io.PrintWriter; +import java.io.Serializable; +import java.io.StringWriter; + +/** + * Outcome of a {@link Subagent} call. + * + *

Sub-agent implementations should capture internal failures into a {@code Result} (via {@link + * #error}) instead of throwing, so callers can inspect {@link #isSuccess()} without try/catch. + * + *

The failure cause is carried as a serializable {@code errorMessage} — the full stack trace of + * the failure — rather than a live exception, so that a {@code Result} can be persisted through + * durable execution. + */ +public class Result implements Serializable { + + private static final long serialVersionUID = 1L; + + private final boolean success; + private final Object result; + private final String errorMessage; + + @JsonCreator + public Result( + @JsonProperty("success") boolean success, + @JsonProperty("result") Object result, + @JsonProperty("errorMessage") String errorMessage) { + this.success = success; + this.result = result; + this.errorMessage = errorMessage; + } + + /** Creates a successful result carrying the given value. */ + public static Result ok(Object result) { + return new Result(true, result, null); + } + + /** Creates a failed result carrying the full stack trace of the given exception. */ + public static Result error(Exception exception) { + return new Result(false, null, exception == null ? null : stackTraceOf(exception)); + } + + /** Creates a failed result carrying the given message. */ + public static Result error(String errorMessage) { + return new Result(false, null, errorMessage); + } + + private static String stackTraceOf(Exception exception) { + StringWriter writer = new StringWriter(); + exception.printStackTrace(new PrintWriter(writer)); + return writer.toString(); + } + + public boolean isSuccess() { + return success; + } + + public Object getResult() { + return result; + } + + public String getErrorMessage() { + return errorMessage; + } + + /** + * Reconstructs an exception carrying the stored stack trace as its message, or null if this + * result is successful. + */ + @JsonIgnore + public Exception getException() { + return success ? null : new RuntimeException(errorMessage); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/subagent/Subagent.java b/api/src/main/java/org/apache/flink/agents/api/subagent/Subagent.java new file mode 100644 index 000000000..a448b2ce3 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/subagent/Subagent.java @@ -0,0 +1,47 @@ +/* + * 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.api.subagent; + +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; + +/** + * Caller-facing interface for all sub-agents (external and internal). + * + *

An invocation is identified by a {@code (sessionId, callId)} pair; the session groups a + * conversation across invocations. Both ids are assigned by the framework ({@link + * RunnerContext#nextSessionId()} and {@link RunnerContext#nextCallId(String)}): callers may supply + * a session id to continue a prior session but never supply a call id. + */ +public interface Subagent { + + /** Synchronously invokes the sub-agent, creating a new session. */ + Result call(RunnerContext ctx, Object prompt) throws Exception; + + /** Synchronously invokes the sub-agent continuing {@code sessionId}. */ + Result call(RunnerContext ctx, Object prompt, String sessionId) throws Exception; + + /** Produces a deferred, durable callable for this sub-agent call, creating a new session. */ + DurableCallable asAsyncCallable(RunnerContext ctx, Object prompt); + + /** + * Produces a deferred, durable callable for this sub-agent call continuing {@code sessionId}. + */ + DurableCallable asAsyncCallable(RunnerContext ctx, Object prompt, String sessionId); +} diff --git a/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java new file mode 100644 index 000000000..9942d9577 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java @@ -0,0 +1,81 @@ +/* + * 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.api.subagent; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.resource.SerializableResource; + +/** + * Base setup for an external sub-agent resource. Serialized into the agent plan as an {@code AGENT} + * resource. + * + *

Hosts the id-resolution chain behind {@link Subagent}: omitted ids are assigned via the + * context before an implementation is ever invoked, and {@link #call} runs the deferred callable + * through durable execution. Implementations only provide the terminal {@code asAsyncCallable}. + */ +public abstract class SubagentSetup extends SerializableResource implements Subagent { + + @Override + @JsonIgnore + public ResourceType getResourceType() { + return ResourceType.AGENT; + } + + @Override + public Result call(RunnerContext ctx, Object prompt) throws Exception { + return call(ctx, prompt, ctx.nextSessionId()); + } + + @Override + public Result call(RunnerContext ctx, Object prompt, String sessionId) throws Exception { + return call(ctx, prompt, sessionId, ctx.nextCallId(sessionId)); + } + + /** + * Synchronously invokes the sub-agent with already-assigned {@code sessionId} and {@code + * callId}, running the deferred callable through durable execution. Framework-facing: the ids + * are assigned by the shorter variants, never supplied by callers. + */ + public Result call(RunnerContext ctx, Object prompt, String sessionId, String callId) + throws Exception { + return ctx.durableExecuteAsync(asAsyncCallable(ctx, prompt, sessionId, callId)); + } + + @Override + public DurableCallable asAsyncCallable(RunnerContext ctx, Object prompt) { + return asAsyncCallable(ctx, prompt, ctx.nextSessionId()); + } + + @Override + public DurableCallable asAsyncCallable( + RunnerContext ctx, Object prompt, String sessionId) { + return asAsyncCallable(ctx, prompt, sessionId, ctx.nextCallId(sessionId)); + } + + /** + * Produces the deferred, durable callable for one invocation; both ids are already assigned. + * The only method implementations must provide. Contract: the returned {@link + * DurableCallable#getId()} MUST be derived solely from the {@code (sessionId, callId)} pair. + */ + protected abstract DurableCallable asAsyncCallable( + RunnerContext ctx, Object prompt, String sessionId, String callId); +} diff --git a/api/src/main/java/org/apache/flink/agents/api/yaml/YamlLoader.java b/api/src/main/java/org/apache/flink/agents/api/yaml/YamlLoader.java index 0a33b4e75..4e5c39486 100644 --- a/api/src/main/java/org/apache/flink/agents/api/yaml/YamlLoader.java +++ b/api/src/main/java/org/apache/flink/agents/api/yaml/YamlLoader.java @@ -275,6 +275,7 @@ public static LoadedFile buildAgents(Path path) { addSharedDescriptors( sharedResources, ResourceType.VECTOR_STORE, doc.getVectorStores(), path); addSharedDescriptors(sharedResources, ResourceType.MCP_SERVER, doc.getMcpServers(), path); + addSharedDescriptors(sharedResources, ResourceType.AGENT, doc.getSubagents(), path); for (ToolSpec t : doc.getTools()) { if (sharedResources.get(ResourceType.TOOL).put(t.getName(), buildTool(t)) != null) { @@ -390,6 +391,7 @@ private static Agent buildAgent(AgentSpec spec) { addAgentDescriptors(agent, ResourceType.EMBEDDING_MODEL, spec.getEmbeddingModelSetups()); addAgentDescriptors(agent, ResourceType.VECTOR_STORE, spec.getVectorStores()); addAgentDescriptors(agent, ResourceType.MCP_SERVER, spec.getMcpServers()); + addAgentDescriptors(agent, ResourceType.AGENT, spec.getSubagents()); for (ToolSpec t : spec.getTools()) { agent.addResource(t.getName(), ResourceType.TOOL, buildTool(t)); diff --git a/api/src/main/java/org/apache/flink/agents/api/yaml/spec/AgentSpec.java b/api/src/main/java/org/apache/flink/agents/api/yaml/spec/AgentSpec.java index 0e37a8aad..1826a61a0 100644 --- a/api/src/main/java/org/apache/flink/agents/api/yaml/spec/AgentSpec.java +++ b/api/src/main/java/org/apache/flink/agents/api/yaml/spec/AgentSpec.java @@ -40,6 +40,7 @@ public final class AgentSpec { private final List embeddingModelSetups; private final List vectorStores; private final List mcpServers; + private final List subagents; @JsonCreator public AgentSpec( @@ -55,7 +56,8 @@ public AgentSpec( List embeddingModelConnections, @JsonProperty("embedding_model_setups") List embeddingModelSetups, @JsonProperty("vector_stores") List vectorStores, - @JsonProperty("mcp_servers") List mcpServers) { + @JsonProperty("mcp_servers") List mcpServers, + @JsonProperty("subagents") List subagents) { this.name = name; this.description = description; this.prompts = orEmpty(prompts); @@ -68,6 +70,7 @@ public AgentSpec( this.embeddingModelSetups = orEmpty(embeddingModelSetups); this.vectorStores = orEmpty(vectorStores); this.mcpServers = orEmpty(mcpServers); + this.subagents = orEmpty(subagents); } private static List orEmpty(List list) { @@ -121,4 +124,8 @@ public List getVectorStores() { public List getMcpServers() { return mcpServers; } + + public List getSubagents() { + return subagents; + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/yaml/spec/YamlAgentsDocument.java b/api/src/main/java/org/apache/flink/agents/api/yaml/spec/YamlAgentsDocument.java index 1b0fbfa99..019c06d4e 100644 --- a/api/src/main/java/org/apache/flink/agents/api/yaml/spec/YamlAgentsDocument.java +++ b/api/src/main/java/org/apache/flink/agents/api/yaml/spec/YamlAgentsDocument.java @@ -39,6 +39,7 @@ public final class YamlAgentsDocument { private final List embeddingModelSetups; private final List vectorStores; private final List mcpServers; + private final List subagents; @JsonCreator public YamlAgentsDocument( @@ -53,7 +54,8 @@ public YamlAgentsDocument( List embeddingModelConnections, @JsonProperty("embedding_model_setups") List embeddingModelSetups, @JsonProperty("vector_stores") List vectorStores, - @JsonProperty("mcp_servers") List mcpServers) { + @JsonProperty("mcp_servers") List mcpServers, + @JsonProperty("subagents") List subagents) { this.agents = orEmpty(agents); this.prompts = orEmpty(prompts); this.tools = orEmpty(tools); @@ -65,6 +67,7 @@ public YamlAgentsDocument( this.embeddingModelSetups = orEmpty(embeddingModelSetups); this.vectorStores = orEmpty(vectorStores); this.mcpServers = orEmpty(mcpServers); + this.subagents = orEmpty(subagents); } private static List orEmpty(List list) { @@ -114,4 +117,8 @@ public List getVectorStores() { public List getMcpServers() { return mcpServers; } + + public List getSubagents() { + return subagents; + } } diff --git a/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentRegisterTest.java b/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentRegisterTest.java new file mode 100644 index 000000000..c5adef6d5 --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentRegisterTest.java @@ -0,0 +1,88 @@ +/* + * 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.api.subagent; + +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.resource.ResourceType; +import org.junit.jupiter.api.Test; + +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Tests registering sub-agents as AGENT resources and the failure-capturing callable base. */ +class SubagentRegisterTest { + + @Test + void registerSubagentSetupAsResource() { + Agent agent = new Agent(); + TestSubagentSetup setup = new TestSubagentSetup(); + agent.addResource("reviewer", ResourceType.AGENT, setup); + + Map agentResources = agent.getResources().get(ResourceType.AGENT); + assertEquals(1, agentResources.size()); + assertSame(setup, agentResources.get("reviewer")); + assertEquals(ResourceType.AGENT, setup.getResourceType()); + } + + @Test + void duplicateNameThrows() { + Agent agent = new Agent(); + agent.addResource("reviewer", ResourceType.AGENT, new TestSubagentSetup()); + assertThrows( + IllegalArgumentException.class, + () -> agent.addResource("reviewer", ResourceType.AGENT, new TestSubagentSetup())); + } + + @Test + void multipleSubagentsRegistered() { + Agent agent = new Agent(); + TestSubagentSetup reviewer = new TestSubagentSetup(); + TestSubagentSetup coder = new TestSubagentSetup(); + agent.addResource("reviewer", ResourceType.AGENT, reviewer); + agent.addResource("coder", ResourceType.AGENT, coder); + + Map agentResources = agent.getResources().get(ResourceType.AGENT); + assertEquals(2, agentResources.size()); + assertSame(reviewer, agentResources.get("reviewer")); + assertSame(coder, agentResources.get("coder")); + } + + @Test + void baseCallableCapturesExceptionIntoErrorResult() { + BaseSubagentCallable failing = + new BaseSubagentCallable("sid-1", "c-1") { + @Override + protected Object callInternal() throws Exception { + throw new IllegalStateException("boom"); + } + }; + + Result result = failing.call(); + assertFalse(result.isSuccess()); + // errorMessage is the full stack trace, which contains the original message. + assertTrue(result.getErrorMessage().contains("boom")); + assertNotNull(result.getException()); + } +} diff --git a/api/src/test/java/org/apache/flink/agents/api/subagent/TestSubagentSetup.java b/api/src/test/java/org/apache/flink/agents/api/subagent/TestSubagentSetup.java new file mode 100644 index 000000000..cc094c7ad --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/subagent/TestSubagentSetup.java @@ -0,0 +1,89 @@ +/* + * 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.api.subagent; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; + +import javax.annotation.Nullable; + +import java.util.List; + +/** + * Shared {@link SubagentSetup} test double, constructible directly or from a {@link + * ResourceDescriptor} (the YAML shape). Echoes the prompt (prefixed with {@code endpoint} when + * set); {@code failOnCall=true} makes every call fail, surfacing through {@link Result}. + */ +public class TestSubagentSetup extends SubagentSetup { + + private static final long serialVersionUID = 1L; + + @Nullable private final String endpoint; + private final boolean failOnCall; + + public TestSubagentSetup() { + this(null, false); + } + + public TestSubagentSetup(@Nullable String endpoint) { + this(endpoint, false); + } + + @JsonCreator + public TestSubagentSetup( + @JsonProperty("endpoint") @Nullable String endpoint, + @JsonProperty("failOnCall") boolean failOnCall) { + this.endpoint = endpoint; + this.failOnCall = failOnCall; + } + + /** Descriptor-based construction, as used by YAML-declared {@code subagents:} entries. */ + public TestSubagentSetup(ResourceDescriptor descriptor, ResourceContext resourceContext) { + this( + (String) descriptor.getArgument("endpoint"), + Boolean.TRUE.equals(descriptor.getArgument("fail_on_call"))); + } + + @Nullable + public String getEndpoint() { + return endpoint; + } + + public boolean isFailOnCall() { + return failOnCall; + } + + @Override + protected DurableCallable asAsyncCallable( + RunnerContext ctx, Object prompt, String sessionId, String callId) { + return new BaseSubagentCallable(sessionId, callId) { + @Override + protected Object callInternal() { + if (failOnCall) { + throw new IllegalStateException("endpoint " + endpoint + " is down"); + } + return List.of(endpoint == null ? prompt : endpoint + ":" + prompt); + } + }; + } +} diff --git a/docs/yaml-schema.json b/docs/yaml-schema.json index 183cc7ac8..27abf5692 100644 --- a/docs/yaml-schema.json +++ b/docs/yaml-schema.json @@ -148,6 +148,13 @@ "title": "Skills", "type": "array" }, + "subagents": { + "items": { + "$ref": "#/$defs/DescriptorSpec" + }, + "title": "Subagents", + "type": "array" + }, "tools": { "items": { "$ref": "#/$defs/ToolSpec" @@ -513,6 +520,13 @@ "title": "Skills", "type": "array" }, + "subagents": { + "items": { + "$ref": "#/$defs/DescriptorSpec" + }, + "title": "Subagents", + "type": "array" + }, "tools": { "items": { "$ref": "#/$defs/ToolSpec" diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/subagent/ExternalSubagentAgent.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/subagent/ExternalSubagentAgent.java new file mode 100644 index 000000000..b28833778 --- /dev/null +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/subagent/ExternalSubagentAgent.java @@ -0,0 +1,65 @@ +/* + * 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.integration.test.subagent; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.subagent.Result; +import org.apache.flink.agents.api.subagent.SubagentSetup; +import org.apache.flink.api.java.functions.KeySelector; + +import java.util.List; + +/** Agent whose input action delegates to the {@code ext-agent} external sub-agent. */ +public class ExternalSubagentAgent extends Agent { + + public static final String SUBAGENT_NAME = "ext-agent"; + + public ExternalSubagentAgent(SubagentSetup setup) throws Exception { + addResource(SUBAGENT_NAME, ResourceType.AGENT, setup); + addAction( + new String[] {InputEvent.EVENT_TYPE}, + ExternalSubagentAgent.class.getMethod( + "callExternal", Event.class, RunnerContext.class)); + } + + /** Calls the sub-agent with an explicit session id and emits its result (or failure). */ + public static void callExternal(Event event, RunnerContext ctx) throws Exception { + SubagentSetup setup = (SubagentSetup) ctx.getResource(SUBAGENT_NAME, ResourceType.AGENT); + Object prompt = InputEvent.fromEvent(event).getInput(); + Result result = setup.call(ctx, prompt, "session-" + prompt); + if (result.isSuccess()) { + ctx.sendEvent(new OutputEvent(((List) result.getResult()).get(0))); + } else { + ctx.sendEvent(new OutputEvent("error:" + result.getErrorMessage())); + } + } + + /** Keys every element by itself. */ + public static class LongKeySelector implements KeySelector { + @Override + public Long getKey(Long value) { + return value; + } + } +} diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/subagent/ExternalSubagentTest.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/subagent/ExternalSubagentTest.java new file mode 100644 index 000000000..ca722d83c --- /dev/null +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/subagent/ExternalSubagentTest.java @@ -0,0 +1,132 @@ +/* + * 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.integration.test.subagent; + +import org.apache.flink.agents.api.AgentsExecutionEnvironment; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.util.CloseableIterator; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.net.URL; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; + +/** + * End-to-end tests for external sub-agents on a real Flink job: registration, invocation, failure + * reporting via Result, and YAML-declared registration. + */ +public class ExternalSubagentTest { + + @Test + public void javaExternalSubagentCallEndToEnd() throws Exception { + List outputs = + runJob(new ExternalSubagentAgent(new MockExternalSubagentSetup("http://ext:8080"))); + + Assertions.assertEquals(2, outputs.size()); + for (long input : new long[] {1L, 2L}) { + Assertions.assertTrue( + outputs.contains("HTTP response for: " + input + " from http://ext:8080"), + "missing output for input " + input + ": " + outputs); + } + } + + @Test + public void javaExternalSubagentFailureSurfacesViaResult() throws Exception { + List outputs = + runJob( + new ExternalSubagentAgent( + new MockExternalSubagentSetup("http://down:8080", true))); + + Assertions.assertEquals(2, outputs.size()); + for (String output : outputs) { + // errorMessage is the full stack trace, which contains the original message. + Assertions.assertTrue( + output.startsWith("error:") + && output.contains("endpoint http://down:8080 is down"), + "unexpected failure output: " + output); + } + } + + @Test + public void yamlExternalSubagentCallEndToEnd() throws Exception { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + DataStream inputStream = env.fromData(1L, 2L); + + AgentsExecutionEnvironment agentsEnv = + AgentsExecutionEnvironment.getExecutionEnvironment(env); + agentsEnv.loadYaml(yamlFixture("external_subagent_agent.yaml")); + DataStream outputStream = + agentsEnv + .fromDataStream(inputStream, new ExternalSubagentAgent.LongKeySelector()) + .apply("external_subagent_yaml_agent") + .toDataStream(); + + List outputs = collect(outputStream, agentsEnv); + + Assertions.assertEquals(2, outputs.size()); + for (long input : new long[] {1L, 2L}) { + Assertions.assertTrue( + outputs.contains( + "HTTP response for: " + input + " from http://yaml-endpoint:8080"), + "missing output for input " + input + ": " + outputs); + } + } + + private static List runJob(ExternalSubagentAgent agent) throws Exception { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + DataStream inputStream = env.fromData(1L, 2L); + + AgentsExecutionEnvironment agentsEnv = + AgentsExecutionEnvironment.getExecutionEnvironment(env); + DataStream outputStream = + agentsEnv + .fromDataStream(inputStream, new ExternalSubagentAgent.LongKeySelector()) + .apply(agent) + .toDataStream(); + + return collect(outputStream, agentsEnv); + } + + private static List collect( + DataStream outputStream, AgentsExecutionEnvironment agentsEnv) + throws Exception { + try (CloseableIterator results = outputStream.collectAsync()) { + agentsEnv.execute(); + List outputs = new ArrayList<>(); + while (results.hasNext()) { + outputs.add(results.next().toString()); + } + return outputs; + } + } + + private static java.nio.file.Path yamlFixture(String fileName) { + URL url = + Objects.requireNonNull( + ExternalSubagentTest.class.getClassLoader().getResource("yaml/" + fileName), + "missing yaml fixture: " + fileName); + return Paths.get(url.getPath()); + } +} diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/subagent/MockExternalSubagentSetup.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/subagent/MockExternalSubagentSetup.java new file mode 100644 index 000000000..c840bc213 --- /dev/null +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/subagent/MockExternalSubagentSetup.java @@ -0,0 +1,94 @@ +/* + * 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.integration.test.subagent; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.subagent.BaseSubagentCallable; +import org.apache.flink.agents.api.subagent.Result; +import org.apache.flink.agents.api.subagent.SubagentSetup; + +import javax.annotation.Nullable; + +import java.util.List; + +/** + * Mock external sub-agent simulating an HTTP-based external agent system, constructible directly or + * from a {@link ResourceDescriptor} (the YAML shape). Internal failures are captured into a {@link + * Result} rather than thrown. + */ +public class MockExternalSubagentSetup extends SubagentSetup { + + private static final long serialVersionUID = 1L; + + private final String endpointUrl; + private final boolean failOnCall; + + public MockExternalSubagentSetup(String endpointUrl) { + this(endpointUrl, false); + } + + @JsonCreator + public MockExternalSubagentSetup( + @JsonProperty("endpointUrl") String endpointUrl, + @JsonProperty("failOnCall") boolean failOnCall) { + this.endpointUrl = endpointUrl; + this.failOnCall = failOnCall; + } + + /** Descriptor-based construction, as used by YAML-declared {@code subagents:} entries. */ + public MockExternalSubagentSetup( + ResourceDescriptor descriptor, @Nullable ResourceContext resourceContext) { + this( + (String) descriptor.getArgument("endpoint"), + Boolean.TRUE.equals(descriptor.getArgument("fail_on_call"))); + } + + public String getEndpointUrl() { + return endpointUrl; + } + + public boolean isFailOnCall() { + return failOnCall; + } + + @Override + protected DurableCallable asAsyncCallable( + RunnerContext ctx, Object prompt, String sessionId, String callId) { + return new BaseSubagentCallable(sessionId, callId) { + @Override + protected Object callInternal() throws Exception { + return simulateHttpCall(prompt); + } + }; + } + + private List simulateHttpCall(Object prompt) throws Exception { + // Token latency standing in for a network round trip. + Thread.sleep(50); + if (failOnCall) { + throw new IllegalStateException("endpoint " + endpointUrl + " is down"); + } + return List.of("HTTP response for: " + prompt + " from " + endpointUrl); + } +} diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/resources/yaml/external_subagent_agent.yaml b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/resources/yaml/external_subagent_agent.yaml new file mode 100644 index 000000000..7c52699a1 --- /dev/null +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/resources/yaml/external_subagent_agent.yaml @@ -0,0 +1,32 @@ +################################################################################ +# 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. +################################################################################ +agents: + - name: external_subagent_yaml_agent + description: E2E agent delegating to a YAML-declared external sub-agent. + + subagents: + - name: ext-agent + type: java + clazz: org.apache.flink.agents.integration.test.subagent.MockExternalSubagentSetup + endpoint: http://yaml-endpoint:8080 + + actions: + - name: call_external + type: java + function: org.apache.flink.agents.integration.test.subagent.ExternalSubagentAgent:callExternal + trigger_conditions: [input] diff --git a/plan/pom.xml b/plan/pom.xml index 02df3c2c3..3fbe9b42c 100644 --- a/plan/pom.xml +++ b/plan/pom.xml @@ -40,6 +40,13 @@ under the License. flink-agents-api ${project.version} + + org.apache.flink + flink-agents-api + ${project.version} + test-jar + test + org.apache.flink flink-agents-integrations-mcp diff --git a/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java b/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java index 3b08bf2cf..084db1627 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java @@ -37,6 +37,7 @@ import org.apache.flink.agents.api.resource.SerializableResource; import org.apache.flink.agents.api.skills.SkillSourceSpec; import org.apache.flink.agents.api.skills.Skills; +import org.apache.flink.agents.api.subagent.SubagentSetup; import org.apache.flink.agents.api.tools.ToolMetadata; import org.apache.flink.agents.api.tools.ToolParameterInjection; import org.apache.flink.agents.api.tools.ToolParameterInjectionValidator; @@ -80,7 +81,14 @@ import static org.apache.flink.agents.api.resource.ResourceType.PROMPT; import static org.apache.flink.agents.api.resource.ResourceType.TOOL; -/** Agent plan compiled from user defined agent. */ +/** + * Agent plan compiled from a user-defined {@link Agent}. + * + *

{@code AgentPlan} is a single-scope structure: each instance carries its own actions, + * actions-by-event dispatch table, resource providers, and optional configuration. Sub-agents are + * represented as {@link SubagentSetup} resources (type {@code AGENT}) within the resource providers + * map. + */ @JsonSerialize(using = AgentPlanJsonSerializer.class) @JsonDeserialize(using = AgentPlanJsonDeserializer.class) public class AgentPlan implements Serializable { @@ -581,6 +589,30 @@ private void extractResourceProvidersFromAgent(Agent agent) throws Exception { + " method on your Agent class so its tools and prompts can be" + " discovered."); } + } else if (type == ResourceType.AGENT) { + for (Map.Entry kv : entry.getValue().entrySet()) { + String name = kv.getKey(); + Object value = kv.getValue(); + if (value instanceof SubagentSetup) { + addResourceProvider( + JavaSerializableResourceProvider.createResourceProvider( + name, ResourceType.AGENT, (SubagentSetup) value)); + } else if (value instanceof ResourceDescriptor) { + // Declared via YAML: the descriptor names a SubagentSetup subclass that is + // instantiated when the resource is first resolved. + addResourceProvider( + createDescriptorResourceProvider( + name, ResourceType.AGENT, (ResourceDescriptor) value)); + } else { + throw new IllegalArgumentException( + "AGENT resource '" + + name + + "' must be a SubagentSetup or a ResourceDescriptor, but" + + " got " + + value.getClass().getName() + + "."); + } + } } else { for (Map.Entry kv : entry.getValue().entrySet()) { ResourceDescriptor descriptor = diff --git a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java new file mode 100644 index 000000000..8c801ae43 --- /dev/null +++ b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java @@ -0,0 +1,87 @@ +/* + * 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; + +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.resource.Resource; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.subagent.SubagentSetup; +import org.apache.flink.agents.api.subagent.TestSubagentSetup; +import org.apache.flink.agents.plan.resourceprovider.ResourceProvider; +import org.junit.jupiter.api.Test; + +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Tests compiling AGENT resources into the agent plan, for both registration shapes: a {@link + * SubagentSetup} instance (programmatic) and a {@link ResourceDescriptor} (the YAML shape). + */ +public class AgentPlanSubagentResourceTest { + + @Test + void subagentSetupInstanceCompilesIntoAgentProvider() throws Exception { + Agent agent = new Agent(); + agent.addResource("reviewer", ResourceType.AGENT, new TestSubagentSetup()); + + AgentPlan plan = new AgentPlan(agent); + + Map agentProviders = + plan.getResourceProviders().get(ResourceType.AGENT); + assertThat(agentProviders).containsKey("reviewer"); + Resource resolved = agentProviders.get("reviewer").provide(null); + assertThat(resolved).isInstanceOf(SubagentSetup.class); + } + + @Test + void agentDescriptorCompilesAndResolvesToSubagentSetup() throws Exception { + Agent agent = new Agent(); + agent.addResource( + "summarizer", + ResourceType.AGENT, + ResourceDescriptor.Builder.newBuilder(TestSubagentSetup.class.getName()) + .addInitialArgument("endpoint", "http://summarizer:8080") + .build()); + + AgentPlan plan = new AgentPlan(agent); + + Map agentProviders = + plan.getResourceProviders().get(ResourceType.AGENT); + assertThat(agentProviders).containsKey("summarizer"); + + Resource resolved = agentProviders.get("summarizer").provide(null); + assertThat(resolved).isInstanceOf(TestSubagentSetup.class); + assertThat(((TestSubagentSetup) resolved).getEndpoint()) + .isEqualTo("http://summarizer:8080"); + assertThat(resolved.getResourceType()).isEqualTo(ResourceType.AGENT); + } + + @Test + void nonSubagentAgentResourceIsRejected() { + Agent agent = new Agent(); + agent.getResources().get(ResourceType.AGENT).put("bad", "not-a-subagent"); + + assertThatThrownBy(() -> new AgentPlan(agent)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("must be a SubagentSetup or a ResourceDescriptor"); + } +} diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java index 750f12adc..0f6ff81f4 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java @@ -164,6 +164,16 @@ public T durableExecuteAsync(DurableCallable callable) throws Exception { return callable.call(); } + @Override + public String nextSessionId() { + throw new UnsupportedOperationException(); + } + + @Override + public String nextCallId(String sessionId) { + throw new UnsupportedOperationException(); + } + @Override public void close() {} } diff --git a/python/flink_agents/api/resource.py b/python/flink_agents/api/resource.py index 21d033429..b46ce1f39 100644 --- a/python/flink_agents/api/resource.py +++ b/python/flink_agents/api/resource.py @@ -32,7 +32,7 @@ class ResourceType(Enum): """Type enum of resource. Currently, support chat_model, chat_model_server, tool, embedding_model, - vector_store, prompt, mcp_server, skills. + vector_store, prompt, mcp_server, skills, agent. """ CHAT_MODEL = "chat_model" @@ -44,6 +44,7 @@ class ResourceType(Enum): PROMPT = "prompt" MCP_SERVER = "mcp_server" SKILLS = "skills" + AGENT = "agent" class Resource(BaseModel, ABC): diff --git a/python/flink_agents/api/runner_context.py b/python/flink_agents/api/runner_context.py index e110bd2f9..1bace0c11 100644 --- a/python/flink_agents/api/runner_context.py +++ b/python/flink_agents/api/runner_context.py @@ -312,6 +312,14 @@ async def my_action(event, ctx): An awaitable object that yields the function result when awaited. """ + @abstractmethod + def next_session_id(self) -> str: + """Create a unique sub-agent session id for the current action.""" + + @abstractmethod + def next_call_id(self, session_id: str) -> str: + """Create a call id for a sub-agent invocation under the given session.""" + @property @abstractmethod def config(self) -> ReadableConfiguration: diff --git a/python/flink_agents/api/subagent.py b/python/flink_agents/api/subagent.py new file mode 100644 index 000000000..9572d55a1 --- /dev/null +++ b/python/flink_agents/api/subagent.py @@ -0,0 +1,202 @@ +################################################################################ +# 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. +################################################################################# +import traceback +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar + +from flink_agents.api.resource import ResourceType, SerializableResource + +if TYPE_CHECKING: + from flink_agents.api.runner_context import RunnerContext + +T = TypeVar("T") + + +@dataclass +class Result: + """Outcome of a sub-agent call. + + A successful call carries a JSON-serializable ``result``; a failed call + carries a serializable ``error_message`` — the full stack trace of the + failure — rather than a live exception, so a result can be persisted + through durable execution and survive failover. + """ + + success: bool + result: Any = None + error_message: str | None = None + + @staticmethod + def ok(result: Any) -> "Result": + """Create a successful result carrying ``result``.""" + return Result(success=True, result=result) + + @staticmethod + def error(error: BaseException | str) -> "Result": + """Create a failed result from an exception or a plain message. + + For an exception the full stack trace is stored as a serializable + string so the result can survive durable execution. + """ + if isinstance(error, BaseException): + message = "".join( + traceback.format_exception(type(error), error, error.__traceback__) + ) + else: + message = error + return Result(success=False, error_message=message) + + @property + def exception(self) -> Exception | None: + """Reconstruct an exception carrying the stored stack trace; None on success.""" + return None if self.success else RuntimeError(self.error_message) + + +@dataclass +class DurableCallable(Generic[T]): + """A callable for durable execution that carries a stable identifier. + + Used with :meth:`RunnerContext.durable_execute` and + :meth:`RunnerContext.durable_execute_async` so each durable call has a + stable id that persists across job restarts. + """ + + id: str + call: Callable[[], T] + reconciler: Callable[[], T] | None = field(default=None) + + +class BaseSubagentCallable(DurableCallable["Result"], ABC): + """Convenience base for the callable returned by ``as_async_callable``. + + Keys the durable call by the framework-assigned identity as + ``session_id#call_id`` (the :class:`SubagentSetup` contract) and captures + exceptions raised by :meth:`call_internal` into :meth:`Result.error`, so + failures are reported through the result rather than raised. + Implementations only provide :meth:`call_internal`. + """ + + def __init__(self, session_id: str, call_id: str) -> None: + """Initialize with the framework-assigned identity.""" + super().__init__(id=f"{session_id}#{call_id}", call=self._invoke) + + def _invoke(self) -> "Result": + try: + return Result.ok(self.call_internal()) + except Exception as e: + return Result.error(e) + + @abstractmethod + def call_internal(self) -> Any: + """Perform the invocation and return the JSON-serializable payload. + + Raised exceptions are captured into a failed :class:`Result`. + """ + + +class Subagent(ABC): + """Caller-facing interface for a sub-agent invocable from within an action. + + An invocation is identified by a ``(session_id, call_id)`` pair; the + session groups a conversation across invocations. Both ids are assigned by + the framework (:meth:`RunnerContext.next_session_id` and + :meth:`RunnerContext.next_call_id`): callers may supply a session id to + continue a prior session but never supply a call id. + """ + + @abstractmethod + def call( + self, + ctx: "RunnerContext", + prompt: Any, + session_id: str | None = None, + ) -> Result: + """Synchronously invoke the sub-agent and return its :class:`Result`. + + An omitted ``session_id`` starts a new session. + """ + + @abstractmethod + def as_async_callable( + self, + ctx: "RunnerContext", + prompt: Any, + session_id: str | None = None, + ) -> DurableCallable[Result]: + """Return the deferred :class:`DurableCallable` for one invocation. + + An omitted ``session_id`` starts a new session. + """ + + +class SubagentSetup(SerializableResource, Subagent, ABC): + """Base setup for a sub-agent resource, registered as an AGENT resource. + + Hosts the id-resolution chain behind :class:`Subagent`: omitted ids are + assigned via :meth:`RunnerContext.next_session_id` and + :meth:`RunnerContext.next_call_id` before an implementation is ever + invoked, and :meth:`call` runs the deferred callable through durable + execution. Implementations only provide the terminal 4-arg + :meth:`as_async_callable` and contribute nothing to identity assignment. + """ + + @classmethod + def resource_type(cls) -> ResourceType: + """Return resource type of class.""" + return ResourceType.AGENT + + def call( + self, + ctx: "RunnerContext", + prompt: Any, + session_id: str | None = None, + call_id: str | None = None, + ) -> Result: + """Synchronously invoke the sub-agent and return its :class:`Result`. + + Omitted ids are assigned from the context (``call_id`` is + framework-facing and never supplied by callers). The deferred callable + runs through durable execution, so the invocation participates in + failover recovery. + """ + if session_id is None: + session_id = ctx.next_session_id() + if call_id is None: + call_id = ctx.next_call_id(session_id) + callable_ = self.as_async_callable(ctx, prompt, session_id, call_id) + return ctx.durable_execute( + callable_.call, reconciler=callable_.reconciler + ) + + @abstractmethod + def as_async_callable( + self, + ctx: "RunnerContext", + prompt: Any, + session_id: str | None = None, + call_id: str | None = None, + ) -> DurableCallable[Result]: + """Return the deferred :class:`DurableCallable` for one invocation. + + Implementations resolve ``None`` ids via ``ctx.next_session_id()`` / + ``ctx.next_call_id(session_id)``; :meth:`call` always passes resolved + ids. Contract: the returned callable's ``id`` MUST be derived solely + from the ``(session_id, call_id)`` pair (extend + :class:`BaseSubagentCallable` to get this for free). + """ diff --git a/python/flink_agents/api/tests/subagent_test_utils.py b/python/flink_agents/api/tests/subagent_test_utils.py new file mode 100644 index 000000000..e7335ad0e --- /dev/null +++ b/python/flink_agents/api/tests/subagent_test_utils.py @@ -0,0 +1,64 @@ +################################################################################ +# 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. +################################################################################ +"""Shared sub-agent test doubles.""" +from typing import Any + +from flink_agents.api.runner_context import RunnerContext +from flink_agents.api.subagent import ( + BaseSubagentCallable, + DurableCallable, + Result, + SubagentSetup, +) + + +class TestSubagentSetup(SubagentSetup): + """Shared ``SubagentSetup`` test double, constructible directly or from a + resource descriptor (the YAML shape). + + Echoes the prompt (prefixed with ``endpoint_url`` when set); + ``fail_on_call=True`` makes every call fail, surfacing through ``Result``. + """ + + endpoint_url: str | None = None + fail_on_call: bool = False + + def as_async_callable( + self, + ctx: RunnerContext, + prompt: Any, + session_id: str | None = None, + call_id: str | None = None, + ) -> DurableCallable[Result]: + """Return a callable echoing (or failing) one invocation.""" + if session_id is None: + session_id = ctx.next_session_id() + if call_id is None: + call_id = ctx.next_call_id(session_id) + setup = self + + class _Call(BaseSubagentCallable): + def call_internal(self) -> Any: + if setup.fail_on_call: + msg = f"endpoint {setup.endpoint_url} is down" + raise RuntimeError(msg) + if setup.endpoint_url is None: + return [prompt] + return [f"{setup.endpoint_url}:{prompt}"] + + return _Call(session_id, call_id) diff --git a/python/flink_agents/api/tests/test_subagent.py b/python/flink_agents/api/tests/test_subagent.py new file mode 100644 index 000000000..cd16d968e --- /dev/null +++ b/python/flink_agents/api/tests/test_subagent.py @@ -0,0 +1,146 @@ +################################################################################ +# 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. +################################################################################ +"""Tests registering sub-agents as AGENT resources and the call routing.""" +from typing import Any + +import pytest + +from flink_agents.api.agents.agent import Agent +from flink_agents.api.resource import ResourceType +from flink_agents.api.subagent import BaseSubagentCallable +from flink_agents.api.tests.subagent_test_utils import TestSubagentSetup + + +class _RecordingContext: + """Fake context recording session minting and durable execution.""" + + def __init__(self, session_id: str = "gen-session") -> None: + self._session_id = session_id + self.next_session_id_called = False + self.next_call_id_calls: list[str] = [] + self.durable_execute_calls: list[tuple] = [] + + def next_session_id(self) -> str: + self.next_session_id_called = True + return self._session_id + + def next_call_id(self, session_id: str) -> str: + self.next_call_id_calls.append(session_id) + return f"{session_id}-{len(self.next_call_id_calls)}" + + def durable_execute( + self, + func: Any, + *args: Any, + reconciler: Any = None, + **kwargs: Any, + ) -> Any: + self.durable_execute_calls.append((func, args, reconciler)) + return func(*args) + + +def test_register_subagent_setup_as_resource() -> None: + """A ``SubagentSetup`` registers under the AGENT resource map.""" + agent = Agent() + setup = TestSubagentSetup() + + agent.add_resource("reviewer", ResourceType.AGENT, setup) + + agent_resources = agent.resources[ResourceType.AGENT] + assert len(agent_resources) == 1 + assert agent_resources["reviewer"] is setup + assert setup.resource_type() == ResourceType.AGENT + + +def test_duplicate_name_throws() -> None: + """Registering a duplicate AGENT name raises.""" + agent = Agent() + agent.add_resource("reviewer", ResourceType.AGENT, TestSubagentSetup()) + + with pytest.raises(ValueError): + agent.add_resource("reviewer", ResourceType.AGENT, TestSubagentSetup()) + + +def test_multiple_subagents_registered() -> None: + """Multiple distinct AGENT resources coexist.""" + agent = Agent() + reviewer = TestSubagentSetup() + coder = TestSubagentSetup() + + agent.add_resource("reviewer", ResourceType.AGENT, reviewer) + agent.add_resource("coder", ResourceType.AGENT, coder) + + agent_resources = agent.resources[ResourceType.AGENT] + assert len(agent_resources) == 2 + assert agent_resources["reviewer"] is reviewer + assert agent_resources["coder"] is coder + + +def test_base_callable_captures_exception_into_error_result() -> None: + """``BaseSubagentCallable`` wraps exceptions into ``Result.error``.""" + + class _Failing(BaseSubagentCallable): + def call_internal(self) -> Any: + msg = "boom" + raise RuntimeError(msg) + + result = _Failing("sid-1", "c-1").call() + assert result.success is False + # error_message is the full stack trace, which contains the original message. + assert "boom" in result.error_message + assert result.exception is not None + + +def test_call_with_explicit_session_id_uses_durable_execute() -> None: + """Explicit-session ``call`` mints only a call id and routes durably.""" + setup = TestSubagentSetup() + ctx = _RecordingContext() + + result = setup.call(ctx, "hello", "explicit-sid") + + assert ctx.next_session_id_called is False + assert ctx.next_call_id_calls == ["explicit-sid"] + assert len(ctx.durable_execute_calls) == 1 + assert result.success is True + assert result.result == ["hello"] + + +def test_call_without_session_id_mints_both_ids() -> None: + """2-arg ``call`` mints a session and a call id, then routes.""" + setup = TestSubagentSetup() + ctx = _RecordingContext(session_id="minted") + + result = setup.call(ctx, "hi") + + assert ctx.next_session_id_called is True + assert ctx.next_call_id_calls == ["minted"] + assert len(ctx.durable_execute_calls) == 1 + assert result.success is True + assert result.result == ["hi"] + + +def test_callable_id_equals_assigned_call_id() -> None: + """The produced callable is keyed by the framework-assigned call id.""" + setup = TestSubagentSetup() + ctx = _RecordingContext() + + callable_ = setup.as_async_callable(ctx, "ping", "sid-1") + + assert ctx.next_call_id_calls == ["sid-1"] + assert callable_.id == "sid-1#sid-1-1" + assert callable_.call().result == ["ping"] diff --git a/python/flink_agents/api/yaml/loader.py b/python/flink_agents/api/yaml/loader.py index 49ba20601..1e9a02e79 100644 --- a/python/flink_agents/api/yaml/loader.py +++ b/python/flink_agents/api/yaml/loader.py @@ -62,6 +62,7 @@ "embedding_model_setups": ResourceType.EMBEDDING_MODEL, "vector_stores": ResourceType.VECTOR_STORE, "mcp_servers": ResourceType.MCP_SERVER, + "subagents": ResourceType.AGENT, } diff --git a/python/flink_agents/api/yaml/specs.py b/python/flink_agents/api/yaml/specs.py index ecfae4e5f..516e28446 100644 --- a/python/flink_agents/api/yaml/specs.py +++ b/python/flink_agents/api/yaml/specs.py @@ -233,6 +233,7 @@ class AgentSpec(BaseModel): embedding_model_setups: List[DescriptorSpec] = Field(default_factory=list) vector_stores: List[DescriptorSpec] = Field(default_factory=list) mcp_servers: List[DescriptorSpec] = Field(default_factory=list) + subagents: List[DescriptorSpec] = Field(default_factory=list) class YamlAgentsDocument(BaseModel): @@ -262,6 +263,7 @@ class YamlAgentsDocument(BaseModel): embedding_model_setups: List[DescriptorSpec] = Field(default_factory=list) vector_stores: List[DescriptorSpec] = Field(default_factory=list) mcp_servers: List[DescriptorSpec] = Field(default_factory=list) + subagents: List[DescriptorSpec] = Field(default_factory=list) def export() -> str: diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/external_subagent_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/external_subagent_agent.py new file mode 100644 index 000000000..c6160d00a --- /dev/null +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/external_subagent_agent.py @@ -0,0 +1,103 @@ +################################################################################ +# 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. +################################################################################ +"""External sub-agent e2e agent: mock setup, agent, and shared actions.""" +import time +from typing import Any + +from pyflink.datastream import KeySelector + +from flink_agents.api.agents.agent import Agent +from flink_agents.api.decorators import action +from flink_agents.api.events.event import Event, InputEvent, OutputEvent +from flink_agents.api.resource import ResourceType +from flink_agents.api.runner_context import RunnerContext +from flink_agents.api.subagent import ( + BaseSubagentCallable, + DurableCallable, + Result, + SubagentSetup, +) + +SUBAGENT_NAME = "ext-agent" + + +class MockExternalSubagentSetup(SubagentSetup): + """Mock external sub-agent simulating an HTTP-based external agent system. + + Constructible directly or from a resource descriptor (the YAML shape). + Internal failures are captured into a ``Result`` rather than raised. + """ + + endpoint_url: str + fail_on_call: bool = False + + def as_async_callable( + self, + ctx: RunnerContext, + prompt: Any, + session_id: str | None = None, + call_id: str | None = None, + ) -> DurableCallable[Result]: + """Return a callable performing one simulated HTTP call.""" + if session_id is None: + session_id = ctx.next_session_id() + if call_id is None: + call_id = ctx.next_call_id(session_id) + setup = self + + class _HttpCall(BaseSubagentCallable): + def call_internal(self) -> Any: + return setup._simulate_http_call(prompt) + + return _HttpCall(session_id, call_id) + + def _simulate_http_call(self, prompt: Any) -> list: + # Token latency standing in for a network round trip. + time.sleep(0.05) + if self.fail_on_call: + msg = f"endpoint {self.endpoint_url} is down" + raise RuntimeError(msg) + return [f"HTTP response for: {prompt} from {self.endpoint_url}"] + + +def call_external(event: Event, ctx: RunnerContext) -> None: + """Call the sub-agent with an explicit session id and emit its result.""" + setup = ctx.get_resource(SUBAGENT_NAME, ResourceType.AGENT) + prompt = InputEvent.from_event(event).input + result = setup.call(ctx, prompt, f"session-{prompt}") + if result.success: + ctx.send_event(OutputEvent(output=result.result[0])) + else: + ctx.send_event(OutputEvent(output=f"error:{result.error_message}")) + + +class ExternalSubagentAgent(Agent): + """Agent whose input action delegates to the external sub-agent.""" + + @action(InputEvent.EVENT_TYPE) + @staticmethod + def call_external(event: Event, ctx: RunnerContext) -> None: + call_external(event, ctx) + + +class InputKeySelector(KeySelector): + """Keys every element by itself.""" + + def get_key(self, value: Any) -> Any: + """Return the element itself as key.""" + return value diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/external_subagent_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/external_subagent_test.py new file mode 100644 index 000000000..ad303ec01 --- /dev/null +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/external_subagent_test.py @@ -0,0 +1,131 @@ +################################################################################ +# 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. +################################################################################ +"""E2E tests for external sub-agents on a real Flink job. + +Covers registration, invocation, failure reporting via ``Result``, and +YAML-declared registration. +""" +import os +import sysconfig +from pathlib import Path + +from pyflink.common import Encoder +from pyflink.common.typeinfo import Types +from pyflink.datastream import StreamExecutionEnvironment +from pyflink.datastream.connectors.file_system import StreamingFileSink + +from flink_agents.api.execution_environment import AgentsExecutionEnvironment +from flink_agents.api.resource import ResourceType +from flink_agents.e2e_tests.e2e_tests_integration.external_subagent_agent import ( + SUBAGENT_NAME, + ExternalSubagentAgent, + InputKeySelector, + MockExternalSubagentSetup, +) + +current_dir = Path(__file__).parent +_RESOURCES = current_dir.parent / "resources" + +os.environ["PYTHONPATH"] = sysconfig.get_paths()["purelib"] + + +def _run_and_collect(agents_env, output_datastream, tmp_path: Path) -> str: + result_dir = tmp_path / "results" + result_dir.mkdir(parents=True, exist_ok=True) + output_datastream.map(str, Types.STRING()).add_sink( + StreamingFileSink.for_row_format( + base_path=str(result_dir.absolute()), + encoder=Encoder.simple_string_encoder(), + ).build() + ) + agents_env.execute() + return "".join(p.read_text() for p in result_dir.rglob("*") if p.is_file()) + + +def test_python_external_subagent_call_end_to_end(tmp_path: Path) -> None: + """Registration -> execution -> sub-agent response in the job output.""" + env = StreamExecutionEnvironment.get_execution_environment() + env.set_parallelism(1) + input_stream = env.from_collection(["a", "b"]) + + agent = ExternalSubagentAgent() + agent.add_resource( + SUBAGENT_NAME, + ResourceType.AGENT, + MockExternalSubagentSetup(endpoint_url="http://ext:8080"), + ) + + agents_env = AgentsExecutionEnvironment.get_execution_environment(env=env) + output = ( + agents_env.from_datastream(input=input_stream, key_selector=InputKeySelector()) + .apply(agent) + .to_datastream() + ) + + contents = _run_and_collect(agents_env, output, tmp_path) + for prompt in ("a", "b"): + assert f"HTTP response for: {prompt} from http://ext:8080" in contents + + +def test_python_external_subagent_failure_surfaces_via_result( + tmp_path: Path, +) -> None: + """A failing endpoint surfaces through Result, not a job failure.""" + env = StreamExecutionEnvironment.get_execution_environment() + env.set_parallelism(1) + input_stream = env.from_collection(["a"]) + + agent = ExternalSubagentAgent() + agent.add_resource( + SUBAGENT_NAME, + ResourceType.AGENT, + MockExternalSubagentSetup(endpoint_url="http://down:8080", fail_on_call=True), + ) + + agents_env = AgentsExecutionEnvironment.get_execution_environment(env=env) + output = ( + agents_env.from_datastream(input=input_stream, key_selector=InputKeySelector()) + .apply(agent) + .to_datastream() + ) + + contents = _run_and_collect(agents_env, output, tmp_path) + # error_message is the full stack trace, which contains the original message. + assert "endpoint http://down:8080 is down" in contents + + +def test_yaml_external_subagent_call_end_to_end(tmp_path: Path) -> None: + """A YAML-declared sub-agent resolves and executes in the job.""" + env = StreamExecutionEnvironment.get_execution_environment() + env.set_parallelism(1) + input_stream = env.from_collection(["a", "b"]) + + agents_env = AgentsExecutionEnvironment.get_execution_environment(env=env) + agents_env.load_yaml(str(_RESOURCES / "external_subagent_agent.yaml")) + + output = ( + agents_env.from_datastream(input=input_stream, key_selector=InputKeySelector()) + .apply("external_subagent_yaml_agent") + .to_datastream() + ) + + contents = _run_and_collect(agents_env, output, tmp_path) + for prompt in ("a", "b"): + assert ( + f"HTTP response for: {prompt} from http://yaml-endpoint:8080" in contents + ) diff --git a/python/flink_agents/e2e_tests/resources/external_subagent_agent.yaml b/python/flink_agents/e2e_tests/resources/external_subagent_agent.yaml new file mode 100644 index 000000000..307f775d5 --- /dev/null +++ b/python/flink_agents/e2e_tests/resources/external_subagent_agent.yaml @@ -0,0 +1,30 @@ +################################################################################ +# 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. +################################################################################ +agents: + - name: external_subagent_yaml_agent + description: E2E agent delegating to a YAML-declared external sub-agent. + + subagents: + - name: ext-agent + clazz: flink_agents.e2e_tests.e2e_tests_integration.external_subagent_agent.MockExternalSubagentSetup + endpoint_url: http://yaml-endpoint:8080 + + actions: + - name: call_external + function: flink_agents.e2e_tests.e2e_tests_integration.external_subagent_agent:call_external + trigger_conditions: [input] diff --git a/python/flink_agents/plan/agent_plan.py b/python/flink_agents/plan/agent_plan.py index 747aa8577..c0f0ecbe3 100644 --- a/python/flink_agents/plan/agent_plan.py +++ b/python/flink_agents/plan/agent_plan.py @@ -33,6 +33,7 @@ LOAD_SKILL_TOOL, Skills, ) +from flink_agents.api.subagent import SubagentSetup from flink_agents.api.tools.function_tool import FunctionTool as ApiFunctionTool from flink_agents.api.tools.tool import Tool from flink_agents.plan.actions.action import Action @@ -424,6 +425,25 @@ def _get_resource_providers( ) _add_skills(all_skills, resource_providers) + # Handle AGENT resources (subagent setups) + for name, value in agent.resources[ResourceType.AGENT].items(): + if isinstance(value, SubagentSetup): + resource_providers.append( + PythonSerializableResourceProvider.from_resource( + name=name, resource=value + ) + ) + elif isinstance(value, ResourceDescriptor): + resource_providers.append( + PythonResourceProvider.get(name=name, descriptor=value) + ) + else: + msg = ( + f"AGENT resource '{name}' must be a SubagentSetup or a " + f"ResourceDescriptor, but got {type(value).__name__}." + ) + raise TypeError(msg) + for resource_type in [ ResourceType.CHAT_MODEL, ResourceType.CHAT_MODEL_CONNECTION, diff --git a/python/flink_agents/plan/tests/test_agent_plan_subagent_resources.py b/python/flink_agents/plan/tests/test_agent_plan_subagent_resources.py new file mode 100644 index 000000000..ed9f5eaab --- /dev/null +++ b/python/flink_agents/plan/tests/test_agent_plan_subagent_resources.py @@ -0,0 +1,74 @@ +################################################################################ +# 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. +################################################################################ +"""Tests for compiling AGENT resources (SubagentSetup) into the agent plan.""" +import pytest + +from flink_agents.api.agents.agent import Agent +from flink_agents.api.resource import ResourceDescriptor, ResourceType +from flink_agents.api.subagent import SubagentSetup +from flink_agents.api.tests.subagent_test_utils import TestSubagentSetup +from flink_agents.plan.agent_plan import AgentPlan +from flink_agents.plan.configuration import AgentConfiguration + + +def test_subagent_setup_compiles_into_agent_provider() -> None: + """A registered SubagentSetup lands in the AGENT provider map and resolves.""" + setup = TestSubagentSetup() + agent = Agent() + agent.add_resource("reviewer", ResourceType.AGENT, setup) + + plan = AgentPlan.from_agent(agent, AgentConfiguration()) + + agents = plan.resource_providers[ResourceType.AGENT] + assert agents is not None + assert "reviewer" in agents + resolved = agents["reviewer"].provide( + resource_context=None, config=AgentConfiguration() + ) + assert isinstance(resolved, SubagentSetup) + result = resolved.as_async_callable(None, "ping", "sid-1", "sid-1-1").call() + assert result.success is True + assert result.result == ["ping"] + + +def test_agent_descriptor_compiles_into_agent_provider() -> None: + """Descriptor-shaped AGENT resources (the YAML path) compile into providers.""" + agent = Agent() + agent.add_resource( + "summarizer", + ResourceType.AGENT, + ResourceDescriptor( + clazz=f"{TestSubagentSetup.__module__}.{TestSubagentSetup.__name__}", + endpoint_url="http://summarizer:8080", + ), + ) + + plan = AgentPlan.from_agent(agent, AgentConfiguration()) + + agents = plan.resource_providers[ResourceType.AGENT] + assert agents is not None + assert "summarizer" in agents + + +def test_non_setup_agent_resource_is_rejected() -> None: + """A bare object registered under AGENT fails plan compilation.""" + agent = Agent() + agent.resources[ResourceType.AGENT]["bad"] = object() + + with pytest.raises(TypeError, match="must be a SubagentSetup"): + AgentPlan.from_agent(agent, AgentConfiguration()) diff --git a/python/flink_agents/runtime/flink_runner_context.py b/python/flink_agents/runtime/flink_runner_context.py index ada5b8e30..d14e324a9 100644 --- a/python/flink_agents/runtime/flink_runner_context.py +++ b/python/flink_agents/runtime/flink_runner_context.py @@ -738,6 +738,16 @@ def durable_execute_async( kwargs, ) + @override + def next_session_id(self) -> str: + """Create a unique sub-agent session id for the current action.""" + return self._j_runner_context.nextSessionId() + + @override + def next_call_id(self, session_id: str) -> str: + """Create a call id for a sub-agent invocation under ``session_id``.""" + return self._j_runner_context.nextCallId(session_id) + @property @override def config(self) -> ReadableConfiguration: diff --git a/python/flink_agents/runtime/tests/test_subagent_identity.py b/python/flink_agents/runtime/tests/test_subagent_identity.py new file mode 100644 index 000000000..21b529d7e --- /dev/null +++ b/python/flink_agents/runtime/tests/test_subagent_identity.py @@ -0,0 +1,79 @@ +################################################################################ +# 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. +################################################################################ +"""Tests for sub-agent identity allocation on the Python side. + +``FlinkRunnerContext`` forwards both allocators to the Java runner context over +pemja; determinism and failover-replay behavior are covered by the Java runtime +tests (``SubagentIdentityContextTest`` / ``SubagentIdentityRecoveryTest``). +""" +from typing import Any, List + +from flink_agents.runtime.flink_runner_context import FlinkRunnerContext + + +class _FakeJavaRunnerContextForIdentity: + """Records nextSessionId()/nextCallId() calls made by FlinkRunnerContext.""" + + def __init__(self) -> None: + self.next_session_id_calls = 0 + self.next_call_id_calls: List[str] = [] + + def nextSessionId(self) -> str: + """Fake Java allocator: return a distinguishable, incrementing id.""" + self.next_session_id_calls += 1 + return f"java-session-{self.next_session_id_calls}" + + def nextCallId(self, session_id: str) -> str: + """Fake Java allocator: record the session id it was given.""" + self.next_call_id_calls.append(session_id) + return f"java-call-{session_id}-{len(self.next_call_id_calls)}" + + +def _create_flink_runner_context(j_runner_context: Any) -> FlinkRunnerContext: + """Bare-construct a FlinkRunnerContext, bypassing __init__. + + Mirrors the ``FlinkRunnerContext.__new__(FlinkRunnerContext)`` pattern + used by ``test_flink_runner_context_reconcilable.py``. + """ + ctx = FlinkRunnerContext.__new__(FlinkRunnerContext) + ctx._j_runner_context = j_runner_context + return ctx + + +def test_flink_runner_context_next_session_id_delegates_to_java() -> None: + """next_session_id() is a pure delegation to the Java allocator.""" + j_ctx = _FakeJavaRunnerContextForIdentity() + ctx = _create_flink_runner_context(j_ctx) + + session_id = ctx.next_session_id() + + assert session_id == "java-session-1" + assert j_ctx.next_session_id_calls == 1 + + +def test_flink_runner_context_next_call_id_delegates_to_java_with_session_id() -> ( + None +): + """next_call_id() forwards the session id to the Java allocator.""" + j_ctx = _FakeJavaRunnerContextForIdentity() + ctx = _create_flink_runner_context(j_ctx) + + call_id = ctx.next_call_id("session-x") + + assert call_id == "java-call-session-x-1" + assert j_ctx.next_call_id_calls == ["session-x"] diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java index 1395bdc56..601e7e8bb 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java @@ -19,7 +19,10 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.MapperFeature; import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.SerializationFeature; +import com.fasterxml.jackson.databind.json.JsonMapper; import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.configuration.ReadableConfiguration; @@ -47,10 +50,12 @@ import javax.annotation.Nullable; import java.util.ArrayList; +import java.util.HashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.UUID; import java.util.concurrent.Callable; /** @@ -108,6 +113,13 @@ public CachedMemoryStore getSensoryMemStore() { /** Context for fine-grained durable execution, may be null if not enabled. */ @Nullable protected DurableExecutionContext durableExecutionContext; + /** + * Per-task context that deterministically assigns sub-agent session/call ids. Attached + * alongside the rest of the per-task state by {@link #switchActionContext} before {@link + * #nextSessionId()} or {@link #nextCallId(String)} are invoked. + */ + @Nullable private SubagentIdentityContext subagentIdentityContext; + public RunnerContextImpl( FlinkAgentsMetricGroupImpl agentMetricGroup, Runnable mailboxThreadChecker, @@ -124,9 +136,14 @@ public void setLongTermMemory(InteranlBaseLongTermMemory ltm) { this.ltm = ltm; } - public void switchActionContext(String actionName, MemoryContext memoryContext, String key) { + public void switchActionContext( + String actionName, + MemoryContext memoryContext, + SubagentIdentityContext subagentIdentityContext, + String key) { this.actionName = actionName; this.memoryContext = memoryContext; + this.subagentIdentityContext = subagentIdentityContext; if (ltm != null) { ltm.switchContext(key); } @@ -265,6 +282,22 @@ public T durableExecuteAsync(DurableCallable callable) throws Exception { return durableExecute(callable); } + @Override + public String nextSessionId() { + Preconditions.checkState( + subagentIdentityContext != null, + "Subagent identity context is not attached to the current action task."); + return subagentIdentityContext.nextSessionId(); + } + + @Override + public String nextCallId(String sessionId) { + Preconditions.checkState( + subagentIdentityContext != null, + "Subagent identity context is not attached to the current action task."); + return subagentIdentityContext.nextCallId(sessionId); + } + /** * Executes a durable call using the completion-only state machine. * @@ -371,6 +404,11 @@ public void clearDurableExecutionContext() { this.durableExecutionContext = null; } + @Nullable + public SubagentIdentityContext getSubagentIdentityContext() { + return subagentIdentityContext; + } + /** * Matches the next call result for recovery, or clears subsequent results if mismatch detected. * @@ -582,6 +620,103 @@ protected static class DurableExecutionRuntimeException extends RuntimeException } } + /** + * Caller-side facts identifying one action execution, used as the namespace for deterministic + * sub-agent id assignment: record key, sequence number, caller action name, and the triggering + * event (represented by its type and attributes, so two replays of the same logical event map + * to the same namespace regardless of the event instance id). + */ + public static final class SubagentIdentityNamespace { + + @JsonProperty("key") + private final String key; + + @JsonProperty("sequenceNumber") + private final long sequenceNumber; + + @JsonProperty("actionName") + private final String actionName; + + @JsonProperty("eventType") + private final String eventType; + + @JsonProperty("eventAttributes") + private final Map eventAttributes; + + public SubagentIdentityNamespace( + Object key, long sequenceNumber, String actionName, Event event) { + this.key = key.toString(); + this.sequenceNumber = sequenceNumber; + this.actionName = actionName; + this.eventType = event.getType(); + this.eventAttributes = event.getAttributes(); + } + } + + /** + * Per-{@code ActionTask} context that deterministically assigns sub-agent session and call ids. + * + *

The namespace is derived purely from caller-side facts, so a failover replay reproduces + * the same digest and therefore the same id sequence. The context is transient per-task heap + * state: continuation resume carries it forward (ordinals continue), failover rebuilds it + * (ordinals restart). The digest is computed lazily on the first allocation. + */ + public static final class SubagentIdentityContext { + + /** + * Sorts map entries and bean properties so the namespace bytes do not depend on map + * iteration order, which is not guaranteed across JVMs. + */ + private static final ObjectMapper DIGEST_MAPPER = + JsonMapper.builder() + .configure(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS, true) + .configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true) + .build(); + + private final SubagentIdentityNamespace namespace; + + /** Computed lazily on the first allocation; mailbox-confined, no synchronization. */ + @Nullable private String namespaceDigest; + + private int sessionOrdinal; + private final Map perSessionCallOrdinals = new HashMap<>(); + + public SubagentIdentityContext( + Object key, long sequenceNumber, String actionName, Event event) { + this.namespace = new SubagentIdentityNamespace(key, sequenceNumber, actionName, event); + } + + /** Creates a new, ordinal-increasing session id scoped to this task's namespace. */ + public String nextSessionId() { + return namespaceDigest() + "-" + (sessionOrdinal++); + } + + /** + * Creates a new call id by appending the per-session ordinal (starting at 1) to the session + * id. Cross-task uniqueness relies on session ids not being shared between action + * executions (see the {@code RunnerContext#nextCallId(String)} contract). + */ + public String nextCallId(String sessionId) { + int ordinal = perSessionCallOrdinals.merge(sessionId, 1, Integer::sum); + return sessionId + "-" + ordinal; + } + + private String namespaceDigest() { + if (namespaceDigest == null) { + try { + namespaceDigest = + String.valueOf( + UUID.nameUUIDFromBytes( + DIGEST_MAPPER.writeValueAsBytes(namespace))); + } catch (JsonProcessingException e) { + throw new IllegalStateException( + "Failed to digest the sub-agent identity namespace", e); + } + } + return namespaceDigest; + } + } + /** * Context for fine-grained durable execution within an action. * diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java index 46d09e0a2..4a0a81583 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java @@ -316,9 +316,11 @@ private void processActionTaskForKey(Object key) throws Exception { } // 2. Invoke the action task. + long sequenceNumber = stateManager.getSequenceNumber(); contextManager.createAndSetRunnerContext( actionTask, key, + sequenceNumber, agentPlan, resourceCache, metricGroup, @@ -329,7 +331,6 @@ private void processActionTaskForKey(Object key) throws Exception { pythonBridge.getPythonRunnerContext(), ltm); - long sequenceNumber = stateManager.getSequenceNumber(); boolean isFinished; List outputEvents; Optional generatedActionTaskOpt = Optional.empty(); @@ -385,6 +386,7 @@ private void processActionTaskForKey(Object key) throws Exception { durableExecManager.removeDurableContext(actionTask); contextManager.removeContinuationContext(actionTask); contextManager.removePythonAwaitableRef(actionTask); + contextManager.removeIdentityContext(actionTask); durableExecManager.maybePersistTaskResult( key, sequenceNumber, diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java index eba25d5d4..6e9e19df9 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java @@ -72,6 +72,8 @@ class ActionTaskContextManager implements AutoCloseable { private final Map actionTaskMemoryContexts; private final Map continuationContexts; private final Map pythonAwaitableRefs; + private final Map + subagentIdentityContexts; private ContinuationActionExecutor continuationActionExecutor; @@ -79,6 +81,7 @@ class ActionTaskContextManager implements AutoCloseable { this.actionTaskMemoryContexts = new HashMap<>(); this.continuationContexts = new HashMap<>(); this.pythonAwaitableRefs = new HashMap<>(); + this.subagentIdentityContexts = new HashMap<>(); this.continuationActionExecutor = new ContinuationActionExecutor(numAsyncThreads); } @@ -170,6 +173,7 @@ RunnerContextImpl createOrGetRunnerContext( void createAndSetRunnerContext( ActionTask actionTask, Object key, + long sequenceNumber, AgentPlan agentPlan, ResourceCache resourceCache, FlinkAgentsMetricGroupImpl metricGroup, @@ -217,8 +221,21 @@ void createAndSetRunnerContext( new CachedMemoryStore(shortTermMemState)); } + RunnerContextImpl.SubagentIdentityContext identityContext = + subagentIdentityContexts.get(actionTask); + if (identityContext == null) { + // Built from caller-side facts only, so a failover replay (task absent from this + // transient map) reproduces an identical namespace and allocation sequence. + identityContext = + new RunnerContextImpl.SubagentIdentityContext( + key, sequenceNumber, actionTask.action.getName(), actionTask.event); + } + context.switchActionContext( - actionTask.action.getName(), memoryContext, String.valueOf(key.hashCode())); + actionTask.action.getName(), + memoryContext, + identityContext, + String.valueOf(key.hashCode())); if (context instanceof JavaRunnerContextImpl) { ContinuationContext continuationContext; @@ -250,14 +267,19 @@ RunnerContextImpl.MemoryContext removeMemoryContext(ActionTask actionTask) { return actionTaskMemoryContexts.remove(actionTask); } + void removeIdentityContext(ActionTask actionTask) { + subagentIdentityContexts.remove(actionTask); + } + /** * Transfers per-task contexts from a finishing action task to the action task it generated. * - *

Always transfers the memory context. For Java tasks, transfers the continuation context. - * For Python tasks, transfers the awaitable reference when present. The durable-execution - * context map lives on {@link DurableExecutionManager}, so that manager is passed in as a - * parameter rather than held as a field — this keeps the no-manager-to-manager-references - * design constraint intact. + *

Always transfers the memory context and, if present, the sub-agent identity context (so a + * same-process continuation resume keeps its id ordinals). For Java tasks, transfers the + * continuation context. For Python tasks, transfers the awaitable reference when present. The + * durable-execution context map lives on {@link DurableExecutionManager}, so that manager is + * passed in as a parameter rather than held as a field — this keeps the + * no-manager-to-manager-references design constraint intact. * * @param fromTask the finishing task whose contexts should be transferred. * @param toTask the newly generated task that will inherit the contexts. @@ -271,6 +293,11 @@ void transferContexts( if (durableContext != null) { durableExecManager.putDurableContext(toTask, durableContext); } + RunnerContextImpl.SubagentIdentityContext identityContext = + fromTask.getRunnerContext().getSubagentIdentityContext(); + if (identityContext != null) { + subagentIdentityContexts.put(toTask, identityContext); + } if (fromTask.getRunnerContext() instanceof JavaRunnerContextImpl) { this.putContinuationContext( toTask, diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/SubagentIdentityContextTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/SubagentIdentityContextTest.java new file mode 100644 index 000000000..69667c8b0 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/SubagentIdentityContextTest.java @@ -0,0 +1,226 @@ +/* + * 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.runtime.context; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.InputEvent; +import org.junit.jupiter.api.Test; + +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Unit tests for {@link RunnerContextImpl.SubagentIdentityContext}: same namespace inputs plus the + * same call sequence must reproduce byte-for-byte identical ids, and any difference in the + * caller-side namespace inputs must separate the id stream. + */ +class SubagentIdentityContextTest { + + private static RunnerContextImpl.SubagentIdentityContext newContext( + Object key, long sequenceNumber, String actionName, Object eventInput) { + return new RunnerContextImpl.SubagentIdentityContext( + key, sequenceNumber, actionName, new InputEvent(eventInput)); + } + + @Test + void eventInstanceIdDoesNotAffectNamespace() { + // The namespace keeps only the event's type and attributes, so a failover replay that + // re-creates the event object (with a new event id) reproduces the same ids. + RunnerContextImpl.SubagentIdentityContext ctx1 = + new RunnerContextImpl.SubagentIdentityContext( + "key-1", 7L, "actionA", new InputEvent("payload")); + RunnerContextImpl.SubagentIdentityContext ctx2 = + new RunnerContextImpl.SubagentIdentityContext( + "key-1", 7L, "actionA", new InputEvent("payload")); + + assertEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void attributeMapIterationOrderDoesNotAffectNamespace() { + // A replay on another JVM may iterate the rebuilt attribute map differently, so the digest + // mapper sorts entries: identical contents in reversed order must yield the same id. + Map forward = new LinkedHashMap<>(); + forward.put("alpha", "1"); + forward.put("beta", "2"); + forward.put("gamma", "3"); + Map reversed = new LinkedHashMap<>(); + reversed.put("gamma", "3"); + reversed.put("beta", "2"); + reversed.put("alpha", "1"); + + RunnerContextImpl.SubagentIdentityContext ctx1 = + new RunnerContextImpl.SubagentIdentityContext( + "key-1", 7L, "actionA", new Event("my.EventType", forward)); + RunnerContextImpl.SubagentIdentityContext ctx2 = + new RunnerContextImpl.SubagentIdentityContext( + "key-1", 7L, "actionA", new Event("my.EventType", reversed)); + + assertEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void deterministicAcrossTwoInstancesWithSameNamespace() { + RunnerContextImpl.SubagentIdentityContext ctx1 = newContext("key-1", 7L, "actionA", "e-1"); + RunnerContextImpl.SubagentIdentityContext ctx2 = newContext("key-1", 7L, "actionA", "e-1"); + + // Same call sequence replayed on two independently constructed instances that share an + // identical namespace must reproduce byte-for-byte identical ids -- this is the core + // property that makes failover replay reproduce the same sub-agent identities. + String session1a = ctx1.nextSessionId(); + String session2a = ctx2.nextSessionId(); + assertEquals(session1a, session2a); + + String call1a = ctx1.nextCallId(session1a); + String call2a = ctx2.nextCallId(session2a); + assertEquals(call1a, call2a); + + String session1b = ctx1.nextSessionId(); + String session2b = ctx2.nextSessionId(); + assertEquals(session1b, session2b); + assertNotEquals(session1a, session1b); + } + + @Test + void namespaceSeparatesOnKey() { + RunnerContextImpl.SubagentIdentityContext ctx1 = newContext("key-1", 7L, "actionA", "e-1"); + RunnerContextImpl.SubagentIdentityContext ctx2 = newContext("key-2", 7L, "actionA", "e-1"); + + // The ordinal is 0 on the first call for both instances, so any difference in the + // returned id must come from the namespace digest alone. + assertNotEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void namespaceSeparatesOnSequenceNumber() { + RunnerContextImpl.SubagentIdentityContext ctx1 = newContext("key-1", 7L, "actionA", "e-1"); + RunnerContextImpl.SubagentIdentityContext ctx2 = newContext("key-1", 8L, "actionA", "e-1"); + + assertNotEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void namespaceSeparatesOnActionName() { + RunnerContextImpl.SubagentIdentityContext ctx1 = newContext("key-1", 7L, "actionA", "e-1"); + RunnerContextImpl.SubagentIdentityContext ctx2 = newContext("key-1", 7L, "actionB", "e-1"); + + assertNotEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void namespaceSeparatesOnTriggeringEventAloneSiblingTaskScenario() { + // Two sibling tasks sharing key/sequenceNumber/actionName but triggered by different + // events: the event attributes alone must keep their identities from colliding. + RunnerContextImpl.SubagentIdentityContext ctx1 = newContext("key-1", 7L, "actionA", "e-1"); + RunnerContextImpl.SubagentIdentityContext ctx2 = newContext("key-1", 7L, "actionA", "e-2"); + + assertNotEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void sessionOrdinalsIncreaseAndAreUnique() { + RunnerContextImpl.SubagentIdentityContext ctx = newContext("key-1", 1L, "actionA", "e-1"); + + String s0 = ctx.nextSessionId(); + String s1 = ctx.nextSessionId(); + String s2 = ctx.nextSessionId(); + + // The namespace digest is a fixed-length UUID string, so "-" unambiguously identifies + // the ordinal suffix. + assertTrue(s0.endsWith("-0")); + assertTrue(s1.endsWith("-1")); + assertTrue(s2.endsWith("-2")); + + Set unique = new HashSet<>(); + unique.add(s0); + unique.add(s1); + unique.add(s2); + assertEquals(3, unique.size()); + } + + @Test + void perSessionCallOrdinalStartsAtOneAndIncrementsPerSession() { + RunnerContextImpl.SubagentIdentityContext ctx = newContext("key-1", 1L, "actionA", "e-1"); + + String sessionA = ctx.nextSessionId(); + String sessionB = ctx.nextSessionId(); + + String callA1 = ctx.nextCallId(sessionA); + String callA2 = ctx.nextCallId(sessionA); + String callB1 = ctx.nextCallId(sessionB); + + assertTrue(callA1.endsWith("-1")); + assertTrue(callA2.endsWith("-2")); + // A different session's ordinal is tracked independently and also starts at 1, rather + // than continuing sessionA's running count. + assertTrue(callB1.endsWith("-1")); + + assertNotEquals(callA1, callA2); + assertNotEquals(callA1, callB1); + } + + @Test + void callIdIsSessionIdPlusOrdinal() { + RunnerContextImpl.SubagentIdentityContext ctx = newContext("key-1", 1L, "actionA", "e-1"); + + // The call id is formed by appending the per-session ordinal to the session id -- the + // session id already carries the namespace digest, so no further hashing is involved. + String sessionId = ctx.nextSessionId(); + assertEquals(sessionId + "-1", ctx.nextCallId(sessionId)); + assertEquals(sessionId + "-2", ctx.nextCallId(sessionId)); + } + + @Test + void explicitSessionIdIsUsedVerbatimNotParsed() { + RunnerContextImpl.SubagentIdentityContext ctx = newContext("key-1", 1L, "actionA", "e-1"); + + // Caller-supplied session ids need not follow the "{digest}-{ordinal}" shape; any string + // is legal and distinct ids never collide (they must not be reused across executions). + String callForExplicit1 = ctx.nextCallId("checkout-session-42"); + String callForExplicit2 = ctx.nextCallId("checkout-session-43"); + assertEquals("checkout-session-42-1", callForExplicit1); + assertEquals("checkout-session-43-1", callForExplicit2); + assertNotEquals(callForExplicit1, callForExplicit2); + + // Determinism holds for explicit session ids too: a second, freshly constructed context + // sharing the same namespace reproduces the same id for the same explicit session id. + RunnerContextImpl.SubagentIdentityContext ctx2 = newContext("key-1", 1L, "actionA", "e-1"); + assertEquals(callForExplicit1, ctx2.nextCallId("checkout-session-42")); + } + + @Test + void idsAreOpaqueAndDoNotLeakPlainKeyOrActionName() { + RunnerContextImpl.SubagentIdentityContext ctx = + newContext("super-secret-key", 1L, "myCustomActionName", "e-1"); + + String sessionId = ctx.nextSessionId(); + String callId = ctx.nextCallId(sessionId); + + assertFalse(sessionId.contains("super-secret-key")); + assertFalse(sessionId.contains("myCustomActionName")); + assertFalse(callId.contains("super-secret-key")); + assertFalse(callId.contains("myCustomActionName")); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryRefTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryRefTest.java index e16361552..3fcdcbbd5 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryRefTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryRefTest.java @@ -128,6 +128,16 @@ public T durableExecuteAsync(DurableCallable callable) throws Exception { return callable.call(); } + @Override + public String nextSessionId() { + throw new UnsupportedOperationException(); + } + + @Override + public String nextCallId(String sessionId) { + throw new UnsupportedOperationException(); + } + @Override public void close() throws Exception {} } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java index 8fec7cb63..e7eb354d8 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java @@ -222,6 +222,69 @@ void closeIsIdempotent() throws Exception { mgr.close(); } + @Test + void createAndSetRunnerContextBuildsFreshIdentityContextOnFirstCall() throws Exception { + try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) { + ActionTask t = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction()); + invokeCreateAndSetRunnerContext(mgr, t); + + assertThat(t.getRunnerContext().getSubagentIdentityContext()).isNotNull(); + } + } + + @Test + void identityContextIsInheritedAcrossTransferWithOrdinalContinuation() throws Exception { + try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) { + Action action = TestActions.noopAction(); + ActionTask from = new JavaActionTask("k", new InputEvent(1L), action); + ActionTask to = new JavaActionTask("k", new InputEvent(2L), action); + + invokeCreateAndSetRunnerContext(mgr, from); + RunnerContextImpl.SubagentIdentityContext fromIdentityCtx = + from.getRunnerContext().getSubagentIdentityContext(); + assertThat(fromIdentityCtx).isNotNull(); + + // Advance the ordinal on `from`'s context so continuation (not reset) is observable + // once `to` inherits it. + String firstSessionId = fromIdentityCtx.nextSessionId(); + + mgr.transferContexts(from, to, new DurableExecutionManager(null)); + invokeCreateAndSetRunnerContext(mgr, to); + + // `to` inherits the exact same instance — a same-process continuation resume keeps + // the ordinal state alive rather than restarting it. + RunnerContextImpl.SubagentIdentityContext toIdentityCtx = + to.getRunnerContext().getSubagentIdentityContext(); + assertThat(toIdentityCtx).isSameAs(fromIdentityCtx); + assertThat(toIdentityCtx.nextSessionId()).isNotEqualTo(firstSessionId); + } + } + + @Test + void removeIdentityContextThenSetupBuildsFreshContext() throws Exception { + try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) { + Action action = TestActions.noopAction(); + ActionTask from = new JavaActionTask("k", new InputEvent(1L), action); + ActionTask to = new JavaActionTask("k", new InputEvent(2L), action); + + invokeCreateAndSetRunnerContext(mgr, from); + RunnerContextImpl.SubagentIdentityContext fromIdentityCtx = + from.getRunnerContext().getSubagentIdentityContext(); + + mgr.transferContexts(from, to, new DurableExecutionManager(null)); + // Simulate the transient per-task map losing `to`'s entry (e.g. a failover restart, + // since this map is heap-only and never checkpointed) before `to` is set up. + mgr.removeIdentityContext(to); + + invokeCreateAndSetRunnerContext(mgr, to); + + RunnerContextImpl.SubagentIdentityContext toIdentityCtx = + to.getRunnerContext().getSubagentIdentityContext(); + assertThat(toIdentityCtx).isNotNull(); + assertThat(toIdentityCtx).isNotSameAs(fromIdentityCtx); + } + } + /** * Shared helper: install a runner context on {@code task} using mocked collaborators. Used by * tests that need a fully wired runner context but do not care about the collaborator details. @@ -238,6 +301,7 @@ private static void invokeCreateAndSetRunnerContext( mgr.createAndSetRunnerContext( task, "k", + 0L, plan, cache, metricGroup, diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/SubagentIdentityRecoveryTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/SubagentIdentityRecoveryTest.java new file mode 100644 index 000000000..af67c86c9 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/SubagentIdentityRecoveryTest.java @@ -0,0 +1,280 @@ +/* + * 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.runtime.operator; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.subagent.Result; +import org.apache.flink.agents.api.subagent.SubagentSetup; +import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.agents.runtime.actionstate.ActionState; +import org.apache.flink.agents.runtime.actionstate.CallResult; +import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore; +import org.apache.flink.agents.runtime.subagent.CapturingSubagentSetup; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Failover-replay tests for deterministic sub-agent identity: replaying a record must recompute + * identical ids from the namespace alone, and a persisted durable {@link CallResult} must turn the + * replayed call into a cache hit instead of a re-execution. + * + *

The crash window under test (durable call persisted, action not yet complete) cannot be + * produced by interrupting the synchronous harness, because a finished single-hop action is + * immediately marked completed and replays without re-entering its body. Each test therefore learns + * the assigned ids from a real stage-1 run, then hand-seeds a fresh store with a not-yet-completed + * {@link ActionState} carrying the learned result (the same technique as {@code + * ActionExecutionOperatorTest}) and re-processes the identical record; the identity context is + * heap-only, so stage 2 recomputes ordinals from scratch. + */ +public class SubagentIdentityRecoveryTest { + + private static final String RESOURCE_NAME = "agent"; + + private static final ObjectMapper OBJECT_MAPPER = + new ObjectMapper().registerModule(new JavaTimeModule()); + + @BeforeEach + void resetCaptures() { + CapturingSubagentSetup.reset(); + } + + // Failover replay reproduces identical ids without duplicating side effects. + + @SuppressWarnings("unused") + public static void singleCall(Event event, RunnerContext ctx) throws Exception { + SubagentSetup setup = (SubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + setup.call(ctx, "recover-me"); + ctx.sendEvent(new OutputEvent("done")); + } + + @Test + void it2FailoverReplayReproducesIdenticalIdsWithoutDuplicateSideEffects() throws Exception { + long key = 11L; + + // Stage 1: learn the deterministically-assigned (sessionId, callId) from a real run + // against a fresh, independent store. + runToCompletion(buildSingleCallPlan(), new InMemoryActionStateStore(false), key); + + assertThat(CapturingSubagentSetup.captures()).hasSize(1); + assertThat(CapturingSubagentSetup.executionCount()).isEqualTo(1); + CapturingSubagentSetup.Capture firstRun = CapturingSubagentSetup.captures().get(0); + + // Stage 2: seed a *different*, fresh store to simulate a crash right after the durable + // sub-agent call persisted its result but before the action as a whole completed. See + // the class-level javadoc for why this must be constructed by hand. + AgentPlan plan2 = buildSingleCallPlan(); + InMemoryActionStateStore store2 = new InMemoryActionStateStore(false); + seedCompletedCallResults( + store2, + key, + plan2, + "singleCall", + key, + new CallResult( + firstRun.sessionId + "#" + firstRun.callId, + "", + OBJECT_MAPPER.writeValueAsBytes(Result.ok("cached")))); + + try (KeyedOneInputStreamOperatorTestHarness harness2 = + newHarness(plan2, store2)) { + harness2.open(); + ActionExecutionOperator operator2 = + (ActionExecutionOperator) harness2.getOperator(); + + harness2.processElement(new StreamRecord<>(key)); + operator2.waitInFlightEventsFinished(); + + List captures = CapturingSubagentSetup.captures(); + assertThat(captures).hasSize(2); + assertThat(captures.get(1).sessionId).isEqualTo(captures.get(0).sessionId); + assertThat(captures.get(1).callId).isEqualTo(captures.get(0).callId); + + assertThat(CapturingSubagentSetup.executionCount()) + .as("Durable call must not be re-executed once its CallResult is persisted") + .isEqualTo(1); + + @SuppressWarnings("unchecked") + List> recordOutput = + (List>) harness2.getRecordOutput(); + assertThat(recordOutput).hasSize(1); + assertThat(recordOutput.get(0).getValue()).isEqualTo("done"); + } + } + + // Fan-out replay: callables replay in creation order without duplicating side effects. + + @SuppressWarnings("unused") + public static void fanOut(Event event, RunnerContext ctx) throws Exception { + SubagentSetup setup = (SubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + DurableCallable c1 = setup.asAsyncCallable(ctx, "fan-a"); + DurableCallable c2 = setup.asAsyncCallable(ctx, "fan-b"); + ctx.durableExecuteAsync(c1); + ctx.durableExecuteAsync(c2); + ctx.sendEvent(new OutputEvent("done")); + } + + @Test + void it4FanOutReplayReproducesIdenticalIdsInOrderWithoutDuplicateSideEffects() + throws Exception { + long key = 6L; + + // Stage 1: learn both callables' deterministically-assigned ids from a real run. + runToCompletion(buildFanOutPlan(), new InMemoryActionStateStore(false), key); + + assertThat(CapturingSubagentSetup.captures()).hasSize(2); + assertThat(CapturingSubagentSetup.executionCount()).isEqualTo(2); + CapturingSubagentSetup.Capture firstRunC1 = CapturingSubagentSetup.captures().get(0); + CapturingSubagentSetup.Capture firstRunC2 = CapturingSubagentSetup.captures().get(1); + + // Stage 2: seed both callables' CallResults as already-persisted successes, simulating a + // crash after both durable calls completed but before the action as a whole finished. + AgentPlan plan2 = buildFanOutPlan(); + InMemoryActionStateStore store2 = new InMemoryActionStateStore(false); + seedCompletedCallResults( + store2, + key, + plan2, + "fanOut", + key, + new CallResult( + firstRunC1.sessionId + "#" + firstRunC1.callId, + "", + OBJECT_MAPPER.writeValueAsBytes(Result.ok("cached-a"))), + new CallResult( + firstRunC2.sessionId + "#" + firstRunC2.callId, + "", + OBJECT_MAPPER.writeValueAsBytes(Result.ok("cached-b")))); + + try (KeyedOneInputStreamOperatorTestHarness harness2 = + newHarness(plan2, store2)) { + harness2.open(); + ActionExecutionOperator operator2 = + (ActionExecutionOperator) harness2.getOperator(); + + harness2.processElement(new StreamRecord<>(key)); + operator2.waitInFlightEventsFinished(); + + List captures = CapturingSubagentSetup.captures(); + assertThat(captures).hasSize(4); + + // The replayed pair must reproduce the same two ids, in the same creation order, as + // the first run. + assertThat(captures.get(2).prompt).isEqualTo("fan-a"); + assertThat(captures.get(3).prompt).isEqualTo("fan-b"); + assertThat(captures.get(2).sessionId).isEqualTo(captures.get(0).sessionId); + assertThat(captures.get(2).callId).isEqualTo(captures.get(0).callId); + assertThat(captures.get(3).sessionId).isEqualTo(captures.get(1).sessionId); + assertThat(captures.get(3).callId).isEqualTo(captures.get(1).callId); + + assertThat(CapturingSubagentSetup.executionCount()) + .as( + "Fan-out durable calls must not be re-executed once their CallResults" + + " are persisted") + .isEqualTo(2); + + @SuppressWarnings("unchecked") + List> recordOutput = + (List>) harness2.getRecordOutput(); + assertThat(recordOutput).hasSize(1); + assertThat(recordOutput.get(0).getValue()).isEqualTo("done"); + } + } + + // Helpers + + private static AgentPlan buildSingleCallPlan() throws Exception { + Agent agent = new Agent(); + agent.addResource(RESOURCE_NAME, ResourceType.AGENT, new CapturingSubagentSetup()); + agent.addAction( + new String[] {InputEvent.EVENT_TYPE}, + SubagentIdentityRecoveryTest.class.getMethod( + "singleCall", Event.class, RunnerContext.class)); + return new AgentPlan(agent); + } + + private static AgentPlan buildFanOutPlan() throws Exception { + Agent agent = new Agent(); + agent.addResource(RESOURCE_NAME, ResourceType.AGENT, new CapturingSubagentSetup()); + agent.addAction( + new String[] {InputEvent.EVENT_TYPE}, + SubagentIdentityRecoveryTest.class.getMethod( + "fanOut", Event.class, RunnerContext.class)); + return new AgentPlan(agent); + } + + private static KeyedOneInputStreamOperatorTestHarness newHarness( + AgentPlan plan, InMemoryActionStateStore actionStateStore) throws Exception { + return new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(plan, true, actionStateStore), + (KeySelector) value -> value, + TypeInformation.of(Long.class)); + } + + /** Runs a single record through a fresh harness to completion, then closes the harness. */ + private static void runToCompletion( + AgentPlan plan, InMemoryActionStateStore actionStateStore, long input) + throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = + newHarness(plan, actionStateStore)) { + harness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) harness.getOperator(); + harness.processElement(new StreamRecord<>(input)); + operator.waitInFlightEventsFinished(); + } + } + + /** + * Seeds {@code actionStateStore} with an {@code ActionState} that is not yet completed but + * already carries the given {@link CallResult}s, in the order the action creates its durable + * callables. Mirrors the private helpers in {@code ActionExecutionOperatorTest}. + */ + private static void seedCompletedCallResults( + InMemoryActionStateStore actionStateStore, + long key, + AgentPlan agentPlan, + String actionName, + long input, + CallResult... callResults) + throws Exception { + InputEvent event = new InputEvent(input); + Action action = agentPlan.getActions().get(actionName); + ActionState actionState = new ActionState(null); + for (CallResult callResult : callResults) { + actionState.addCallResult(callResult); + } + actionStateStore.put(key, 0L, action, event, actionState); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/CapturingSubagentSetup.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/CapturingSubagentSetup.java new file mode 100644 index 000000000..26441a5d0 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/CapturingSubagentSetup.java @@ -0,0 +1,108 @@ +/* + * 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.runtime.subagent; + +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.subagent.Result; +import org.apache.flink.agents.api.subagent.SubagentSetup; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Test double that records every {@code (sessionId, callId)} pair assigned to it. + * + *

Capture happens at callable-creation time rather than in {@code call()}, so tests can assert + * on assigned ids even when the durable call is served from cache; {@link #executionCount()} + * separately tracks real executions. State is static because one setup instance is shared by every + * resolving task — call {@link #reset()} before each independent scenario. + */ +public class CapturingSubagentSetup extends SubagentSetup { + + /** One {@code (sessionId, callId)} assignment captured at callable-creation time. */ + public static final class Capture { + public final String sessionId; + public final String callId; + public final Object prompt; + + Capture(String sessionId, String callId, Object prompt) { + this.sessionId = sessionId; + this.callId = callId; + this.prompt = prompt; + } + + @Override + public String toString() { + return "Capture{sessionId=" + + sessionId + + ", callId=" + + callId + + ", prompt=" + + prompt + + "}"; + } + } + + private static final List CAPTURES = Collections.synchronizedList(new ArrayList<>()); + private static final AtomicInteger EXECUTION_COUNT = new AtomicInteger(); + + /** Clears all captures and the execution counter. Call before each independent scenario. */ + public static void reset() { + CAPTURES.clear(); + EXECUTION_COUNT.set(0); + } + + /** Snapshot of every assignment captured since the last {@link #reset()}, in creation order. */ + public static List captures() { + synchronized (CAPTURES) { + return new ArrayList<>(CAPTURES); + } + } + + /** Number of times a produced callable's {@code call()} body actually ran. */ + public static int executionCount() { + return EXECUTION_COUNT.get(); + } + + @Override + public DurableCallable asAsyncCallable( + RunnerContext ctx, Object prompt, String sessionId, String callId) { + CAPTURES.add(new Capture(sessionId, callId, prompt)); + return new DurableCallable() { + @Override + public String getId() { + return sessionId + "#" + callId; + } + + @Override + public Class getResultClass() { + return Result.class; + } + + @Override + public Result call() { + EXECUTION_COUNT.incrementAndGet(); + return Result.ok(sessionId + "|" + callId + "|" + prompt); + } + }; + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/SubagentIdentityIntegrationTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/SubagentIdentityIntegrationTest.java new file mode 100644 index 000000000..7ff510395 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/SubagentIdentityIntegrationTest.java @@ -0,0 +1,262 @@ +/* + * 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.runtime.subagent; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.agents.Agent; +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.subagent.Result; +import org.apache.flink.agents.api.subagent.SubagentSetup; +import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.runtime.operator.ActionExecutionOperator; +import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorFactory; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Operator-harness tests for deterministic sub-agent identity assignment: same-round determinism + * and uniqueness, sibling-task separation, and fan-out ordering. {@link CapturingSubagentSetup} + * records the {@code (sessionId, callId)} pair handed to every produced callable. The + * failover-replay counterpart lives in {@code SubagentIdentityRecoveryTest}. + */ +public class SubagentIdentityIntegrationTest { + + private static final String RESOURCE_NAME = "agent"; + private static final String EXPLICIT_SESSION_ID = "explicit-session-checkout-1"; + + @BeforeEach + void resetCaptures() { + CapturingSubagentSetup.reset(); + } + + // Same-round determinism and uniqueness. + + @SuppressWarnings("unused") + public static void mixedCalls(Event event, RunnerContext ctx) throws Exception { + SubagentSetup setup = (SubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + setup.call(ctx, "p1"); + setup.call(ctx, "p2"); + setup.call(ctx, "p3a", EXPLICIT_SESSION_ID); + setup.call(ctx, "p3b", EXPLICIT_SESSION_ID); + ctx.sendEvent(new OutputEvent("done")); + } + + @Test + void it1MixedCallsProduceUniqueDeterministicIds() throws Exception { + List captures = runMixedCallsScenario(1L); + assertThat(captures).hasSize(4); + + String autoSession1 = captures.get(0).sessionId; + String autoSession2 = captures.get(1).sessionId; + String explicitSessionA = captures.get(2).sessionId; + String explicitSessionB = captures.get(3).sessionId; + + // The two auto-assigned sessions follow the runtime's session id shape (a fixed-length + // name-UUID namespace plus the session ordinal) and are distinct. + assertThat(autoSession1) + .matches("[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}-0"); + assertThat(autoSession2) + .matches("[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}-1"); + + // An explicit session id is used verbatim, not reshaped into the auto-assigned form. + assertThat(explicitSessionA).isEqualTo(EXPLICIT_SESSION_ID); + assertThat(explicitSessionB).isEqualTo(EXPLICIT_SESSION_ID); + + // All four call ids are globally distinct. + List callIds = captures.stream().map(c -> c.callId).collect(Collectors.toList()); + assertThat(callIds).doesNotHaveDuplicates(); + + // A call id is its session id plus the per-session ordinal (1, then 2 for the two + // explicit-session calls; 1 for each auto session's single call). + assertThat(captures.get(0).callId).isEqualTo(autoSession1 + "-1"); + assertThat(captures.get(1).callId).isEqualTo(autoSession2 + "-1"); + assertThat(captures.get(2).callId).isEqualTo(EXPLICIT_SESSION_ID + "-1"); + assertThat(captures.get(3).callId).isEqualTo(EXPLICIT_SESSION_ID + "-2"); + } + + @Test + void it1RerunningIdenticalInputReproducesIdenticalCaptureSequence() throws Exception { + List firstRun = runMixedCallsScenario(2L); + List secondRun = runMixedCallsScenario(2L); + + assertThat(toIdPairs(secondRun)).isEqualTo(toIdPairs(firstRun)); + } + + private static List> toIdPairs(List captures) { + return captures.stream() + .map(c -> List.of(c.sessionId, c.callId)) + .collect(Collectors.toList()); + } + + /** Builds a fresh plan/harness, runs {@link #mixedCalls}, and returns its capture sequence. */ + private static List runMixedCallsScenario(long inputValue) + throws Exception { + CapturingSubagentSetup.reset(); + Agent agent = new Agent(); + agent.addResource(RESOURCE_NAME, ResourceType.AGENT, new CapturingSubagentSetup()); + agent.addAction( + new String[] {InputEvent.EVENT_TYPE}, + SubagentIdentityIntegrationTest.class.getMethod( + "mixedCalls", Event.class, RunnerContext.class)); + AgentPlan plan = new AgentPlan(agent); + + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(plan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(inputValue)); + operator.waitInFlightEventsFinished(); + } + return CapturingSubagentSetup.captures(); + } + + // Sibling tasks fired from the same record must not collide. + + /** Trigger event whose attributes vary per branch, used to separate sibling task namespaces. */ + public static class SiblingTriggerEvent extends Event { + public static final String EVENT_TYPE = "SiblingTriggerEvent"; + + public SiblingTriggerEvent(String branch) { + super(EVENT_TYPE); + setAttr("branch", branch); + } + + public String getBranch() { + return (String) getAttr("branch"); + } + } + + @SuppressWarnings("unused") + public static void siblingProducer(Event event, RunnerContext ctx) { + // Same key, same sequenceNumber, same downstream action name for both -- only the + // triggering event differs between the two siblings. + ctx.sendEvent(new SiblingTriggerEvent("branch-a")); + ctx.sendEvent(new SiblingTriggerEvent("branch-b")); + } + + @SuppressWarnings("unused") + public static void siblingConsumer(SiblingTriggerEvent event, RunnerContext ctx) + throws Exception { + SubagentSetup setup = (SubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + setup.call(ctx, "prompt-" + event.getBranch()); + ctx.sendEvent(new OutputEvent(event.getBranch())); + } + + @Test + void it3SiblingTasksFromSameRecordGetSeparateIdentityNamespaces() throws Exception { + Agent agent = new Agent(); + agent.addResource(RESOURCE_NAME, ResourceType.AGENT, new CapturingSubagentSetup()); + agent.addAction( + new String[] {InputEvent.EVENT_TYPE}, + SubagentIdentityIntegrationTest.class.getMethod( + "siblingProducer", Event.class, RunnerContext.class)); + agent.addAction( + new String[] {SiblingTriggerEvent.EVENT_TYPE}, + SubagentIdentityIntegrationTest.class.getMethod( + "siblingConsumer", SiblingTriggerEvent.class, RunnerContext.class)); + AgentPlan plan = new AgentPlan(agent); + + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(plan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(9L)); + operator.waitInFlightEventsFinished(); + + @SuppressWarnings("unchecked") + List> output = + (List>) testHarness.getRecordOutput(); + assertThat(output).hasSize(2); + } + + List captures = CapturingSubagentSetup.captures(); + assertThat(captures).hasSize(2); + + // Both siblings share (key, sequenceNumber, actionName); only the triggering event's + // digest differs. Their identities must not collide. + assertThat(captures.get(0).sessionId).isNotEqualTo(captures.get(1).sessionId); + assertThat(captures.get(0).callId).isNotEqualTo(captures.get(1).callId); + } + + // Async fan-out assigns distinct ids in creation order. + + @SuppressWarnings("unused") + public static void fanOut(Event event, RunnerContext ctx) throws Exception { + SubagentSetup setup = (SubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + DurableCallable c1 = setup.asAsyncCallable(ctx, "fan-a"); + DurableCallable c2 = setup.asAsyncCallable(ctx, "fan-b"); + ctx.durableExecuteAsync(c1); + ctx.durableExecuteAsync(c2); + ctx.sendEvent(new OutputEvent("done")); + } + + @Test + void it4FanOutCallablesGetDistinctIdsInCreationOrder() throws Exception { + Agent agent = new Agent(); + agent.addResource(RESOURCE_NAME, ResourceType.AGENT, new CapturingSubagentSetup()); + agent.addAction( + new String[] {InputEvent.EVENT_TYPE}, + SubagentIdentityIntegrationTest.class.getMethod( + "fanOut", Event.class, RunnerContext.class)); + AgentPlan plan = new AgentPlan(agent); + + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(plan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(4L)); + operator.waitInFlightEventsFinished(); + } + + List captures = CapturingSubagentSetup.captures(); + assertThat(captures).hasSize(2); + assertThat(captures.get(0).prompt).isEqualTo("fan-a"); + assertThat(captures.get(1).prompt).isEqualTo("fan-b"); + assertThat(captures.get(0).sessionId).isNotEqualTo(captures.get(1).sessionId); + assertThat(captures.get(0).callId).isNotEqualTo(captures.get(1).callId); + } +}