diff --git a/flocks/agent/registry.py b/flocks/agent/registry.py index e1643b541..6280f4c70 100644 --- a/flocks/agent/registry.py +++ b/flocks/agent/registry.py @@ -139,6 +139,62 @@ def _build_available_agents(agents: Dict[str, AgentInfo]) -> List[AvailableAgent return available +def _storage_custom_agent_to_info(agent_data: Dict[str, Any]) -> Optional[AgentInfo]: + """Convert a Storage-backed custom agent record into AgentInfo.""" + name = agent_data.get("name") + if not name: + return None + + model_raw = agent_data.get("model") + model: Optional[AgentModel] = None + if isinstance(model_raw, dict): + model_id = model_raw.get("model_id") or model_raw.get("modelID") + provider_id = model_raw.get("provider_id") or model_raw.get("providerID") + if model_id and provider_id: + model = AgentModel(model_id=model_id, provider_id=provider_id) + + return AgentInfo( + name=name, + description=agent_data.get("description"), + description_cn=agent_data.get("description_cn") or agent_data.get("descriptionCn"), + prompt=agent_data.get("prompt"), + temperature=agent_data.get("temperature"), + color=agent_data.get("color"), + mode=agent_data.get("mode", "primary"), + model=model, + native=False, + hidden=agent_data.get("hidden", False), + tools=agent_data.get("tools", []), + tags=agent_data.get("tags", []), + ) + + +async def _load_storage_custom_agents(existing_names: Set[str]) -> Dict[str, AgentInfo]: + """Load Storage-backed custom agents created by POST /api/agent.""" + try: + from flocks.storage.storage import Storage + entries = await Storage.list_entries("agent/custom/") + except Exception as exc: + log.warn("agent.registry.storage_custom_load_error", {"error": str(exc)}) + return {} + + loaded: Dict[str, AgentInfo] = {} + for key, agent_data in entries: + if not isinstance(agent_data, dict) or not agent_data.get("name"): + continue + agent = _storage_custom_agent_to_info(agent_data) + if agent is None: + continue + if agent.name in existing_names or agent.name in loaded: + log.warn("agent.registry.storage_custom_name_conflict", { + "name": agent.name, + "key": key, + }) + continue + loaded[agent.name] = agent + return loaded + + # --------------------------------------------------------------------------- # Agent registry # --------------------------------------------------------------------------- @@ -339,6 +395,9 @@ def _permission_dict_to_tools(permission_cfg: Dict[str, Any]) -> List[str]: if isinstance(value.permission, dict): item.tools = _permission_dict_to_tools(value.permission) + storage_custom_agents = await _load_storage_custom_agents(set(result.keys())) + result.update(storage_custom_agents) + # enabled_agents whitelist filter if cfg.enabled_agents is not None: allowed = set(cfg.enabled_agents) diff --git a/tests/server/routes/test_agent_routes.py b/tests/server/routes/test_agent_routes.py index 9f276c4d9..0fa1f98b2 100644 --- a/tests/server/routes/test_agent_routes.py +++ b/tests/server/routes/test_agent_routes.py @@ -106,11 +106,30 @@ async def test_created_agent_retrievable_by_name(self, client: AsyncClient): resp = await client.post("/api/agent", json=_AGENT_PAYLOAD) assert resp.status_code == status.HTTP_200_OK - # Direct GET by name always reads Storage, so the new agent is visible + # Direct GET by name uses the refreshed agent registry, so the new agent is visible get_resp = await client.get("/api/agent/test-agent") assert get_resp.status_code == status.HTTP_200_OK assert get_resp.json()["name"] == "test-agent" + @pytest.mark.asyncio + async def test_created_agent_survives_registry_reload(self, client: AsyncClient): + """Storage-backed custom agents remain visible after process cache reload.""" + from flocks.agent.registry import Agent + + resp = await client.post("/api/agent", json=_AGENT_PAYLOAD) + assert resp.status_code == status.HTTP_200_OK + + Agent._custom_agents.clear() + Agent.invalidate_cache() + + get_resp = await client.get("/api/agent/test-agent") + assert get_resp.status_code == status.HTTP_200_OK + assert get_resp.json()["name"] == "test-agent" + + list_resp = await client.get("/api/agent") + assert list_resp.status_code == status.HTTP_200_OK + assert "test-agent" in [agent["name"] for agent in list_resp.json()] + # =========================================================================== # Update