智能体配置

2026-05-13 20:04:06 +08:00
1 parent 83d7fa26fe
commit 31abad86d5
@@ -89,6 +89,12 @@ def _current_user_id(request: Request) -> str:
    return get_effective_user_id()


+def _is_admin_user(request: Request) -> bool:
+    """Return True when the request is authenticated as a system admin."""
+    user = getattr(request.state, "user", None)
+    return getattr(user, "system_role", None) == "admin"
+
+
 def _validate_display_name(name: str) -> str:
    normalized = name.strip()
    if not normalized:
@@ -203,6 +209,29 @@ async def _get_owned_or_raise(store: AgentStore, agent_id: str, user_id: str) ->
    raise HTTPException(status_code=403, detail="You can only modify agents that you created")


+async def _get_editable_or_raise(
+    store: AgentStore,
+    agent_id: str,
+    user_id: str,
+    *,
+    is_admin: bool,
+) -> tuple[dict[str, Any], str]:
+    """Return (record, edit_mode) where edit_mode is 'owner' or 'admin'.
+
+    'admin' is granted only when the agent is built-in (user_id IS NULL) and the
+    caller is a system admin. Custom agents owned by other users still 403.
+    """
+    owned = await store.get_owned(agent_id, user_id)
+    if owned is not None:
+        return owned, "owner"
+    visible = await store.get_visible(agent_id, user_id)
+    if visible is None:
+        raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found")
+    if is_admin and visible.get("user_id") is None:
+        return visible, "admin"
+    raise HTTPException(status_code=403, detail="You can only modify agents that you created")
+
+
@router.get(
    "/agents",
    response_model=AgentsListResponse,
@@ -322,7 +351,10 @@ async def update_agent(agent_id: str, request: AgentUpdateRequest, http_request:
    try:
        store = get_agent_store(http_request)
        user_id = _current_user_id(http_request)
-        current = await _get_owned_or_raise(store, agent_id, user_id)
+        is_admin = _is_admin_user(http_request)
+        current, edit_mode = await _get_editable_or_raise(
+            store, agent_id, user_id, is_admin=is_admin
+        )

        fields_set = request.model_fields_set
        update_data: dict[str, Any] = {}
@@ -354,7 +386,12 @@ async def update_agent(agent_id: str, request: AgentUpdateRequest, http_request:
        if request.soul is not None:
            _write_agent_soul(agent_id, request.soul)

-        record = await store.update(agent_id, user_id, update_data) if update_data else current
+        if not update_data:
+            record = current
+        elif edit_mode == "admin":
+            record = await store.update_builtin(agent_id, update_data)
+        else:
+            record = await store.update(agent_id, user_id, update_data)
        if record is None:
            raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found")

@@ -8,8 +8,9 @@ from deerflow.agents.lead_agent.prompt import apply_prompt_template
 from deerflow.agents.memory.summarization_hook import memory_flush_hook
 from deerflow.agents.middlewares.clarification_middleware import ClarificationMiddleware
 from deerflow.agents.middlewares.loop_detection_middleware import LoopDetectionMiddleware
-from deerflow.agents.middlewares.memory_middleware import MemoryMiddleware
-from deerflow.agents.middlewares.subagent_limit_middleware import SubagentLimitMiddleware
+from deerflow.agents.middlewares.memory_middleware import MemoryMiddleware
+from deerflow.agents.middlewares.setup_agent_approval_middleware import SetupAgentApprovalMiddleware
+from deerflow.agents.middlewares.subagent_limit_middleware import SubagentLimitMiddleware
 from deerflow.agents.middlewares.summarization_middleware import BeforeSummarizationHook, DeerFlowSummarizationMiddleware
 from deerflow.agents.middlewares.title_middleware import TitleMiddleware
 from deerflow.agents.middlewares.todo_middleware import TodoMiddleware
@@ -303,8 +304,10 @@ def _build_middlewares(
    if custom_middlewares:
        middlewares.extend(custom_middlewares)

-    # ClarificationMiddleware should always be last
-    middlewares.append(ClarificationMiddleware())
+    middlewares.append(SetupAgentApprovalMiddleware())
+
+    # ClarificationMiddleware should always be last
+    middlewares.append(ClarificationMiddleware())
    return middlewares


@@ -0,0 +1,90 @@
+"""Approval middleware for the custom-agent setup tool."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from hashlib import sha256
+from typing import Any, override
+
+from langchain.agents import AgentState
+from langchain.agents.middleware import AgentMiddleware
+from langchain_core.messages import ToolMessage
+from langgraph.graph import END
+from langgraph.prebuilt.tool_node import ToolCallRequest
+from langgraph.types import Command
+
+logger = logging.getLogger(__name__)
+
+
+class SetupAgentApprovalMiddlewareState(AgentState):
+    """Compatible with the DeerFlow thread state schema."""
+
+    pass
+
+
+class SetupAgentApprovalMiddleware(AgentMiddleware[SetupAgentApprovalMiddlewareState]):
+    """Stop before ``setup_agent`` writes files and ask the user to approve.
+
+    The first ``setup_agent`` call is converted into a structured tool message
+    and the run ends. A subsequent request can set ``setup_agent_approved`` in
+    runtime context to allow the tool to execute.
+    """
+
+    state_schema = SetupAgentApprovalMiddlewareState
+
+    def _is_approved(self, request: ToolCallRequest) -> bool:
+        context = getattr(request.runtime, "context", None) or {}
+        if not isinstance(context, dict):
+            return False
+        return context.get("setup_agent_approved") is True
+
+    def _stable_message_id(self, tool_call_id: str, args: dict[str, Any]) -> str:
+        if tool_call_id:
+            return f"setup-agent-approval:{tool_call_id}"
+        digest = sha256(repr(sorted(args.items())).encode("utf-8")).hexdigest()[:16]
+        return f"setup-agent-approval:{digest}"
+
+    def _build_payload(self, args: dict[str, Any]) -> dict[str, Any]:
+        skills = args.get("skills")
+        return {
+            "name": str(args.get("name") or "").strip(),
+            "description": str(args.get("description") or "").strip(),
+            "soul": str(args.get("soul") or "").strip(),
+            "skills": skills if isinstance(skills, list) else [],
+        }
+
+    def _handle_setup_agent(self, request: ToolCallRequest) -> Command:
+        args = request.tool_call.get("args", {}) or {}
+        tool_call_id = request.tool_call.get("id", "")
+        payload = self._build_payload(args)
+        logger.info("Setup agent approval requested before writing agent files")
+
+        message = ToolMessage(
+            id=self._stable_message_id(tool_call_id, args),
+            content="请确认是否创建/更新这个专家配置。",
+            tool_call_id=tool_call_id,
+            name="setup_agent",
+            additional_kwargs={"setup_agent_approval": payload},
+        )
+        return Command(update={"messages": [message]}, goto=END)
+
+    @override
+    def wrap_tool_call(
+        self,
+        request: ToolCallRequest,
+        handler: Callable[[ToolCallRequest], ToolMessage | Command],
+    ) -> ToolMessage | Command:
+        if request.tool_call.get("name") != "setup_agent" or self._is_approved(request):
+            return handler(request)
+        return self._handle_setup_agent(request)
+
+    @override
+    async def awrap_tool_call(
+        self,
+        request: ToolCallRequest,
+        handler: Callable[[ToolCallRequest], ToolMessage | Command],
+    ) -> ToolMessage | Command:
+        if request.tool_call.get("name") != "setup_agent" or self._is_approved(request):
+            return await handler(request)
+        return self._handle_setup_agent(request)
@@ -35,6 +35,15 @@ class AgentStore(ABC):
    async def update(self, agent_id: str, user_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
        raise NotImplementedError

+    @abstractmethod
+    async def update_builtin(self, agent_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
+        """Update a built-in agent (user_id IS NULL).
+
+        Caller is responsible for authorization — this method does NOT check
+        any user identity and returns ``None`` if the target is not a built-in.
+        """
+        raise NotImplementedError
+
    @abstractmethod
    async def delete(self, agent_id: str, user_id: str) -> bool:
        raise NotImplementedError
@@ -91,6 +91,16 @@ class MemoryAgentStore(AgentStore):
        row["updated_at"] = _now()
        return _copy(row)

+    async def update_builtin(self, agent_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
+        row = self._agents.get(agent_id)
+        if row is None or row.get("user_id") is not None:
+            return None
+        for key in ("name", "description", "published"):
+            if key in data:
+                row[key] = data[key]
+        row["updated_at"] = _now()
+        return _copy(row)
+
    async def delete(self, agent_id: str, user_id: str) -> bool:
        row = self._agents.get(agent_id)
        if row is None or row.get("user_id") != user_id:
@@ -107,6 +107,19 @@ class AgentRepository(AgentStore):
            await session.refresh(row)
            return _row_to_dict(row)

+    async def update_builtin(self, agent_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
+        async with self._sf() as session:
+            row = await session.get(AgentRow, agent_id)
+            if row is None or row.user_id is not None:
+                return None
+            for key in ("name", "description", "published"):
+                if key in data:
+                    setattr(row, key, data[key])
+            row.updated_at = datetime.now(UTC)
+            await session.commit()
+            await session.refresh(row)
+            return _row_to_dict(row)
+
    async def delete(self, agent_id: str, user_id: str) -> bool:
        async with self._sf() as session:
            row = await session.get(AgentRow, agent_id)
@@ -17,6 +17,7 @@ def setup_agent(
    soul: str,
    description: str,
    runtime: ToolRuntime,
+    name: str | None = None,
    skills: list[str] | None = None,
 ) -> Command:
    """Setup the custom DeerFlow agent.
@@ -24,6 +25,8 @@ def setup_agent(
    Args:
        soul: Full SOUL.md content defining the agent's personality and behavior.
        description: One-line description of what the agent does.
+        name: Optional human-readable display name for the agent. If omitted,
+            the agent's stable id is used as the display name.
        skills: Optional list of skill names this agent should use. None or empty means no skills.
    """

@@ -40,7 +43,8 @@ def setup_agent(

        if agent_name:
            # If agent_name is provided, we are creating a custom agent in the agents/ directory
-            config_data: dict = {"id": agent_name, "name": agent_name}
+            display_name = (name or "").strip() or agent_name
+            config_data: dict = {"id": agent_name, "name": display_name}
            if description:
                config_data["description"] = description
            if skills is not None:
@@ -52,6 +52,63 @@ async def test_agent_store_only_owner_can_update_or_delete() -> None:
    assert await store.delete("private-a", "user-a") is True


+@pytest.mark.asyncio
+async def test_update_builtin_only_targets_builtins() -> None:
+    store = MemoryAgentStore()
+
+    await store.ensure_builtin({"id": "builtin-agent", "name": "Built in"})
+    await store.create({"id": "owned-agent", "name": "Owned", "user_id": "user-a"})
+
+    # update_builtin succeeds on a built-in agent
+    updated = await store.update_builtin("builtin-agent", {"name": "Renamed by admin"})
+    assert updated is not None
+    assert updated["name"] == "Renamed by admin"
+
+    # update_builtin refuses to touch a user-owned agent (returns None)
+    assert await store.update_builtin("owned-agent", {"name": "Hijacked"}) is None
+
+    # And on a missing agent
+    assert await store.update_builtin("does-not-exist", {"name": "Ghost"}) is None
+
+
+@pytest.mark.asyncio
+async def test_get_editable_or_raise_admin_can_edit_builtin() -> None:
+    """Route helper: admin gets edit access to built-ins; non-admin still 403."""
+    from fastapi import HTTPException
+
+    from app.gateway.routers.agents import _get_editable_or_raise
+
+    store = MemoryAgentStore()
+    await store.ensure_builtin({"id": "builtin-agent", "name": "Built in"})
+    await store.create({"id": "owned-agent", "name": "Owned", "user_id": "user-a"})
+
+    # Owner edits their own agent → returns ('owner', record)
+    record, mode = await _get_editable_or_raise(
+        store, "owned-agent", "user-a", is_admin=False
+    )
+    assert mode == "owner" and record["id"] == "owned-agent"
+
+    # Admin edits a built-in agent → returns ('admin', record)
+    record, mode = await _get_editable_or_raise(
+        store, "builtin-agent", "user-a", is_admin=True
+    )
+    assert mode == "admin" and record["id"] == "builtin-agent"
+
+    # Non-admin tries to edit a built-in → 403
+    with pytest.raises(HTTPException) as exc:
+        await _get_editable_or_raise(store, "builtin-agent", "user-a", is_admin=False)
+    assert exc.value.status_code == 403
+
+    # Admin tries to edit someone else's PUBLISHED custom agent → still 403
+    # (admin override only applies to built-ins, not to others' custom agents).
+    await store.create(
+        {"id": "published-other", "name": "Pub", "user_id": "user-x", "published": True}
+    )
+    with pytest.raises(HTTPException) as exc:
+        await _get_editable_or_raise(store, "published-other", "user-a", is_admin=True)
+    assert exc.value.status_code == 403
+
+
 def test_build_run_config_injects_custom_assistant_as_agent_id() -> None:
    config = build_run_config("thread-1", None, None, assistant_id="agent-123")