diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index add4df1f..6f6b0875 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -523,6 +523,9 @@ class ToolExecutor: # get_tools() resolves EXACTLY these ids — builtin synthetic ids and # user_tools rows alike — with no defaults mixed in. None = unscoped. self.allowed_tool_ids: Optional[List[str]] = None + # Who an explicit tool-id scope resolves as: the workflow owner for a + # node, so whoever runs it gets the owner's tools. None = ``user``. + self.tool_owner: Optional[str] = None # Tool id -> the user to resolve it as, for a workflow node's tools # sponsored by an editor (see resource_access.active_sponsor). self.tool_principals: Dict[str, str] = {} @@ -582,17 +585,21 @@ class ToolExecutor: """Resolve an explicit tool-id scope — exactly these ids, no defaults. Used by workflow agent nodes: the node's configured tools (builtin - synthetic ids like Artifact/Code Executor/Read Document, or the user's - ``user_tools`` rows) are the node's WHOLE toolset. An unresolvable id - is dropped with a warning rather than failing the node. + synthetic ids like Artifact/Code Executor/Read Document, or the + ``user_tools`` rows of ``tool_owner``) are the node's WHOLE toolset. + Rows resolve as the workflow owner, then as the editor who attached + them, never as whoever runs the workflow — the same rule as an agent's + own tools. An unresolvable id is dropped with a warning rather than + failing the node. """ if not tool_ids: return {} + principal = self.tool_owner or self.user with db_readonly() as conn: tools_repo = UserToolsRepository(conn) tools: List[Dict] = [] for tid in tool_ids: - row = resolve_tool_by_id(tid, self.user, user_tools_repo=tools_repo) + row = resolve_tool_by_id(tid, principal, user_tools_repo=tools_repo) if row is None and str(tid) in self.tool_principals: row = resolve_tool_by_id(tid, self.tool_principals[str(tid)], user_tools_repo=tools_repo) if row is None: diff --git a/docsgpt/agents/workflows/node_agent.py b/docsgpt/agents/workflows/node_agent.py index 45215317..a2927f14 100644 --- a/docsgpt/agents/workflows/node_agent.py +++ b/docsgpt/agents/workflows/node_agent.py @@ -20,6 +20,7 @@ class _WorkflowNodeMixin: api_key: str, tool_ids: Optional[List[str]] = None, tool_principals: Optional[Dict[str, str]] = None, + tool_owner: Optional[str] = None, **kwargs, ): super().__init__( @@ -35,7 +36,9 @@ class _WorkflowNodeMixin: # (Artifact / Code Executor / Read Document) and ``user_tools`` rows # alike, and an empty list means the node's LLM gets no tools. self.tool_executor.allowed_tool_ids = [str(t) for t in (tool_ids or [])] - # Tools the owner can't use resolve as the editor who attached them. + # The node's tools are the workflow owner's whoever runs it; tools the + # owner can't use resolve as the editor who attached them. + self.tool_executor.tool_owner = tool_owner self.tool_executor.tool_principals = dict(tool_principals or {}) diff --git a/docsgpt/agents/workflows/workflow_engine.py b/docsgpt/agents/workflows/workflow_engine.py index 22017726..15a8dcb4 100644 --- a/docsgpt/agents/workflows/workflow_engine.py +++ b/docsgpt/agents/workflows/workflow_engine.py @@ -405,6 +405,7 @@ class WorkflowEngine: "model_user_id": getattr(self.agent, "model_user_id", None), "api_key": node_api_key, "tool_ids": node_config.tools, + "tool_owner": self._workflow_owner_id(), "tool_principals": self._node_tool_principals(node_config.tools), "prompt": node_prompt, "chat_history": self.agent.chat_history, @@ -1322,6 +1323,15 @@ class WorkflowEngine: docs_together = "\n\n".join(docs_together_parts) if docs_together_parts else None return docs, docs_together + def _workflow_owner_id(self) -> Optional[str]: + """The workflow's owner, whom node tools and sources run as. + + Returns: + The owner's user id, or None when the run has none. + """ + resolve_owner = getattr(self.agent, "_resolve_owner_id", None) + return (resolve_owner() if callable(resolve_owner) else None) or self._resolve_user_id() + def _node_tool_principals(self, tool_ids) -> Dict[str, str]: """Node tool id -> the editor to resolve it as, for sponsored tools. @@ -1373,10 +1383,7 @@ class WorkflowEngine: if not sources: return [] ids = sources if isinstance(sources, list) else [sources] - resolve_owner = getattr(self.agent, "_resolve_owner_id", None) - owner = (resolve_owner() if callable(resolve_owner) else None) or ( - self._resolve_user_id() - ) + owner = self._workflow_owner_id() if not owner: logger.warning("Workflow node sources dropped: no owner to authorize.") return [] diff --git a/tests/api/user/test_resource_sponsors.py b/tests/api/user/test_resource_sponsors.py index 7dff59f6..5ce2ec45 100644 --- a/tests/api/user/test_resource_sponsors.py +++ b/tests/api/user/test_resource_sponsors.py @@ -45,6 +45,8 @@ _DB_MODULES = ( "docsgpt.agents.tool_executor", "docsgpt.api.answer.services.stream_processor", "docsgpt.agents.workflows.workflow_engine", + "docsgpt.connectors.resolve", + "docsgpt.connectors.service", "docsgpt.storage.db.session", ) @@ -389,8 +391,10 @@ class TestRunTime: # --------------------------------------------------------------------------- -def _wf_body(tool=None, source=None): - agent_cfg = {"agent_type": "classic", "system_prompt": "s", "tools": [tool] if tool else [], +def _wf_body(tool=None, source=None, tools=None): + if tools is None: + tools = [tool] if tool else [] + agent_cfg = {"agent_type": "classic", "system_prompt": "s", "tools": tools, "sources": [source] if source else []} return { "name": "WF", @@ -457,3 +461,94 @@ class TestWorkflows: TeamResourceGrantsRepository(pg_conn).revoke(team_id, "agent", agent_id, target_user_id=EDITOR) assert engine._node_tool_principals([tool]) == {} assert engine._authorized_node_sources([source]) == [] + + def _node_executor(self, monkeypatch, pg_conn, wid, caller, tools): + """Run one agent node of ``wid`` for ``caller``; return its agent's tool executor.""" + from docsgpt.agents.tool_executor import ToolExecutor + from docsgpt.agents.workflows.node_agent import WorkflowNodeAgentFactory, _WorkflowNodeMixin + from docsgpt.agents.workflows.schemas import NodeType, Workflow, WorkflowGraph, WorkflowNode + from docsgpt.agents.workflows.workflow_engine import WorkflowEngine + + class _Base: + # Stands in for BaseAgent: the executor is built for the token's user. + def __init__(self, decoded_token=None, **_kwargs): + sub = (decoded_token or {}).get("sub") + self.tool_executor = ToolExecutor(user=sub, decoded_token=decoded_token) + + class _NodeAgent(_WorkflowNodeMixin, _Base): + def gen(self, _prompt): + yield {"answer": "ok"} + + built = [] + + def _create(agent_type, **kwargs): + built.append(_NodeAgent(**kwargs)) + return built[-1] + + monkeypatch.setattr(WorkflowNodeAgentFactory, "create", staticmethod(_create)) + monkeypatch.setattr("docsgpt.core.model_utils.get_api_key_for_provider", lambda _name: None) + agent = SimpleNamespace( + endpoint="stream", llm_name="openai", model_id="gpt-4o-mini", api_key="k", chat_history=[], + decoded_token={"sub": caller}, user=caller, workflow_owner=OWNER, + _resolve_owner_id=lambda: OWNER, workflow_row=WorkflowsRepository(pg_conn).get_by_id(wid), + ) + engine = WorkflowEngine(WorkflowGraph(workflow=Workflow(name="wf"), nodes=[], edges=[]), agent) + engine.state["query"] = "q" + node = WorkflowNode( + id="a1", workflow_id=wid, type=NodeType.AGENT, title="A", position={"x": 0, "y": 0}, + config={"agent_type": "classic", "system_prompt": "s", "tools": tools}, + ) + with _patch_db(pg_conn): + list(engine._execute_agent_node(node)) + return built[0].tool_executor + + @pytest.mark.parametrize("caller", [OWNER, VIEWER, OTHER]) + def test_node_tools_resolve_as_the_owner_whoever_runs_it(self, app, pg_conn, monkeypatch, caller): + """A teammate or public-link user gets the node tools the owner gets.""" + wid, _ = self._setup(pg_conn) + owner_tool = str(UserToolsRepository(pg_conn).create(OWNER, "api_tool")["id"]) + editor_tool, _, _ = _editor_resources(pg_conn) + assert _status(self._put(app, pg_conn, wid, EDITOR, _wf_body(editor_tool))) == 200 + tools = [owner_tool, editor_tool] + assert _status(self._put(app, pg_conn, wid, OWNER, _wf_body(tools=tools))) == 200 + + executor = self._node_executor(monkeypatch, pg_conn, wid, caller, tools) + with _patch_db(pg_conn): + resolved = executor.get_tools() + assert {tid: row["user_id"] for tid, row in resolved.items()} == {owner_tool: OWNER, editor_tool: EDITOR} + # Whoever runs the node stays the invoker: member-mode connections are theirs. + assert executor.user == caller + + def test_caller_only_tool_is_not_in_the_node(self, app, pg_conn, monkeypatch): + """A tool only the caller can use never resolves: nodes are the owner's.""" + wid, _ = self._setup(pg_conn) + viewer_tool = str(UserToolsRepository(pg_conn).create(VIEWER, "api_tool")["id"]) + executor = self._node_executor(monkeypatch, pg_conn, wid, VIEWER, [viewer_tool]) + with _patch_db(pg_conn): + assert executor.get_tools() == {} + + @pytest.mark.parametrize(("mode", "account", "delegated"), [("member", VIEWER, False), ("owner", OWNER, True)]) + def test_node_connection_follows_the_credential_mode(self, app, pg_conn, monkeypatch, mode, account, delegated): + from sqlalchemy import text + + def _connection(user): + return str(pg_conn.execute( + text( + "INSERT INTO connector_sessions (user_id, provider, connector_key, auth_kind, status, " + "account_label) VALUES (:u, 'telegram', 'telegram', 'api_key', 'connected', :u) RETURNING id" + ), + {"u": user}, + ).scalar()) + + wid, _ = self._setup(pg_conn) + _connection(VIEWER) + tool = str(UserToolsRepository(pg_conn).create( + OWNER, "telegram", connection_id=_connection(OWNER), credential_mode=mode, + )["id"]) + assert _status(self._put(app, pg_conn, wid, OWNER, _wf_body(tool))) == 200 + + executor = self._node_executor(monkeypatch, pg_conn, wid, VIEWER, [tool]) + with _patch_db(pg_conn): + resolved = executor._resolve_connection(executor.get_tools()[tool]) + assert resolved.row["user_id"] == account + assert resolved.delegated is delegated