mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-11 03:12:55 +00:00
Give workflow node tools the owner's toolset whoever runs the workflow
A node's tools resolved as the person running the workflow, so a teammate or public-link user lost every owner tool they could not use themselves. They now resolve as the workflow owner, then as the editor who attached them, like an agent's own tools. The runner stays the invoker, so a member-mode connection still uses their own account.
This commit is contained in:
1 parent
23e7145cee
commit
5a503e9b3c
4 files changed
+123
-11
No files matched your search
@@ -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:
|
||||
|
||||
@@ -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 {})
|
||||
|
||||
|
||||
|
||||
@@ -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 []
|
||||
|
||||
@@ -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
|
||||
Reference in new issue
Block a user