This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new eca1969fd [runtime][java][python] Release subagent handles after await
failures (#1320)
eca1969fd is described below
commit eca1969fde826bfbfc41f146e51b5b1b7786217b
Author: Wenjin Xie <[email protected]>
AuthorDate: Fri Oct 23 16:04:50 2026 +0800
[runtime][java][python] Release subagent handles after await failures
(#1320)
Untrack deferred and async subagent handles in finally blocks so handled
model-driven await failures can reach action finalization without leaving
pending calls. Preserve caller-driven exceptions, cancellation and
interruption.
Document the await failure contract and cover model dispatch, lifecycle
finalization and direct exception propagation in both languages.
Fixes #1279
Generated-by: Codex 0.153.4 (GPT-6)
Co-authored-by: Codex <[email protected]>
---
.../flink/agents/api/subagent/SubagentFuture.java | 11 +-
python/flink_agents/api/subagent.py | 10 +-
python/flink_agents/runtime/async_subagent.py | 28 +--
python/flink_agents/runtime/deferred_subagent.py | 21 +-
.../runtime/tests/test_subagent_await_cleanup.py | 163 ++++++++++++++
.../runtime/subagent/AsyncSubagentFuture.java | 16 +-
.../runtime/subagent/DeferredSubagentFuture.java | 14 +-
.../runtime/subagent/SubagentAwaitCleanupTest.java | 247 +++++++++++++++++++++
8 files changed, 473 insertions(+), 37 deletions(-)
diff --git
a/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFuture.java
b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFuture.java
index 38e2eb8cd..005c2c2be 100644
--- a/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFuture.java
+++ b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFuture.java
@@ -44,8 +44,15 @@ public abstract class SubagentFuture {
public abstract boolean isDone();
/**
- * Resolves the invocation, waiting until it reaches a terminal state.
Failures converge into a
- * failed {@link SubagentResult} rather than a separately reported
exceptional completion.
+ * Resolves the invocation, waiting until it reaches a terminal state. A
comprehensible
+ * sub-agent failure is returned as a failed {@link SubagentResult};
system-level exceptions,
+ * cancellation and interruption propagate to the caller.
+ *
+ * <p>Do not retry an await that throws. Any retry belongs inside its
await implementation.
+ * Successful repeated awaits may return the cached result. The built-in
deferred and async
+ * handles release their pending-call registration when the await exits,
including
+ * exceptionally; callers need no additional cleanup to finish an action
after handling an await
+ * failure.
*/
public abstract SubagentResult await() throws Exception;
diff --git a/python/flink_agents/api/subagent.py
b/python/flink_agents/api/subagent.py
index 10b62f9e9..d08d56179 100644
--- a/python/flink_agents/api/subagent.py
+++ b/python/flink_agents/api/subagent.py
@@ -176,8 +176,14 @@ class SubagentFuture(ABC):
def __await__(self) -> Any:
"""Resolve the invocation, waiting until it reaches a terminal state.
- Failures converge into a failed :class:`SubagentResult` rather than a
- separately raised exception.
+ Comprehensible sub-agent failures return a failed
:class:`SubagentResult`;
+ system-level exceptions, cancellation and interruption propagate.
+
+ Do not retry an await that raises. Any retry belongs inside its
implementation.
+ Successful repeated awaits may return the cached result. The built-in
deferred
+ and async handles release their pending-call registration when the
await
+ exits, including exceptionally, so
+ callers need no extra cleanup after handling an await failure.
"""
diff --git a/python/flink_agents/runtime/async_subagent.py
b/python/flink_agents/runtime/async_subagent.py
index 585de4919..3a88d4c0c 100644
--- a/python/flink_agents/runtime/async_subagent.py
+++ b/python/flink_agents/runtime/async_subagent.py
@@ -143,19 +143,21 @@ class AsyncSubagentFuture(SubagentFuture):
msg = f"Sub-agent call cancelled: {self.identity}"
raise CancelledError(msg)
if not self._consumed:
- # Build the durable awaitable first, then yield from it, so the
- # await composition and its durable execution cannot be misread
- # as one serial call.
- awaitable = self._ctx.durable_execute_async(
- self._setup._await_until_terminal,
- self.session_id,
- self.call_id,
- durable_id=f"{self.identity}#await",
- )
- self._value = yield from awaitable.__await__()
- self._consumed = True
- if self._registry is not None:
- self._registry.untrack_pending_subagent_call(self.identity)
+ try:
+ # Build the durable awaitable first, then yield from it, so the
+ # await composition and its durable execution cannot be misread
+ # as one serial call.
+ awaitable = self._ctx.durable_execute_async(
+ self._setup._await_until_terminal,
+ self.session_id,
+ self.call_id,
+ durable_id=f"{self.identity}#await",
+ )
+ self._value = yield from awaitable.__await__()
+ self._consumed = True
+ finally:
+ if self._registry is not None:
+ self._registry.untrack_pending_subagent_call(self.identity)
return self._value
def cancel(self) -> None:
diff --git a/python/flink_agents/runtime/deferred_subagent.py
b/python/flink_agents/runtime/deferred_subagent.py
index 68024b604..43432cd28 100644
--- a/python/flink_agents/runtime/deferred_subagent.py
+++ b/python/flink_agents/runtime/deferred_subagent.py
@@ -96,13 +96,18 @@ class DeferredSubagentFuture(SubagentFuture):
A system-level failure escaping durable execution propagates and fails
the action instead of being folded into an error result.
"""
- durable_id, call, reconcile = self.prepare()
- value = yield from self._ctx.durable_execute_async(
- call,
- reconciler=reconcile,
- durable_id=durable_id,
- ).__await__()
- self._resolve(value)
+ try:
+ durable_id, call, reconcile = self.prepare()
+ value = yield from self._ctx.durable_execute_async(
+ call,
+ reconciler=reconcile,
+ durable_id=durable_id,
+ ).__await__()
+ self._resolve(value)
+ finally:
+ # The dispatcher may handle the exception and finish the action.
+ if self._registry is not None:
+ self._registry.untrack_pending_subagent_call(self.identity)
def combine(self, *others: SubagentFuture) -> SubagentFutures:
"""Group this handle with others for a batched resolve."""
@@ -120,8 +125,6 @@ class DeferredSubagentFuture(SubagentFuture):
def _resolve(self, value: SubagentResult) -> None:
self._value = value
self._done = True
- if self._registry is not None:
- self._registry.untrack_pending_subagent_call(self.identity)
class DeferredSubagentSetup(BaseSubagentSetup, ABC):
diff --git a/python/flink_agents/runtime/tests/test_subagent_await_cleanup.py
b/python/flink_agents/runtime/tests/test_subagent_await_cleanup.py
new file mode 100644
index 000000000..ff2c92bc2
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_subagent_await_cleanup.py
@@ -0,0 +1,163 @@
+################################################################################
+# 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.
+################################################################################
+"""Await failures must not leave handles pending at action finalization."""
+
+import asyncio
+from concurrent.futures import CancelledError
+from typing import Any
+
+import pytest
+from pydantic import PrivateAttr
+
+from flink_agents.api.subagent import CALLABLE_NAME_PREFIX, SubagentResult
+from flink_agents.plan.actions.tool_call_action import process_tool_request
+from flink_agents.plan.tests.actions.test_tool_call_action_subagent import (
+ _Context,
+ _request,
+)
+from flink_agents.runtime.async_subagent import BaseAsyncSubagentSetup,
RunStatus
+from flink_agents.runtime.deferred_subagent import DeferredSubagentSetup
+from flink_agents.runtime.tests.test_base_subagent import _FakeTask, _run
+
+
+class _AwaitContext(_Context):
+ def durable_execute_async(self, func, *args: Any, **kwargs: Any) -> Any:
+ kwargs.pop("durable_id", None)
+ kwargs.pop("reconciler", None)
+
+ async def execute() -> Any:
+ return func(*args, **kwargs)
+
+ return execute()
+
+
+class _DeferredSetup(DeferredSubagentSetup):
+ _failure: BaseException | None = PrivateAttr(default=None)
+ _stage: str = PrivateAttr(default="call")
+
+ def prepare(self, ctx, prompt, session_id, call_id) -> tuple:
+ if self._stage == "prepare":
+ raise self._failure
+
+ def call() -> SubagentResult:
+ if self._failure is not None:
+ raise self._failure
+ return SubagentResult.error("remote declined")
+
+ return f"{session_id}#{call_id}", call, None
+
+
+class _AsyncSetup(BaseAsyncSubagentSetup):
+ _failure: BaseException | None = PrivateAttr(default=None)
+ _stage: str = PrivateAttr(default="status")
+
+ def call_submit_request(self, session_id, call_id, prompt) -> None:
+ pass
+
+ def call_query_status(self, session_id, call_id) -> RunStatus:
+ if self._failure is not None and self._stage == "status":
+ raise self._failure
+ return RunStatus.completed()
+
+ def call_fetch_result(self, session_id, call_id) -> SubagentResult:
+ if self._failure is not None:
+ raise self._failure
+ return SubagentResult.error("remote declined")
+
+
+def _setup(mode, stage, failure) -> _DeferredSetup | _AsyncSetup:
+ setup = _DeferredSetup() if mode == "deferred" else _AsyncSetup()
+ setup._stage = stage
+ setup._failure = failure
+ setup.set_subagent_name("reviewer")
+ return setup
+
+
[email protected]("parallel", [False, True])
[email protected](
+ ("mode", "stage"),
+ [
+ ("deferred", "prepare"),
+ ("deferred", "call"),
+ ("async", "status"),
+ ("async", "fetch"),
+ ],
+)
[email protected]("raises", [False, True])
+def test_model_handled_failure_allows_action_finalization(
+ mode, stage, parallel, raises
+):
+ failure = RuntimeError("remote unavailable") if raises else None
+ # A prepare failure has no result-returning variant.
+ setup = _setup(mode, stage if raises else "call", failure)
+ task = _FakeTask(action_name="tool_call_action")
+ setup.on_action_prepared(task)
+ ctx = _AwaitContext().with_agent("reviewer", setup)
+ if parallel:
+ ctx.with_parallel_tool_calls()
+ request = _request(
+ "model",
+ [
+ {
+ "id": f"call-{i}",
+ "function": {
+ "name": CALLABLE_NAME_PREFIX + "reviewer",
+ "arguments": {"code": "hello"},
+ },
+ }
+ for i in range(2)
+ ],
+ )
+
+ asyncio.run(process_tool_request(request, ctx))
+
+ response = ctx.sent_events[0]
+ for i in range(2):
+ assert response.success[f"call-{i}"] is False
+ assert (
+ "remote unavailable" if raises else "remote declined"
+ ) in response.error[f"call-{i}"]
+ # This is the lifecycle check that used to fail after dispatch returned.
+ setup.on_action_finishing(task)
+
+
[email protected](
+ ("mode", "stage"),
+ [
+ ("deferred", "prepare"),
+ ("deferred", "call"),
+ ("async", "status"),
+ ("async", "fetch"),
+ ],
+)
[email protected](
+ "failure_type",
+ [RuntimeError, CancelledError, asyncio.CancelledError, InterruptedError],
+)
+def test_direct_await_preserves_exception_and_releases_handle(
+ mode, stage, failure_type
+):
+ failure = failure_type("await failed")
+ setup = _setup(mode, stage, failure)
+ task = _FakeTask()
+ setup.on_action_prepared(task)
+ future = _run(setup.submit(_AwaitContext(), "hello"))
+ with pytest.raises(failure_type) as caught:
+ _run(future)
+ assert caught.value is failure
+ setup.on_action_finishing(task)
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/AsyncSubagentFuture.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/AsyncSubagentFuture.java
index 254e11896..707a612a5 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/AsyncSubagentFuture.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/AsyncSubagentFuture.java
@@ -84,12 +84,16 @@ final class AsyncSubagentFuture extends SubagentFuture {
"Sub-agent call cancelled: " + getSessionId() + "#" +
getCallId());
}
if (!consumed) {
- DurableCallable<SubagentResult> awaitCall =
- setup.awaitResult(ctx, getSessionId(), getCallId());
- value = ctx.durableExecuteAsync(awaitCall).await();
- consumed = true;
- if (registry != null) {
- registry.untrackPendingSubagentCall(identity());
+ try {
+ DurableCallable<SubagentResult> awaitCall =
+ setup.awaitResult(ctx, getSessionId(), getCallId());
+ value = ctx.durableExecuteAsync(awaitCall).await();
+ consumed = true;
+ } finally {
+ // Preserve the exception while releasing the action's
pending-handle obligation.
+ if (registry != null) {
+ registry.untrackPendingSubagentCall(identity());
+ }
}
}
return value;
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentFuture.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentFuture.java
index a6e62fed8..64b30308b 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentFuture.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentFuture.java
@@ -76,8 +76,15 @@ public final class DeferredSubagentFuture extends
SubagentFuture implements Phas
*/
@Override
public void execute() throws Exception {
- prepare();
- complete(ctx.durableExecuteAsync(prepared).await());
+ try {
+ prepare();
+ complete(ctx.durableExecuteAsync(prepared).await());
+ } finally {
+ // Await failures may be handled by ToolCallAction; they still
resolve this obligation.
+ if (registry != null) {
+ registry.untrackPendingSubagentCall(identity());
+ }
+ }
}
/**
@@ -103,9 +110,6 @@ public final class DeferredSubagentFuture extends
SubagentFuture implements Phas
private void complete(SubagentResult outcome) {
this.value = outcome;
this.done = true;
- if (registry != null) {
- registry.untrackPendingSubagentCall(identity());
- }
}
@Override
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/SubagentAwaitCleanupTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/SubagentAwaitCleanupTest.java
new file mode 100644
index 000000000..fe83e787b
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/SubagentAwaitCleanupTest.java
@@ -0,0 +1,247 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.agents.runtime.subagent;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.OutputEvent;
+import org.apache.flink.agents.api.agents.Agent;
+import org.apache.flink.agents.api.agents.AgentExecutionOptions;
+import org.apache.flink.agents.api.chat.messages.ToolCallBlock;
+import org.apache.flink.agents.api.context.DurableCallable;
+import org.apache.flink.agents.api.context.DurableFuture;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.event.ToolRequestEvent;
+import org.apache.flink.agents.api.event.ToolResponseEvent;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.agents.api.subagent.SubagentFuture;
+import org.apache.flink.agents.api.subagent.SubagentResult;
+import org.apache.flink.agents.plan.AgentConfiguration;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.runtime.operator.ActionExecutionOperator;
+import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorFactory;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CancellationException;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/** Exercises handled await failures through dispatch and operator action
finalization. */
+public class SubagentAwaitCleanupTest {
+ @ParameterizedTest
+ @CsvSource({
+ "deferred,prepare,false", "deferred,prepare,true",
+ "deferred,call,false", "deferred,call,true",
+ "async,status,false", "async,status,true",
+ "async,fetch,false", "async,fetch,true",
+ "deferred,result,false", "deferred,result,true",
+ "async,result,false", "async,result,true"
+ })
+ @Timeout(30)
+ void modelHandledFailureAllowsActionToFinish(String mode, String stage,
boolean parallel)
+ throws Exception {
+ Agent agent = new Agent();
+ agent.addResource(
+ "child",
+ ResourceType.AGENT,
+ ResourceDescriptor.Builder.newBuilder(
+ (mode.equals("deferred") ? DeferredSetup.class
: AsyncSetup.class)
+ .getName())
+ .addInitialArgument("stage", stage)
+ .build());
+ agent.addAction(
+ new String[] {InputEvent.EVENT_TYPE},
+ SubagentAwaitCleanupTest.class.getMethod(
+ "request", Event.class, RunnerContext.class));
+ agent.addAction(
+ new String[] {ToolResponseEvent.EVENT_TYPE},
+ SubagentAwaitCleanupTest.class.getMethod(
+ "receive", Event.class, RunnerContext.class));
+ AgentConfiguration config = new AgentConfiguration();
+ config.set(AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED, false);
+ config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, parallel);
+ config.set(AgentExecutionOptions.ASYNC_BATCH_PARALLELISM, 2);
+ AgentPlan compiled = new AgentPlan(agent, config);
+ Map<String, org.apache.flink.agents.plan.actions.Action> actions =
+ new LinkedHashMap<>(compiled.getActions());
+ actions.remove("chat_model_action");
+ AgentPlan plan = new AgentPlan(actions,
compiled.getResourceProviders(), config);
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
harness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<Long, Object>(plan,
true),
+ (KeySelector<Long, Long>) value -> 0L,
+ TypeInformation.of(Long.class))) {
+ harness.open();
+ harness.processElement(new StreamRecord<>(1L));
+ ((ActionExecutionOperator<?, ?>)
harness.getOperator()).waitInFlightEventsFinished();
+ assertThat(harness.getRecordOutput())
+ .extracting(StreamRecord::getValue)
+ .containsExactly("handled");
+ }
+ }
+
+ @ParameterizedTest
+ @CsvSource({
+ "deferred,error",
+ "deferred,cancel",
+ "deferred,interrupt",
+ "async,error",
+ "async,cancel",
+ "async,interrupt"
+ })
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ void directAwaitPreservesFailureAndReleasesHandle(String mode, String
failureKind)
+ throws Exception {
+ Exception failure =
+ failureKind.equals("cancel")
+ ? new CancellationException("await failed")
+ : failureKind.equals("interrupt")
+ ? new InterruptedException("await failed")
+ : new IllegalStateException("await failed");
+ RunnerContext ctx = mock(RunnerContext.class);
+ when(ctx.durableExecuteAsync(any(DurableCallable.class)))
+ .thenReturn(
+ (DurableFuture)
+ () -> {
+ throw failure;
+ });
+ PendingSubagentCallRegistry registry = new
PendingSubagentCallRegistry("caller");
+ ResourceDescriptor descriptor =
+ ResourceDescriptor.Builder.newBuilder(
+ (mode.equals("deferred") ? DeferredSetup.class
: AsyncSetup.class)
+ .getName())
+ .addInitialArgument("stage", "result")
+ .build();
+ SubagentFuture future =
+ mode.equals("deferred")
+ ? new DeferredSubagentFuture(
+ "s",
+ "c",
+ ctx,
+ registry,
+ () ->
+ new DeferredSetup(descriptor, null)
+ .prepare(ctx, "hello", "s",
"c"))
+ : new AsyncSubagentFuture(
+ new AsyncSetup(descriptor, null), ctx, "s",
"c", registry);
+ assertThatThrownBy(future::await).isSameAs(failure);
+ registry.checkEmpty();
+ }
+
+ public static void request(Event event, RunnerContext ctx) {
+ ctx.sendEvent(
+ new ToolRequestEvent(
+ "model",
+ List.of(
+ new ToolCallBlock(
+ "call-1", "_subagent_child",
Map.of("question", "hello")),
+ new ToolCallBlock(
+ "call-2",
+ "_subagent_child",
+ Map.of("question", "hello")))));
+ }
+
+ public static void receive(Event event, RunnerContext ctx) {
+ ToolResponseEvent response = ToolResponseEvent.fromEvent(event);
+ for (String id : List.of("call-1", "call-2")) {
+ assertThat(response.getSuccess().get(id)).isFalse();
+ assertThat(response.getError().get(id)).contains("remote
unavailable");
+ }
+ ctx.sendEvent(new OutputEvent("handled"));
+ }
+
+ public static class DeferredSetup extends BaseDeferredSubagentSetup {
+ private final String stage;
+
+ public DeferredSetup(ResourceDescriptor descriptor, ResourceContext
context) {
+ super(descriptor, context);
+ stage = descriptor.getArgument("stage");
+ }
+
+ @Override
+ protected DurableCallable<SubagentResult> prepare(
+ RunnerContext ctx, Object prompt, String sessionId, String
callId) {
+ if (stage.equals("prepare")) {
+ throw new IllegalStateException("remote unavailable");
+ }
+ return new DurableCallable<SubagentResult>() {
+ @Override
+ public String getId() {
+ return sessionId + "#" + callId;
+ }
+
+ @Override
+ public Class<SubagentResult> getResultClass() {
+ return SubagentResult.class;
+ }
+
+ @Override
+ public SubagentResult call() {
+ if (stage.equals("call")) {
+ throw new IllegalStateException("remote unavailable");
+ }
+ return SubagentResult.error("remote unavailable");
+ }
+ };
+ }
+ }
+
+ public static class AsyncSetup extends BaseAsyncSubagentSetup {
+ private final String stage;
+
+ public AsyncSetup(ResourceDescriptor descriptor, ResourceContext
context) {
+ super(descriptor, context);
+ stage = descriptor.getArgument("stage");
+ }
+
+ @Override
+ protected void callSubmitRequest(String sessionId, String callId,
Object prompt) {}
+
+ @Override
+ protected RunStatus callQueryStatus(String sessionId, String callId) {
+ if (stage.equals("status")) {
+ throw new IllegalStateException("remote unavailable");
+ }
+ return RunStatus.completed();
+ }
+
+ @Override
+ protected SubagentResult callFetchResult(String sessionId, String
callId) {
+ if (stage.equals("fetch")) {
+ throw new IllegalStateException("remote unavailable");
+ }
+ return SubagentResult.error("remote unavailable");
+ }
+ }
+}
