diff --git a/examples/runtime-configuration/agent.py b/examples/runtime-configuration/agent.py new file mode 100644 index 00000000..f40fbb19 --- /dev/null +++ b/examples/runtime-configuration/agent.py @@ -0,0 +1,41 @@ +import asyncio + +from fast_agent import FastAgent +from fast_agent.core import Core +from fast_agent.mcp_server_registry import ServerRegistry + + +def change_env_of_server(server_name: str, env_name: str, new_env: str): + def change_server_registry(server_registry: ServerRegistry): + server_config = server_registry.registry[server_name] + if server_config.env is None: + server_config.env = {} + server_config.env[env_name] = new_env + return server_registry + + def inner(core: Core): + if core._context is None: + raise ValueError("Context is not initialized") + registry = core._context.server_registry + if registry is None: + raise ValueError("Server registry is not initialized") + registry = change_server_registry(registry) + core._context.server_registry = registry + return core + + return inner + + +fast = FastAgent("RAG Application") + + +@fast.agent(servers=["env_get_server"]) +async def main(token: str) -> None: + change_fn = change_env_of_server("env_get_server", "TOKEN", token) + async with fast.run(core_modifiers=[change_fn]) as agent: + await agent.default.generate("What is the value of 'TOKEN' env variable?") + + +if __name__ == "__main__": + asyncio.run(main("CUSTOM_TOKEN_VALUE_1")) + asyncio.run(main("CUSTOM_TOKEN_VALUE_2")) diff --git a/examples/runtime-configuration/fastagent.config.yaml b/examples/runtime-configuration/fastagent.config.yaml new file mode 100644 index 00000000..bc447fb4 --- /dev/null +++ b/examples/runtime-configuration/fastagent.config.yaml @@ -0,0 +1,15 @@ +execution_engine: asyncio +default_model: gpt-4.1 + +logger: + type: file + level: error + truncate_tools: true + +mcp: + servers: + env_get_server: + command: "uv" + args: ["run", "mcp_server.py"] + env: + TOKEN: "DEFAULT_TOKEN" diff --git a/examples/runtime-configuration/mcp_server.py b/examples/runtime-configuration/mcp_server.py new file mode 100644 index 00000000..91f8e918 --- /dev/null +++ b/examples/runtime-configuration/mcp_server.py @@ -0,0 +1,14 @@ +import os + +from mcp.server.fastmcp import FastMCP + +app = FastMCP(name="ENV-GET server") + + +@app.tool(description="Returns ENV variable specified") +def get_env_var(key: str) -> str: + return os.environ.get(key, "") + + +if __name__ == "__main__": + app.run() diff --git a/src/fast_agent/core/fastagent.py b/src/fast_agent/core/fastagent.py index c9402937..9bb67de1 100644 --- a/src/fast_agent/core/fastagent.py +++ b/src/fast_agent/core/fastagent.py @@ -17,6 +17,7 @@ Awaitable, Callable, Dict, + Iterable, List, Literal, Optional, @@ -85,6 +86,12 @@ from fast_agent.interfaces import AgentProtocol from fast_agent.types import PromptMessageExtended +CORE_MODIFIER_T = Callable[ + [ + Core, + ], + Core, +] F = TypeVar("F", bound=Callable[..., Any]) # For decorated functions logger = get_logger(__name__) @@ -409,7 +416,10 @@ def evaluator_optimizer( evaluator_optimizer = evaluator_optimizer_decorator @asynccontextmanager - async def run(self) -> AsyncIterator["AgentApp"]: + async def run( + self, + core_modifiers: Iterable[CORE_MODIFIER_T] | None = None, + ) -> AsyncIterator["AgentApp"]: """ Context manager for running the application. Initializes all registered agents. @@ -417,6 +427,10 @@ async def run(self) -> AsyncIterator["AgentApp"]: active_agents: Dict[str, AgentProtocol] = {} had_error = False await self.app.initialize() + app_core_copy = self.app + for modifier in core_modifiers or []: + app_core_copy = modifier(app_core_copy) + self.app = app_core_copy # Handle quiet mode and CLI model override safely # Define these *before* they are used, checking if self.args exists and has the attributes diff --git a/tests/integration/mcp_filtering/test_mcp_filtering.py b/tests/integration/mcp_filtering/test_mcp_filtering.py index 0259bbfe..3c699696 100644 --- a/tests/integration/mcp_filtering/test_mcp_filtering.py +++ b/tests/integration/mcp_filtering/test_mcp_filtering.py @@ -11,7 +11,6 @@ @pytest.mark.integration @pytest.mark.asyncio -@pytest.mark.e2e async def test_tool_filtering_basic_agent(fast_agent): """Test tool filtering with basic agent - no filtering vs with filtering""" fast = fast_agent @@ -83,7 +82,6 @@ async def agent_with_filter(): @pytest.mark.integration @pytest.mark.asyncio -@pytest.mark.e2e async def test_resource_filtering_basic_agent(fast_agent): """Test resource filtering with basic agent - no filtering vs with filtering""" fast = fast_agent @@ -151,7 +149,6 @@ async def agent_with_filter(): @pytest.mark.integration @pytest.mark.asyncio -@pytest.mark.e2e async def test_prompt_filtering_basic_agent(fast_agent): """Test prompt filtering with basic agent - no filtering vs with filtering""" fast = fast_agent @@ -214,10 +211,8 @@ async def agent_with_filter(): await agent_with_filter() -@pytest.mark.integration @pytest.mark.integration @pytest.mark.asyncio -@pytest.mark.e2e async def test_tool_filtering_custom_agent(fast_agent): """Test tool filtering with custom agent""" fast = fast_agent @@ -262,7 +257,6 @@ async def custom_string_agent(): @pytest.mark.integration @pytest.mark.asyncio -@pytest.mark.e2e async def test_combined_filtering(fast_agent): """Test combined tool, resource, and prompt filtering""" fast = fast_agent diff --git a/tests/integration/runtime-configuration/fastagent.config.yaml b/tests/integration/runtime-configuration/fastagent.config.yaml new file mode 100644 index 00000000..bc447fb4 --- /dev/null +++ b/tests/integration/runtime-configuration/fastagent.config.yaml @@ -0,0 +1,15 @@ +execution_engine: asyncio +default_model: gpt-4.1 + +logger: + type: file + level: error + truncate_tools: true + +mcp: + servers: + env_get_server: + command: "uv" + args: ["run", "mcp_server.py"] + env: + TOKEN: "DEFAULT_TOKEN" diff --git a/tests/integration/runtime-configuration/mcp_server.py b/tests/integration/runtime-configuration/mcp_server.py new file mode 100644 index 00000000..91f8e918 --- /dev/null +++ b/tests/integration/runtime-configuration/mcp_server.py @@ -0,0 +1,14 @@ +import os + +from mcp.server.fastmcp import FastMCP + +app = FastMCP(name="ENV-GET server") + + +@app.tool(description="Returns ENV variable specified") +def get_env_var(key: str) -> str: + return os.environ.get(key, "") + + +if __name__ == "__main__": + app.run() diff --git a/tests/integration/runtime-configuration/test_runtime_configuration.py b/tests/integration/runtime-configuration/test_runtime_configuration.py new file mode 100644 index 00000000..948cfd89 --- /dev/null +++ b/tests/integration/runtime-configuration/test_runtime_configuration.py @@ -0,0 +1,97 @@ +import pytest + +from fast_agent.core import Core +from fast_agent.mcp_server_registry import ServerRegistry + + +def change_env_of_server(server_name: str, env_name: str, new_env: str): + def change_server_registry(server_registry: ServerRegistry): + server_config = server_registry.registry[server_name] + if server_config.env is None: + server_config.env = {} + server_config.env[env_name] = new_env + return server_registry + + def inner(core: Core): + if core._context is None: + raise ValueError("Context is not initialized") + registry = core._context.server_registry + if registry is None: + raise ValueError("Server registry is not initialized") + registry = change_server_registry(registry) + core._context.server_registry = registry + return core + + return inner + + +@pytest.mark.integration +@pytest.mark.asyncio +@pytest.mark.xfail(reason="Environment variables are not set when we initialize the agent app") +async def test_runtime_configuration_no_change(fast_agent): + """Test if environment variables are correctly set when we do nothing""" + fast = fast_agent + + @fast.agent(servers=["env_get_server"]) + async def agent_no_change(): + async with fast.run(): + env = fast.app.server_registry.registry["env_get_server"].env + assert env == {"TOKEN": "DEFAULT_TOKEN"} + + await agent_no_change() + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_runtime_configuration_change_env(fast_agent): + """Test if environment variables are correctly set when we change them at runtime""" + fast = fast_agent + + @fast.agent(servers=["env_get_server"]) + async def agent_change_env(): + new_value = "NEW_TOKEN" + change_fn = change_env_of_server("env_get_server", "TOKEN", new_value) + + async with fast.run(core_modifiers=[change_fn]): + env = fast.app.server_registry.registry["env_get_server"].env + assert env == {"TOKEN": new_value} + + await agent_change_env() + + +@pytest.mark.integration +@pytest.mark.asyncio +@pytest.mark.xfail(reason="Environment variables are not set when we initialize the agent app") +async def test_runtime_configuration_add_env(fast_agent): + """Test if environment variables are correctly set when we add a new one""" + fast = fast_agent + + @fast.agent(servers=["env_get_server"]) + async def agent_add_env(): + new_value = "NEW_TOKEN" + change_fn = change_env_of_server("env_get_server", "OTHER_TOKEN", new_value) + + async with fast.run(core_modifiers=[change_fn]): + env = fast.app.server_registry.registry["env_get_server"].env + assert env == {"TOKEN": "DEFAULT_TOKEN", "OTHER_TOKEN": new_value} + + await agent_add_env() + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_runtime_configuration_many_apps(fast_agent): + """Test if environment variables are correctly set when we have more than one application""" + fast = fast_agent + + @fast.agent(servers=["env_get_server"]) + async def agent_app(env_name, env_value): + change_fn = change_env_of_server("env_get_server", env_name, env_value) + + async with fast.run(core_modifiers=[change_fn]): + env = fast.app.server_registry.registry["env_get_server"].env + expected_env = {**env, env_name: env_value} + assert env == expected_env + + await agent_app("NEW_ENV1", "VALUE_ONE") + await agent_app("NEW_ENV2", "VALUE_TWO")