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
commit a30eab5503d6baf1983ffa795ef22373a4f7d2a7 Author: WenjinXie <[email protected]> AuthorDate: Tue Oct 20 11:12:12 2026 +0800 [runtime][python] Support root resource fallback for internal sub-agents Resolve missing child resources from the root while preserving child overrides and root lifecycle ownership. Cover scope creation, shared resources, nested calls and ReAct prompt isolation, and document the shared lookup rule. Generated-by: Codex CLI 0.153.4 (GPT-6) Co-authored-by: Codex <[email protected]> --- .../docs/development/multi-agent/subagents.md | 7 +- .../react_agent_subagent_test.py | 61 ++++++--- .../flink_agents/runtime/flink_runner_context.py | 8 +- python/flink_agents/runtime/resource_cache.py | 7 + .../runtime/tests/test_eager_materialize.py | 22 ++++ .../tests/test_resource_cache_inheritance.py | 144 +++++++++++++++++++++ 6 files changed, 229 insertions(+), 20 deletions(-) diff --git a/docs/content/docs/development/multi-agent/subagents.md b/docs/content/docs/development/multi-agent/subagents.md index cbd696f6b..aa3e2cd13 100644 --- a/docs/content/docs/development/multi-agent/subagents.md +++ b/docs/content/docs/development/multi-agent/subagents.md @@ -135,9 +135,10 @@ While a call runs, the child behaves like a scoped agent of its own: - Its `OutputEvent`s are accumulated, in order, into the call's result instead of landing on the parent's context. See [the result](#the-result) for the shape. -- It resolves resources against its own plan. Register the resources a child uses, - such as its chat model connection, on the child itself. In Python a resource - registered solely on the parent is not visible to the child. +- The child first looks for resources registered on itself, then falls back to + resources registered on the root agent. You can share a chat model connection + by registering it on the root, or override it on a child using the same name + and resource type. This also applies to nested sub-agents. - Its sensory and short-term memory start empty for each call, including calls continuing a session, and are shared by the child's actions within that call. Long-term memory, when configured, is shared with the caller. diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_subagent_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_subagent_test.py index c7e8681ed..317204106 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_subagent_test.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_subagent_test.py @@ -39,6 +39,7 @@ import sysconfig from pathlib import Path from typing import Any +import pytest from pyflink.common import Encoder from pyflink.common.typeinfo import Types from pyflink.datastream import KeySelector, StreamExecutionEnvironment @@ -61,6 +62,7 @@ from flink_agents.api.chat_result import ChatResult from flink_agents.api.decorators import action from flink_agents.api.events.event import Event, InputEvent, OutputEvent from flink_agents.api.execution_environment import AgentsExecutionEnvironment +from flink_agents.api.prompts.prompt import Prompt from flink_agents.api.resource import ResourceDescriptor, ResourceType from flink_agents.api.runner_context import RunnerContext from flink_agents.api.subagent import InternalSubagentResourceDescriptorBuilder @@ -172,13 +174,8 @@ class ToolCallingChildConnection(BaseChatModelConnection): ) -def _child_react_agent() -> ReActAgent: - """A ReActAgent with neither a prompt nor an output schema. - - The connection is registered on the child itself, keeping the child plan - self-contained: a child scope resolves resources against the child plan - only, so the connection must live there. - """ +def _child_react_agent(*, local_connection: bool = True) -> ReActAgent: + """A bare ReActAgent using either its own or the root's connection.""" child = ReActAgent( chat_model=ResourceDescriptor( clazz=f"{ScriptedModelSetup.__module__}.{ScriptedModelSetup.__name__}", @@ -186,14 +183,19 @@ def _child_react_agent() -> ReActAgent: model="mock-model", ), ) - child.add_resource( + if local_connection: + _register_child_connection(child) + return child + + +def _register_child_connection(agent: Agent) -> None: + agent.add_resource( "child_connection", ResourceType.CHAT_MODEL_CONNECTION, ResourceDescriptor( clazz=f"{ChildScriptedConnection.__module__}.{ChildScriptedConnection.__name__}", ), ) - return child def _tool_calling_child_react_agent() -> ReActAgent: @@ -234,8 +236,8 @@ class CallerRootAgent(Agent): ctx.send_event(OutputEvent(output=msg)) -def _delegating_parent() -> ReActAgent: - """A bare ReActAgent parent that declares the child to its model. +def _delegating_parent(*, shared_connection: bool = False) -> ReActAgent: + """A prompted ReActAgent parent that declares the child to its model. Sub-classing ReActAgent would trip the plan compiler, which rejects inherited ``@action`` methods on a concrete agent, so the parent is a @@ -248,7 +250,10 @@ def _delegating_parent() -> ReActAgent: model="mock-model", subagents=[CHILD_SCOPE], ), + prompt=Prompt.from_text("PARENT ONLY: {input}"), ) + if shared_connection: + _register_child_connection(parent) parent.add_resource( "parent_connection", ResourceType.CHAT_MODEL_CONNECTION, @@ -261,7 +266,7 @@ def _delegating_parent() -> ReActAgent: CHILD_SCOPE, ResourceType.AGENT, InternalSubagentResourceDescriptorBuilder() - .agent(_child_react_agent()) + .agent(_child_react_agent(local_connection=not shared_connection)) .description("Reviews a code diff.") .input_schema( '{"type":"object","properties":{"input":{"type":"string"}},' @@ -297,10 +302,28 @@ def _run_and_collect(root: Agent, tmp_path: Path) -> str: return "".join(p.read_text() for p in result_dir.rglob("*") if p.is_file()) -def test_caller_driven_react_subagent(tmp_path: Path) -> None: [email protected]("connection_scope", ["child", "root", "nested_root", "override"]) +def test_caller_driven_react_subagent(tmp_path: Path, connection_scope: str) -> None: """A submitted string runs the child's ReAct loop and returns its answer.""" root = CallerRootAgent() - root.add_resource(CHILD_SCOPE, ResourceType.AGENT, _child_react_agent()) + child = _child_react_agent(local_connection=connection_scope in {"child", "override"}) + if connection_scope in {"root", "nested_root"}: + _register_child_connection(root) + if connection_scope == "override": + # This connection would ask for an unavailable tool if the child failed + # to prefer its own connection registered under the same name and type. + root.add_resource( + "child_connection", + ResourceType.CHAT_MODEL_CONNECTION, + ResourceDescriptor( + clazz=f"{ToolCallingChildConnection.__module__}.{ToolCallingChildConnection.__name__}", + ), + ) + if connection_scope == "nested_root": + middle = CallerRootAgent() + middle.add_resource(CHILD_SCOPE, ResourceType.AGENT, child) + child = middle + root.add_resource(CHILD_SCOPE, ResourceType.AGENT, child) contents = _run_and_collect(root, tmp_path) @@ -311,9 +334,14 @@ def test_caller_driven_react_subagent(tmp_path: Path) -> None: ) -def test_model_driven_react_subagent_delegation(tmp_path: Path) -> None: [email protected]("shared_connection", [False, True]) +def test_model_driven_react_subagent_delegation( + tmp_path: Path, shared_connection: bool +) -> None: """A model-issued delegation reaches the declared child and closes the loop.""" - contents = _run_and_collect(_delegating_parent(), tmp_path) + contents = _run_and_collect( + _delegating_parent(shared_connection=shared_connection), tmp_path + ) # The final answer echoes the tool observation, proving the whole loop: the # model's {"input": ...} arguments became the child's user message through @@ -322,6 +350,7 @@ def test_model_driven_react_subagent_delegation(tmp_path: Path) -> None: # this parent answer does -- so both fragments here mean the round-trip # completed; the assertion stays split because the observation carries the # result in the framework's string form, not verbatim. + assert "PARENT ONLY" not in contents assert "observed:" in contents, f"no observation echoed; collected: {contents!r}" assert "child saw: review the diff" in contents, ( f"child output missing from the loop; collected: {contents!r}" diff --git a/python/flink_agents/runtime/flink_runner_context.py b/python/flink_agents/runtime/flink_runner_context.py index e1ee0f9f4..efd88b1da 100644 --- a/python/flink_agents/runtime/flink_runner_context.py +++ b/python/flink_agents/runtime/flink_runner_context.py @@ -749,7 +749,13 @@ class FlinkRunnerContext(RunnerContext, ExecutionReporter): from flink_agents.plan.agent_plan import AgentPlan scoped_plan = AgentPlan.model_validate_json(plan_json) - cache = ResourceCache(scoped_plan.resource_providers, scoped_plan.config) + # Match Java's child scopes: local resources override resources + # shared by the root, including for nested sub-agents. + cache = ResourceCache( + scoped_plan.resource_providers, + scoped_plan.config, + parent=self.__resource_cache, + ) cache.set_java_resource_adapter(self.__j_resource_adapter) entry = (scoped_plan, cache) self.__scoped_resource_caches[plan_json] = entry diff --git a/python/flink_agents/runtime/resource_cache.py b/python/flink_agents/runtime/resource_cache.py index 4dfc9c6e6..08ad891b4 100644 --- a/python/flink_agents/runtime/resource_cache.py +++ b/python/flink_agents/runtime/resource_cache.py @@ -41,6 +41,7 @@ class ResourceCache: self, resource_providers: Dict[ResourceType, Dict[str, ResourceProvider]], config: AgentConfiguration | None = None, + parent: "ResourceCache | None" = None, ) -> None: """Create a ResourceCache from the given resource providers and config. @@ -50,9 +51,13 @@ class ResourceCache: Two-level mapping of resource type to resource name to provider. config : AgentConfiguration | None Agent configuration passed to providers during resource creation. + parent : ResourceCache | None + Fallback cache for resources not declared in this scope. Resources + resolved there remain owned by that cache, including their lifecycle. """ self._resource_providers = resource_providers or {} self._config = config + self._parent = parent self._cache: Dict[ResourceType, Dict[str, Resource]] = {} self._j_resource_adapter: Any = None self._resource_context = ResourceContextImpl(self) @@ -86,6 +91,8 @@ class ResourceCache: return cached providers = self._resource_providers.get(type) if providers is None or name not in providers: + if self._parent is not None: + return self._parent.get_resource(name, type) msg = f"Resource not found: '{name}' of type {type}" raise KeyError(msg) resource_provider = providers[name] diff --git a/python/flink_agents/runtime/tests/test_eager_materialize.py b/python/flink_agents/runtime/tests/test_eager_materialize.py index a54562f6a..35dfa02ea 100644 --- a/python/flink_agents/runtime/tests/test_eager_materialize.py +++ b/python/flink_agents/runtime/tests/test_eager_materialize.py @@ -24,14 +24,17 @@ without a live interpreter. """ from typing import Any +from unittest.mock import Mock from flink_agents.api.resource import ResourceType +from flink_agents.plan.agent_plan import AgentPlan from flink_agents.plan.resource_provider import ( JavaResourceProvider, PythonResourceProvider, PythonSerializableResourceProvider, ) from flink_agents.runtime.flink_runner_context import FlinkRunnerContext +from flink_agents.runtime.resource_cache import ResourceCache class _StubResourceCache: @@ -205,3 +208,22 @@ def test_close_closes_the_root_and_every_scoped_cache() -> None: assert scoped_a.closed is True assert scoped_b.closed is True assert ctx._FlinkRunnerContext__scoped_resource_caches == {} + + +def test_new_child_scopes_can_resolve_root_resources() -> None: + """Lazy lookup and eager materialization share the root-backed child cache.""" + ctx, root_cache = _context({}) + ctx._FlinkRunnerContext__scoped_resource_caches = {} + ctx._FlinkRunnerContext__j_resource_adapter = None + ctx._j_runner_context = Mock() + + for name in ("child", "nested_child"): + plan_json = AgentPlan(actions={}, agent_name=name).model_dump_json() + # Inherited resources must not be materialized or owned by the child. + assert ctx.eager_materialize(ResourceType.CHAT_MODEL.value, plan_json) == {} + ctx._j_runner_context.getActiveScopePlanJson.return_value = plan_json + cache = ctx._FlinkRunnerContext__active_resource_cache() + assert isinstance(cache, ResourceCache) + assert cache.get_resource("shared", ResourceType.CHAT_MODEL) == "resource:shared" + + assert root_cache.resolved == [("shared", ResourceType.CHAT_MODEL)] * 2 diff --git a/python/flink_agents/runtime/tests/test_resource_cache_inheritance.py b/python/flink_agents/runtime/tests/test_resource_cache_inheritance.py new file mode 100644 index 000000000..323ef544b --- /dev/null +++ b/python/flink_agents/runtime/tests/test_resource_cache_inheritance.py @@ -0,0 +1,144 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################ +"""Shared resource lookup, overriding, and lifecycle ownership.""" + +from unittest.mock import Mock + +import pytest + +from flink_agents.api.resource import Resource, ResourceType +from flink_agents.plan.resource_provider import ResourceProvider +from flink_agents.runtime.resource_cache import ResourceCache + + +def _provider(resource: Resource) -> Mock: + provider = Mock(spec=ResourceProvider) + provider.provide.return_value = resource + return provider + + +def test_siblings_share_one_parent_instance_without_owning_it() -> None: + resource = Mock(spec=Resource) + provider = _provider(resource) + root = ResourceCache({ResourceType.TOOL: {"shared": provider}}) + children = [ResourceCache({}, parent=root) for _ in range(2)] + + for child in children: + assert child.get_resource("shared", ResourceType.TOOL) is resource + assert child.get_resource("shared", ResourceType.TOOL) is resource + child.close() + + provider.provide.assert_called_once() + resource.open.assert_called_once() + resource.close.assert_not_called() + assert root.get_resource("shared", ResourceType.TOOL) is resource + root.close() + resource.close.assert_called_once() + + +def test_child_override_and_dependency_lookup_use_child_scope() -> None: + root_resource, child_resource, dependent = [Mock(spec=Resource) for _ in range(3)] + root_provider = _provider(root_resource) + child_provider = _provider(child_resource) + dependent_provider = _provider(dependent) + + def build_dependent(*, resource_context, config) -> Resource: + assert ( + resource_context.get_resource("shared", ResourceType.TOOL) is child_resource + ) + return dependent + + dependent_provider.provide.side_effect = build_dependent + root = ResourceCache({ResourceType.TOOL: {"shared": root_provider}}) + child = ResourceCache( + { + ResourceType.TOOL: { + "shared": child_provider, + "dependent": dependent_provider, + } + }, + parent=root, + ) + # Even an already cached parent instance must not hide a child's provider. + assert root.get_resource("shared", ResourceType.TOOL) is root_resource + assert child.get_resource("dependent", ResourceType.TOOL) is dependent + assert child.get_resource("shared", ResourceType.TOOL) is child_resource + child.close() + child_resource.close.assert_called_once() + dependent.close.assert_called_once() + root_resource.close.assert_not_called() + root.close() + + +def test_missing_name_in_existing_type_falls_back_but_wrong_type_does_not() -> None: + resource = Mock(spec=Resource) + root = ResourceCache({ResourceType.TOOL: {"shared": _provider(resource)}}) + child = ResourceCache( + {ResourceType.TOOL: {"other": _provider(Mock())}}, parent=root + ) + assert child.get_resource("shared", ResourceType.TOOL) is resource + with pytest.raises(KeyError, match="Resource not found"): + child.get_resource("shared", ResourceType.CHAT_MODEL) + with pytest.raises(KeyError, match="Resource not found"): + root.get_resource("other", ResourceType.TOOL) + child.close() + root.close() + + +def test_shared_resource_dependencies_remain_in_the_root_scope() -> None: + root_dependency, child_dependency, shared = [Mock(spec=Resource) for _ in range(3)] + provider = _provider(shared) + + def build_shared(*, resource_context, config) -> Resource: + assert ( + resource_context.get_resource("dependency", ResourceType.TOOL) + is root_dependency + ) + return shared + + provider.provide.side_effect = build_shared + root = ResourceCache( + { + ResourceType.TOOL: { + "shared": provider, + "dependency": _provider(root_dependency), + } + } + ) + child = ResourceCache( + {ResourceType.TOOL: {"dependency": _provider(child_dependency)}}, parent=root + ) + assert child.get_resource("shared", ResourceType.TOOL) is shared + child_dependency.open.assert_not_called() + root.close() + child.close() + shared.close.assert_called_once() + root_dependency.close.assert_called_once() + + +def test_child_provider_failure_does_not_silently_fall_back() -> None: + root_provider = _provider(Mock(spec=Resource)) + broken_provider = Mock(spec=ResourceProvider) + broken_provider.provide.side_effect = RuntimeError("broken child resource") + root = ResourceCache({ResourceType.TOOL: {"shared": root_provider}}) + child = ResourceCache({ResourceType.TOOL: {"shared": broken_provider}}, parent=root) + with pytest.raises(RuntimeError, match="broken child resource"): + child.get_resource("shared", ResourceType.TOOL) + root_provider.provide.assert_not_called() + child.close() + root.close()
