From 5a686524f94959a020c0d01b61bde31a4d6eb982 Mon Sep 17 00:00:00 2001 From: daken Date: Sun, 19 Jul 2026 22:37:35 +0800 Subject: [PATCH 01/13] [feature] add_durableExecuteAllAsync_java --- .../api/agents/AgentExecutionOptions.java | 17 ++ .../flink/agents/api/context/Outcome.java | 53 +++++ .../agents/api/context/RunnerContext.java | 26 +++ .../async/ContinuationActionExecutor.java | 30 +++ .../context/JavaRunnerContextImpl.java | 200 +++++++++++++++++- .../runtime/context/RunnerContextImpl.java | 152 +++++++++++-- .../operator/ActionExecutionOperator.java | 5 +- .../operator/ActionTaskContextManager.java | 24 ++- .../async/ContinuationActionExecutor.java | 90 +++++++- .../runtime/async/ContinuationContext.java | 18 ++ 10 files changed, 581 insertions(+), 34 deletions(-) create mode 100644 api/src/main/java/org/apache/flink/agents/api/context/Outcome.java diff --git a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java index a0b9269e4..5f8d9d0bd 100644 --- a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java +++ b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java @@ -20,6 +20,8 @@ import org.apache.flink.agents.api.configuration.ConfigOption; +import java.time.Duration; + public class AgentExecutionOptions { public static final ConfigOption ERROR_HANDLING_STRATEGY = new ConfigOption<>( @@ -45,6 +47,21 @@ public class AgentExecutionOptions { public static final ConfigOption TOOL_CALL_ASYNC = new ConfigOption<>("tool-call.async", Boolean.class, true); + public static final ConfigOption TOOL_CALL_PARALLEL = + new ConfigOption<>("tool-call.parallel", Boolean.class, true); + + public static final ConfigOption TOOL_CALL_NUM_ASYNC_THREADS = + new ConfigOption<>( + "tool-call.num-async-threads", + Integer.class, + Runtime.getRuntime().availableProcessors() * 2); + + public static final ConfigOption TOOL_CALL_MAX_PARALLELISM = + new ConfigOption<>("tool-call.max-parallelism", Integer.class, Integer.MAX_VALUE); + + public static final ConfigOption TOOL_CALL_BATCH_TIMEOUT = + new ConfigOption<>("tool-call.batch.timeout", Duration.class, Duration.ofMillis(-1)); + public static final ConfigOption RAG_ASYNC = new ConfigOption<>("rag.async", Boolean.class, true); diff --git a/api/src/main/java/org/apache/flink/agents/api/context/Outcome.java b/api/src/main/java/org/apache/flink/agents/api/context/Outcome.java new file mode 100644 index 000000000..0b16ee92c --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/context/Outcome.java @@ -0,0 +1,53 @@ +/* + * 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.context; + +/** Result of a durable batch execution slot. */ +public class Outcome { + private final T value; + private final Exception error; + + private Outcome(T value, Exception error) { + this.value = value; + this.error = error; + } + + public static Outcome success(T value) { + return new Outcome<>(value, null); + } + + public static Outcome failure(Exception error) { + return new Outcome<>(null, error); + } + + public boolean isSuccess() { + return error == null; + } + + public boolean isFailure() { + return error != null; + } + + public T getValue() { + return value; + } + + public Exception getError() { + return error; + } +} 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..cf35e6905 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 @@ -24,6 +24,7 @@ import org.apache.flink.agents.api.resource.Resource; import org.apache.flink.agents.api.resource.ResourceType; +import java.util.List; import java.util.Map; /** @@ -146,6 +147,31 @@ public interface RunnerContext { */ T durableExecuteAsync(DurableCallable callable) throws Exception; + /** + * Executes multiple durable callables as one batch and returns outcomes in input order. + * + *

On JDK 21+, implementations may submit uncached calls concurrently and yield the current + * action execution until the batch completes. On JDK < 21, this falls back to serial durable + * execution. + * + *

The callable list must be deterministic across recovery: same order and same {@link + * DurableCallable#getId()} values. + * + *

Access to memory and sendEvent are prohibited within the callables. + */ + default List> durableExecuteAllAsync(List> callables) + throws Exception { + java.util.ArrayList> outcomes = new java.util.ArrayList<>(callables.size()); + for (DurableCallable callable : callables) { + try { + outcomes.add(Outcome.success(durableExecute(callable))); + } catch (Exception e) { + outcomes.add(Outcome.failure(e)); + } + } + return outcomes; + } + /** Clean up the resource. */ void close() throws Exception; } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java index 7c1a526d2..54ccbbd0d 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java @@ -17,7 +17,13 @@ */ package org.apache.flink.agents.runtime.async; +import org.apache.flink.agents.api.context.Outcome; + +import java.time.Duration; +import java.util.List; +import java.util.concurrent.Callable; import java.util.function.Supplier; +import java.util.stream.Collectors; /** * Executor for Java actions that supports asynchronous execution. @@ -56,6 +62,30 @@ public T executeAsync(ContinuationContext context, Supplier supplier) { return supplier.get(); } + /** + * Executes all suppliers as one batch. In JDK 11, this falls back to serial execution and + * captures each supplier's success or failure as an {@link Outcome}. + * + * @param context the continuation context + * @param suppliers the suppliers to execute + * @param timeout ignored in the JDK 11 fallback + * @param the result type + * @return outcomes in supplier order + */ + public List> executeAllAsync( + ContinuationContext context, List> suppliers, Duration timeout) { + return suppliers.stream() + .map( + supplier -> { + try { + return Outcome.success(supplier.call()); + } catch (Exception e) { + return Outcome.failure(e); + } + }) + .collect(Collectors.toList()); + } + public void close() {} /** diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java index ae0661ead..b365e010d 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java @@ -17,15 +17,22 @@ */ package org.apache.flink.agents.runtime.context; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.runtime.ResourceCache; +import org.apache.flink.agents.runtime.actionstate.CallResult; import org.apache.flink.agents.runtime.async.ContinuationActionExecutor; import org.apache.flink.agents.runtime.async.ContinuationContext; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; import java.util.concurrent.Callable; import java.util.function.Supplier; +import java.util.stream.Collectors; /** * Java-specific implementation of RunnerContext that includes ContinuationActionExecutor for async @@ -33,6 +40,7 @@ */ public class JavaRunnerContextImpl extends RunnerContextImpl { private final ContinuationActionExecutor continuationExecutor; + private final ContinuationActionExecutor toolCallContinuationExecutor; private ContinuationContext continuationContext; public JavaRunnerContextImpl( @@ -41,9 +49,28 @@ public JavaRunnerContextImpl( AgentPlan agentPlan, ResourceCache resourceCache, String jobIdentifier, - ContinuationActionExecutor continuationExecutor) { + ContinuationActionExecutor continuationExecutor, + ContinuationActionExecutor toolCallContinuationExecutor) { super(agentMetricGroup, mailboxThreadChecker, agentPlan, resourceCache, jobIdentifier); this.continuationExecutor = continuationExecutor; + this.toolCallContinuationExecutor = toolCallContinuationExecutor; + } + + public JavaRunnerContextImpl( + FlinkAgentsMetricGroupImpl agentMetricGroup, + Runnable mailboxThreadChecker, + AgentPlan agentPlan, + ResourceCache resourceCache, + String jobIdentifier, + ContinuationActionExecutor continuationExecutor) { + this( + agentMetricGroup, + mailboxThreadChecker, + agentPlan, + resourceCache, + jobIdentifier, + continuationExecutor, + continuationExecutor); } public ContinuationActionExecutor getContinuationExecutor() { @@ -75,6 +102,172 @@ private T durableExecuteAsyncWithReconcile( callable, reconcileCallable, () -> executeAsyncCallable(callable)); } + @Override + public List> durableExecuteAllAsync(List> callables) + throws Exception { + if (callables.isEmpty()) { + return List.of(); + } + if (durableExecutionContext == null) { + return executeAllWithoutDurableState(callables); + } + + String argsDigest = ""; + int base = durableExecutionContext.getCurrentCallIndex(); + BatchExecutionPlan plan = buildBatchExecutionPlan(callables, base, argsDigest); + + reservePendingBatchIfNeeded(callables, argsDigest, plan); + + List> executed = executeOutcomeSuppliers(plan.suppliers); + finalizeExecutedOutcomes(callables, base, argsDigest, plan, executed); + + advanceCallIndexBy(callables.size()); + return plan.outcomes; + } + + private BatchExecutionPlan buildBatchExecutionPlan( + List> callables, int base, String argsDigest) throws Exception { + BatchExecutionPlan plan = new BatchExecutionPlan<>(callables.size()); + for (int i = 0; i < callables.size(); i++) { + DurableCallable callable = callables.get(i); + CallResult current = getCallResultAt(base + i); + if (current == null) { + markReservationStart(plan, i); + addExecutableCall(plan, i, callable::call); + continue; + } + if (!current.matches(callable.getId(), argsDigest)) { + clearCallResultsFromAndPersist(base + i); + plan.needsReservation = true; + plan.executionStart = i; + addExecutableCall(plan, i, callable::call); + appendRemainingExecutions(callables, plan, i + 1); + break; + } + if (current.isPending()) { + Callable reconcileCallable = callable.reconciler(); + Callable executionCallable = + reconcileCallable != null ? reconcileCallable : callable::call; + addExecutableCall(plan, i, executionCallable); + } else { + plan.outcomes.add( + readTerminalOutcomeAt( + base + i, callable.getId(), argsDigest, callable.getResultClass())); + } + } + return plan; + } + + private void markReservationStart(BatchExecutionPlan plan, int callIndex) { + plan.needsReservation = true; + if (plan.executionStart < 0) { + plan.executionStart = callIndex; + } + } + + private void addExecutableCall( + BatchExecutionPlan plan, int callIndex, Callable executionCallable) { + plan.outcomes.add(null); + plan.suppliers.add(executionCallable); + plan.executableCallIndexes.add(callIndex); + } + + private void appendRemainingExecutions( + List> callables, BatchExecutionPlan plan, int startIndex) { + for (int i = startIndex; i < callables.size(); i++) { + DurableCallable remaining = callables.get(i); + addExecutableCall(plan, i, remaining::call); + } + } + + private void reservePendingBatchIfNeeded( + List> callables, String argsDigest, BatchExecutionPlan plan) { + if (!plan.needsReservation) { + return; + } + List ids = new ArrayList<>(); + for (DurableCallable callable : + callables.subList(plan.executionStart, callables.size())) { + ids.add(callable.getId()); + } + reservePendingBatch(ids, argsDigest); + } + + private void finalizeExecutedOutcomes( + List> callables, + int base, + String argsDigest, + BatchExecutionPlan plan, + List> executed) + throws Exception { + for (int i = 0; i < executed.size(); i++) { + int callIndex = plan.executableCallIndexes.get(i); + Outcome outcome = executed.get(i); + DurableCallable callable = callables.get(callIndex); + finalizeCallAt( + base + callIndex, + callable.getId(), + argsDigest, + serializeDurableResult(outcome.getValue()), + serializeDurableException(outcome.getError())); + plan.outcomes.set(callIndex, outcome); + } + } + + private static class BatchExecutionPlan { + private final List> outcomes; + private final List> suppliers = new ArrayList<>(); + private final List executableCallIndexes = new ArrayList<>(); + private boolean needsReservation; + private int executionStart = -1; + + private BatchExecutionPlan(int size) { + this.outcomes = new ArrayList<>(size); + } + } + + private List> executeAllWithoutDurableState(List> callables) + throws Exception { + List> suppliers = new ArrayList<>(); + for (DurableCallable callable : callables) { + suppliers.add(callable::call); + } + return executeOutcomeSuppliers(suppliers); + } + + private List> executeOutcomeSuppliers(List> suppliers) + throws Exception { + if (suppliers.isEmpty()) { + return List.of(); + } + if (continuationExecutor == null || continuationContext == null) { + return suppliers.stream() + .map( + supplier -> { + try { + return Outcome.success(supplier.call()); + } catch (Exception e) { + return Outcome.failure(e); + } + }) + .collect(Collectors.toList()); + } + Duration timeout = getConfig().get(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT); + int maxParallelism = getConfig().get(AgentExecutionOptions.TOOL_CALL_MAX_PARALLELISM); + if (maxParallelism <= 0 || maxParallelism >= suppliers.size()) { + return toolCallContinuationExecutor.executeAllAsync( + continuationContext, suppliers, timeout); + } + List> results = new ArrayList<>(suppliers.size()); + for (int i = 0; i < suppliers.size(); i += maxParallelism) { + int end = Math.min(i + maxParallelism, suppliers.size()); + results.addAll( + toolCallContinuationExecutor.executeAllAsync( + continuationContext, suppliers.subList(i, end), timeout)); + } + return results; + } + private T executeAsyncCallable(DurableCallable callable) throws Exception { Supplier wrappedSupplier = @@ -103,4 +296,9 @@ private T executeAsyncCallable(DurableCallable callable) throws Exception throw (Exception) e.getCause(); } } + + @Override + public void close() throws Exception { + super.close(); + } } 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..5844da838 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 @@ -26,6 +26,7 @@ import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.MemoryUpdate; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.memory.BaseLongTermMemory; import org.apache.flink.agents.api.resource.Resource; @@ -59,7 +60,7 @@ */ public class RunnerContextImpl implements RunnerContext { - private static final ObjectMapper OBJECT_MAPPER = + protected static final ObjectMapper OBJECT_MAPPER = new ObjectMapper().registerModule(new JavaTimeModule()); public static class MemoryContext { @@ -265,6 +266,20 @@ public T durableExecuteAsync(DurableCallable callable) throws Exception { return durableExecute(callable); } + @Override + public List> durableExecuteAllAsync(List> callables) + throws Exception { + List> outcomes = new ArrayList<>(callables.size()); + for (DurableCallable callable : callables) { + try { + outcomes.add(Outcome.success(durableExecute(callable))); + } catch (Exception e) { + outcomes.add(Outcome.failure(e)); + } + } + return outcomes; + } + /** * Executes a durable call using the completion-only state machine. * @@ -417,6 +432,13 @@ public void appendPendingCall(String functionId, String argsDigest) { } } + public void reservePendingBatch(List functionIds, String argsDigest) { + mailboxThreadChecker.run(); + if (durableExecutionContext != null && !functionIds.isEmpty()) { + durableExecutionContext.reservePendingBatch(functionIds, argsDigest); + } + } + /** Finalizes the pending durable call slot at the current call index. */ public void finalizeCurrentCall( String functionId, String argsDigest, byte[] resultPayload, byte[] exceptionPayload) { @@ -427,6 +449,26 @@ public void finalizeCurrentCall( } } + public void finalizeCallAt( + int index, + String functionId, + String argsDigest, + byte[] resultPayload, + byte[] exceptionPayload) { + mailboxThreadChecker.run(); + if (durableExecutionContext != null) { + durableExecutionContext.finalizeCallAt( + index, functionId, argsDigest, resultPayload, exceptionPayload); + } + } + + public void advanceCallIndexBy(int count) { + mailboxThreadChecker.run(); + if (durableExecutionContext != null) { + durableExecutionContext.advanceCallIndexBy(count); + } + } + /** * Clears persisted call results from the current call index onward and persists immediately. */ @@ -437,6 +479,13 @@ public void clearCallResultsFromCurrentIndexAndPersist() { } } + public void clearCallResultsFromAndPersist(int index) { + mailboxThreadChecker.run(); + if (durableExecutionContext != null) { + durableExecutionContext.clearCallResultsFromAndPersist(index); + } + } + /** * Returns the current durable call result as an array of fields for bridge consumers, or null * if no persisted slot exists at the current call index. @@ -455,6 +504,29 @@ public Object[] getCurrentCallResultFields() { }; } + protected Outcome readTerminalOutcomeAt( + int index, String functionId, String argsDigest, Class resultClass) + throws Exception { + CallResult callResult = getCallResultAt(index); + if (callResult == null || callResult.isPending()) { + throw new IllegalStateException( + String.format( + "Expected a terminal durable call result at index %s for " + + "functionId=%s, argsDigest=%s", + index, functionId, argsDigest)); + } + if (callResult.getExceptionPayload() != null) { + DurableExecutionException exception = + OBJECT_MAPPER.readValue( + callResult.getExceptionPayload(), DurableExecutionException.class); + return Outcome.failure(exception.toException()); + } + if (callResult.getResultPayload() == null) { + return Outcome.success(null); + } + return Outcome.success(OBJECT_MAPPER.readValue(callResult.getResultPayload(), resultClass)); + } + protected CallResult getCurrentCallResult() { mailboxThreadChecker.run(); if (durableExecutionContext != null) { @@ -463,6 +535,14 @@ protected CallResult getCurrentCallResult() { return null; } + protected CallResult getCallResultAt(int index) { + mailboxThreadChecker.run(); + if (durableExecutionContext != null) { + return durableExecutionContext.getCallResultAt(index); + } + return null; + } + protected Optional tryGetCachedResult( String functionId, String argsDigest, Class resultClass) throws Exception { Object[] cached = matchNextOrClearSubsequentCallResult(functionId, argsDigest); @@ -637,8 +717,12 @@ public ActionState getActionState() { * yet have a persisted slot. */ public CallResult getCurrentCallResult() { - if (currentCallIndex < recoveryCallResults.size()) { - return recoveryCallResults.get(currentCallIndex); + return getCallResultAt(currentCallIndex); + } + + public CallResult getCallResultAt(int index) { + if (index < recoveryCallResults.size()) { + return recoveryCallResults.get(index); } return null; } @@ -738,6 +822,15 @@ public void appendPendingCall(String functionId, String argsDigest) { argsDigest); } + public void reservePendingBatch(List functionIds, String argsDigest) { + for (String functionId : functionIds) { + CallResult pending = CallResult.pending(functionId, argsDigest); + actionState.addCallResult(pending); + recoveryCallResults.add(pending); + } + persistActionState(); + } + /** * Replaces the current persisted slot with a terminal call result and advances the current * call index. @@ -747,42 +840,52 @@ public void finalizeCurrentCall( String argsDigest, byte[] resultPayload, byte[] exceptionPayload) { - CallResult current = getCurrentCallResult(); + finalizeCallAt( + currentCallIndex, functionId, argsDigest, resultPayload, exceptionPayload); + currentCallIndex++; + } + + public void finalizeCallAt( + int index, + String functionId, + String argsDigest, + byte[] resultPayload, + byte[] exceptionPayload) { + CallResult current = getCallResultAt(index); if (current == null) { throw new IllegalStateException( String.format( - "Cannot finalize current call at index %s because no persisted " - + "slot exists", - currentCallIndex)); + "Cannot finalize call at index %s because no persisted slot exists", + index)); } if (!current.matches(functionId, argsDigest)) { throw new IllegalStateException( String.format( - "Cannot finalize current call at index %s because the persisted " - + "slot does not match functionId=%s, argsDigest=%s", - currentCallIndex, functionId, argsDigest)); + "Cannot finalize call at index %s because the persisted slot does not match functionId=%s, argsDigest=%s", + index, functionId, argsDigest)); } if (!current.isPending()) { throw new IllegalStateException( String.format( - "Cannot finalize current call at index %s because the persisted " - + "slot is not pending", - currentCallIndex)); + "Cannot finalize call at index %s because the persisted slot is not pending", + index)); } CallResult terminal = new CallResult(functionId, argsDigest, resultPayload, exceptionPayload); - actionState.replaceCallResult(currentCallIndex, terminal); - recoveryCallResults.set(currentCallIndex, terminal); + actionState.replaceCallResult(index, terminal); + recoveryCallResults.set(index, terminal); persistActionState(); LOG.debug( "Finalized and persisted CallResult at index {}: functionId={}, argsDigest={}", - currentCallIndex, + index, functionId, argsDigest); + } - currentCallIndex++; + public void advanceCallIndexBy(int count) { + currentCallIndex += count; } /** @@ -794,12 +897,21 @@ public void clearCallResultsFromCurrentIndexAndPersist() { persistActionState(); } - private void clearCallResultsFromCurrentIndex() { - actionState.clearCallResultsFrom(currentCallIndex); + public void clearCallResultsFromAndPersist(int index) { + clearCallResultsFrom(index); + persistActionState(); + } + + public void clearCallResultsFrom(int index) { + actionState.clearCallResultsFrom(index); recoveryCallResults = new ArrayList<>( recoveryCallResults.subList( - 0, Math.min(currentCallIndex, recoveryCallResults.size()))); + 0, Math.min(index, recoveryCallResults.size()))); + } + + private void clearCallResultsFromCurrentIndex() { + clearCallResultsFrom(currentCallIndex); } private void persistActionState() { 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..bfa5d7b4d 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 @@ -188,7 +188,10 @@ public void open() throws Exception { // init context manager for runner context creation and memory contexts contextManager = new ActionTaskContextManager( - agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS)); + agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS), + agentPlan + .getConfig() + .get(AgentExecutionOptions.TOOL_CALL_NUM_ASYNC_THREADS)); mailboxProcessor = getMailboxProcessor(); 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..286f038e5 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 @@ -73,13 +73,19 @@ class ActionTaskContextManager implements AutoCloseable { private final Map continuationContexts; private final Map pythonAwaitableRefs; - private ContinuationActionExecutor continuationActionExecutor; + private final ContinuationActionExecutor actionContinuationExecutor; + private final ContinuationActionExecutor toolCallContinuationExecutor; - ActionTaskContextManager(int numAsyncThreads) { + ActionTaskContextManager(int numAsyncThreads, int numToolCallAsyncThreads) { this.actionTaskMemoryContexts = new HashMap<>(); this.continuationContexts = new HashMap<>(); this.pythonAwaitableRefs = new HashMap<>(); - this.continuationActionExecutor = new ContinuationActionExecutor(numAsyncThreads); + this.actionContinuationExecutor = new ContinuationActionExecutor(numAsyncThreads); + this.toolCallContinuationExecutor = new ContinuationActionExecutor(numToolCallAsyncThreads); + } + + ActionTaskContextManager(int numAsyncThreads) { + this(numAsyncThreads, numAsyncThreads); } /** @@ -111,7 +117,7 @@ RunnerContextImpl createOrGetRunnerContext( @Nullable InteranlBaseLongTermMemory longTermMemory) { if (isJava) { if (runnerContext == null) { - if (continuationActionExecutor == null) { + if (actionContinuationExecutor == null || toolCallContinuationExecutor == null) { throw new IllegalStateException( "ContinuationActionExecutor has not been initialized."); } @@ -122,7 +128,8 @@ RunnerContextImpl createOrGetRunnerContext( agentPlan, resourceCache, jobIdentifier, - continuationActionExecutor); + actionContinuationExecutor, + toolCallContinuationExecutor); if (longTermMemory != null) { runnerContext.setLongTermMemory(longTermMemory); } @@ -324,8 +331,11 @@ public void close() throws Exception { runnerContext = null; } } - if (continuationActionExecutor != null) { - continuationActionExecutor.close(); + if (actionContinuationExecutor != null) { + actionContinuationExecutor.close(); + } + if (toolCallContinuationExecutor != null) { + toolCallContinuationExecutor.close(); } } } diff --git a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java index 4dc30d817..a721dbcaf 100644 --- a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java +++ b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java @@ -17,12 +17,20 @@ */ package org.apache.flink.agents.runtime.async; +import org.apache.flink.agents.api.context.Outcome; + import jdk.internal.vm.Continuation; import jdk.internal.vm.ContinuationScope; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.TimeoutException; import java.util.function.Supplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -59,16 +67,20 @@ public ContinuationActionExecutor(int numAsyncThreads) { */ public boolean executeAction(ContinuationContext context, Runnable action) { // Check if we have a pending async Future from previous yield + if (context.hasPendingAsync()) { + return false; + } + Future pending = context.getPendingFuture(); if (pending != null) { - if (!pending.isDone()) { - // Async task not done yet, return false to wait - return false; - } - // Async task done, clear the pending future and resume LOG.debug("Async task done..."); context.setPendingFuture(null); } + Future pendingBatch = context.getPendingBatchFuture(); + if (pendingBatch != null) { + LOG.debug("Async batch done..."); + context.setPendingBatchFuture(null); + } Continuation currentContinuation = context.getCurrentContinuation(); if (currentContinuation == null) { @@ -148,6 +160,74 @@ public T executeAsync(ContinuationContext context, Supplier supplier) thr return (T) context.getAsyncResultRef().get(); } + /** + * Executes all suppliers as one async batch and returns one {@link Outcome} per supplier. + * Supplier failures are captured in their own outcome so one failed supplier does not abort the + * whole batch. + * + * @param context the continuation context for this action + * @param suppliers the suppliers to execute + * @param timeout the timeout for the whole batch; null or non-positive means no timeout + * @param the result type + * @return outcomes in supplier order + */ + public List> executeAllAsync( + ContinuationContext context, List> suppliers, Duration timeout) + throws Exception { + context.clearAsyncState(); + + List>> futures = new ArrayList<>(suppliers.size()); + for (Callable supplier : suppliers) { + futures.add( + CompletableFuture.supplyAsync( + () -> { + try { + return Outcome.success(supplier.call()); + } catch (Exception e) { + return Outcome.failure(e); + } + }, + asyncExecutor)); + } + + CompletableFuture barrier = + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])); + context.setPendingBatchFuture(barrier); + + long deadlineNanos = getDeadlineNanos(timeout); + while (!barrier.isDone()) { + if (System.nanoTime() >= deadlineNanos) { + TimeoutException exception = + new TimeoutException( + "Async durable batch execution timed out after " + timeout); + for (CompletableFuture> future : futures) { + future.cancel(true); + } + barrier.cancel(true); + context.setPendingBatchFuture(null); + List> results = new ArrayList<>(suppliers.size()); + for (int i = 0; i < suppliers.size(); i++) { + results.add(Outcome.failure(exception)); + } + return results; + } + Continuation.yield(SCOPE); + } + + context.setPendingBatchFuture(null); + List> results = new ArrayList<>(futures.size()); + for (CompletableFuture> future : futures) { + results.add(future.join()); + } + return results; + } + + private long getDeadlineNanos(Duration timeout) { + return timeout == null || timeout.isZero() || timeout.isNegative() + ? Long.MAX_VALUE + : System.nanoTime() + timeout.toNanos(); + } + public void close() { asyncExecutor.shutdownNow(); } diff --git a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java index fbff2b118..a6453a9d6 100644 --- a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java +++ b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java @@ -27,6 +27,7 @@ public class ContinuationContext { private Continuation currentContinuation; private volatile Future pendingFuture; + private volatile Future pendingBatchFuture; private final AtomicReference asyncResult = new AtomicReference<>(); private final AtomicReference asyncException = new AtomicReference<>(); @@ -46,6 +47,22 @@ public void setPendingFuture(Future pendingFuture) { this.pendingFuture = pendingFuture; } + public Future getPendingBatchFuture() { + return pendingBatchFuture; + } + + public void setPendingBatchFuture(Future pendingBatchFuture) { + this.pendingBatchFuture = pendingBatchFuture; + } + + public boolean hasPendingAsync() { + return isPending(pendingFuture) || isPending(pendingBatchFuture); + } + + private static boolean isPending(Future future) { + return future != null && !future.isDone(); + } + public AtomicReference getAsyncResultRef() { return asyncResult; } @@ -56,6 +73,7 @@ public AtomicReference getAsyncExceptionRef() { public void clearAsyncState() { pendingFuture = null; + pendingBatchFuture = null; asyncResult.set(null); asyncException.set(null); } From b01c949688e51191f5e7b692c270ff1e41da8631 Mon Sep 17 00:00:00 2001 From: daken Date: Mon, 20 Jul 2026 00:14:29 +0800 Subject: [PATCH 02/13] refactor(plan): Refactor tool-call execution logic for parallel processing --- .../agents/plan/actions/ToolCallAction.java | 244 ++++++++++++++---- 1 file changed, 197 insertions(+), 47 deletions(-) diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java index bc82e1b61..e1ce84477 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java @@ -22,6 +22,7 @@ import org.apache.flink.agents.api.configuration.ConfigOption; import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.event.ToolRequestEvent; import org.apache.flink.agents.api.event.ToolResponseEvent; @@ -34,6 +35,7 @@ import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.tools.FunctionTool; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -54,11 +56,39 @@ public static Action getToolCallAction() throws Exception { public static void processToolRequest(Event event, RunnerContext ctx) { ToolRequestEvent toolRequest = ToolRequestEvent.fromEvent(event); boolean toolCallAsync = ctx.getConfig().get(AgentExecutionOptions.TOOL_CALL_ASYNC); + boolean toolCallParallel = ctx.getConfig().get(AgentExecutionOptions.TOOL_CALL_PARALLEL); Map success = new HashMap<>(); Map error = new HashMap<>(); Map responses = new HashMap<>(); Map externalIds = new HashMap<>(); + List executions = buildToolCallExecutions(toolRequest, ctx, externalIds); + List callableExecutions = new ArrayList<>(); + List> callables = new ArrayList<>(); + + for (ToolCallExecution execution : executions) { + if (execution.response != null) { + applyInlineResponse(execution, success, error, responses); + } else { + callableExecutions.add(execution); + callables.add(execution.callable); + } + } + + if (toolCallAsync && toolCallParallel && callables.size() > 1) { + executeParallel(callableExecutions, callables, ctx, success, error, responses); + } else { + executeSequentially(callableExecutions, toolCallAsync, ctx, success, error, responses); + } + + ctx.sendEvent( + new ToolResponseEvent(toolRequest.getId(), responses, success, error, externalIds)); + } + + @SuppressWarnings("unchecked") + private static List buildToolCallExecutions( + ToolRequestEvent toolRequest, RunnerContext ctx, Map externalIds) { + List executions = new ArrayList<>(); for (Map toolCall : toolRequest.getToolCalls()) { String id = String.valueOf(toolCall.get("id")); Map function = (Map) toolCall.get("function"); @@ -79,55 +109,175 @@ public static void processToolRequest(Event event, RunnerContext ctx) { diagnosticError = e.getMessage(); } - if (tool != null) { - try { - // Framework-owned injected args must win over model-provided values so hidden - // context such as tenant ids cannot be spoofed by a tool call payload. - mergedArguments.putAll(resolveInjectedArguments(tool, ctx)); - ToolResponse response; - final Tool toolRef = tool; - final Map callArguments = mergedArguments; - DurableCallable callable = - new DurableCallable<>() { - @Override - public String getId() { - return "tool-call"; - } - - @Override - public Class getResultClass() { - return ToolResponse.class; - } - - @Override - public ToolResponse call() throws Exception { - return toolRef.call(new ToolParameters(callArguments)); - } - }; - response = - toolCallAsync - ? ctx.durableExecuteAsync(callable) - : ctx.durableExecute(callable); - success.put(id, response.isSuccess()); - responses.put(id, response); - if (!response.isSuccess() && response.getError() != null) { - error.put(id, response.getError()); - } - } catch (Exception e) { - success.put(id, false); - responses.put( - id, ToolResponse.error(String.format("Tool %s execute failed.", name))); - error.put(id, e.getMessage()); - } - } else { - success.put(id, false); - responses.put( - id, ToolResponse.error(String.format("Tool %s does not exist.", name))); - error.put(id, diagnosticError != null ? diagnosticError : "Tool does not exist."); + if (tool == null) { + executions.add( + ToolCallExecution.withResponse( + id, + name, + ToolResponse.error(String.format("Tool %s does not exist.", name)), + diagnosticError != null + ? diagnosticError + : "Tool does not exist.")); + continue; } + + try { + // Framework-owned injected args must win over model-provided values so hidden + // context such as tenant ids cannot be spoofed by a tool call payload. + mergedArguments.putAll(resolveInjectedArguments(tool, ctx)); + } catch (Exception e) { + executions.add( + ToolCallExecution.withResponse( + id, + name, + ToolResponse.error(String.format("Tool %s execute failed.", name)), + e.getMessage())); + continue; + } + + final Tool toolRef = tool; + final Map callArguments = mergedArguments; + DurableCallable callable = + new DurableCallable<>() { + @Override + public String getId() { + return "tool-call-" + id; + } + + @Override + public Class getResultClass() { + return ToolResponse.class; + } + + @Override + public ToolResponse call() throws Exception { + return toolRef.call(new ToolParameters(callArguments)); + } + }; + executions.add(ToolCallExecution.withCallable(id, name, callable)); } - ctx.sendEvent( - new ToolResponseEvent(toolRequest.getId(), responses, success, error, externalIds)); + return executions; + } + + private static void executeParallel( + List executions, + List> callables, + RunnerContext ctx, + Map success, + Map error, + Map responses) { + try { + List> outcomes = ctx.durableExecuteAllAsync(callables); + for (int i = 0; i < outcomes.size(); i++) { + applyOutcome(executions.get(i), outcomes.get(i), success, error, responses); + } + } catch (Exception e) { + for (ToolCallExecution execution : executions) { + applyExecutionException(execution, e, success, error, responses); + } + } + } + + private static void executeSequentially( + List executions, + boolean toolCallAsync, + RunnerContext ctx, + Map success, + Map error, + Map responses) { + for (ToolCallExecution execution : executions) { + try { + ToolResponse response = + toolCallAsync + ? ctx.durableExecuteAsync(execution.callable) + : ctx.durableExecute(execution.callable); + applyToolResponse(execution.id, response, success, error, responses); + } catch (Exception e) { + applyExecutionException(execution, e, success, error, responses); + } + } + } + + private static void applyOutcome( + ToolCallExecution execution, + Outcome outcome, + Map success, + Map error, + Map responses) { + if (outcome.isFailure()) { + applyExecutionException(execution, outcome.getError(), success, error, responses); + } else { + applyToolResponse(execution.id, outcome.getValue(), success, error, responses); + } + } + + private static void applyInlineResponse( + ToolCallExecution execution, + Map success, + Map error, + Map responses) { + applyToolResponse(execution.id, execution.response, success, error, responses); + if (execution.diagnosticError != null) { + error.put(execution.id, execution.diagnosticError); + } + } + + private static void applyExecutionException( + ToolCallExecution execution, + Exception exception, + Map success, + Map error, + Map responses) { + success.put(execution.id, false); + responses.put( + execution.id, + ToolResponse.error(String.format("Tool %s execute failed.", execution.name))); + error.put(execution.id, exception.getMessage()); + } + + private static void applyToolResponse( + String id, + ToolResponse response, + Map success, + Map error, + Map responses) { + success.put(id, response.isSuccess()); + responses.put(id, response); + if (!response.isSuccess() && response.getError() != null) { + error.put(id, response.getError()); + } + } + + private static class ToolCallExecution { + private final String id; + private final String name; + private final DurableCallable callable; + private final ToolResponse response; + + private ToolCallExecution( + String id, + String name, + DurableCallable callable, + ToolResponse response) { + this.id = id; + this.name = name; + this.callable = callable; + this.response = response; + } + + private static ToolCallExecution withCallable( + String id, String name, DurableCallable callable) { + return new ToolCallExecution(id, name, callable, null); + } + + private static ToolCallExecution withResponse( + String id, String name, ToolResponse response, String diagnosticError) { + ToolCallExecution execution = new ToolCallExecution(id, name, null, response); + execution.diagnosticError = diagnosticError; + return execution; + } + + private String diagnosticError; } private static Map resolveInjectedArguments(Tool tool, RunnerContext ctx) From 1cdfbfb1eedab57ad7432ea440dc489d44d121be Mon Sep 17 00:00:00 2001 From: daken Date: Mon, 20 Jul 2026 18:10:29 +0800 Subject: [PATCH 03/13] fix bug --- .../api/agents/AgentExecutionOptions.java | 3 -- .../context/JavaRunnerContextImpl.java | 15 +------ .../async/ContinuationActionExecutor.java | 44 +++++++++++++++---- 3 files changed, 38 insertions(+), 24 deletions(-) diff --git a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java index 5f8d9d0bd..509d3644a 100644 --- a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java +++ b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java @@ -56,9 +56,6 @@ public class AgentExecutionOptions { Integer.class, Runtime.getRuntime().availableProcessors() * 2); - public static final ConfigOption TOOL_CALL_MAX_PARALLELISM = - new ConfigOption<>("tool-call.max-parallelism", Integer.class, Integer.MAX_VALUE); - public static final ConfigOption TOOL_CALL_BATCH_TIMEOUT = new ConfigOption<>("tool-call.batch.timeout", Duration.class, Duration.ofMillis(-1)); diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java index b365e010d..18d75aa36 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java @@ -253,19 +253,8 @@ private List> executeOutcomeSuppliers(List> suppliers .collect(Collectors.toList()); } Duration timeout = getConfig().get(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT); - int maxParallelism = getConfig().get(AgentExecutionOptions.TOOL_CALL_MAX_PARALLELISM); - if (maxParallelism <= 0 || maxParallelism >= suppliers.size()) { - return toolCallContinuationExecutor.executeAllAsync( - continuationContext, suppliers, timeout); - } - List> results = new ArrayList<>(suppliers.size()); - for (int i = 0; i < suppliers.size(); i += maxParallelism) { - int end = Math.min(i + maxParallelism, suppliers.size()); - results.addAll( - toolCallContinuationExecutor.executeAllAsync( - continuationContext, suppliers.subList(i, end), timeout)); - } - return results; + return toolCallContinuationExecutor.executeAllAsync( + continuationContext, suppliers, timeout); } private T executeAsyncCallable(DurableCallable callable) throws Exception { diff --git a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java index a721dbcaf..6035bf23a 100644 --- a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java +++ b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java @@ -200,21 +200,25 @@ public List> executeAllAsync( TimeoutException exception = new TimeoutException( "Async durable batch execution timed out after " + timeout); - for (CompletableFuture> future : futures) { - future.cancel(true); - } barrier.cancel(true); context.setPendingBatchFuture(null); - List> results = new ArrayList<>(suppliers.size()); - for (int i = 0; i < suppliers.size(); i++) { - results.add(Outcome.failure(exception)); - } - return results; + return collectBatchOutcomesOnTimeout(futures, exception); } Continuation.yield(SCOPE); } context.setPendingBatchFuture(null); + return collectBatchOutcomes(futures); + } + + /** + * Collects per-slot outcomes after the batch barrier completes normally. + * + *

Each supplier already wraps success and failure into an {@link Outcome}, so {@code join()} + * returns that outcome rather than throwing for ordinary tool exceptions. + */ + private static List> collectBatchOutcomes( + List>> futures) { List> results = new ArrayList<>(futures.size()); for (CompletableFuture> future : futures) { results.add(future.join()); @@ -222,6 +226,30 @@ public List> executeAllAsync( return results; } + /** + * Collects per-slot outcomes when the batch deadline elapses. + * + *

Completed slots keep their success or failure outcome. Only slots that are still running + * (or become cancelled) are finalized as timeout failures. {@code cancel(true)} is attempted + * only for unfinished futures; a future that completes between the check and cancel stays + * non-cancelled and is collected as a normal outcome. + */ + private static List> collectBatchOutcomesOnTimeout( + List>> futures, TimeoutException timeoutException) { + List> results = new ArrayList<>(futures.size()); + for (CompletableFuture> future : futures) { + if (!future.isDone()) { + future.cancel(true); + } + if (future.isDone() && !future.isCancelled()) { + results.add(future.join()); + } else { + results.add(Outcome.failure(timeoutException)); + } + } + return results; + } + private long getDeadlineNanos(Duration timeout) { return timeout == null || timeout.isZero() || timeout.isNegative() ? Long.MAX_VALUE From 87a01ec875ff7361b5f6e955088fd153948a3b97 Mon Sep 17 00:00:00 2001 From: daken Date: Tue, 21 Jul 2026 18:11:44 +0800 Subject: [PATCH 04/13] add user docs --- .../api/agents/AgentExecutionOptions.java | 22 +++++++++++++++++++ docs/content/docs/operations/configuration.md | 5 ++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java index 509d3644a..32f248d8b 100644 --- a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java +++ b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java @@ -44,18 +44,40 @@ public class AgentExecutionOptions { public static final ConfigOption CHAT_ASYNC = new ConfigOption<>("chat.async", Boolean.class, true); + /** Whether the built-in tool-call action runs each tool via durable async execution. */ public static final ConfigOption TOOL_CALL_ASYNC = new ConfigOption<>("tool-call.async", Boolean.class, true); + /** + * Whether multiple tool calls from one {@code ToolRequestEvent} run as one parallel durable + * batch when {@link #TOOL_CALL_ASYNC} is also enabled (JDK 21+). + * + *

Default is {@code true}. A parallel batch raises the number of in-flight external calls; + * after failover, tools whose results were not yet persisted may be submitted again. + * Side-effecting tools should be idempotent or provide a {@code reconciler()}. Set to {@code + * false} to keep serial async or sync tool execution. + */ public static final ConfigOption TOOL_CALL_PARALLEL = new ConfigOption<>("tool-call.parallel", Boolean.class, true); + /** + * Size of the dedicated thread pool used for tool-call async and parallel batch execution. + * + *

Separate from {@link #NUM_ASYNC_THREADS} so a large tool batch does not exhaust the global + * async pool. + */ public static final ConfigOption TOOL_CALL_NUM_ASYNC_THREADS = new ConfigOption<>( "tool-call.num-async-threads", Integer.class, Runtime.getRuntime().availableProcessors() * 2); + /** + * Overall timeout for one parallel tool-call batch. + * + *

Non-positive values disable the timeout. When the deadline elapses, unfinished slots are + * failed; slots that already completed keep their success or failure outcome. + */ public static final ConfigOption TOOL_CALL_BATCH_TIMEOUT = new ConfigOption<>("tool-call.batch.timeout", Duration.class, Duration.ofMillis(-1)); diff --git a/docs/content/docs/operations/configuration.md b/docs/content/docs/operations/configuration.md index 9ba9e200e..98f2bf39f 100644 --- a/docs/content/docs/operations/configuration.md +++ b/docs/content/docs/operations/configuration.md @@ -131,7 +131,10 @@ Here is the list of all built-in core configuration options. | `max-retries` | 3 | int | Number of retries when using `ErrorHandlingStrategy.RETRY`. | | `retry-wait-interval` | 1 | int | Base wait interval in seconds between retries when using `ErrorHandlingStrategy.RETRY`. Uses exponential backoff: the actual wait time for the Nth retry is `retry-wait-interval * 2^(N-1)` seconds. For example, with default 1s, waits are 1s, 2s, 4s, etc. Retry count and total wait time are reported in `ChatResponseEvent` and recorded as metrics (`retryCount`, `retryWaitSec`) under the connection name. | | `chat.async` | true | boolean | Whether chat asynchronously for built-in chat action. | -| `tool-call.async` | true | boolean | Whether process tool call for built-in tool call action. | +| `tool-call.async` | true | boolean | Whether the built-in tool-call action runs each tool via durable async execution. | +| `tool-call.parallel` | true | boolean | When `tool-call.async` is also true (JDK 21+), run multiple tool calls from one `ToolRequestEvent` as one parallel durable batch. Increases in-flight external calls; after failover, unfinished tools may be submitted again — side-effecting tools should be idempotent or provide a reconciler. Set to `false` for serial tool execution. | +| `tool-call.num-async-threads` | os cpu count * 2 | int | Dedicated thread pool size for tool-call async / parallel batch execution. Separate from `num-async-threads` so a large tool batch does not exhaust the global async pool. | +| `tool-call.batch.timeout` | -1ms (disabled) | Duration | Overall timeout for one parallel tool-call batch. Non-positive disables it. On timeout, unfinished slots fail; already completed slots keep their outcome. | | `rag.async` | true | boolean | Whether retrieve context asynchronously for built-in context retrieval action. | | `num-async-threads` | os cpu count * 2 | int | The thread pool size for async executor. | | `job-identifier` | none | String | The unique identifier of job, remaining consistent after restoring from a savepoint. If not set, uses flink job id. | From c21cf3d70d15a383e17f681bc2ccbe309d125ea2 Mon Sep 17 00:00:00 2001 From: daken Date: Fri, 24 Jul 2026 13:57:35 +0800 Subject: [PATCH 05/13] Update toolCallAction --- .../agents/plan/actions/ToolCallAction.java | 117 ++++++++---------- 1 file changed, 49 insertions(+), 68 deletions(-) diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java index e1ce84477..e2f5c512e 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java @@ -52,7 +52,6 @@ public static Action getToolCallAction() throws Exception { List.of(ToolRequestEvent.EVENT_TYPE)); } - @SuppressWarnings("unchecked") public static void processToolRequest(Event event, RunnerContext ctx) { ToolRequestEvent toolRequest = ToolRequestEvent.fromEvent(event); boolean toolCallAsync = ctx.getConfig().get(AgentExecutionOptions.TOOL_CALL_ASYNC); @@ -62,23 +61,14 @@ public static void processToolRequest(Event event, RunnerContext ctx) { Map error = new HashMap<>(); Map responses = new HashMap<>(); Map externalIds = new HashMap<>(); - List executions = buildToolCallExecutions(toolRequest, ctx, externalIds); - List callableExecutions = new ArrayList<>(); - List> callables = new ArrayList<>(); - - for (ToolCallExecution execution : executions) { - if (execution.response != null) { - applyInlineResponse(execution, success, error, responses); - } else { - callableExecutions.add(execution); - callables.add(execution.callable); - } - } + List executions = + buildToolCallExecutions( + toolRequest, ctx, externalIds, success, error, responses); - if (toolCallAsync && toolCallParallel && callables.size() > 1) { - executeParallel(callableExecutions, callables, ctx, success, error, responses); + if (toolCallAsync && toolCallParallel && executions.size() > 1) { + executeParallel(executions, ctx, success, error, responses); } else { - executeSequentially(callableExecutions, toolCallAsync, ctx, success, error, responses); + executeSequentially(executions, toolCallAsync, ctx, success, error, responses); } ctx.sendEvent( @@ -87,7 +77,12 @@ public static void processToolRequest(Event event, RunnerContext ctx) { @SuppressWarnings("unchecked") private static List buildToolCallExecutions( - ToolRequestEvent toolRequest, RunnerContext ctx, Map externalIds) { + ToolRequestEvent toolRequest, + RunnerContext ctx, + Map externalIds, + Map success, + Map error, + Map responses) { List executions = new ArrayList<>(); for (Map toolCall : toolRequest.getToolCalls()) { String id = String.valueOf(toolCall.get("id")); @@ -110,14 +105,13 @@ private static List buildToolCallExecutions( } if (tool == null) { - executions.add( - ToolCallExecution.withResponse( - id, - name, - ToolResponse.error(String.format("Tool %s does not exist.", name)), - diagnosticError != null - ? diagnosticError - : "Tool does not exist.")); + recordInlineResponse( + id, + ToolResponse.error(String.format("Tool %s does not exist.", name)), + diagnosticError != null ? diagnosticError : "Tool does not exist.", + success, + error, + responses); continue; } @@ -126,12 +120,13 @@ private static List buildToolCallExecutions( // context such as tenant ids cannot be spoofed by a tool call payload. mergedArguments.putAll(resolveInjectedArguments(tool, ctx)); } catch (Exception e) { - executions.add( - ToolCallExecution.withResponse( - id, - name, - ToolResponse.error(String.format("Tool %s execute failed.", name)), - e.getMessage())); + recordInlineResponse( + id, + ToolResponse.error(String.format("Tool %s execute failed.", name)), + e.getMessage(), + success, + error, + responses); continue; } @@ -154,26 +149,29 @@ public ToolResponse call() throws Exception { return toolRef.call(new ToolParameters(callArguments)); } }; - executions.add(ToolCallExecution.withCallable(id, name, callable)); + executions.add(new ToolCallExecution(id, name, callable)); } return executions; } private static void executeParallel( List executions, - List> callables, RunnerContext ctx, Map success, Map error, Map responses) { + List> callables = new ArrayList<>(executions.size()); + for (ToolCallExecution execution : executions) { + callables.add(execution.callable); + } try { List> outcomes = ctx.durableExecuteAllAsync(callables); for (int i = 0; i < outcomes.size(); i++) { - applyOutcome(executions.get(i), outcomes.get(i), success, error, responses); + recordOutcome(executions.get(i), outcomes.get(i), success, error, responses); } } catch (Exception e) { for (ToolCallExecution execution : executions) { - applyExecutionException(execution, e, success, error, responses); + recordExecutionException(execution, e, success, error, responses); } } } @@ -191,38 +189,40 @@ private static void executeSequentially( toolCallAsync ? ctx.durableExecuteAsync(execution.callable) : ctx.durableExecute(execution.callable); - applyToolResponse(execution.id, response, success, error, responses); + recordToolResponse(execution.id, response, success, error, responses); } catch (Exception e) { - applyExecutionException(execution, e, success, error, responses); + recordExecutionException(execution, e, success, error, responses); } } } - private static void applyOutcome( + private static void recordOutcome( ToolCallExecution execution, Outcome outcome, Map success, Map error, Map responses) { if (outcome.isFailure()) { - applyExecutionException(execution, outcome.getError(), success, error, responses); + recordExecutionException(execution, outcome.getError(), success, error, responses); } else { - applyToolResponse(execution.id, outcome.getValue(), success, error, responses); + recordToolResponse(execution.id, outcome.getValue(), success, error, responses); } } - private static void applyInlineResponse( - ToolCallExecution execution, + private static void recordInlineResponse( + String id, + ToolResponse response, + String diagnosticError, Map success, Map error, Map responses) { - applyToolResponse(execution.id, execution.response, success, error, responses); - if (execution.diagnosticError != null) { - error.put(execution.id, execution.diagnosticError); + recordToolResponse(id, response, success, error, responses); + if (diagnosticError != null) { + error.put(id, diagnosticError); } } - private static void applyExecutionException( + private static void recordExecutionException( ToolCallExecution execution, Exception exception, Map success, @@ -235,7 +235,7 @@ private static void applyExecutionException( error.put(execution.id, exception.getMessage()); } - private static void applyToolResponse( + private static void recordToolResponse( String id, ToolResponse response, Map success, @@ -248,36 +248,17 @@ private static void applyToolResponse( } } - private static class ToolCallExecution { + private static final class ToolCallExecution { private final String id; private final String name; private final DurableCallable callable; - private final ToolResponse response; private ToolCallExecution( - String id, - String name, - DurableCallable callable, - ToolResponse response) { + String id, String name, DurableCallable callable) { this.id = id; this.name = name; this.callable = callable; - this.response = response; - } - - private static ToolCallExecution withCallable( - String id, String name, DurableCallable callable) { - return new ToolCallExecution(id, name, callable, null); } - - private static ToolCallExecution withResponse( - String id, String name, ToolResponse response, String diagnosticError) { - ToolCallExecution execution = new ToolCallExecution(id, name, null, response); - execution.diagnosticError = diagnosticError; - return execution; - } - - private String diagnosticError; } private static Map resolveInjectedArguments(Tool tool, RunnerContext ctx) From 97b9421fe971e0e32cd0af3c11069b77b17796b8 Mon Sep 17 00:00:00 2001 From: daken Date: Fri, 24 Jul 2026 15:07:35 +0800 Subject: [PATCH 06/13] Update ContinuationActionExecutor --- .../async/ContinuationActionExecutor.java | 9 ++++---- .../runtime/async/ContinuationContext.java | 22 +++++++++++++++++++ 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java index 6035bf23a..284ff441c 100644 --- a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java +++ b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java @@ -66,8 +66,9 @@ public ContinuationActionExecutor(int numAsyncThreads) { * @return true if the action completed, false if waiting for async execution */ public boolean executeAction(ContinuationContext context, Runnable action) { - // Check if we have a pending async Future from previous yield - if (context.hasPendingAsync()) { + // Wait while async work is still pending, unless the batch deadline elapsed — then resume + // so executeAllAsync can finalize timed-out slots. + if (context.hasPendingAsync() && !context.isBatchDeadlineElapsed()) { return false; } @@ -192,9 +193,9 @@ public List> executeAllAsync( CompletableFuture barrier = CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])); - context.setPendingBatchFuture(barrier); - long deadlineNanos = getDeadlineNanos(timeout); + context.setPendingBatchFuture(barrier, deadlineNanos); + while (!barrier.isDone()) { if (System.nanoTime() >= deadlineNanos) { TimeoutException exception = diff --git a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java index a6453a9d6..0523abb5b 100644 --- a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java +++ b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java @@ -28,6 +28,9 @@ public class ContinuationContext { private Continuation currentContinuation; private volatile Future pendingFuture; private volatile Future pendingBatchFuture; + /** {@link Long#MAX_VALUE} means the pending batch has no deadline. */ + private volatile long pendingBatchDeadlineNanos = Long.MAX_VALUE; + private final AtomicReference asyncResult = new AtomicReference<>(); private final AtomicReference asyncException = new AtomicReference<>(); @@ -52,13 +55,31 @@ public Future getPendingBatchFuture() { } public void setPendingBatchFuture(Future pendingBatchFuture) { + setPendingBatchFuture(pendingBatchFuture, Long.MAX_VALUE); + } + + /** + * Records a pending batch barrier and its absolute nanoTime deadline. + * + *

{@link #hasPendingAsync()} only reflects whether the barrier future is still running. + * {@link ContinuationActionExecutor#executeAction} consults {@link #isBatchDeadlineElapsed()} + * separately so a timed-out batch can still resume and finalize unfinished slots. + */ + public void setPendingBatchFuture(Future pendingBatchFuture, long deadlineNanos) { this.pendingBatchFuture = pendingBatchFuture; + this.pendingBatchDeadlineNanos = + pendingBatchFuture == null ? Long.MAX_VALUE : deadlineNanos; } public boolean hasPendingAsync() { return isPending(pendingFuture) || isPending(pendingBatchFuture); } + /** Whether the pending batch deadline has elapsed. No deadline means {@link Long#MAX_VALUE}. */ + public boolean isBatchDeadlineElapsed() { + return System.nanoTime() >= pendingBatchDeadlineNanos; + } + private static boolean isPending(Future future) { return future != null && !future.isDone(); } @@ -74,6 +95,7 @@ public AtomicReference getAsyncExceptionRef() { public void clearAsyncState() { pendingFuture = null; pendingBatchFuture = null; + pendingBatchDeadlineNanos = Long.MAX_VALUE; asyncResult.set(null); asyncException.set(null); } From e3e178c18731f7b5606e930266cdcde1a8e84c78 Mon Sep 17 00:00:00 2001 From: daken Date: Sun, 26 Jul 2026 00:06:29 +0800 Subject: [PATCH 07/13] feat(runtime): add java test --- .../agents/api/context/RunnerContext.java | 14 +- .../agents/plan/actions/ToolCallAction.java | 6 +- .../plan/actions/ToolCallActionTest.java | 184 ++++++++++++++++-- .../context/JavaRunnerContextImpl.java | 6 +- ...nerContextImplDurableExecuteAsyncTest.java | 146 ++++++++++++++ .../agents/runtime/memory/MemoryRefTest.java | 11 ++ 6 files changed, 332 insertions(+), 35 deletions(-) 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 cf35e6905..a6d816300 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 @@ -159,18 +159,8 @@ public interface RunnerContext { * *

Access to memory and sendEvent are prohibited within the callables. */ - default List> durableExecuteAllAsync(List> callables) - throws Exception { - java.util.ArrayList> outcomes = new java.util.ArrayList<>(callables.size()); - for (DurableCallable callable : callables) { - try { - outcomes.add(Outcome.success(durableExecute(callable))); - } catch (Exception e) { - outcomes.add(Outcome.failure(e)); - } - } - return outcomes; - } + List> durableExecuteAllAsync(List> callables) + throws Exception; /** Clean up the resource. */ void close() throws Exception; diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java index e2f5c512e..412595125 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java @@ -62,8 +62,7 @@ public static void processToolRequest(Event event, RunnerContext ctx) { Map responses = new HashMap<>(); Map externalIds = new HashMap<>(); List executions = - buildToolCallExecutions( - toolRequest, ctx, externalIds, success, error, responses); + buildToolCallExecutions(toolRequest, ctx, externalIds, success, error, responses); if (toolCallAsync && toolCallParallel && executions.size() > 1) { executeParallel(executions, ctx, success, error, responses); @@ -253,8 +252,7 @@ private static final class ToolCallExecution { private final String name; private final DurableCallable callable; - private ToolCallExecution( - String id, String name, DurableCallable callable) { + private ToolCallExecution(String id, String name, DurableCallable callable) { this.id = id; this.name = name; this.callable = callable; 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..f8accc967 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 @@ -18,11 +18,13 @@ package org.apache.flink.agents.plan.actions; import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; import org.apache.flink.agents.api.annotation.ToolParam; import org.apache.flink.agents.api.configuration.ReadableConfiguration; import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.MemoryRef; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.event.ToolRequestEvent; import org.apache.flink.agents.api.event.ToolResponseEvent; @@ -46,6 +48,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; @@ -76,7 +79,19 @@ public ToolResponse call(ToolParameters parameters) { static class FakeRunnerContext implements RunnerContext { private final List sentEvents = new ArrayList<>(); private final AgentConfiguration config = - new AgentConfiguration(Map.of("tenant_id", "tenant-1")); + new AgentConfiguration( + Map.of( + "tenant_id", + "tenant-1", + AgentExecutionOptions.TOOL_CALL_ASYNC.getKey(), + true, + AgentExecutionOptions.TOOL_CALL_PARALLEL.getKey(), + true)); + private final List durableExecuteIds = new ArrayList<>(); + private final List durableExecuteAsyncIds = new ArrayList<>(); + private final List> durableExecuteAllAsyncIds = new ArrayList<>(); + private final AtomicInteger queryOrderCalls = new AtomicInteger(); + private List> durableExecuteAllAsyncOutcomes; private ToolParameterInjection injection = ToolParameterInjection.fromConfig("tenant_id"); private MemoryObject sensoryMemory; private MemoryObject shortTermMemory; @@ -86,6 +101,21 @@ FakeRunnerContext withInjection(ToolParameterInjection injection) { return this; } + FakeRunnerContext withToolCallAsync(boolean enabled) { + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, enabled); + return this; + } + + FakeRunnerContext withToolCallParallel(boolean enabled) { + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, enabled); + return this; + } + + FakeRunnerContext withDurableExecuteAllAsyncOutcomes(List> outcomes) { + this.durableExecuteAllAsyncOutcomes = outcomes; + return this; + } + FakeRunnerContext withSensoryMemory(Map values) { this.sensoryMemory = new FakeMemoryObject(values); return this; @@ -156,14 +186,35 @@ public Object getActionConfigValue(String key) { @Override public T durableExecute(DurableCallable callable) throws Exception { + durableExecuteIds.add(callable.getId()); return callable.call(); } @Override public T durableExecuteAsync(DurableCallable callable) throws Exception { + durableExecuteAsyncIds.add(callable.getId()); return callable.call(); } + @Override + @SuppressWarnings("unchecked") + public List> durableExecuteAllAsync(List> callables) + throws Exception { + List ids = new ArrayList<>(); + for (DurableCallable callable : callables) { + ids.add(callable.getId()); + } + durableExecuteAllAsyncIds.add(ids); + if (durableExecuteAllAsyncOutcomes != null) { + return (List>) (List) durableExecuteAllAsyncOutcomes; + } + List> outcomes = new ArrayList<>(callables.size()); + for (DurableCallable callable : callables) { + outcomes.add(Outcome.success(callable.call())); + } + return outcomes; + } + @Override public void close() {} } @@ -378,20 +429,123 @@ public Resource getResource(String name, ResourceType type) { assertThat(response.getError()).containsEntry("call-1", "missing resource"); } + @Test + void processToolRequestUsesParallelBatchForMultipleTools() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext(); + + ToolCallAction.processToolRequest(toolRequest("queryOrder", "call-1", "call-2"), ctx); + + assertThat(ctx.durableExecuteAllAsyncIds) + .containsExactly(List.of("tool-call-call-1", "tool-call-call-2")); + assertThat(ctx.durableExecuteAsyncIds).isEmpty(); + assertThat(ctx.durableExecuteIds).isEmpty(); + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getResponses().get("call-1").getResult()).isEqualTo("tenant-1:order-1"); + assertThat(response.getResponses().get("call-2").getResult()).isEqualTo("tenant-1:order-2"); + } + + @Test + void processToolRequestUsesSerialAsyncWhenParallelDisabled() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext().withToolCallParallel(false); + + ToolCallAction.processToolRequest(toolRequest("queryOrder", "call-1", "call-2"), ctx); + + assertThat(ctx.durableExecuteAllAsyncIds).isEmpty(); + assertThat(ctx.durableExecuteAsyncIds) + .containsExactly("tool-call-call-1", "tool-call-call-2"); + assertThat(ctx.durableExecuteIds).isEmpty(); + } + + @Test + void processToolRequestUsesSyncWhenAsyncDisabled() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext().withToolCallAsync(false); + + ToolCallAction.processToolRequest(toolRequest("queryOrder", "call-1", "call-2"), ctx); + + assertThat(ctx.durableExecuteAllAsyncIds).isEmpty(); + assertThat(ctx.durableExecuteAsyncIds).isEmpty(); + assertThat(ctx.durableExecuteIds).containsExactly("tool-call-call-1", "tool-call-call-2"); + } + + @Test + void processToolRequestDoesNotBatchSingleTool() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext(); + + ToolCallAction.processToolRequest(toolRequest("queryOrder"), ctx); + + assertThat(ctx.durableExecuteAllAsyncIds).isEmpty(); + assertThat(ctx.durableExecuteAsyncIds).containsExactly("tool-call-call-1"); + assertThat(ctx.durableExecuteIds).isEmpty(); + } + + @Test + void processToolRequestExcludesMissingToolFromParallelBatch() throws Exception { + FakeRunnerContext ctx = + new FakeRunnerContext() { + @Override + public Resource getResource(String name, ResourceType type) throws Exception { + if ("missingTool".equals(name)) { + throw new RuntimeException("missing resource"); + } + return super.getResource(name, type); + } + }; + + ToolCallAction.processToolRequest( + new ToolRequestEvent( + "model", + List.of( + toolCall("missingTool", "missing-call", "order-0"), + toolCall("queryOrder", "call-1", "order-1"), + toolCall("queryOrder", "call-2", "order-2"))), + ctx); + + assertThat(ctx.durableExecuteAllAsyncIds) + .containsExactly(List.of("tool-call-call-1", "tool-call-call-2")); + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("missing-call", false); + assertThat(response.getError()).containsEntry("missing-call", "missing resource"); + } + + @Test + void processToolRequestRecordsOutcomeFailureAsToolError() throws Exception { + FakeRunnerContext ctx = + new FakeRunnerContext() + .withDurableExecuteAllAsyncOutcomes( + List.of( + Outcome.success(ToolResponse.success("ok")), + Outcome.failure(new RuntimeException("boom")))); + + ToolCallAction.processToolRequest(toolRequest("queryOrder", "call-1", "call-2"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", true); + assertThat(response.getResponses().get("call-1").getResult()).isEqualTo("ok"); + assertThat(response.getSuccess()).containsEntry("call-2", false); + assertThat(response.getResponses().get("call-2").getError()) + .isEqualTo("Tool queryOrder execute failed."); + assertThat(response.getError()).containsEntry("call-2", "boom"); + } + private static ToolRequestEvent toolRequest(String toolName) { - return new ToolRequestEvent( - "model", - List.of( - Map.of( - "id", - "call-1", - "type", - "function", - "function", - Map.of( - "name", - toolName, - "arguments", - Map.of("orderId", "order-1"))))); + return new ToolRequestEvent("model", List.of(toolCall(toolName, "call-1", "order-1"))); + } + + private static ToolRequestEvent toolRequest(String toolName, String... ids) { + List> toolCalls = new ArrayList<>(); + for (int i = 0; i < ids.length; i++) { + toolCalls.add(toolCall(toolName, ids[i], "order-" + (i + 1))); + } + return new ToolRequestEvent("model", toolCalls); + } + + private static Map toolCall(String toolName, String id, String orderId) { + return Map.of( + "id", + id, + "type", + "function", + "function", + Map.of("name", toolName, "arguments", Map.of("orderId", orderId))); } } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java index 18d75aa36..f5e0070c1 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java @@ -226,8 +226,7 @@ private BatchExecutionPlan(int size) { } } - private List> executeAllWithoutDurableState(List> callables) - throws Exception { + private List> executeAllWithoutDurableState(List> callables) { List> suppliers = new ArrayList<>(); for (DurableCallable callable : callables) { suppliers.add(callable::call); @@ -235,8 +234,7 @@ private List> executeAllWithoutDurableState(List List> executeOutcomeSuppliers(List> suppliers) - throws Exception { + private List> executeOutcomeSuppliers(List> suppliers) { if (suppliers.isEmpty()) { return List.of(); } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java index ff378fb6a..3dfd77017 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java @@ -20,6 +20,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.runtime.actionstate.ActionState; @@ -32,6 +33,8 @@ import org.junit.jupiter.api.Test; import java.util.HashMap; +import java.util.List; +import java.util.concurrent.Callable; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Supplier; @@ -195,6 +198,121 @@ void testDurableExecuteAsyncReconcilableReconcileExceptionPersistsFailure() thro assertEquals(1, context.getDurableExecutionContext().getCurrentCallIndex()); } + @Test + void testDurableExecuteAllAsyncInitialBatchPersistsOutcomes() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + JavaRunnerContextImpl context = createContext(new ActionState(null), executor); + TestDurableCallable first = + new TestDurableCallable<>("batch-1", String.class, () -> "one"); + TestDurableCallable second = + new TestDurableCallable<>("batch-2", String.class, () -> "two"); + + List> outcomes = context.durableExecuteAllAsync(List.of(first, second)); + + assertEquals("one", outcomes.get(0).getValue()); + assertEquals("two", outcomes.get(1).getValue()); + assertEquals(1, executor.getExecuteAllAsyncCallCount()); + assertEquals(List.of(2), executor.getExecuteAllAsyncBatchSizes()); + assertEquals(1, first.getCallCount()); + assertEquals(1, second.getCallCount()); + assertEquals(3, persistCallCount.get()); + assertEquals(2, context.getDurableExecutionContext().getCurrentCallIndex()); + List persisted = + context.getDurableExecutionContext().getActionState().getCallResults(); + assertEquals(2, persisted.size()); + assertEquals("batch-1", persisted.get(0).getFunctionId()); + assertTrue(persisted.get(0).isSuccess()); + assertEquals("batch-2", persisted.get(1).getFunctionId()); + assertTrue(persisted.get(1).isSuccess()); + } + + @Test + void testDurableExecuteAllAsyncReconcilesPendingSlot() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + ActionState actionState = new ActionState(null); + actionState.addCallResult(CallResult.pending("batch-1", "")); + JavaRunnerContextImpl context = createContext(actionState, executor); + TestReconcilableCallable callable = + new TestReconcilableCallable<>( + "batch-1", + String.class, + () -> fail("call should not be executed"), + () -> "recovered"); + + List> outcomes = context.durableExecuteAllAsync(List.of(callable)); + + assertEquals("recovered", outcomes.get(0).getValue()); + assertEquals(0, callable.getCallCount()); + assertEquals(1, callable.getReconcileCount()); + assertEquals(1, executor.getExecuteAllAsyncCallCount()); + assertEquals(1, persistCallCount.get()); + assertTrue(actionState.getCallResults().get(0).isSuccess()); + assertEquals(1, context.getDurableExecutionContext().getCurrentCallIndex()); + } + + @Test + void testDurableExecuteAllAsyncRecoversPartialFinalizedBatch() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + ActionState actionState = new ActionState(null); + actionState.addCallResult( + new CallResult("batch-1", "", OBJECT_MAPPER.writeValueAsBytes("cached-one"))); + actionState.addCallResult( + new CallResult("batch-2", "", OBJECT_MAPPER.writeValueAsBytes("cached-two"))); + actionState.addCallResult(CallResult.pending("batch-3", "")); + JavaRunnerContextImpl context = createContext(actionState, executor); + TestDurableCallable first = + new TestDurableCallable<>( + "batch-1", String.class, () -> fail("cached slot should not execute")); + TestDurableCallable second = + new TestDurableCallable<>( + "batch-2", String.class, () -> fail("cached slot should not execute")); + TestDurableCallable third = + new TestDurableCallable<>("batch-3", String.class, () -> "fresh-three"); + + List> outcomes = + context.durableExecuteAllAsync(List.of(first, second, third)); + + assertEquals("cached-one", outcomes.get(0).getValue()); + assertEquals("cached-two", outcomes.get(1).getValue()); + assertEquals("fresh-three", outcomes.get(2).getValue()); + assertEquals(0, first.getCallCount()); + assertEquals(0, second.getCallCount()); + assertEquals(1, third.getCallCount()); + assertEquals(1, executor.getExecuteAllAsyncCallCount()); + assertEquals(List.of(1), executor.getExecuteAllAsyncBatchSizes()); + assertEquals("batch-3", actionState.getCallResults().get(2).getFunctionId()); + assertTrue(actionState.getCallResults().get(2).isSuccess()); + assertEquals(1, persistCallCount.get()); + assertEquals(3, context.getDurableExecutionContext().getCurrentCallIndex()); + } + + @Test + void testDurableExecuteAllAsyncReturnsCachedFailureOutcome() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + ActionState actionState = new ActionState(null); + actionState.addCallResult( + new CallResult( + "batch-1", + "", + null, + OBJECT_MAPPER.writeValueAsBytes( + RunnerContextImpl.DurableExecutionException.fromException( + new IllegalStateException("cached failure"))))); + JavaRunnerContextImpl context = createContext(actionState, executor); + TestDurableCallable callable = + new TestDurableCallable<>( + "batch-1", String.class, () -> fail("cached slot should not execute")); + + List> outcomes = context.durableExecuteAllAsync(List.of(callable)); + + assertTrue(outcomes.get(0).isFailure()); + assertTrue(outcomes.get(0).getError().getMessage().contains("cached failure")); + assertEquals(0, callable.getCallCount()); + assertEquals(0, executor.getExecuteAllAsyncCallCount()); + assertEquals(0, persistCallCount.get()); + assertEquals(1, context.getDurableExecutionContext().getCurrentCallIndex()); + } + private JavaRunnerContextImpl createContext( ActionState actionState, ContinuationActionExecutor executor) { JavaRunnerContextImpl context = @@ -226,6 +344,8 @@ private static final class InspectingContinuationActionExecutor extends ContinuationActionExecutor { private Runnable beforeExecute; private int executeAsyncCallCount; + private int executeAllAsyncCallCount; + private final List executeAllAsyncBatchSizes = new java.util.ArrayList<>(); private InspectingContinuationActionExecutor() { super(1); @@ -240,6 +360,24 @@ public T executeAsync(ContinuationContext context, Supplier supplier) { return supplier.get(); } + @Override + public List> executeAllAsync( + ContinuationContext context, + List> suppliers, + java.time.Duration timeout) { + executeAllAsyncCallCount++; + executeAllAsyncBatchSizes.add(suppliers.size()); + List> outcomes = new java.util.ArrayList<>(suppliers.size()); + for (Callable supplier : suppliers) { + try { + outcomes.add(Outcome.success(supplier.call())); + } catch (Exception e) { + outcomes.add(Outcome.failure(e)); + } + } + return outcomes; + } + private void setBeforeExecute(Runnable beforeExecute) { this.beforeExecute = beforeExecute; } @@ -247,5 +385,13 @@ private void setBeforeExecute(Runnable beforeExecute) { private int getExecuteAsyncCallCount() { return executeAsyncCallCount; } + + private int getExecuteAllAsyncCallCount() { + return executeAllAsyncCallCount; + } + + private List getExecuteAllAsyncBatchSizes() { + return executeAllAsyncBatchSizes; + } } } 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..615c6c748 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 @@ -21,6 +21,7 @@ import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.MemoryRef; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.memory.BaseLongTermMemory; import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; @@ -128,6 +129,16 @@ public T durableExecuteAsync(DurableCallable callable) throws Exception { return callable.call(); } + @Override + public List> durableExecuteAllAsync(List> callables) + throws Exception { + List> outcomes = new ArrayList<>(callables.size()); + for (DurableCallable callable : callables) { + outcomes.add(Outcome.success(callable.call())); + } + return outcomes; + } + @Override public void close() throws Exception {} } From ffbe1c5ad47e61444a9e871eab23e81a1126ada0 Mon Sep 17 00:00:00 2001 From: daken Date: Sun, 26 Jul 2026 15:47:04 +0800 Subject: [PATCH 08/13] [feature] add_durableExecuteAllAsync_python --- python/flink_agents/api/core_options.py | 12 + python/flink_agents/api/runner_context.py | 52 +++- .../api/tests/test_core_options.py | 12 + .../plan/actions/tool_call_action.py | 177 ++++++++++--- .../tests/actions/test_tool_call_action.py | 136 +++++++++- .../runtime/flink_runner_context.py | 174 +++++++++++++ .../test_flink_runner_context_reconcilable.py | 233 +++++++++++++++++- .../runtime/context/RunnerContextImpl.java | 27 +- 8 files changed, 770 insertions(+), 53 deletions(-) diff --git a/python/flink_agents/api/core_options.py b/python/flink_agents/api/core_options.py index d196777c4..164585d12 100644 --- a/python/flink_agents/api/core_options.py +++ b/python/flink_agents/api/core_options.py @@ -254,6 +254,18 @@ class AgentExecutionOptions: default=True, ) + TOOL_CALL_PARALLEL = ConfigOption( + key="tool-call.parallel", + config_type=bool, + default=True, + ) + + TOOL_CALL_BATCH_TIMEOUT = ConfigOption( + key="tool-call.batch.timeout", + config_type=int, + default=-1, + ) + RAG_ASYNC = ConfigOption( key="rag.async", config_type=bool, diff --git a/python/flink_agents/api/runner_context.py b/python/flink_agents/api/runner_context.py index e110bd2f9..70796de79 100644 --- a/python/flink_agents/api/runner_context.py +++ b/python/flink_agents/api/runner_context.py @@ -16,6 +16,7 @@ # limitations under the License. ################################################################################# from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable, Dict from flink_agents.api.configuration import ReadableConfiguration @@ -24,12 +25,49 @@ from flink_agents.api.metric_group import MetricGroup from flink_agents.api.resource import Resource, ResourceType -__all__ = ["AsyncExecutionResult", "RunnerContext"] +__all__ = ["AsyncExecutionResult", "DurableCall", "Outcome", "RunnerContext"] if TYPE_CHECKING: from flink_agents.api.memory_object import MemoryObject +@dataclass(frozen=True) +class DurableCall: + """A deterministic durable call entry for batch execution.""" + + id: str + func: Callable[..., Any] + args: tuple[Any, ...] = () + kwargs: dict[str, Any] | None = None + reconciler: Callable[[], Any] | None = None + + +@dataclass(frozen=True) +class Outcome: + """Result or failure for one durable batch slot.""" + + value: Any = None + error: BaseException | None = None + + @classmethod + def success(cls, value: Any) -> "Outcome": + """Create a successful outcome.""" + return cls(value=value) + + @classmethod + def failure(cls, error: BaseException) -> "Outcome": + """Create a failed outcome.""" + return cls(error=error) + + def is_success(self) -> bool: + """Return whether this outcome is successful.""" + return self.error is None + + def is_failure(self) -> bool: + """Return whether this outcome is failed.""" + return self.error is not None + + class AsyncExecutionResult: """This class wraps an asynchronous task that will be submitted to a thread pool only when awaited. This ensures lazy submission and serial execution semantics. @@ -312,6 +350,18 @@ async def my_action(event, ctx): An awaitable object that yields the function result when awaited. """ + @abstractmethod + def durable_execute_all_async( + self, + callables: list[DurableCall], + ) -> "AsyncExecutionResult": + """Execute multiple durable callables as one asynchronous durable batch. + + The returned awaitable resolves to a list of ``Outcome`` values in the + same order as the input calls. Individual callable exceptions are returned + as failure outcomes instead of failing the whole batch. + """ + @property @abstractmethod def config(self) -> ReadableConfiguration: diff --git a/python/flink_agents/api/tests/test_core_options.py b/python/flink_agents/api/tests/test_core_options.py index f1e773a83..ed6ad37f0 100644 --- a/python/flink_agents/api/tests/test_core_options.py +++ b/python/flink_agents/api/tests/test_core_options.py @@ -107,6 +107,18 @@ def test_agent_config_options_are_explicitly_declared() -> None: assert options["EVENT_LOG_LEVEL"].get_default_value() is EventLogLevel.STANDARD +def test_agent_execution_options_include_parallel_tool_call_options() -> None: + from flink_agents.api.core_options import AgentExecutionOptions + + options = _collect_config_options(AgentExecutionOptions) + assert options["TOOL_CALL_ASYNC"].get_key() == "tool-call.async" + assert options["TOOL_CALL_ASYNC"].get_default_value() is True + assert options["TOOL_CALL_PARALLEL"].get_key() == "tool-call.parallel" + assert options["TOOL_CALL_PARALLEL"].get_default_value() is True + assert options["TOOL_CALL_BATCH_TIMEOUT"].get_key() == "tool-call.batch.timeout" + assert options["TOOL_CALL_BATCH_TIMEOUT"].get_default_value() == -1 + + def test_unknown_agent_config_option_raises_attribute_error() -> None: from flink_agents.api.core_options import AgentConfigOptions diff --git a/python/flink_agents/plan/actions/tool_call_action.py b/python/flink_agents/plan/actions/tool_call_action.py index f605c4807..84bf33543 100644 --- a/python/flink_agents/plan/actions/tool_call_action.py +++ b/python/flink_agents/plan/actions/tool_call_action.py @@ -16,13 +16,14 @@ # limitations under the License. ################################################################################# import logging +from dataclasses import dataclass from flink_agents.api.core_options import AgentExecutionOptions from flink_agents.api.events.event import Event from flink_agents.api.events.tool_event import ToolRequestEvent, ToolResponseEvent from flink_agents.api.memory_object import MemoryObject from flink_agents.api.resource import ResourceType -from flink_agents.api.runner_context import RunnerContext +from flink_agents.api.runner_context import DurableCall, Outcome, RunnerContext from flink_agents.api.tools.tool_parameter_injection import ( InjectedArg, ToolParameterSource, @@ -34,10 +35,18 @@ _logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class _ToolCallExecution: + id: str + name: str + durable_call: DurableCall + + async def process_tool_request(event: Event, ctx: RunnerContext) -> None: """Built-in action for processing tool call requests.""" event = ToolRequestEvent.from_event(event) tool_call_async = ctx.config.get(AgentExecutionOptions.TOOL_CALL_ASYNC) + tool_call_parallel = ctx.config.get(AgentExecutionOptions.TOOL_CALL_PARALLEL) if tool_call_async: # To avoid https://github.com/alibaba/pemja/issues/88, we log a message here. @@ -47,11 +56,53 @@ async def process_tool_request(event: Event, ctx: RunnerContext) -> None: success = {} error = {} external_ids = {} + executions = _build_tool_call_executions( + event, + ctx, + responses, + success, + error, + external_ids, + ) + + if tool_call_async and tool_call_parallel and len(executions) > 1: + await _execute_parallel(executions, ctx, responses, success, error) + else: + await _execute_sequentially( + executions, + tool_call_async=tool_call_async, + ctx=ctx, + responses=responses, + success=success, + error=error, + ) + + ctx.send_event( + ToolResponseEvent( + request_id=event.id, + responses=responses, + external_ids=external_ids, + success=success, + error=error, + ) + ) + + +def _build_tool_call_executions( + event: ToolRequestEvent, + ctx: RunnerContext, + responses: dict, + success: dict, + error: dict, + external_ids: dict, +) -> list[_ToolCallExecution]: + executions = [] for tool_call in event.tool_calls: call_id = tool_call["id"] name = tool_call["function"]["name"] kwargs = tool_call["function"]["arguments"] external_id = tool_call.get("original_id") + external_ids[call_id] = external_id try: tool = ctx.get_resource(name, ResourceType.TOOL) @@ -62,37 +113,105 @@ async def process_tool_request(event: Event, ctx: RunnerContext) -> None: responses[call_id] = f"Tool `{name}` does not exist." success[call_id] = False error.setdefault(call_id, f"Tool `{name}` does not exist.") - external_ids[call_id] = external_id continue - else: - try: - call_kwargs = dict(kwargs or {}) - # Framework-owned injected args must win over model-provided values so - # hidden context such as tenant ids cannot be spoofed by tool calls. - call_kwargs.update(_resolve_injected_arguments(tool, ctx)) - if tool_call_async: - response = await ctx.durable_execute_async( - tool.call, **call_kwargs - ) - else: - response = ctx.durable_execute(tool.call, **call_kwargs) - responses[call_id] = response - success[call_id] = True - except Exception as e: - responses[call_id] = f"Tool `{name}` execute failed." - success[call_id] = False - error[call_id] = str(e) - external_ids[call_id] = external_id - ctx.send_event( - ToolResponseEvent( - request_id=event.id, - responses=responses, - external_ids=external_ids, - success=success, - error=error, + try: + call_kwargs = dict(kwargs or {}) + # Framework-owned injected args must win over model-provided values so + # hidden context such as tenant ids cannot be spoofed by tool calls. + call_kwargs.update(_resolve_injected_arguments(tool, ctx)) + except Exception as e: + responses[call_id] = f"Tool `{name}` execute failed." + success[call_id] = False + error[call_id] = str(e) + continue + + executions.append( + _ToolCallExecution( + id=call_id, + name=name, + durable_call=DurableCall( + id=f"tool-call:{call_id}", + func=tool.call, + kwargs=call_kwargs, + ), + ) ) - ) + return executions + + +async def _execute_parallel( + executions: list[_ToolCallExecution], + ctx: RunnerContext, + responses: dict, + success: dict, + error: dict, +) -> None: + try: + outcomes = await ctx.durable_execute_all_async( + [execution.durable_call for execution in executions] + ) + for execution, outcome in zip(executions, outcomes, strict=True): + _record_outcome(execution, outcome, responses, success, error) + except Exception as e: + for execution in executions: + _record_execution_exception(execution, e, responses, success, error) + + +async def _execute_sequentially( + executions: list[_ToolCallExecution], + *, + tool_call_async: bool, + ctx: RunnerContext, + responses: dict, + success: dict, + error: dict, +) -> None: + for execution in executions: + try: + call = execution.durable_call + if tool_call_async: + response = await ctx.durable_execute_async( + call.func, + *call.args, + **(call.kwargs or {}), + ) + else: + response = ctx.durable_execute( + call.func, + *call.args, + **(call.kwargs or {}), + ) + responses[execution.id] = response + success[execution.id] = True + except Exception as e: # noqa: PERF203 + _record_execution_exception(execution, e, responses, success, error) + + +def _record_outcome( + execution: _ToolCallExecution, + outcome: Outcome, + responses: dict, + success: dict, + error: dict, +) -> None: + if outcome.is_failure(): + _record_execution_exception(execution, outcome.error, responses, success, error) + else: + responses[execution.id] = outcome.value + success[execution.id] = True + + +def _record_execution_exception( + execution: _ToolCallExecution, + exception: BaseException, + responses: dict, + success: dict, + error: dict, +) -> None: + responses[execution.id] = f"Tool `{execution.name}` execute failed." + success[execution.id] = False + error[execution.id] = str(exception) def _resolve_injected_arguments(tool: object, ctx: RunnerContext) -> dict: diff --git a/python/flink_agents/plan/tests/actions/test_tool_call_action.py b/python/flink_agents/plan/tests/actions/test_tool_call_action.py index d9690981a..ce5d11f53 100644 --- a/python/flink_agents/plan/tests/actions/test_tool_call_action.py +++ b/python/flink_agents/plan/tests/actions/test_tool_call_action.py @@ -22,6 +22,7 @@ from flink_agents.api.events.tool_event import ToolRequestEvent, ToolResponseEvent from flink_agents.api.memory_object import MemoryObject from flink_agents.api.resource import ResourceType +from flink_agents.api.runner_context import Outcome from flink_agents.api.tools import InjectedArg from flink_agents.plan.actions.tool_call_action import process_tool_request from flink_agents.plan.configuration import AgentConfiguration @@ -41,29 +42,48 @@ def __init__( sensory_memory: Any | None = None, short_term_memory: Any | None = None, ) -> None: - self.config = config or AgentConfiguration({"tenant_id": "tenant-1"}) - if isinstance(self.config, AgentConfiguration): + if config is None: + self.config = AgentConfiguration({"tenant_id": "tenant-1"}) self.config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, False) + else: + self.config = config self.injected_args = injected_args or { "tenant_id": InjectedArg.from_config("tenant_id") } self.sensory_memory = sensory_memory self.short_term_memory = short_term_memory self.sent_events = [] + self.durable_execute_calls = [] + self.durable_execute_async_calls = [] + self.durable_execute_all_async_calls = [] + self.durable_execute_all_async_outcomes = None def get_resource(self, name: str, type: ResourceType) -> FunctionTool: - assert name == "query_order" assert type == ResourceType.TOOL + if name != "query_order": + msg = f"Tool `{name}` does not exist." + raise ValueError(msg) return FunctionTool( func=PythonFunction.from_callable(query_order), injected_args=self.injected_args, ) - def durable_execute(self, func: Any, **kwargs: Any) -> Any: - return func(**kwargs) + def durable_execute(self, func: Any, *args: Any, **kwargs: Any) -> Any: + self.durable_execute_calls.append((func, args, kwargs)) + return func(*args, **kwargs) - async def durable_execute_async(self, func: Any, **kwargs: Any) -> Any: - return func(**kwargs) + async def durable_execute_async(self, func: Any, *args: Any, **kwargs: Any) -> Any: + self.durable_execute_async_calls.append((func, args, kwargs)) + return func(*args, **kwargs) + + async def durable_execute_all_async(self, callables: list[Any]) -> list[Outcome]: + self.durable_execute_all_async_calls.append(callables) + if self.durable_execute_all_async_outcomes is not None: + return self.durable_execute_all_async_outcomes + return [ + Outcome.success(call.func(*call.args, **(call.kwargs or {}))) + for call in callables + ] def send_event(self, event: Any) -> None: self.sent_events.append(event) @@ -102,7 +122,10 @@ def get_fields(self) -> dict[str, Any]: class _WrongConfig: def get(self, option: Any) -> bool: - assert option == AgentExecutionOptions.TOOL_CALL_ASYNC + assert option in ( + AgentExecutionOptions.TOOL_CALL_ASYNC, + AgentExecutionOptions.TOOL_CALL_PARALLEL, + ) return False @@ -277,14 +300,105 @@ def test_tool_call_action_uses_sync_execution_in_test_context() -> None: assert ctx.config.get(AgentExecutionOptions.TOOL_CALL_ASYNC) is False -def tool_request() -> ToolRequestEvent: +def test_tool_call_action_uses_parallel_batch_for_multiple_tools() -> None: + config = AgentConfiguration({"tenant_id": "tenant-1"}) + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, True) + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, True) + ctx = _Context(config=config) + + asyncio.run(process_tool_request(tool_request("call-1", "call-2"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.responses == { + "call-1": "tenant-1:order-call-1", + "call-2": "tenant-1:order-call-2", + } + assert response.success == {"call-1": True, "call-2": True} + assert len(ctx.durable_execute_all_async_calls) == 1 + assert [call.id for call in ctx.durable_execute_all_async_calls[0]] == [ + "tool-call:call-1", + "tool-call:call-2", + ] + assert ctx.durable_execute_async_calls == [] + + +def test_tool_call_action_uses_serial_async_when_parallel_disabled() -> None: + config = AgentConfiguration({"tenant_id": "tenant-1"}) + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, True) + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, False) + ctx = _Context(config=config) + + asyncio.run(process_tool_request(tool_request("call-1", "call-2"), ctx)) + + assert ctx.durable_execute_all_async_calls == [] + assert len(ctx.durable_execute_async_calls) == 2 + + +def test_tool_call_action_does_not_batch_single_tool() -> None: + config = AgentConfiguration({"tenant_id": "tenant-1"}) + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, True) + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, True) + ctx = _Context(config=config) + + asyncio.run(process_tool_request(tool_request("call-1"), ctx)) + + assert ctx.durable_execute_all_async_calls == [] + assert len(ctx.durable_execute_async_calls) == 1 + + +def test_tool_call_action_excludes_missing_tool_from_parallel_batch() -> None: + config = AgentConfiguration({"tenant_id": "tenant-1"}) + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, True) + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, True) + ctx = _Context(config=config) + + asyncio.run(process_tool_request(tool_request("call-1", "missing"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.success["call-1"] is True + assert response.success["missing"] is False + assert len(ctx.durable_execute_async_calls) == 1 + assert ctx.durable_execute_all_async_calls == [] + + +def test_tool_call_action_records_parallel_outcome_failure() -> None: + config = AgentConfiguration({"tenant_id": "tenant-1"}) + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, True) + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, True) + ctx = _Context(config=config) + ctx.durable_execute_all_async_outcomes = [ + Outcome.success("ok"), + Outcome.failure(ValueError("boom")), + ] + + asyncio.run(process_tool_request(tool_request("call-1", "call-2"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.responses["call-1"] == "ok" + assert response.success["call-1"] is True + assert response.responses["call-2"] == "Tool `query_order` execute failed." + assert response.success["call-2"] is False + assert response.error["call-2"] == "boom" + + +def tool_request(*call_ids: str) -> ToolRequestEvent: + if not call_ids: + call_ids = ("call-1",) return ToolRequestEvent( model="model", tool_calls=[ { - "id": "call-1", + "id": call_id, "type": "function", - "function": {"name": "query_order", "arguments": {"order_id": "order-1"}}, + "function": { + "name": "missing_tool" if call_id == "missing" else "query_order", + "arguments": { + "order_id": "order-1" + if len(call_ids) == 1 + else f"order-{call_id}" + }, + }, } + for call_id in call_ids ], ) diff --git a/python/flink_agents/runtime/flink_runner_context.py b/python/flink_agents/runtime/flink_runner_context.py index ada5b8e30..9b40ae257 100644 --- a/python/flink_agents/runtime/flink_runner_context.py +++ b/python/flink_agents/runtime/flink_runner_context.py @@ -17,6 +17,7 @@ ################################################################################# import logging import os +import time from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from functools import partial @@ -26,6 +27,7 @@ from typing_extensions import override from flink_agents.api.configuration import ReadableConfiguration +from flink_agents.api.core_options import AgentExecutionOptions from flink_agents.api.events.event import Event from flink_agents.api.memory.long_term_memory import ( BaseLongTermMemory, @@ -36,6 +38,8 @@ from flink_agents.api.resource import Resource, ResourceType from flink_agents.api.runner_context import ( AsyncExecutionResult, + DurableCall, + Outcome, RunnerContext, ) from flink_agents.runtime.durable_execution import ( @@ -73,6 +77,14 @@ class _ReconcilerExecutionPlan: needs_append_pending: bool = False +@dataclass(frozen=True) +class _BatchExecutionPlan: + outcomes: list[Outcome] + suppliers: list[tuple[int, Callable[[], Any]]] + needs_reservation: bool = False + execution_start: int = -1 + + class _DurableExecutionResult: """Wrapper that holds result and triggers recording when unwrapped.""" @@ -241,6 +253,43 @@ def __await__(self) -> Any: return result +class _DurableBatchAsyncExecutionResult(AsyncExecutionResult): + def __init__(self, ctx: "FlinkRunnerContext", calls: list[DurableCall]) -> None: + self._ctx = ctx + self._calls = calls + + def __await__(self) -> Any: + plan = self._ctx._prepare_batch_execution(self._calls) + futures = [ + self._ctx.executor.submit(supplier) for _, supplier in plan.suppliers + ] + timeout_ms = self._ctx.config.get(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT) + deadline = time.monotonic() + timeout_ms / 1000 if timeout_ms >= 0 else None + while any(not future.done() for future in futures): + if deadline is not None and time.monotonic() >= deadline: + exception = TimeoutError( + f"Async durable batch execution timed out after {timeout_ms} ms" + ) + for future in futures: + future.cancel() + executed = [Outcome.failure(exception) for _ in futures] + return self._ctx._finalize_batch_execution(self._calls, plan, executed) + yield + + executed = _collect_outcomes(futures) + return self._ctx._finalize_batch_execution(self._calls, plan, executed) + + +def _collect_outcomes(futures: list[Any]) -> list[Outcome]: + outcomes = [] + for future in futures: + try: + outcomes.append(Outcome.success(future.result())) + except Exception as e: # noqa: PERF203 + outcomes.append(Outcome.failure(e)) + return outcomes + + class FlinkRunnerContext(RunnerContext): """Providing context for agent execution in Flink Environment. @@ -273,6 +322,7 @@ def __init__( self.__agent_plan.resource_providers, self.__agent_plan.config ) self.__resource_cache.set_java_resource_adapter(j_resource_adapter) + self.__config = self.__agent_plan.config self.executor = executor def set_long_term_memory(self, ltm: InternalBaseLongTermMemory) -> None: @@ -490,6 +540,24 @@ def _serialize_call_payloads( exception_payload = cloudpickle.dumps(exception) if exception else None return result_payload, exception_payload + def _read_call_result_at(self, index: int) -> _PersistedCallResult | None: + current = self._j_runner_context.getCallResultFieldsAt(index) + if current is None: + return None + + function_id, args_digest, status, result_payload, exception_payload = current + return _PersistedCallResult( + function_id=function_id, + args_digest=args_digest, + status=status, + result_payload=bytes(result_payload) + if result_payload is not None + else None, + exception_payload=( + bytes(exception_payload) if exception_payload is not None else None + ), + ) + def _peek_current_call_result(self) -> _PersistedCallResult | None: current = self._j_runner_context.getCurrentCallResultFields() if current is None: @@ -636,6 +704,110 @@ def wrapped_func(*a: Any, **kw: Any) -> Any: return wrapped_func + def _call_matches( + self, current: _PersistedCallResult, call: DurableCall, args_digest: str + ) -> bool: + return current.function_id == call.id and current.args_digest == args_digest + + def _read_terminal_outcome(self, current: _PersistedCallResult) -> Outcome: + if current.exception_payload is not None: + return Outcome.failure(cloudpickle.loads(current.exception_payload)) + if current.result_payload is None: + return Outcome.success(None) + return Outcome.success(cloudpickle.loads(current.result_payload)) + + def _callable_for_durable_call(self, call: DurableCall) -> Callable[[], Any]: + kwargs = call.kwargs or {} + return partial(call.func, *call.args, **kwargs) + + def _prepare_batch_execution(self, calls: list[DurableCall]) -> _BatchExecutionPlan: + args_digest = "" + base = self._j_runner_context.getCurrentCallIndex() + outcomes: list[Outcome | None] = [] + suppliers: list[tuple[int, Callable[[], Any]]] = [] + needs_reservation = False + execution_start = -1 + + for index, call in enumerate(calls): + current = self._read_call_result_at(base + index) + if current is None: + needs_reservation = True + if execution_start < 0: + execution_start = index + outcomes.append(None) + suppliers.append((index, self._callable_for_durable_call(call))) + continue + + if not self._call_matches(current, call, args_digest): + self._j_runner_context.clearCallResultsFromAndPersist(base + index) + needs_reservation = True + execution_start = index + outcomes.append(None) + suppliers.append((index, self._callable_for_durable_call(call))) + for remaining_index in range(index + 1, len(calls)): + outcomes.append(None) + suppliers.append( + ( + remaining_index, + self._callable_for_durable_call(calls[remaining_index]), + ) + ) + break + + if current.status == "PENDING": + outcomes.append(None) + suppliers.append( + ( + index, + call.reconciler or self._callable_for_durable_call(call), + ) + ) + else: + outcomes.append(self._read_terminal_outcome(current)) + + if needs_reservation: + function_ids = [call.id for call in calls[execution_start:]] + self._j_runner_context.reservePendingBatch(function_ids, args_digest) + + return _BatchExecutionPlan( + outcomes=outcomes, + suppliers=suppliers, + needs_reservation=needs_reservation, + execution_start=execution_start, + ) + + def _finalize_batch_execution( + self, + calls: list[DurableCall], + plan: _BatchExecutionPlan, + executed: list[Outcome], + ) -> list[Outcome]: + args_digest = "" + base = self._j_runner_context.getCurrentCallIndex() + outcomes = list(plan.outcomes) + for (call_index, _), outcome in zip(plan.suppliers, executed, strict=True): + result_payload, exception_payload = self._serialize_call_payloads( + outcome.value, + outcome.error, + ) + self._j_runner_context.finalizeCallAt( + base + call_index, + calls[call_index].id, + args_digest, + result_payload, + exception_payload, + ) + outcomes[call_index] = outcome + self._j_runner_context.advanceCallIndexBy(len(calls)) + return outcomes + + @override + def durable_execute_all_async( + self, + callables: list[DurableCall], + ) -> AsyncExecutionResult: + return _DurableBatchAsyncExecutionResult(self, callables) + @override def durable_execute( self, @@ -748,6 +920,8 @@ def config(self) -> ReadableConfiguration: ReadableConfiguration The configuration for flink agents. """ + if hasattr(self, "_FlinkRunnerContext__config"): + return self.__config return self.__agent_plan.config @override diff --git a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py index a41a33c7b..e7d7c4a2f 100644 --- a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py +++ b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py @@ -15,7 +15,7 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################# -import asyncio +import time from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from typing import Any, Callable @@ -23,6 +23,8 @@ import cloudpickle import pytest +from flink_agents.api.runner_context import DurableCall +from flink_agents.plan.configuration import AgentConfiguration from flink_agents.runtime.durable_execution import ( _compute_args_digest, _compute_function_id, @@ -45,6 +47,22 @@ def __init__(self) -> None: self.current_call_index = 0 self.operations: list[str] = [] + def getCurrentCallIndex(self) -> int: + return self.current_call_index + + def getCallResultFieldsAt(self, index: int) -> list[Any] | None: + self.operations.append(f"read_at:{index}") + if index < len(self.call_results): + current = self.call_results[index] + return [ + current.function_id, + current.args_digest, + current.status, + current.result_payload, + current.exception_payload, + ] + return None + def getCurrentCallResultFields(self) -> list[Any] | None: self.operations.append("peek") if self.current_call_index < len(self.call_results): @@ -128,15 +146,57 @@ def clearCallResultsFromCurrentIndexAndPersist(self) -> None: self.operations.append("clear") self.call_results = self.call_results[: self.current_call_index] + def reservePendingBatch(self, function_ids: list[str], args_digest: str) -> None: + self.operations.append(f"reserve:{len(function_ids)}") + for function_id in function_ids: + self.call_results.append( + _StoredCallResult( + function_id=function_id, + args_digest=args_digest, + status="PENDING", + ) + ) + + def finalizeCallAt( + self, + index: int, + function_id: str, + args_digest: str, + result_payload: bytes | None, + exception_payload: bytes | None, + ) -> None: + self.operations.append(f"finalize_at:{index}") + current = self.call_results[index] + assert current.status == "PENDING" + assert current.function_id == function_id + assert current.args_digest == args_digest + self.call_results[index] = _StoredCallResult( + function_id=function_id, + args_digest=args_digest, + status="FAILED" if exception_payload is not None else "SUCCEEDED", + result_payload=result_payload, + exception_payload=exception_payload, + ) + + def clearCallResultsFromAndPersist(self, index: int) -> None: + self.operations.append(f"clear_at:{index}") + self.call_results = self.call_results[:index] + + def advanceCallIndexBy(self, count: int) -> None: + self.operations.append(f"advance:{count}") + self.current_call_index += count + def _create_runner_context( j_runner_context: _FakeJavaRunnerContext, + config: AgentConfiguration | None = None, ) -> FlinkRunnerContext: ctx = FlinkRunnerContext.__new__(FlinkRunnerContext) ctx._j_runner_context = j_runner_context - ctx.executor = ThreadPoolExecutor(max_workers=1) + ctx.executor = ThreadPoolExecutor(max_workers=2) ctx._FlinkRunnerContext__agent_plan = None ctx._FlinkRunnerContext__ltm = None + ctx._FlinkRunnerContext__config = config or AgentConfiguration({}) return ctx @@ -145,10 +205,14 @@ def _close_runner_context(ctx: FlinkRunnerContext) -> None: def _run_async(result: Any) -> object: - async def _await_result() -> Any: - return await result - - return asyncio.run(_await_result()) + iterator = result.__await__() + value = None + while True: + try: + iterator.send(value) + value = None + except StopIteration as e: # noqa: PERF203 + return e.value def _preload_pending( @@ -415,3 +479,160 @@ def collect_kwargs(**kwargs: Any) -> dict[str, Any]: _close_runner_context(ctx) assert result == {} + + +def test_flink_runner_context_durable_execute_all_async_initial_batch() -> None: + j_runner_context = _FakeJavaRunnerContext() + ctx = _create_runner_context(j_runner_context) + + try: + outcomes = _run_async( + ctx.durable_execute_all_async( + [ + DurableCall("tool-call:1", _call_value, args=("one",)), + DurableCall("tool-call:2", _call_value, args=("two",)), + ] + ) + ) + finally: + _close_runner_context(ctx) + + assert [outcome.value for outcome in outcomes] == ["call:one", "call:two"] + assert [result.status for result in j_runner_context.call_results] == [ + "SUCCEEDED", + "SUCCEEDED", + ] + assert j_runner_context.current_call_index == 2 + assert "reserve:2" in j_runner_context.operations + assert "finalize_at:0" in j_runner_context.operations + assert "finalize_at:1" in j_runner_context.operations + + +def test_flink_runner_context_durable_execute_all_async_recovers_partial_batch() -> None: + j_runner_context = _FakeJavaRunnerContext() + j_runner_context.call_results.extend( + [ + _StoredCallResult( + function_id="tool-call:1", + args_digest="", + status="SUCCEEDED", + result_payload=cloudpickle.dumps("cached-one"), + ), + _StoredCallResult( + function_id="tool-call:2", + args_digest="", + status="PENDING", + ), + ] + ) + call_count = 0 + + def tracked_call(value: str) -> str: + nonlocal call_count + call_count += 1 + return _call_value(value) + + ctx = _create_runner_context(j_runner_context) + try: + outcomes = _run_async( + ctx.durable_execute_all_async( + [ + DurableCall("tool-call:1", tracked_call, args=("one",)), + DurableCall("tool-call:2", tracked_call, args=("two",)), + ] + ) + ) + finally: + _close_runner_context(ctx) + + assert [outcome.value for outcome in outcomes] == ["cached-one", "call:two"] + assert call_count == 1 + assert j_runner_context.call_results[1].status == "SUCCEEDED" + assert j_runner_context.current_call_index == 2 + + +def test_flink_runner_context_durable_execute_all_async_returns_cached_failure() -> None: + j_runner_context = _FakeJavaRunnerContext() + j_runner_context.call_results.append( + _StoredCallResult( + function_id="tool-call:1", + args_digest="", + status="FAILED", + exception_payload=cloudpickle.dumps(ValueError("cached failure")), + ) + ) + ctx = _create_runner_context(j_runner_context) + + try: + outcomes = _run_async( + ctx.durable_execute_all_async( + [DurableCall("tool-call:1", _call_value, args=("one",))] + ) + ) + finally: + _close_runner_context(ctx) + + assert outcomes[0].is_failure() + assert isinstance(outcomes[0].error, ValueError) + assert str(outcomes[0].error) == "cached failure" + assert j_runner_context.current_call_index == 1 + + +def test_flink_runner_context_durable_execute_all_async_collects_failures() -> None: + j_runner_context = _FakeJavaRunnerContext() + ctx = _create_runner_context(j_runner_context) + + def fail_call() -> str: + msg = "failed" + raise ValueError(msg) + + try: + outcomes = _run_async( + ctx.durable_execute_all_async( + [ + DurableCall("tool-call:1", _call_value, args=("one",)), + DurableCall("tool-call:2", fail_call), + ] + ) + ) + finally: + _close_runner_context(ctx) + + assert outcomes[0].value == "call:one" + assert outcomes[1].is_failure() + assert isinstance(outcomes[1].error, ValueError) + assert [result.status for result in j_runner_context.call_results] == [ + "SUCCEEDED", + "FAILED", + ] + assert j_runner_context.current_call_index == 2 + + +def test_flink_runner_context_durable_execute_all_async_timeout_returns_failures() -> None: + j_runner_context = _FakeJavaRunnerContext() + config = AgentConfiguration({"tool-call.batch.timeout": 1}) + ctx = _create_runner_context(j_runner_context, config=config) + + def slow_call() -> str: + time.sleep(0.05) + return "slow" + + try: + outcomes = _run_async( + ctx.durable_execute_all_async( + [ + DurableCall("tool-call:1", slow_call), + DurableCall("tool-call:2", slow_call), + ] + ) + ) + finally: + _close_runner_context(ctx) + + assert all(outcome.is_failure() for outcome in outcomes) + assert all(isinstance(outcome.error, TimeoutError) for outcome in outcomes) + assert [result.status for result in j_runner_context.call_results] == [ + "FAILED", + "FAILED", + ] + assert j_runner_context.current_call_index == 2 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 5844da838..be36a4e8b 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 @@ -486,12 +486,16 @@ public void clearCallResultsFromAndPersist(int index) { } } - /** - * Returns the current durable call result as an array of fields for bridge consumers, or null - * if no persisted slot exists at the current call index. - */ - public Object[] getCurrentCallResultFields() { - CallResult current = getCurrentCallResult(); + public int getCurrentCallIndex() { + mailboxThreadChecker.run(); + if (durableExecutionContext == null) { + return 0; + } + return durableExecutionContext.getCurrentCallIndex(); + } + + public Object[] getCallResultFieldsAt(int index) { + CallResult current = getCallResultAt(index); if (current == null) { return null; } @@ -504,6 +508,17 @@ public Object[] getCurrentCallResultFields() { }; } + /** + * Returns the current durable call result as an array of fields for bridge consumers, or null + * if no persisted slot exists at the current call index. + */ + public Object[] getCurrentCallResultFields() { + if (durableExecutionContext == null) { + return null; + } + return getCallResultFieldsAt(durableExecutionContext.getCurrentCallIndex()); + } + protected Outcome readTerminalOutcomeAt( int index, String functionId, String argsDigest, Class resultClass) throws Exception { From 3d3a427a0e83215e1d49ca220a0bce6152d40871 Mon Sep 17 00:00:00 2001 From: daken Date: Sun, 26 Jul 2026 17:05:31 +0800 Subject: [PATCH 09/13] add batch-timeout test --- docs/content/docs/operations/configuration.md | 2 +- python/flink_agents/api/core_options.py | 6 ++ .../api/tests/test_core_options.py | 3 + .../plan/actions/tool_call_action.py | 2 +- .../tests/actions/test_tool_call_action.py | 4 +- .../runtime/flink_runner_context.py | 28 +++++++-- .../test_flink_runner_context_reconcilable.py | 37 ++++++------ .../python/utils/PythonActionExecutor.java | 12 ++++ ...nerContextImplDurableExecuteAsyncTest.java | 60 +++++++++++++++++++ 9 files changed, 128 insertions(+), 26 deletions(-) diff --git a/docs/content/docs/operations/configuration.md b/docs/content/docs/operations/configuration.md index 98f2bf39f..f1f70a7e3 100644 --- a/docs/content/docs/operations/configuration.md +++ b/docs/content/docs/operations/configuration.md @@ -134,7 +134,7 @@ Here is the list of all built-in core configuration options. | `tool-call.async` | true | boolean | Whether the built-in tool-call action runs each tool via durable async execution. | | `tool-call.parallel` | true | boolean | When `tool-call.async` is also true (JDK 21+), run multiple tool calls from one `ToolRequestEvent` as one parallel durable batch. Increases in-flight external calls; after failover, unfinished tools may be submitted again — side-effecting tools should be idempotent or provide a reconciler. Set to `false` for serial tool execution. | | `tool-call.num-async-threads` | os cpu count * 2 | int | Dedicated thread pool size for tool-call async / parallel batch execution. Separate from `num-async-threads` so a large tool batch does not exhaust the global async pool. | -| `tool-call.batch.timeout` | -1ms (disabled) | Duration | Overall timeout for one parallel tool-call batch. Non-positive disables it. On timeout, unfinished slots fail; already completed slots keep their outcome. | +| `tool-call.batch.timeout` | -1ms (disabled) | Duration | Overall timeout for one parallel tool-call batch. Non-positive disables it. On timeout, completed slots keep their outcome and unfinished slots fail. Timeout cancellation is best-effort; external side effects from unfinished tool calls may still complete, so side-effecting tools should be idempotent or provide a reconciler. | | `rag.async` | true | boolean | Whether retrieve context asynchronously for built-in context retrieval action. | | `num-async-threads` | os cpu count * 2 | int | The thread pool size for async executor. | | `job-identifier` | none | String | The unique identifier of job, remaining consistent after restoring from a savepoint. If not set, uses flink job id. | diff --git a/python/flink_agents/api/core_options.py b/python/flink_agents/api/core_options.py index 164585d12..142f0c383 100644 --- a/python/flink_agents/api/core_options.py +++ b/python/flink_agents/api/core_options.py @@ -260,6 +260,12 @@ class AgentExecutionOptions: default=True, ) + TOOL_CALL_NUM_ASYNC_THREADS = ConfigOption( + key="tool-call.num-async-threads", + config_type=int, + default=os.cpu_count() * 2, + ) + TOOL_CALL_BATCH_TIMEOUT = ConfigOption( key="tool-call.batch.timeout", config_type=int, diff --git a/python/flink_agents/api/tests/test_core_options.py b/python/flink_agents/api/tests/test_core_options.py index ed6ad37f0..2abbedf92 100644 --- a/python/flink_agents/api/tests/test_core_options.py +++ b/python/flink_agents/api/tests/test_core_options.py @@ -17,6 +17,7 @@ ################################################################################# import ast import importlib +import os import sys import types from pathlib import Path @@ -115,6 +116,8 @@ def test_agent_execution_options_include_parallel_tool_call_options() -> None: assert options["TOOL_CALL_ASYNC"].get_default_value() is True assert options["TOOL_CALL_PARALLEL"].get_key() == "tool-call.parallel" assert options["TOOL_CALL_PARALLEL"].get_default_value() is True + assert options["TOOL_CALL_NUM_ASYNC_THREADS"].get_key() == "tool-call.num-async-threads" + assert options["TOOL_CALL_NUM_ASYNC_THREADS"].get_default_value() == os.cpu_count() * 2 assert options["TOOL_CALL_BATCH_TIMEOUT"].get_key() == "tool-call.batch.timeout" assert options["TOOL_CALL_BATCH_TIMEOUT"].get_default_value() == -1 diff --git a/python/flink_agents/plan/actions/tool_call_action.py b/python/flink_agents/plan/actions/tool_call_action.py index 84bf33543..5651f453f 100644 --- a/python/flink_agents/plan/actions/tool_call_action.py +++ b/python/flink_agents/plan/actions/tool_call_action.py @@ -131,7 +131,7 @@ def _build_tool_call_executions( id=call_id, name=name, durable_call=DurableCall( - id=f"tool-call:{call_id}", + id=f"tool-call-{call_id}", func=tool.call, kwargs=call_kwargs, ), diff --git a/python/flink_agents/plan/tests/actions/test_tool_call_action.py b/python/flink_agents/plan/tests/actions/test_tool_call_action.py index ce5d11f53..b329e853f 100644 --- a/python/flink_agents/plan/tests/actions/test_tool_call_action.py +++ b/python/flink_agents/plan/tests/actions/test_tool_call_action.py @@ -316,8 +316,8 @@ def test_tool_call_action_uses_parallel_batch_for_multiple_tools() -> None: assert response.success == {"call-1": True, "call-2": True} assert len(ctx.durable_execute_all_async_calls) == 1 assert [call.id for call in ctx.durable_execute_all_async_calls[0]] == [ - "tool-call:call-1", - "tool-call:call-2", + "tool-call-call-1", + "tool-call-call-2", ] assert ctx.durable_execute_async_calls == [] diff --git a/python/flink_agents/runtime/flink_runner_context.py b/python/flink_agents/runtime/flink_runner_context.py index 9b40ae257..b3ff7bea4 100644 --- a/python/flink_agents/runtime/flink_runner_context.py +++ b/python/flink_agents/runtime/flink_runner_context.py @@ -261,7 +261,7 @@ def __init__(self, ctx: "FlinkRunnerContext", calls: list[DurableCall]) -> None: def __await__(self) -> Any: plan = self._ctx._prepare_batch_execution(self._calls) futures = [ - self._ctx.executor.submit(supplier) for _, supplier in plan.suppliers + self._ctx.tool_call_executor.submit(supplier) for _, supplier in plan.suppliers ] timeout_ms = self._ctx.config.get(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT) deadline = time.monotonic() + timeout_ms / 1000 if timeout_ms >= 0 else None @@ -270,9 +270,7 @@ def __await__(self) -> Any: exception = TimeoutError( f"Async durable batch execution timed out after {timeout_ms} ms" ) - for future in futures: - future.cancel() - executed = [Outcome.failure(exception) for _ in futures] + executed = _collect_outcomes_on_timeout(futures, exception) return self._ctx._finalize_batch_execution(self._calls, plan, executed) yield @@ -290,6 +288,23 @@ def _collect_outcomes(futures: list[Any]) -> list[Outcome]: return outcomes +def _collect_outcomes_on_timeout( + futures: list[Any], timeout_exception: TimeoutError +) -> list[Outcome]: + outcomes = [] + for future in futures: + if not future.done(): + future.cancel() + if future.done() and not future.cancelled(): + try: + outcomes.append(Outcome.success(future.result())) + except Exception as e: + outcomes.append(Outcome.failure(e)) + else: + outcomes.append(Outcome.failure(timeout_exception)) + return outcomes + + class FlinkRunnerContext(RunnerContext): """Providing context for agent execution in Flink Environment. @@ -305,6 +320,7 @@ def __init__( j_runner_context: Any, agent_plan_json: str, executor: ThreadPoolExecutor, + tool_call_executor: ThreadPoolExecutor, j_resource_adapter: Any, ) -> None: """Initialize a flink runner context with the given java runner context. @@ -324,6 +340,7 @@ def __init__( self.__resource_cache.set_java_resource_adapter(j_resource_adapter) self.__config = self.__agent_plan.config self.executor = executor + self.tool_call_executor = tool_call_executor def set_long_term_memory(self, ltm: InternalBaseLongTermMemory) -> None: """Set long term memory instance to this context. @@ -940,12 +957,13 @@ def create_flink_runner_context( j_runner_context: Any, agent_plan_json: str, executor: ThreadPoolExecutor, + tool_call_executor: ThreadPoolExecutor, j_resource_adapter: Any, job_identifier: str, ) -> FlinkRunnerContext: """Used to create a FlinkRunnerContext Python object in Pemja environment.""" ctx = FlinkRunnerContext( - j_runner_context, agent_plan_json, executor, j_resource_adapter + j_runner_context, agent_plan_json, executor, tool_call_executor, j_resource_adapter ) ltm = _init_long_term_memory(ctx, job_identifier) if ltm is not None: diff --git a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py index e7d7c4a2f..2755752b4 100644 --- a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py +++ b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py @@ -194,6 +194,7 @@ def _create_runner_context( ctx = FlinkRunnerContext.__new__(FlinkRunnerContext) ctx._j_runner_context = j_runner_context ctx.executor = ThreadPoolExecutor(max_workers=2) + ctx.tool_call_executor = ThreadPoolExecutor(max_workers=2) ctx._FlinkRunnerContext__agent_plan = None ctx._FlinkRunnerContext__ltm = None ctx._FlinkRunnerContext__config = config or AgentConfiguration({}) @@ -202,6 +203,7 @@ def _create_runner_context( def _close_runner_context(ctx: FlinkRunnerContext) -> None: ctx.executor.shutdown(wait=True) + ctx.tool_call_executor.shutdown(wait=True) def _run_async(result: Any) -> object: @@ -489,8 +491,8 @@ def test_flink_runner_context_durable_execute_all_async_initial_batch() -> None: outcomes = _run_async( ctx.durable_execute_all_async( [ - DurableCall("tool-call:1", _call_value, args=("one",)), - DurableCall("tool-call:2", _call_value, args=("two",)), + DurableCall("tool-call-1", _call_value, args=("one",)), + DurableCall("tool-call-2", _call_value, args=("two",)), ] ) ) @@ -513,13 +515,13 @@ def test_flink_runner_context_durable_execute_all_async_recovers_partial_batch() j_runner_context.call_results.extend( [ _StoredCallResult( - function_id="tool-call:1", + function_id="tool-call-1", args_digest="", status="SUCCEEDED", result_payload=cloudpickle.dumps("cached-one"), ), _StoredCallResult( - function_id="tool-call:2", + function_id="tool-call-2", args_digest="", status="PENDING", ), @@ -537,8 +539,8 @@ def tracked_call(value: str) -> str: outcomes = _run_async( ctx.durable_execute_all_async( [ - DurableCall("tool-call:1", tracked_call, args=("one",)), - DurableCall("tool-call:2", tracked_call, args=("two",)), + DurableCall("tool-call-1", tracked_call, args=("one",)), + DurableCall("tool-call-2", tracked_call, args=("two",)), ] ) ) @@ -555,7 +557,7 @@ def test_flink_runner_context_durable_execute_all_async_returns_cached_failure() j_runner_context = _FakeJavaRunnerContext() j_runner_context.call_results.append( _StoredCallResult( - function_id="tool-call:1", + function_id="tool-call-1", args_digest="", status="FAILED", exception_payload=cloudpickle.dumps(ValueError("cached failure")), @@ -566,7 +568,7 @@ def test_flink_runner_context_durable_execute_all_async_returns_cached_failure() try: outcomes = _run_async( ctx.durable_execute_all_async( - [DurableCall("tool-call:1", _call_value, args=("one",))] + [DurableCall("tool-call-1", _call_value, args=("one",))] ) ) finally: @@ -590,8 +592,8 @@ def fail_call() -> str: outcomes = _run_async( ctx.durable_execute_all_async( [ - DurableCall("tool-call:1", _call_value, args=("one",)), - DurableCall("tool-call:2", fail_call), + DurableCall("tool-call-1", _call_value, args=("one",)), + DurableCall("tool-call-2", fail_call), ] ) ) @@ -608,9 +610,9 @@ def fail_call() -> str: assert j_runner_context.current_call_index == 2 -def test_flink_runner_context_durable_execute_all_async_timeout_returns_failures() -> None: +def test_flink_runner_context_durable_execute_all_async_timeout_keeps_completed_results() -> None: j_runner_context = _FakeJavaRunnerContext() - config = AgentConfiguration({"tool-call.batch.timeout": 1}) + config = AgentConfiguration({"tool-call.batch.timeout": 10}) ctx = _create_runner_context(j_runner_context, config=config) def slow_call() -> str: @@ -621,18 +623,19 @@ def slow_call() -> str: outcomes = _run_async( ctx.durable_execute_all_async( [ - DurableCall("tool-call:1", slow_call), - DurableCall("tool-call:2", slow_call), + DurableCall("tool-call-1", lambda: "fast"), + DurableCall("tool-call-2", slow_call), ] ) ) finally: _close_runner_context(ctx) - assert all(outcome.is_failure() for outcome in outcomes) - assert all(isinstance(outcome.error, TimeoutError) for outcome in outcomes) + assert outcomes[0].value == "fast" + assert outcomes[1].is_failure() + assert isinstance(outcomes[1].error, TimeoutError) assert [result.status for result in j_runner_context.call_results] == [ - "FAILED", + "SUCCEEDED", "FAILED", ] assert j_runner_context.current_call_index == 2 diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java index 67c80f38d..29bbbe110 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java @@ -74,6 +74,7 @@ public class PythonActionExecutor { private final JavaResourceAdapter javaResourceAdapter; private final String jobIdentifier; private PyObject pythonAsyncThreadPool; + private PyObject pythonToolCallAsyncThreadPool; private PyObject pythonRunnerContext; public PythonActionExecutor( @@ -102,6 +103,13 @@ public void open() throws Exception { interpreter.invoke( CREATE_ASYNC_THREAD_POOL, agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS)); + pythonToolCallAsyncThreadPool = + (PyObject) + interpreter.invoke( + CREATE_ASYNC_THREAD_POOL, + agentPlan + .getConfig() + .get(AgentExecutionOptions.TOOL_CALL_NUM_ASYNC_THREADS)); pythonRunnerContext = (PyObject) @@ -110,6 +118,7 @@ public void open() throws Exception { runnerContext, new ObjectMapper().writeValueAsString(agentPlan), pythonAsyncThreadPool, + pythonToolCallAsyncThreadPool, javaResourceAdapter, jobIdentifier); } @@ -192,6 +201,9 @@ public void close() throws Exception { if (pythonAsyncThreadPool != null) { interpreter.invoke(CLOSE_ASYNC_THREAD_POOL, pythonAsyncThreadPool); } + if (pythonToolCallAsyncThreadPool != null) { + interpreter.invoke(CLOSE_ASYNC_THREAD_POOL, pythonToolCallAsyncThreadPool); + } if (pythonRunnerContext != null) { try { diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java index 3dfd77017..62676e939 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java @@ -20,6 +20,8 @@ import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.configuration.Configuration; import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.actions.Action; @@ -35,6 +37,7 @@ import java.util.HashMap; import java.util.List; import java.util.concurrent.Callable; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Supplier; @@ -313,6 +316,39 @@ void testDurableExecuteAllAsyncReturnsCachedFailureOutcome() throws Exception { assertEquals(1, context.getDurableExecutionContext().getCurrentCallIndex()); } + @Test + void testDurableExecuteAllAsyncTimeoutKeepsCompletedOutcomes() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + executor.setUseTimeoutCollection(true); + JavaRunnerContextImpl context = createContext(new ActionState(null), executor); + ((Configuration) context.getConfig()) + .set( + AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT, + java.time.Duration.ofMillis(10)); + TestDurableCallable first = + new TestDurableCallable<>("batch-1", String.class, () -> "fast"); + TestDurableCallable second = + new TestDurableCallable<>( + "batch-2", + String.class, + () -> { + Thread.sleep(100); + return "slow"; + }); + + List> outcomes = context.durableExecuteAllAsync(List.of(first, second)); + + assertEquals("fast", outcomes.get(0).getValue()); + assertTrue(outcomes.get(1).isFailure()); + assertInstanceOf(TimeoutException.class, outcomes.get(1).getError()); + List persisted = + context.getDurableExecutionContext().getActionState().getCallResults(); + assertTrue(persisted.get(0).isSuccess()); + assertTrue(persisted.get(1).isFailure()); + assertEquals(2, context.getDurableExecutionContext().getCurrentCallIndex()); + executor.close(); + } + private JavaRunnerContextImpl createContext( ActionState actionState, ContinuationActionExecutor executor) { JavaRunnerContextImpl context = @@ -343,6 +379,7 @@ private JavaRunnerContextImpl createContext( private static final class InspectingContinuationActionExecutor extends ContinuationActionExecutor { private Runnable beforeExecute; + private boolean useTimeoutCollection; private int executeAsyncCallCount; private int executeAllAsyncCallCount; private final List executeAllAsyncBatchSizes = new java.util.ArrayList<>(); @@ -367,6 +404,9 @@ public List> executeAllAsync( java.time.Duration timeout) { executeAllAsyncCallCount++; executeAllAsyncBatchSizes.add(suppliers.size()); + if (useTimeoutCollection) { + return collectTimedOutOutcomes(suppliers); + } List> outcomes = new java.util.ArrayList<>(suppliers.size()); for (Callable supplier : suppliers) { try { @@ -378,10 +418,30 @@ public List> executeAllAsync( return outcomes; } + private List> collectTimedOutOutcomes(List> suppliers) { + List> outcomes = new java.util.ArrayList<>(suppliers.size()); + for (int i = 0; i < suppliers.size(); i++) { + if (i == 0) { + try { + outcomes.add(Outcome.success(suppliers.get(i).call())); + } catch (Exception e) { + outcomes.add(Outcome.failure(e)); + } + } else { + outcomes.add(Outcome.failure(new TimeoutException("batch timeout"))); + } + } + return outcomes; + } + private void setBeforeExecute(Runnable beforeExecute) { this.beforeExecute = beforeExecute; } + private void setUseTimeoutCollection(boolean useTimeoutCollection) { + this.useTimeoutCollection = useTimeoutCollection; + } + private int getExecuteAsyncCallCount() { return executeAsyncCallCount; } From 05e93636994a1d6952c3b1a6c453bcbbf6559be0 Mon Sep 17 00:00:00 2001 From: daken Date: Sun, 26 Jul 2026 17:42:51 +0800 Subject: [PATCH 10/13] update cross-language test --- .../check_java_python_config_options_parity.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py b/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py index 0e2762d73..8a7eb712a 100644 --- a/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py +++ b/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py @@ -42,6 +42,7 @@ "java.lang.Float": float, "java.lang.Double": float, "java.util.List": list, + "java.time.Duration": int, } @@ -84,11 +85,20 @@ def _java_type_matches_python(java_type_name: str, python_config_type: type) -> return python_config_type.__name__ == java_simple_name -def normalize_java_default(java_default: Any, python_config_type: type) -> Any: +def _java_duration_to_millis(java_duration: Any) -> int: + return int(java_duration.toMillis()) + + +def normalize_java_default( + java_default: Any, java_type_name: str, python_config_type: type +) -> Any: """Convert a Java default value into a Python-comparable form.""" if java_default is None: return None + if java_type_name == "java.time.Duration": + return _java_duration_to_millis(java_default) + if hasattr(java_default, "name") and callable(java_default.name): enum_name = java_default.name() if isinstance(python_config_type, type) and issubclass(python_config_type, Enum): @@ -122,6 +132,7 @@ def assert_python_option_matches_java( python_default = python_option.get_default_value() java_default = normalize_java_default( java_option.getDefaultValue(), + java_type_name, python_config_type, ) assert python_default == java_default, ( From 9169bb2ab921c211e4972ad2996d5e5ef110d991 Mon Sep 17 00:00:00 2001 From: daken Date: Mon, 27 Jul 2026 23:47:13 +0800 Subject: [PATCH 11/13] add e2e test --- .../integration/test/AsyncExecutionAgent.java | 155 ++++++++++++++++++ .../integration/test/AsyncExecutionTest.java | 72 ++++++++ .../test_flink_runner_context_reconcilable.py | 38 ++++- 3 files changed, 264 insertions(+), 1 deletion(-) diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java index c8df332f5..ac9ed762d 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java @@ -23,13 +23,33 @@ import org.apache.flink.agents.api.OutputEvent; import org.apache.flink.agents.api.agents.Agent; import org.apache.flink.agents.api.annotation.Action; +import org.apache.flink.agents.api.annotation.ChatModelConnection; +import org.apache.flink.agents.api.annotation.ChatModelSetup; +import org.apache.flink.agents.api.annotation.ToolParam; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.chat.model.BaseChatModelConnection; +import org.apache.flink.agents.api.chat.model.BaseChatModelSetup; import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.ChatRequestEvent; +import org.apache.flink.agents.api.event.ChatResponseEvent; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolMetadata; +import org.apache.flink.agents.api.tools.ToolParameters; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.tools.ToolType; import org.apache.flink.api.java.functions.KeySelector; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + /** * Agent definition for testing async execution functionality. * @@ -71,6 +91,141 @@ public Integer getKey(AsyncRequest request) { } } + /** Chat connection that emits one tool request containing multiple slow tool calls. */ + public static class ToolBatchChatConnection extends BaseChatModelConnection { + public ToolBatchChatConnection( + ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + } + + @Override + public ChatMessage chat( + List messages, List tools, Map modelParams) { + ChatMessage lastMessage = messages.get(messages.size() - 1); + if (lastMessage.getRole() == MessageRole.TOOL) { + return new ChatMessage(MessageRole.ASSISTANT, lastMessage.getContent()); + } + + String requestId = lastMessage.getContent(); + return new ChatMessage( + MessageRole.ASSISTANT, + "", + List.of( + toolCall("call-1", requestId, 1), + toolCall("call-2", requestId, 2), + toolCall("call-3", requestId, 3))); + } + } + + /** Chat model setup bound to the timed tool. */ + public static class ToolBatchChatModel extends BaseChatModelSetup { + public ToolBatchChatModel(ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + } + + @Override + public Map getParameters() { + return new HashMap<>(); + } + } + + private static Map toolCall(String id, String requestId, int index) { + return Map.of( + "id", + id, + "type", + "function", + "function", + Map.of( + "name", + "timed_tool", + "arguments", + Map.of("request_id", requestId, "call_index", index))); + } + + /** Agent that requests one chat turn that produces multiple slow tool calls. */ + public static class ToolBatchAgent extends Agent { + public ToolBatchAgent(int sleepTimeMs) {} + + @ChatModelConnection + public static ResourceDescriptor toolBatchChatConnection() { + return ResourceDescriptor.Builder.newBuilder(ToolBatchChatConnection.class.getName()) + .build(); + } + + @ChatModelSetup + public static ResourceDescriptor toolBatchChatModel() { + return ResourceDescriptor.Builder.newBuilder(ToolBatchChatModel.class.getName()) + .addInitialArgument("connection", "toolBatchChatConnection") + .addInitialArgument("model", "test-model") + .addInitialArgument("tools", List.of("timed_tool")) + .build(); + } + + @org.apache.flink.agents.api.annotation.Tool( + description = "Records timing for a slow tool call.") + public static String timed_tool( + @ToolParam(name = "request_id") String requestId, + @ToolParam(name = "call_index") String callIndex) { + long start = System.currentTimeMillis(); + try { + Thread.sleep(500); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + long end = System.currentTimeMillis(); + return String.format( + "request=%s,call=%s,start=%d,end=%d", requestId, callIndex, start, end); + } + + @Action(EventType.InputEvent) + public static void requestTools(Event event, RunnerContext ctx) { + InputEvent inputEvent = InputEvent.fromEvent(event); + AsyncRequest request = (AsyncRequest) inputEvent.getInput(); + ctx.sendEvent( + new ChatRequestEvent( + "toolBatchChatModel", + List.of( + new ChatMessage( + MessageRole.USER, String.valueOf(request.id))))); + } + + @Action(EventType.ChatResponseEvent) + public static void emitToolTimings(Event event, RunnerContext ctx) { + ChatResponseEvent responseEvent = ChatResponseEvent.fromEvent(event); + ctx.sendEvent(new OutputEvent(responseEvent.getResponse().getContent())); + } + } + + /** Tool that records its execution time range in the response result. */ + public static class TimedTool extends Tool { + private final int sleepTimeMs; + + public TimedTool(int sleepTimeMs) { + super(new ToolMetadata("timed_tool", "Records timing for a slow tool call.", "{}")); + this.sleepTimeMs = sleepTimeMs; + } + + @Override + public ToolType getToolType() { + return ToolType.FUNCTION; + } + + @Override + public ToolResponse call(ToolParameters parameters) { + int callIndex = parameters.getParameter("call_index", Integer.class); + long start = System.currentTimeMillis(); + try { + Thread.sleep(sleepTimeMs); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + long end = System.currentTimeMillis(); + return ToolResponse.success( + String.format("call=%d,start=%d,end=%d", callIndex, start, end)); + } + } + /** Custom event type for internal agent communication. */ public static class AsyncProcessedEvent extends Event { public static final String EVENT_TYPE = "AsyncProcessedEvent"; diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java index 271b067fa..1c4a62afe 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java @@ -18,6 +18,7 @@ package org.apache.flink.agents.integration.test; import org.apache.flink.agents.api.AgentsExecutionEnvironment; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; import org.apache.flink.agents.runtime.async.ContinuationActionExecutor; import org.apache.flink.streaming.api.datastream.DataStream; import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; @@ -27,6 +28,8 @@ import java.util.ArrayList; import java.util.List; +import java.util.regex.Matcher; +import java.util.regex.Pattern; /** * End-to-end tests for Java async execution functionality. @@ -333,6 +336,61 @@ public void testAsyncExecutionIsActuallyParallel() throws Exception { System.out.println("=== Test Passed ==="); } + @Test + public void testToolCallBatchExecutionIsActuallyParallel() throws Exception { + boolean continuationSupported = ContinuationActionExecutor.isContinuationSupported(); + int javaVersion = Runtime.version().feature(); + + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + + int sleepTimeMs = 500; + DataStream inputStream = + env.fromElements( + new AsyncExecutionAgent.AsyncRequest(1, "tool-batch", sleepTimeMs)); + + AgentsExecutionEnvironment agentsEnv = + AgentsExecutionEnvironment.getExecutionEnvironment(env); + agentsEnv.getConfig().set(AgentExecutionOptions.TOOL_CALL_ASYNC, true); + agentsEnv.getConfig().set(AgentExecutionOptions.TOOL_CALL_PARALLEL, true); + agentsEnv.getConfig().set(AgentExecutionOptions.TOOL_CALL_NUM_ASYNC_THREADS, 3); + + DataStream outputStream = + agentsEnv + .fromDataStream( + inputStream, new AsyncExecutionAgent.AsyncRequestKeySelector()) + .apply(new AsyncExecutionAgent.ToolBatchAgent(sleepTimeMs)) + .toDataStream(); + + CloseableIterator results = outputStream.collectAsync(); + agentsEnv.execute(); + + List executionRanges = new ArrayList<>(); + while (results.hasNext()) { + String output = results.next().toString(); + Matcher matcher = Pattern.compile("start=(\\d+),end=(\\d+)").matcher(output); + while (matcher.find()) { + long start = Long.parseLong(matcher.group(1)); + long end = Long.parseLong(matcher.group(2)); + executionRanges.add(new long[] {start, end}); + } + } + results.close(); + + Assertions.assertEquals(3, executionRanges.size()); + int overlapCount = countOverlaps(executionRanges); + if (continuationSupported && javaVersion >= 21) { + Assertions.assertTrue( + overlapCount >= 2, + "On JDK 21+, tool calls in one ToolRequestEvent should run in parallel."); + } else { + Assertions.assertEquals( + 0, + overlapCount, + "On JDK < 21, tool-call batch execution should use the sequential fallback."); + } + } + /** * Tests that durableExecute (sync) works correctly. * @@ -387,4 +445,18 @@ public void testDurableExecuteSync() throws Exception { "Output should contain processed data: " + output); } } + + private static int countOverlaps(List executionRanges) { + int overlapCount = 0; + for (int i = 0; i < executionRanges.size(); i++) { + for (int j = i + 1; j < executionRanges.size(); j++) { + long[] range1 = executionRanges.get(i); + long[] range2 = executionRanges.get(j); + if (range1[0] < range2[1] && range2[0] < range1[1]) { + overlapCount++; + } + } + } + return overlapCount; + } } diff --git a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py index 2755752b4..59f5283c1 100644 --- a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py +++ b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py @@ -190,11 +190,12 @@ def advanceCallIndexBy(self, count: int) -> None: def _create_runner_context( j_runner_context: _FakeJavaRunnerContext, config: AgentConfiguration | None = None, + tool_call_workers: int = 2, ) -> FlinkRunnerContext: ctx = FlinkRunnerContext.__new__(FlinkRunnerContext) ctx._j_runner_context = j_runner_context ctx.executor = ThreadPoolExecutor(max_workers=2) - ctx.tool_call_executor = ThreadPoolExecutor(max_workers=2) + ctx.tool_call_executor = ThreadPoolExecutor(max_workers=tool_call_workers) ctx._FlinkRunnerContext__agent_plan = None ctx._FlinkRunnerContext__ltm = None ctx._FlinkRunnerContext__config = config or AgentConfiguration({}) @@ -483,6 +484,41 @@ def collect_kwargs(**kwargs: Any) -> dict[str, Any]: assert result == {} +def test_flink_runner_context_durable_execute_all_async_runs_calls_in_parallel() -> None: + j_runner_context = _FakeJavaRunnerContext() + config = AgentConfiguration({"tool-call.batch.timeout": -1}) + ctx = _create_runner_context(j_runner_context, config=config, tool_call_workers=3) + sleep_seconds = 0.2 + + def slow_call(value: str) -> str: + time.sleep(sleep_seconds) + return value + + try: + start = time.perf_counter() + outcomes = _run_async( + ctx.durable_execute_all_async( + [ + DurableCall("tool-call-1", slow_call, args=("one",)), + DurableCall("tool-call-2", slow_call, args=("two",)), + DurableCall("tool-call-3", slow_call, args=("three",)), + ] + ) + ) + elapsed = time.perf_counter() - start + finally: + _close_runner_context(ctx) + + assert [outcome.value for outcome in outcomes] == ["one", "two", "three"] + assert elapsed < sleep_seconds * 2 + assert [result.status for result in j_runner_context.call_results] == [ + "SUCCEEDED", + "SUCCEEDED", + "SUCCEEDED", + ] + assert j_runner_context.current_call_index == 3 + + def test_flink_runner_context_durable_execute_all_async_initial_batch() -> None: j_runner_context = _FakeJavaRunnerContext() ctx = _create_runner_context(j_runner_context) From f1173547e5314bb019b1d9f2dd7e971afbf47d61 Mon Sep 17 00:00:00 2001 From: daken Date: Thu, 30 Jul 2026 18:33:45 +0800 Subject: [PATCH 12/13] fix cr: - Change timeout configuration to Long type and improve timeout e2e tests - Fix pending-slot error for non-reconcilable calls under serial execution - Consolidate all exception handling inside durableExecuteAllAsync - fix py timeout > 0 --- .../api/agents/AgentExecutionOptions.java | 8 +- docs/content/docs/operations/configuration.md | 2 +- .../pom.xml | 56 ++++-- .../integration/test/AsyncExecutionAgent.java | 142 ++++++++++++- .../integration/test/AsyncExecutionTest.java | 57 ++++++ .../agents/plan/actions/ToolCallAction.java | 4 +- .../plan/actions/ToolCallActionTest.java | 10 +- python/flink_agents/api/core_options.py | 5 +- python/flink_agents/api/runner_context.py | 10 + .../api/tests/test_core_options.py | 4 +- .../plan/actions/tool_call_action.py | 19 +- .../tests/actions/test_tool_call_action.py | 35 +++- ...check_java_python_config_options_parity.py | 8 - .../flink_agents/runtime/durable_execution.py | 22 ++ .../runtime/flink_runner_context.py | 100 ++++++--- .../test_flink_runner_context_reconcilable.py | 134 +++++++++--- .../context/JavaRunnerContextImpl.java | 13 +- .../runtime/context/RunnerContextImpl.java | 31 ++- .../context/DurableExecutionContextTest.java | 13 ++ ...nerContextImplDurableExecuteAsyncTest.java | 190 ++++++++++++++++-- .../RunnerContextImplDurableExecuteTest.java | 54 +++++ 21 files changed, 772 insertions(+), 145 deletions(-) diff --git a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java index 32f248d8b..f1e27ba77 100644 --- a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java +++ b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java @@ -20,8 +20,6 @@ import org.apache.flink.agents.api.configuration.ConfigOption; -import java.time.Duration; - public class AgentExecutionOptions { public static final ConfigOption ERROR_HANDLING_STRATEGY = new ConfigOption<>( @@ -73,13 +71,13 @@ public class AgentExecutionOptions { Runtime.getRuntime().availableProcessors() * 2); /** - * Overall timeout for one parallel tool-call batch. + * Overall timeout for one parallel tool-call batch, in milliseconds. * *

Non-positive values disable the timeout. When the deadline elapses, unfinished slots are * failed; slots that already completed keep their success or failure outcome. */ - public static final ConfigOption TOOL_CALL_BATCH_TIMEOUT = - new ConfigOption<>("tool-call.batch.timeout", Duration.class, Duration.ofMillis(-1)); + public static final ConfigOption TOOL_CALL_BATCH_TIMEOUT_MS = + new ConfigOption<>("tool-call.batch.timeout.ms", Long.class, -1L); public static final ConfigOption RAG_ASYNC = new ConfigOption<>("rag.async", Boolean.class, true); diff --git a/docs/content/docs/operations/configuration.md b/docs/content/docs/operations/configuration.md index f1f70a7e3..5691e584f 100644 --- a/docs/content/docs/operations/configuration.md +++ b/docs/content/docs/operations/configuration.md @@ -134,7 +134,7 @@ Here is the list of all built-in core configuration options. | `tool-call.async` | true | boolean | Whether the built-in tool-call action runs each tool via durable async execution. | | `tool-call.parallel` | true | boolean | When `tool-call.async` is also true (JDK 21+), run multiple tool calls from one `ToolRequestEvent` as one parallel durable batch. Increases in-flight external calls; after failover, unfinished tools may be submitted again — side-effecting tools should be idempotent or provide a reconciler. Set to `false` for serial tool execution. | | `tool-call.num-async-threads` | os cpu count * 2 | int | Dedicated thread pool size for tool-call async / parallel batch execution. Separate from `num-async-threads` so a large tool batch does not exhaust the global async pool. | -| `tool-call.batch.timeout` | -1ms (disabled) | Duration | Overall timeout for one parallel tool-call batch. Non-positive disables it. On timeout, completed slots keep their outcome and unfinished slots fail. Timeout cancellation is best-effort; external side effects from unfinished tool calls may still complete, so side-effecting tools should be idempotent or provide a reconciler. | +| `tool-call.batch.timeout.ms` | -1 (disabled) | long (milliseconds) | Overall timeout for one parallel tool-call batch. Non-positive disables it. On timeout, completed slots keep their outcome and unfinished slots fail. Timeout cancellation is best-effort; external side effects from unfinished tool calls may still complete, so side-effecting tools should be idempotent or provide a reconciler. | | `rag.async` | true | boolean | Whether retrieve context asynchronously for built-in context retrieval action. | | `num-async-threads` | os cpu count * 2 | int | The thread pool size for async executor. | | `job-identifier` | none | String | The unique identifier of job, remaining consistent after restoring from a savepoint. If not set, uses flink job id. | diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/pom.xml b/e2e-test/flink-agents-end-to-end-tests-integration/pom.xml index 479a47ba5..df126555a 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/pom.xml +++ b/e2e-test/flink-agents-end-to-end-tests-integration/pom.xml @@ -122,6 +122,40 @@ under the License. flink-test-utils ${flink.version} test + + + + org.apache.logging.log4j + log4j-api + + + org.apache.logging.log4j + log4j-core + + + org.apache.logging.log4j + log4j-slf4j-impl + + + + + + org.apache.logging.log4j + log4j-api + ${log4j2.version} + test + + + org.apache.logging.log4j + log4j-core + ${log4j2.version} + test + + + org.apache.logging.log4j + log4j-slf4j-impl + ${log4j2.version} + test com.squareup.okhttp3 @@ -159,28 +193,6 @@ under the License. ${flink.1.20.version} flink-agents-dist-flink-1.20 - - - - org.apache.logging.log4j - log4j-api - ${log4j2.version} - test - - - org.apache.logging.log4j - log4j-core - ${log4j2.version} - test - - - org.apache.logging.log4j - log4j-slf4j-impl - ${log4j2.version} - test - - diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java index ac9ed762d..dde280a70 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java @@ -129,7 +129,25 @@ public Map getParameters() { } } + /** Chat model setup for batch timeout e2e tests. */ + public static class ToolBatchTimeoutChatModel extends BaseChatModelSetup { + public ToolBatchTimeoutChatModel( + ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + } + + @Override + public Map getParameters() { + return new HashMap<>(); + } + } + private static Map toolCall(String id, String requestId, int index) { + return toolCallWithSleep(id, requestId, index, 500); + } + + private static Map toolCallWithSleep( + String id, String requestId, int index, int sleepMs) { return Map.of( "id", id, @@ -140,7 +158,129 @@ private static Map toolCall(String id, String requestId, int ind "name", "timed_tool", "arguments", - Map.of("request_id", requestId, "call_index", index))); + Map.of( + "request_id", + requestId, + "call_index", + String.valueOf(index), + "sleep_ms", + sleepMs))); + } + + private static Map timeoutToolCallWithSleep( + String id, String requestId, int index, int sleepMs) { + return Map.of( + "id", + id, + "type", + "function", + "function", + Map.of( + "name", + "timed_tool_with_sleep", + "arguments", + Map.of( + "request_id", + requestId, + "call_index", + String.valueOf(index), + "sleep_ms", + sleepMs))); + } + + /** + * Chat connection that issues one fast and one slow tool call, then aggregates all tool + * responses into the assistant reply so timeout e2e tests can observe partial batch outcomes. + */ + public static class ToolBatchTimeoutChatConnection extends BaseChatModelConnection { + public ToolBatchTimeoutChatConnection( + ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + } + + @Override + public ChatMessage chat( + List messages, List tools, Map modelParams) { + ChatMessage lastMessage = messages.get(messages.size() - 1); + if (lastMessage.getRole() == MessageRole.TOOL) { + StringBuilder aggregated = new StringBuilder(); + for (ChatMessage message : messages) { + if (message.getRole() == MessageRole.TOOL) { + if (aggregated.length() > 0) { + aggregated.append('|'); + } + aggregated.append(message.getContent()); + } + } + return new ChatMessage(MessageRole.ASSISTANT, aggregated.toString()); + } + + String requestId = lastMessage.getContent(); + return new ChatMessage( + MessageRole.ASSISTANT, + "", + List.of( + timeoutToolCallWithSleep("call-1", requestId, 1, 0), + timeoutToolCallWithSleep("call-2", requestId, 2, 800))); + } + } + + /** Agent that drives a two-tool batch used to exercise batch timeout behavior. */ + public static class ToolBatchTimeoutAgent extends Agent { + @ChatModelConnection + public static ResourceDescriptor toolBatchTimeoutChatConnection() { + return ResourceDescriptor.Builder.newBuilder( + ToolBatchTimeoutChatConnection.class.getName()) + .build(); + } + + @ChatModelSetup + public static ResourceDescriptor toolBatchTimeoutChatModel() { + return ResourceDescriptor.Builder.newBuilder(ToolBatchTimeoutChatModel.class.getName()) + .addInitialArgument("connection", "toolBatchTimeoutChatConnection") + .addInitialArgument("model", "test-model") + .addInitialArgument("tools", List.of("timed_tool_with_sleep")) + .build(); + } + + @org.apache.flink.agents.api.annotation.Tool( + description = "Records timing for a tool call with configurable sleep.") + public static String timed_tool_with_sleep( + @ToolParam(name = "request_id") String requestId, + @ToolParam(name = "call_index") String callIndex, + @ToolParam(name = "sleep_ms") Integer sleepMs) { + long start = System.currentTimeMillis(); + int sleepMillis = sleepMs != null ? sleepMs : 0; + if (sleepMillis > 0) { + try { + Thread.sleep(sleepMillis); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + long end = System.currentTimeMillis(); + return String.format( + "request=%s,call=%s,sleep_ms=%d,start=%d,end=%d", + requestId, callIndex, sleepMillis, start, end); + } + + @Action(EventType.InputEvent) + public static void requestTools(Event event, RunnerContext ctx) { + InputEvent inputEvent = InputEvent.fromEvent(event); + AsyncRequest request = (AsyncRequest) inputEvent.getInput(); + ctx.sendEvent( + new ChatRequestEvent( + "toolBatchTimeoutChatModel", + List.of( + new ChatMessage( + MessageRole.USER, String.valueOf(request.id))))); + } + + @Action(EventType.ChatResponseEvent) + public static void emitToolTimings(Event event, RunnerContext ctx) { + ChatResponseEvent responseEvent = ChatResponseEvent.fromEvent(event); + ctx.sendEvent(new OutputEvent(responseEvent.getResponse().getContent())); + } } /** Agent that requests one chat turn that produces multiple slow tool calls. */ diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java index 1c4a62afe..cf15fe0d6 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java @@ -391,6 +391,63 @@ public void testToolCallBatchExecutionIsActuallyParallel() throws Exception { } } + @Test + public void testToolCallBatchTimeoutKeepsCompletedOutcomes() throws Exception { + boolean continuationSupported = ContinuationActionExecutor.isContinuationSupported(); + int javaVersion = Runtime.version().feature(); + if (!continuationSupported || javaVersion < 21) { + System.out.println( + "Skipping batch timeout e2e: requires JDK 21+ Continuation execution"); + return; + } + + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + + DataStream inputStream = + env.fromElements(new AsyncExecutionAgent.AsyncRequest(1, "tool-batch-timeout")); + + AgentsExecutionEnvironment agentsEnv = + AgentsExecutionEnvironment.getExecutionEnvironment(env); + agentsEnv.getConfig().set(AgentExecutionOptions.TOOL_CALL_ASYNC, true); + agentsEnv.getConfig().set(AgentExecutionOptions.TOOL_CALL_PARALLEL, true); + agentsEnv.getConfig().set(AgentExecutionOptions.TOOL_CALL_NUM_ASYNC_THREADS, 2); + agentsEnv.getConfig().set(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT_MS, 200L); + + DataStream outputStream = + agentsEnv + .fromDataStream( + inputStream, new AsyncExecutionAgent.AsyncRequestKeySelector()) + .apply(new AsyncExecutionAgent.ToolBatchTimeoutAgent()) + .toDataStream(); + + CloseableIterator results = outputStream.collectAsync(); + agentsEnv.execute(); + + List outputList = new ArrayList<>(); + while (results.hasNext()) { + outputList.add(results.next().toString()); + } + results.close(); + + Assertions.assertEquals(1, outputList.size()); + String output = outputList.get(0); + + Pattern fastToolPattern = + Pattern.compile("call=1.*sleep_ms=0.*start=\\d+,end=\\d+", Pattern.DOTALL); + Assertions.assertTrue( + fastToolPattern.matcher(output).find(), + "Fast tool should complete before the batch deadline: " + output); + Assertions.assertTrue( + output.contains("execute failed") || output.toLowerCase().contains("timed out"), + "Slow tool should fail when the batch deadline elapses: " + output); + Assertions.assertFalse( + Pattern.compile("call=2.*sleep_ms=800.*start=\\d+,end=\\d+", Pattern.DOTALL) + .matcher(output) + .find(), + "Slow tool should not report a successful timed result: " + output); + } + /** * Tests that durableExecute (sync) works correctly. * diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java index 412595125..9a6cebced 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java @@ -42,6 +42,8 @@ /** Built-in action for processing tool call. */ public class ToolCallAction { + static final String TOOL_CALL_DURABLE_ID = "tool-call"; + public static Action getToolCallAction() throws Exception { return new Action( "tool_call_action", @@ -135,7 +137,7 @@ private static List buildToolCallExecutions( new DurableCallable<>() { @Override public String getId() { - return "tool-call-" + id; + return TOOL_CALL_DURABLE_ID; } @Override 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 f8accc967..90815747f 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 @@ -436,7 +436,7 @@ void processToolRequestUsesParallelBatchForMultipleTools() throws Exception { ToolCallAction.processToolRequest(toolRequest("queryOrder", "call-1", "call-2"), ctx); assertThat(ctx.durableExecuteAllAsyncIds) - .containsExactly(List.of("tool-call-call-1", "tool-call-call-2")); + .containsExactly(List.of("tool-call", "tool-call")); assertThat(ctx.durableExecuteAsyncIds).isEmpty(); assertThat(ctx.durableExecuteIds).isEmpty(); ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); @@ -452,7 +452,7 @@ void processToolRequestUsesSerialAsyncWhenParallelDisabled() throws Exception { assertThat(ctx.durableExecuteAllAsyncIds).isEmpty(); assertThat(ctx.durableExecuteAsyncIds) - .containsExactly("tool-call-call-1", "tool-call-call-2"); + .containsExactly("tool-call", "tool-call"); assertThat(ctx.durableExecuteIds).isEmpty(); } @@ -464,7 +464,7 @@ void processToolRequestUsesSyncWhenAsyncDisabled() throws Exception { assertThat(ctx.durableExecuteAllAsyncIds).isEmpty(); assertThat(ctx.durableExecuteAsyncIds).isEmpty(); - assertThat(ctx.durableExecuteIds).containsExactly("tool-call-call-1", "tool-call-call-2"); + assertThat(ctx.durableExecuteIds).containsExactly("tool-call", "tool-call"); } @Test @@ -474,7 +474,7 @@ void processToolRequestDoesNotBatchSingleTool() throws Exception { ToolCallAction.processToolRequest(toolRequest("queryOrder"), ctx); assertThat(ctx.durableExecuteAllAsyncIds).isEmpty(); - assertThat(ctx.durableExecuteAsyncIds).containsExactly("tool-call-call-1"); + assertThat(ctx.durableExecuteAsyncIds).containsExactly("tool-call"); assertThat(ctx.durableExecuteIds).isEmpty(); } @@ -501,7 +501,7 @@ public Resource getResource(String name, ResourceType type) throws Exception { ctx); assertThat(ctx.durableExecuteAllAsyncIds) - .containsExactly(List.of("tool-call-call-1", "tool-call-call-2")); + .containsExactly(List.of("tool-call", "tool-call")); ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); assertThat(response.getSuccess()).containsEntry("missing-call", false); assertThat(response.getError()).containsEntry("missing-call", "missing resource"); diff --git a/python/flink_agents/api/core_options.py b/python/flink_agents/api/core_options.py index 142f0c383..c2b0a913c 100644 --- a/python/flink_agents/api/core_options.py +++ b/python/flink_agents/api/core_options.py @@ -266,8 +266,9 @@ class AgentExecutionOptions: default=os.cpu_count() * 2, ) - TOOL_CALL_BATCH_TIMEOUT = ConfigOption( - key="tool-call.batch.timeout", + # Non-positive disables the batch timeout; positive values are milliseconds. + TOOL_CALL_BATCH_TIMEOUT_MS = ConfigOption( + key="tool-call.batch.timeout.ms", config_type=int, default=-1, ) diff --git a/python/flink_agents/api/runner_context.py b/python/flink_agents/api/runner_context.py index 70796de79..07b41bf8f 100644 --- a/python/flink_agents/api/runner_context.py +++ b/python/flink_agents/api/runner_context.py @@ -241,6 +241,7 @@ def durable_execute( func: Callable[[Any], Any], *args: Any, reconciler: Callable[[], Any] | None = None, + durable_id: str | None = None, **kwargs: Any, ) -> Any: """Synchronously execute the provided function with durable execution support. @@ -283,6 +284,10 @@ def my_action(event, ctx): Optional zero-argument reconciler callable used only during recovery. This is a reserved keyword-only parameter and is not forwarded to `func`. + durable_id : str | None + Optional durable journal identity. When provided, the runtime keys + recovery on this id with an empty args digest instead of deriving + the identity from the callable and arguments. **kwargs : Any Keyword arguments to pass to the function. @@ -298,6 +303,7 @@ def durable_execute_async( func: Callable[[Any], Any], *args: Any, reconciler: Callable[[], Any] | None = None, + durable_id: str | None = None, **kwargs: Any, ) -> "AsyncExecutionResult": """Asynchronously execute the provided function with durable execution support. @@ -341,6 +347,10 @@ async def my_action(event, ctx): Optional zero-argument reconciler callable used only during recovery. This is a reserved keyword-only parameter and is not forwarded to `func`. + durable_id : str | None + Optional durable journal identity. When provided, the runtime keys + recovery on this id with an empty args digest instead of deriving + the identity from the callable and arguments. **kwargs : Any Keyword arguments to pass to the function. diff --git a/python/flink_agents/api/tests/test_core_options.py b/python/flink_agents/api/tests/test_core_options.py index 2abbedf92..bb8933b74 100644 --- a/python/flink_agents/api/tests/test_core_options.py +++ b/python/flink_agents/api/tests/test_core_options.py @@ -118,8 +118,8 @@ def test_agent_execution_options_include_parallel_tool_call_options() -> None: assert options["TOOL_CALL_PARALLEL"].get_default_value() is True assert options["TOOL_CALL_NUM_ASYNC_THREADS"].get_key() == "tool-call.num-async-threads" assert options["TOOL_CALL_NUM_ASYNC_THREADS"].get_default_value() == os.cpu_count() * 2 - assert options["TOOL_CALL_BATCH_TIMEOUT"].get_key() == "tool-call.batch.timeout" - assert options["TOOL_CALL_BATCH_TIMEOUT"].get_default_value() == -1 + assert options["TOOL_CALL_BATCH_TIMEOUT_MS"].get_key() == "tool-call.batch.timeout.ms" + assert options["TOOL_CALL_BATCH_TIMEOUT_MS"].get_default_value() == -1 def test_unknown_agent_config_option_raises_attribute_error() -> None: diff --git a/python/flink_agents/plan/actions/tool_call_action.py b/python/flink_agents/plan/actions/tool_call_action.py index 5651f453f..df42c35b4 100644 --- a/python/flink_agents/plan/actions/tool_call_action.py +++ b/python/flink_agents/plan/actions/tool_call_action.py @@ -32,6 +32,8 @@ from flink_agents.plan.function import PythonFunction from flink_agents.plan.tools.function_tool import FunctionTool +from flink_agents.runtime.durable_execution import durable_identity_for_call + _logger = logging.getLogger(__name__) @@ -126,12 +128,13 @@ def _build_tool_call_executions( error[call_id] = str(e) continue + function_id, _ = durable_identity_for_call(tool.call, (), call_kwargs) executions.append( _ToolCallExecution( id=call_id, name=name, durable_call=DurableCall( - id=f"tool-call-{call_id}", + id=function_id, func=tool.call, kwargs=call_kwargs, ), @@ -147,15 +150,11 @@ async def _execute_parallel( success: dict, error: dict, ) -> None: - try: - outcomes = await ctx.durable_execute_all_async( - [execution.durable_call for execution in executions] - ) - for execution, outcome in zip(executions, outcomes, strict=True): - _record_outcome(execution, outcome, responses, success, error) - except Exception as e: - for execution in executions: - _record_execution_exception(execution, e, responses, success, error) + outcomes = await ctx.durable_execute_all_async( + [execution.durable_call for execution in executions] + ) + for execution, outcome in zip(executions, outcomes, strict=True): + _record_outcome(execution, outcome, responses, success, error) async def _execute_sequentially( diff --git a/python/flink_agents/plan/tests/actions/test_tool_call_action.py b/python/flink_agents/plan/tests/actions/test_tool_call_action.py index b329e853f..92befd89c 100644 --- a/python/flink_agents/plan/tests/actions/test_tool_call_action.py +++ b/python/flink_agents/plan/tests/actions/test_tool_call_action.py @@ -28,12 +28,33 @@ from flink_agents.plan.configuration import AgentConfiguration from flink_agents.plan.function import PythonFunction from flink_agents.plan.tools.function_tool import FunctionTool +from flink_agents.runtime.durable_execution import durable_identity_for_call def query_order(order_id: str, tenant_id: str) -> str: return f"{tenant_id}:{order_id}" +def _query_order_tool(injected_args: dict[str, InjectedArg] | None = None) -> FunctionTool: + return FunctionTool( + func=PythonFunction.from_callable(query_order), + injected_args=injected_args or {"tenant_id": InjectedArg.from_config("tenant_id")}, + ) + + +def _expected_durable_function_id( + order_id: str, + tenant_id: str = "tenant-1", +) -> str: + tool = _query_order_tool() + function_id, _ = durable_identity_for_call( + tool.call, + (), + {"order_id": order_id, "tenant_id": tenant_id}, + ) + return function_id + + class _Context: def __init__( self, @@ -70,11 +91,15 @@ def get_resource(self, name: str, type: ResourceType) -> FunctionTool: def durable_execute(self, func: Any, *args: Any, **kwargs: Any) -> Any: self.durable_execute_calls.append((func, args, kwargs)) - return func(*args, **kwargs) + call_kwargs = dict(kwargs) + call_kwargs.pop("durable_id", None) + return func(*args, **call_kwargs) async def durable_execute_async(self, func: Any, *args: Any, **kwargs: Any) -> Any: self.durable_execute_async_calls.append((func, args, kwargs)) - return func(*args, **kwargs) + call_kwargs = dict(kwargs) + call_kwargs.pop("durable_id", None) + return func(*args, **call_kwargs) async def durable_execute_all_async(self, callables: list[Any]) -> list[Outcome]: self.durable_execute_all_async_calls.append(callables) @@ -315,9 +340,10 @@ def test_tool_call_action_uses_parallel_batch_for_multiple_tools() -> None: } assert response.success == {"call-1": True, "call-2": True} assert len(ctx.durable_execute_all_async_calls) == 1 + expected_id = _expected_durable_function_id("order-call-1") assert [call.id for call in ctx.durable_execute_all_async_calls[0]] == [ - "tool-call-call-1", - "tool-call-call-2", + expected_id, + expected_id, ] assert ctx.durable_execute_async_calls == [] @@ -332,6 +358,7 @@ def test_tool_call_action_uses_serial_async_when_parallel_disabled() -> None: assert ctx.durable_execute_all_async_calls == [] assert len(ctx.durable_execute_async_calls) == 2 + assert all("durable_id" not in call_kwargs for _, _, call_kwargs in ctx.durable_execute_async_calls) def test_tool_call_action_does_not_batch_single_tool() -> None: diff --git a/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py b/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py index 8a7eb712a..dfeae20fe 100644 --- a/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py +++ b/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py @@ -42,7 +42,6 @@ "java.lang.Float": float, "java.lang.Double": float, "java.util.List": list, - "java.time.Duration": int, } @@ -85,10 +84,6 @@ def _java_type_matches_python(java_type_name: str, python_config_type: type) -> return python_config_type.__name__ == java_simple_name -def _java_duration_to_millis(java_duration: Any) -> int: - return int(java_duration.toMillis()) - - def normalize_java_default( java_default: Any, java_type_name: str, python_config_type: type ) -> Any: @@ -96,9 +91,6 @@ def normalize_java_default( if java_default is None: return None - if java_type_name == "java.time.Duration": - return _java_duration_to_millis(java_default) - if hasattr(java_default, "name") and callable(java_default.name): enum_name = java_default.name() if isinstance(python_config_type, type) and issubclass(python_config_type, Enum): diff --git a/python/flink_agents/runtime/durable_execution.py b/python/flink_agents/runtime/durable_execution.py index 9db22bfa9..e66f1c719 100644 --- a/python/flink_agents/runtime/durable_execution.py +++ b/python/flink_agents/runtime/durable_execution.py @@ -22,6 +22,28 @@ import cloudpickle +def _resolve_durable_identity( + func: Callable, + args: tuple, + kwargs: dict, + durable_id: str | None = None, +) -> tuple[str, str]: + """Resolve the durable journal identity for a single-call execution.""" + if durable_id is not None: + return durable_id, "" + return _compute_function_id(func), _compute_args_digest(args, kwargs) + + +def durable_identity_for_call( + func: Callable, + args: tuple, + kwargs: dict | None, +) -> tuple[str, str]: + """Return the durable journal identity for a single callable invocation.""" + call_kwargs = kwargs or {} + return _compute_function_id(func), _compute_args_digest(args, call_kwargs) + + def _compute_function_id(func: Callable) -> str: """Compute a stable function identifier from a callable.""" module_obj = inspect.getmodule(func) diff --git a/python/flink_agents/runtime/flink_runner_context.py b/python/flink_agents/runtime/flink_runner_context.py index b3ff7bea4..e584fac5a 100644 --- a/python/flink_agents/runtime/flink_runner_context.py +++ b/python/flink_agents/runtime/flink_runner_context.py @@ -45,7 +45,9 @@ from flink_agents.runtime.durable_execution import ( _compute_args_digest, _compute_function_id, + _resolve_durable_identity, _validate_reconciler_callable, + durable_identity_for_call, ) from flink_agents.runtime.flink_memory_object import FlinkMemoryObject from flink_agents.runtime.flink_metric_group import FlinkMetricGroup @@ -263,8 +265,8 @@ def __await__(self) -> Any: futures = [ self._ctx.tool_call_executor.submit(supplier) for _, supplier in plan.suppliers ] - timeout_ms = self._ctx.config.get(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT) - deadline = time.monotonic() + timeout_ms / 1000 if timeout_ms >= 0 else None + timeout_ms = self._ctx.config.get(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT_MS) + deadline = time.monotonic() + timeout_ms / 1000 if timeout_ms > 0 else None while any(not future.done() for future in futures): if deadline is not None and time.monotonic() >= deadline: exception = TimeoutError( @@ -469,7 +471,11 @@ def action_metric_group(self) -> FlinkMetricGroup: return FlinkMetricGroup(self._j_runner_context.getActionMetricGroup()) def _try_get_cached_result( - self, func: Callable, args: tuple, kwargs: dict + self, + func: Callable, + args: tuple, + kwargs: dict, + durable_id: str | None = None, ) -> tuple[bool, Any]: """Try to get a cached result from a previous execution. @@ -479,8 +485,9 @@ def _try_get_cached_result( A tuple of (is_hit, result_or_exception). If is_hit is True, the second element is the cached result or an exception to re-raise. """ - function_id = _compute_function_id(func) - args_digest = _compute_args_digest(args, kwargs) + function_id, args_digest = _resolve_durable_identity( + func, args, kwargs, durable_id + ) cached_exception: BaseException | None = None try: @@ -517,6 +524,7 @@ def _record_call_completion( kwargs: dict, result: Any, exception: BaseException | None, + durable_id: str | None = None, ) -> None: """Record the completion of a call for durable execution. @@ -533,8 +541,9 @@ def _record_call_completion( exception : BaseException | None The exception raised by the function (None if successful). """ - function_id = _compute_function_id(func) - args_digest = _compute_args_digest(args, kwargs) + function_id, args_digest = _resolve_durable_identity( + func, args, kwargs, durable_id + ) try: result_payload = None if exception else cloudpickle.dumps(result) @@ -606,9 +615,11 @@ def _finalize_current_call( kwargs: dict, result: Any, exception: BaseException | None, + durable_id: str | None = None, ) -> None: - function_id = _compute_function_id(func) - args_digest = _compute_args_digest(args, kwargs) + function_id, args_digest = _resolve_durable_identity( + func, args, kwargs, durable_id + ) result_payload, exception_payload = self._serialize_call_payloads( result, exception, @@ -637,8 +648,7 @@ def _plan_reconciler_execution( reconciler: Callable[[], Any], kwargs: dict, ) -> _ReconcilerExecutionPlan: - function_id = _compute_function_id(func) - args_digest = _compute_args_digest(args, kwargs) + function_id, args_digest = _resolve_durable_identity(func, args, kwargs, None) current = self._peek_current_call_result() durable_call = partial(func, *args, **kwargs) @@ -702,7 +712,24 @@ def _wrap_completion_only_func( func: Callable, args: tuple, kwargs: dict, + durable_id: str | None = None, ) -> Callable[..., Any]: + def record_call_completion( + call_func: Callable, + call_args: tuple, + call_kwargs: dict, + result: Any, + exception: BaseException | None, + ) -> None: + self._record_call_completion( + call_func, + call_args, + call_kwargs, + result, + exception, + durable_id=durable_id, + ) + def wrapped_func(*a: Any, **kw: Any) -> Any: exception = None result = None @@ -713,18 +740,23 @@ def wrapped_func(*a: Any, **kw: Any) -> Any: if exception: raise _DurableExecutionException( - func, args, kwargs, result, exception, self._record_call_completion + func, args, kwargs, result, exception, record_call_completion ) return _DurableExecutionResult( - func, args, kwargs, result, self._record_call_completion + func, args, kwargs, result, record_call_completion ) return wrapped_func - def _call_matches( - self, current: _PersistedCallResult, call: DurableCall, args_digest: str - ) -> bool: - return current.function_id == call.id and current.args_digest == args_digest + def _durable_identity(self, call: DurableCall) -> tuple[str, str]: + return durable_identity_for_call(call.func, call.args, call.kwargs) + + def _call_matches(self, current: _PersistedCallResult, call: DurableCall) -> bool: + function_id, args_digest = self._durable_identity(call) + return ( + current.function_id == function_id + and current.args_digest == args_digest + ) def _read_terminal_outcome(self, current: _PersistedCallResult) -> Outcome: if current.exception_payload is not None: @@ -738,7 +770,6 @@ def _callable_for_durable_call(self, call: DurableCall) -> Callable[[], Any]: return partial(call.func, *call.args, **kwargs) def _prepare_batch_execution(self, calls: list[DurableCall]) -> _BatchExecutionPlan: - args_digest = "" base = self._j_runner_context.getCurrentCallIndex() outcomes: list[Outcome | None] = [] suppliers: list[tuple[int, Callable[[], Any]]] = [] @@ -746,6 +777,7 @@ def _prepare_batch_execution(self, calls: list[DurableCall]) -> _BatchExecutionP execution_start = -1 for index, call in enumerate(calls): + function_id, args_digest = self._durable_identity(call) current = self._read_call_result_at(base + index) if current is None: needs_reservation = True @@ -755,7 +787,7 @@ def _prepare_batch_execution(self, calls: list[DurableCall]) -> _BatchExecutionP suppliers.append((index, self._callable_for_durable_call(call))) continue - if not self._call_matches(current, call, args_digest): + if not self._call_matches(current, call): self._j_runner_context.clearCallResultsFromAndPersist(base + index) needs_reservation = True execution_start = index @@ -783,8 +815,13 @@ def _prepare_batch_execution(self, calls: list[DurableCall]) -> _BatchExecutionP outcomes.append(self._read_terminal_outcome(current)) if needs_reservation: - function_ids = [call.id for call in calls[execution_start:]] - self._j_runner_context.reservePendingBatch(function_ids, args_digest) + function_ids = [] + args_digests = [] + for call in calls[execution_start:]: + function_id, args_digest = self._durable_identity(call) + function_ids.append(function_id) + args_digests.append(args_digest) + self._j_runner_context.reservePendingBatch(function_ids, args_digests) return _BatchExecutionPlan( outcomes=outcomes, @@ -799,17 +836,18 @@ def _finalize_batch_execution( plan: _BatchExecutionPlan, executed: list[Outcome], ) -> list[Outcome]: - args_digest = "" base = self._j_runner_context.getCurrentCallIndex() outcomes = list(plan.outcomes) for (call_index, _), outcome in zip(plan.suppliers, executed, strict=True): + call = calls[call_index] + function_id, args_digest = self._durable_identity(call) result_payload, exception_payload = self._serialize_call_payloads( outcome.value, outcome.error, ) self._j_runner_context.finalizeCallAt( base + call_index, - calls[call_index].id, + function_id, args_digest, result_payload, exception_payload, @@ -831,6 +869,7 @@ def durable_execute( func: Callable[[Any], Any], *args: Any, reconciler: Callable[[], Any] | None = None, + durable_id: str | None = None, **kwargs: Any, ) -> Any: """Synchronously execute the provided function with durable execution support. @@ -864,7 +903,9 @@ def durable_execute( ) # Try to get cached result for recovery - is_hit, cached_result = self._try_get_cached_result(func, args, kwargs) + is_hit, cached_result = self._try_get_cached_result( + func, args, kwargs, durable_id + ) if is_hit: return cached_result @@ -877,7 +918,9 @@ def durable_execute( exception = e # Record the completion - self._record_call_completion(func, args, kwargs, result, exception) + self._record_call_completion( + func, args, kwargs, result, exception, durable_id=durable_id + ) if exception: raise exception @@ -889,6 +932,7 @@ def durable_execute_async( func: Callable[[Any], Any], *args: Any, reconciler: Callable[[], Any] | None = None, + durable_id: str | None = None, **kwargs: Any, ) -> AsyncExecutionResult: """Asynchronously execute the provided function with durable execution support. @@ -915,14 +959,16 @@ def durable_execute_async( ) # Try to get cached result for recovery - is_hit, cached_result = self._try_get_cached_result(func, args, kwargs) + is_hit, cached_result = self._try_get_cached_result( + func, args, kwargs, durable_id + ) if is_hit: # Return a pre-completed AsyncExecutionResult return _CachedAsyncExecutionResult(cached_result) return _DurableAsyncExecutionResult( self.executor, - self._wrap_completion_only_func(func, args, kwargs), + self._wrap_completion_only_func(func, args, kwargs, durable_id), args, kwargs, ) diff --git a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py index 59f5283c1..5b04c493f 100644 --- a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py +++ b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py @@ -28,6 +28,7 @@ from flink_agents.runtime.durable_execution import ( _compute_args_digest, _compute_function_id, + durable_identity_for_call, ) from flink_agents.runtime.flink_runner_context import FlinkRunnerContext @@ -41,6 +42,29 @@ class _StoredCallResult: exception_payload: bytes | None = None +def _durable_call( + func: Callable[..., Any], + *args: Any, + **kwargs: Any, +) -> DurableCall: + function_id, _ = durable_identity_for_call(func, args, kwargs or None) + return DurableCall(id=function_id, func=func, args=args, kwargs=kwargs or None) + + +def _stored_call( + func: Callable[..., Any], + *args: Any, + status: str, + **kwargs: Any, +) -> _StoredCallResult: + function_id, args_digest = durable_identity_for_call(func, args, kwargs or None) + return _StoredCallResult( + function_id=function_id, + args_digest=args_digest, + status=status, + ) + + class _FakeJavaRunnerContext: def __init__(self) -> None: self.call_results: list[_StoredCallResult] = [] @@ -146,13 +170,15 @@ def clearCallResultsFromCurrentIndexAndPersist(self) -> None: self.operations.append("clear") self.call_results = self.call_results[: self.current_call_index] - def reservePendingBatch(self, function_ids: list[str], args_digest: str) -> None: + def reservePendingBatch( + self, function_ids: list[str], args_digests: list[str] + ) -> None: self.operations.append(f"reserve:{len(function_ids)}") - for function_id in function_ids: + for function_id, digest in zip(function_ids, args_digests, strict=True): self.call_results.append( _StoredCallResult( function_id=function_id, - args_digest=args_digest, + args_digest=digest, status="PENDING", ) ) @@ -486,7 +512,7 @@ def collect_kwargs(**kwargs: Any) -> dict[str, Any]: def test_flink_runner_context_durable_execute_all_async_runs_calls_in_parallel() -> None: j_runner_context = _FakeJavaRunnerContext() - config = AgentConfiguration({"tool-call.batch.timeout": -1}) + config = AgentConfiguration({"tool-call.batch.timeout.ms": -1}) ctx = _create_runner_context(j_runner_context, config=config, tool_call_workers=3) sleep_seconds = 0.2 @@ -499,9 +525,9 @@ def slow_call(value: str) -> str: outcomes = _run_async( ctx.durable_execute_all_async( [ - DurableCall("tool-call-1", slow_call, args=("one",)), - DurableCall("tool-call-2", slow_call, args=("two",)), - DurableCall("tool-call-3", slow_call, args=("three",)), + _durable_call(slow_call, "one"), + _durable_call(slow_call, "two"), + _durable_call(slow_call, "three"), ] ) ) @@ -527,8 +553,8 @@ def test_flink_runner_context_durable_execute_all_async_initial_batch() -> None: outcomes = _run_async( ctx.durable_execute_all_async( [ - DurableCall("tool-call-1", _call_value, args=("one",)), - DurableCall("tool-call-2", _call_value, args=("two",)), + _durable_call(_call_value, "one"), + _durable_call(_call_value, "two"), ] ) ) @@ -541,42 +567,80 @@ def test_flink_runner_context_durable_execute_all_async_initial_batch() -> None: "SUCCEEDED", ] assert j_runner_context.current_call_index == 2 + + +def test_flink_runner_context_durable_execute_all_async_finalize_failure_aborts_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + j_runner_context = _FakeJavaRunnerContext() + ctx = _create_runner_context(j_runner_context) + original = FlinkRunnerContext._serialize_call_payloads + + def failing_serialize( + result: Any, exception: BaseException | None + ) -> tuple[bytes | None, bytes | None]: + if result == "two": + msg = "serialize failed" + raise RuntimeError(msg) + return original(result, exception) + + monkeypatch.setattr( + FlinkRunnerContext, + "_serialize_call_payloads", + staticmethod(failing_serialize), + ) + + try: + with pytest.raises(RuntimeError, match="serialize failed"): + _run_async( + ctx.durable_execute_all_async( + [ + _durable_call(lambda: "one"), + _durable_call(lambda: "two"), + ] + ) + ) + finally: + _close_runner_context(ctx) + + assert j_runner_context.call_results[0].status == "SUCCEEDED" + assert j_runner_context.current_call_index == 0 assert "reserve:2" in j_runner_context.operations assert "finalize_at:0" in j_runner_context.operations - assert "finalize_at:1" in j_runner_context.operations def test_flink_runner_context_durable_execute_all_async_recovers_partial_batch() -> None: j_runner_context = _FakeJavaRunnerContext() + call_count = 0 + + def tracked_call(value: str) -> str: + nonlocal call_count + call_count += 1 + return _call_value(value) + + cached_call = _durable_call(tracked_call, "one") + _, cached_digest = durable_identity_for_call( + cached_call.func, cached_call.args, cached_call.kwargs + ) j_runner_context.call_results.extend( [ _StoredCallResult( - function_id="tool-call-1", - args_digest="", + function_id=cached_call.id, + args_digest=cached_digest, status="SUCCEEDED", result_payload=cloudpickle.dumps("cached-one"), ), - _StoredCallResult( - function_id="tool-call-2", - args_digest="", - status="PENDING", - ), + _stored_call(tracked_call, "two", status="PENDING"), ] ) - call_count = 0 - - def tracked_call(value: str) -> str: - nonlocal call_count - call_count += 1 - return _call_value(value) ctx = _create_runner_context(j_runner_context) try: outcomes = _run_async( ctx.durable_execute_all_async( [ - DurableCall("tool-call-1", tracked_call, args=("one",)), - DurableCall("tool-call-2", tracked_call, args=("two",)), + _durable_call(tracked_call, "one"), + _durable_call(tracked_call, "two"), ] ) ) @@ -591,10 +655,14 @@ def tracked_call(value: str) -> str: def test_flink_runner_context_durable_execute_all_async_returns_cached_failure() -> None: j_runner_context = _FakeJavaRunnerContext() + cached_call = _durable_call(_call_value, "one") + _, cached_digest = durable_identity_for_call( + cached_call.func, cached_call.args, cached_call.kwargs + ) j_runner_context.call_results.append( _StoredCallResult( - function_id="tool-call-1", - args_digest="", + function_id=cached_call.id, + args_digest=cached_digest, status="FAILED", exception_payload=cloudpickle.dumps(ValueError("cached failure")), ) @@ -604,7 +672,7 @@ def test_flink_runner_context_durable_execute_all_async_returns_cached_failure() try: outcomes = _run_async( ctx.durable_execute_all_async( - [DurableCall("tool-call-1", _call_value, args=("one",))] + [_durable_call(_call_value, "one")] ) ) finally: @@ -628,8 +696,8 @@ def fail_call() -> str: outcomes = _run_async( ctx.durable_execute_all_async( [ - DurableCall("tool-call-1", _call_value, args=("one",)), - DurableCall("tool-call-2", fail_call), + _durable_call(_call_value, "one"), + _durable_call(fail_call), ] ) ) @@ -648,7 +716,7 @@ def fail_call() -> str: def test_flink_runner_context_durable_execute_all_async_timeout_keeps_completed_results() -> None: j_runner_context = _FakeJavaRunnerContext() - config = AgentConfiguration({"tool-call.batch.timeout": 10}) + config = AgentConfiguration({"tool-call.batch.timeout.ms": 10}) ctx = _create_runner_context(j_runner_context, config=config) def slow_call() -> str: @@ -659,8 +727,8 @@ def slow_call() -> str: outcomes = _run_async( ctx.durable_execute_all_async( [ - DurableCall("tool-call-1", lambda: "fast"), - DurableCall("tool-call-2", slow_call), + _durable_call(lambda: "fast"), + _durable_call(slow_call), ] ) ) diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java index f5e0070c1..2c65afdac 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java @@ -152,7 +152,10 @@ private BatchExecutionPlan buildBatchExecutionPlan( } else { plan.outcomes.add( readTerminalOutcomeAt( - base + i, callable.getId(), argsDigest, callable.getResultClass())); + base + i, + callable.getId(), + argsDigest, + callable.getResultClass())); } } return plan; @@ -186,11 +189,13 @@ private void reservePendingBatchIfNeeded( return; } List ids = new ArrayList<>(); + List argsDigests = new ArrayList<>(); for (DurableCallable callable : callables.subList(plan.executionStart, callables.size())) { ids.add(callable.getId()); + argsDigests.add(argsDigest); } - reservePendingBatch(ids, argsDigest); + reservePendingBatch(ids, argsDigests); } private void finalizeExecutedOutcomes( @@ -250,7 +255,9 @@ private List> executeOutcomeSuppliers(List> suppliers }) .collect(Collectors.toList()); } - Duration timeout = getConfig().get(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT); + Long timeoutMs = getConfig().get(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT_MS); + Duration timeout = + timeoutMs == null || timeoutMs <= 0 ? null : Duration.ofMillis(timeoutMs); return toolCallContinuationExecutor.executeAllAsync( continuationContext, suppliers, timeout); } 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 be36a4e8b..176264f87 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 @@ -294,6 +294,11 @@ protected T durableExecuteCompletionOnly( // argsDigest is empty because DurableCallable encapsulates all arguments internally String argsDigest = ""; + CallResult current = getCurrentCallResult(); + if (current != null && current.matches(functionId, argsDigest) && current.isPending()) { + return executeAndFinalizeCurrentCall(functionId, argsDigest, executionCallable); + } + Optional cachedResult = tryGetCachedResult(functionId, argsDigest, durableCallable.getResultClass()); if (cachedResult.isPresent()) { @@ -432,10 +437,10 @@ public void appendPendingCall(String functionId, String argsDigest) { } } - public void reservePendingBatch(List functionIds, String argsDigest) { + public void reservePendingBatch(List functionIds, List argsDigests) { mailboxThreadChecker.run(); if (durableExecutionContext != null && !functionIds.isEmpty()) { - durableExecutionContext.reservePendingBatch(functionIds, argsDigest); + durableExecutionContext.reservePendingBatch(functionIds, argsDigests); } } @@ -755,6 +760,15 @@ public Object[] matchNextOrClearSubsequentCallResult(String functionId, String a CallResult result = recoveryCallResults.get(currentCallIndex); if (result.matches(functionId, argsDigest)) { + if (result.isPending()) { + LOG.debug( + "Pending CallResult at index {} treated as cache miss: " + + "functionId={}, argsDigest={}", + currentCallIndex, + functionId, + argsDigest); + return null; + } LOG.debug( "CallResult hit at index {}: functionId={}, argsDigest={}", currentCallIndex, @@ -837,9 +851,16 @@ public void appendPendingCall(String functionId, String argsDigest) { argsDigest); } - public void reservePendingBatch(List functionIds, String argsDigest) { - for (String functionId : functionIds) { - CallResult pending = CallResult.pending(functionId, argsDigest); + public void reservePendingBatch(List functionIds, List argsDigests) { + if (functionIds.size() != argsDigests.size()) { + throw new IllegalArgumentException( + String.format( + "functionIds size (%s) must match argsDigests size (%s)", + functionIds.size(), argsDigests.size())); + } + for (int i = 0; i < functionIds.size(); i++) { + CallResult pending = + CallResult.pending(functionIds.get(i), argsDigests.get(i)); actionState.addCallResult(pending); recoveryCallResults.add(pending); } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/DurableExecutionContextTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/DurableExecutionContextTest.java index b659a77ee..cf7f6cf1b 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/DurableExecutionContextTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/DurableExecutionContextTest.java @@ -105,6 +105,19 @@ void testMatchNextOrClearSubsequentCallResultMiss() { assertEquals(0, context.getCurrentCallIndex()); } + @Test + void testMatchNextOrClearSubsequentCallResultPendingTreatedAsMiss() { + actionState.addCallResult(CallResult.pending("funcA", "digestA")); + + RunnerContextImpl.DurableExecutionContext context = createContext(); + + Object[] result = context.matchNextOrClearSubsequentCallResult("funcA", "digestA"); + + assertNull(result); + assertEquals(0, context.getCurrentCallIndex()); + assertTrue(actionState.getCallResults().get(0).isPending()); + } + @Test void testMatchNextOrClearSubsequentCallResultMismatch() { actionState.addCallResult(new CallResult("funcA", "digestA", "result".getBytes())); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java index 62676e939..d699497ec 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java @@ -18,6 +18,7 @@ package org.apache.flink.agents.runtime.context; +import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.agents.AgentExecutionOptions; @@ -34,9 +35,11 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import java.time.Duration; import java.util.HashMap; import java.util.List; import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Supplier; @@ -201,6 +204,27 @@ void testDurableExecuteAsyncReconcilableReconcileExceptionPersistsFailure() thro assertEquals(1, context.getDurableExecutionContext().getCurrentCallIndex()); } + @Test + void testDurableExecuteAsyncCompletionOnlyReExecutesPendingSlot() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + ActionState actionState = new ActionState(null); + actionState.addCallResult(CallResult.pending("tool-call", "")); + JavaRunnerContextImpl context = createContext(actionState, executor); + TestDurableCallable callable = + new TestDurableCallable<>("tool-call", String.class, () -> "recovered"); + + String result = context.durableExecuteAsync(callable); + + assertEquals("recovered", result); + assertEquals(1, callable.getCallCount()); + assertEquals(1, executor.getExecuteAsyncCallCount()); + assertEquals(1, persistCallCount.get()); + assertEquals(1, context.getDurableExecutionContext().getCurrentCallIndex()); + CallResult persisted = + context.getDurableExecutionContext().getActionState().getCallResults().get(0); + assertTrue(persisted.isSuccess()); + } + @Test void testDurableExecuteAllAsyncInitialBatchPersistsOutcomes() throws Exception { InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); @@ -322,9 +346,7 @@ void testDurableExecuteAllAsyncTimeoutKeepsCompletedOutcomes() throws Exception executor.setUseTimeoutCollection(true); JavaRunnerContextImpl context = createContext(new ActionState(null), executor); ((Configuration) context.getConfig()) - .set( - AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT, - java.time.Duration.ofMillis(10)); + .set(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT_MS, 10L); TestDurableCallable first = new TestDurableCallable<>("batch-1", String.class, () -> "fast"); TestDurableCallable second = @@ -341,6 +363,9 @@ void testDurableExecuteAllAsyncTimeoutKeepsCompletedOutcomes() throws Exception assertEquals("fast", outcomes.get(0).getValue()); assertTrue(outcomes.get(1).isFailure()); assertInstanceOf(TimeoutException.class, outcomes.get(1).getError()); + assertEquals(Duration.ofMillis(10), executor.getLastExecuteAllAsyncTimeout()); + assertEquals(1, first.getCallCount()); + assertEquals(1, second.getCallCount()); List persisted = context.getDurableExecutionContext().getActionState().getCallResults(); assertTrue(persisted.get(0).isSuccess()); @@ -349,6 +374,49 @@ void testDurableExecuteAllAsyncTimeoutKeepsCompletedOutcomes() throws Exception executor.close(); } + @Test + void testDurableExecuteAllAsyncFinalizeFailureAbortsBatch() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + FailingSerializeOnValueContext context = + new FailingSerializeOnValueContext( + metricGroup, + () -> {}, + new AgentPlan(new HashMap<>(), new HashMap<>()), + null, + "test-job", + executor, + "two"); + context.setContinuationContext(new ContinuationContext()); + ActionStatePersister persister = + (key, sequenceNumber, action, event, state) -> { + persistCallCount.incrementAndGet(); + lastPersistedState = state; + }; + context.setDurableExecutionContext( + new RunnerContextImpl.DurableExecutionContext( + "test-key", + 1L, + mock(Action.class), + mock(Event.class), + new ActionState(null), + persister)); + TestDurableCallable first = + new TestDurableCallable<>("batch-1", String.class, () -> "one"); + TestDurableCallable second = + new TestDurableCallable<>("batch-2", String.class, () -> "two"); + + JsonProcessingException thrown = + assertThrows( + JsonProcessingException.class, + () -> context.durableExecuteAllAsync(List.of(first, second))); + + assertTrue(thrown.getMessage().contains("serialize failed")); + List persisted = + context.getDurableExecutionContext().getActionState().getCallResults(); + assertTrue(persisted.get(0).isSuccess()); + assertEquals(0, context.getDurableExecutionContext().getCurrentCallIndex()); + } + private JavaRunnerContextImpl createContext( ActionState actionState, ContinuationActionExecutor executor) { JavaRunnerContextImpl context = @@ -380,6 +448,7 @@ private static final class InspectingContinuationActionExecutor extends ContinuationActionExecutor { private Runnable beforeExecute; private boolean useTimeoutCollection; + private Duration lastExecuteAllAsyncTimeout; private int executeAsyncCallCount; private int executeAllAsyncCallCount; private final List executeAllAsyncBatchSizes = new java.util.ArrayList<>(); @@ -401,11 +470,12 @@ public T executeAsync(ContinuationContext context, Supplier supplier) { public List> executeAllAsync( ContinuationContext context, List> suppliers, - java.time.Duration timeout) { + Duration timeout) { executeAllAsyncCallCount++; executeAllAsyncBatchSizes.add(suppliers.size()); + lastExecuteAllAsyncTimeout = timeout; if (useTimeoutCollection) { - return collectTimedOutOutcomes(suppliers); + return executeAllAsyncWithDeadline(suppliers, timeout); } List> outcomes = new java.util.ArrayList<>(suppliers.size()); for (Callable supplier : suppliers) { @@ -418,20 +488,74 @@ public List> executeAllAsync( return outcomes; } - private List> collectTimedOutOutcomes(List> suppliers) { - List> outcomes = new java.util.ArrayList<>(suppliers.size()); - for (int i = 0; i < suppliers.size(); i++) { - if (i == 0) { - try { - outcomes.add(Outcome.success(suppliers.get(i).call())); - } catch (Exception e) { - outcomes.add(Outcome.failure(e)); - } + private List> executeAllAsyncWithDeadline( + List> suppliers, Duration timeout) { + List>> futures = new java.util.ArrayList<>(suppliers.size()); + for (Callable supplier : suppliers) { + futures.add( + CompletableFuture.supplyAsync( + () -> { + try { + return Outcome.success(supplier.call()); + } catch (Exception e) { + return Outcome.failure(e); + } + })); + } + + CompletableFuture barrier = + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])); + long deadlineNanos = getDeadlineNanos(timeout); + while (!barrier.isDone()) { + if (System.nanoTime() >= deadlineNanos) { + TimeoutException exception = + new TimeoutException( + "Async durable batch execution timed out after " + timeout); + barrier.cancel(true); + return collectBatchOutcomesOnTimeout(futures, exception); + } + try { + Thread.sleep(1); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + TimeoutException exception = + new TimeoutException("Async durable batch execution interrupted"); + return collectBatchOutcomesOnTimeout(futures, exception); + } + } + return collectBatchOutcomes(futures); + } + + private static long getDeadlineNanos(Duration timeout) { + return timeout == null || timeout.isZero() || timeout.isNegative() + ? Long.MAX_VALUE + : System.nanoTime() + timeout.toNanos(); + } + + private static List> collectBatchOutcomes( + List>> futures) { + List> results = new java.util.ArrayList<>(futures.size()); + for (CompletableFuture> future : futures) { + results.add(future.join()); + } + return results; + } + + private static List> collectBatchOutcomesOnTimeout( + List>> futures, + TimeoutException timeoutException) { + List> results = new java.util.ArrayList<>(futures.size()); + for (CompletableFuture> future : futures) { + if (!future.isDone()) { + future.cancel(true); + } + if (future.isDone() && !future.isCancelled()) { + results.add(future.join()); } else { - outcomes.add(Outcome.failure(new TimeoutException("batch timeout"))); + results.add(Outcome.failure(timeoutException)); } } - return outcomes; + return results; } private void setBeforeExecute(Runnable beforeExecute) { @@ -453,5 +577,39 @@ private int getExecuteAllAsyncCallCount() { private List getExecuteAllAsyncBatchSizes() { return executeAllAsyncBatchSizes; } + + private Duration getLastExecuteAllAsyncTimeout() { + return lastExecuteAllAsyncTimeout; + } + } + + private static final class FailingSerializeOnValueContext extends JavaRunnerContextImpl { + private final String failOnResult; + + private FailingSerializeOnValueContext( + FlinkAgentsMetricGroupImpl agentMetricGroup, + Runnable mailboxThreadChecker, + AgentPlan agentPlan, + org.apache.flink.agents.runtime.ResourceCache resourceCache, + String jobIdentifier, + ContinuationActionExecutor continuationExecutor, + String failOnResult) { + super( + agentMetricGroup, + mailboxThreadChecker, + agentPlan, + resourceCache, + jobIdentifier, + continuationExecutor); + this.failOnResult = failOnResult; + } + + @Override + protected byte[] serializeDurableResult(Object result) throws JsonProcessingException { + if (failOnResult != null && failOnResult.equals(result)) { + throw new JsonProcessingException("serialize failed") {}; + } + return super.serializeDurableResult(result); + } } } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java index 10b42decd..9e60fab7a 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplDurableExecuteTest.java @@ -244,6 +244,60 @@ void testDurableExecuteReconcilableReconcileExceptionPersistsFailure() throws Ex assertEquals(1, context.getDurableExecutionContext().getCurrentCallIndex()); } + @Test + void testDurableExecuteCompletionOnlyReExecutesPendingSlot() throws Exception { + ActionState actionState = new ActionState(null); + actionState.addCallResult(CallResult.pending("tool-call", "")); + RunnerContextImpl context = createContext(actionState); + TestDurableCallable callable = + new TestDurableCallable<>("tool-call", String.class, () -> "recovered"); + + String result = context.durableExecute(callable); + + assertEquals("recovered", result); + assertEquals(1, callable.getCallCount()); + assertEquals(1, persistCallCount.get()); + assertEquals(1, context.getDurableExecutionContext().getCurrentCallIndex()); + assertEquals(1, context.getDurableExecutionContext().getActionState().getCallResultCount()); + CallResult persisted = + context.getDurableExecutionContext().getActionState().getCallResults().get(0); + assertTrue(persisted.isSuccess()); + } + + @Test + void testDurableExecuteCompletionOnlyReExecutesPendingSlotAfterBatchReservation() + throws Exception { + ActionState actionState = new ActionState(null); + actionState.addCallResult(CallResult.pending("tool-call", "")); + actionState.addCallResult(CallResult.pending("tool-call", "")); + RunnerContextImpl context = createContext(actionState); + TestDurableCallable first = + new TestDurableCallable<>("tool-call", String.class, () -> "one"); + TestDurableCallable second = + new TestDurableCallable<>("tool-call", String.class, () -> "two"); + + assertEquals("one", context.durableExecute(first)); + assertEquals("two", context.durableExecute(second)); + + assertEquals(1, first.getCallCount()); + assertEquals(1, second.getCallCount()); + assertEquals(2, persistCallCount.get()); + assertEquals(2, context.getDurableExecutionContext().getCurrentCallIndex()); + assertEquals(2, context.getDurableExecutionContext().getActionState().getCallResultCount()); + assertTrue( + context.getDurableExecutionContext() + .getActionState() + .getCallResults() + .get(0) + .isSuccess()); + assertTrue( + context.getDurableExecutionContext() + .getActionState() + .getCallResults() + .get(1) + .isSuccess()); + } + @Test void testDurableExecuteReconcilableMismatchStartsNewCall() throws Exception { ActionState actionState = new ActionState(null); From 8ebeadd9bd616928b4c459a1b0af70ab4d2e437e Mon Sep 17 00:00:00 2001 From: daken Date: Thu, 30 Jul 2026 20:36:00 +0800 Subject: [PATCH 13/13] fix test --- .../agents/integration/test/AsyncExecutionAgent.java | 11 ++++++++++- python/flink_agents/plan/actions/tool_call_action.py | 1 - 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java index dde280a70..200b17d3f 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java @@ -103,7 +103,16 @@ public ChatMessage chat( List messages, List tools, Map modelParams) { ChatMessage lastMessage = messages.get(messages.size() - 1); if (lastMessage.getRole() == MessageRole.TOOL) { - return new ChatMessage(MessageRole.ASSISTANT, lastMessage.getContent()); + StringBuilder aggregated = new StringBuilder(); + for (ChatMessage message : messages) { + if (message.getRole() == MessageRole.TOOL) { + if (aggregated.length() > 0) { + aggregated.append('|'); + } + aggregated.append(message.getContent()); + } + } + return new ChatMessage(MessageRole.ASSISTANT, aggregated.toString()); } String requestId = lastMessage.getContent(); diff --git a/python/flink_agents/plan/actions/tool_call_action.py b/python/flink_agents/plan/actions/tool_call_action.py index df42c35b4..3b432d6bb 100644 --- a/python/flink_agents/plan/actions/tool_call_action.py +++ b/python/flink_agents/plan/actions/tool_call_action.py @@ -31,7 +31,6 @@ from flink_agents.plan.actions.action import Action from flink_agents.plan.function import PythonFunction from flink_agents.plan.tools.function_tool import FunctionTool - from flink_agents.runtime.durable_execution import durable_identity_for_call _logger = logging.getLogger(__name__)