mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 14:12:58 +00:00
156 lines
6.8 KiB
Python
156 lines
6.8 KiB
Python
"""Runtime resolution of team-shared tools (owner context, live grant checks)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from contextlib import contextmanager
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
from docsgpt.agents.default_tools import resolve_tool_by_id
|
|
from docsgpt.agents.tool_executor import ToolExecutor
|
|
from docsgpt.api.user.resource_access import set_settings
|
|
from docsgpt.security.encryption import encrypt_credentials
|
|
from docsgpt.storage.db.repositories.agents import AgentsRepository
|
|
from docsgpt.storage.db.repositories.team_members import TeamMembersRepository
|
|
from docsgpt.storage.db.repositories.team_resource_grants import (
|
|
TeamResourceGrantsRepository,
|
|
)
|
|
from docsgpt.storage.db.repositories.teams import TeamsRepository
|
|
from docsgpt.storage.db.repositories.user_tool_preferences import (
|
|
UserToolPreferencesRepository,
|
|
)
|
|
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
|
|
|
OWNER = "alice"
|
|
|
|
|
|
def _tool(conn, name="read_webpage", config=None, status=True):
|
|
return UserToolsRepository(conn).create(
|
|
OWNER, name, config=config or {}, display_name=name, description="",
|
|
actions=[{"name": "act", "active": True}], status=status,
|
|
)
|
|
|
|
|
|
def _share(conn, tool_id, member, level="viewer"):
|
|
team = TeamsRepository(conn).create("Acme", f"t-{uuid.uuid4().hex[:8]}", OWNER)
|
|
TeamMembersRepository(conn).add_member(str(team["id"]), member)
|
|
TeamResourceGrantsRepository(conn).grant(
|
|
str(team["id"]), "tool", str(tool_id), OWNER, OWNER, access_level=level
|
|
)
|
|
return team
|
|
|
|
|
|
def _revoke(conn, team, tool_id):
|
|
TeamResourceGrantsRepository(conn).revoke(str(team["id"]), "tool", str(tool_id))
|
|
|
|
|
|
@pytest.fixture
|
|
def use_conn(monkeypatch, pg_conn):
|
|
@contextmanager
|
|
def _yield():
|
|
yield pg_conn
|
|
|
|
monkeypatch.setattr("docsgpt.agents.tool_executor.db_readonly", _yield)
|
|
return pg_conn
|
|
|
|
|
|
class TestResolveToolById:
|
|
def test_owner_resolves(self, pg_conn):
|
|
tool = _tool(pg_conn)
|
|
row = resolve_tool_by_id(str(tool["id"]), OWNER, user_tools_repo=UserToolsRepository(pg_conn))
|
|
assert row is not None and row["user_id"] == OWNER
|
|
|
|
def test_grantee_with_use_in_own_resolves_owner_row(self, pg_conn):
|
|
tool = _tool(pg_conn)
|
|
_share(pg_conn, tool["id"], "bob")
|
|
row = resolve_tool_by_id(str(tool["id"]), "bob", user_tools_repo=UserToolsRepository(pg_conn))
|
|
assert row is not None and row["user_id"] == OWNER
|
|
|
|
def test_viewer_without_use_in_own_does_not_resolve(self, pg_conn):
|
|
tool = _tool(pg_conn)
|
|
_share(pg_conn, tool["id"], "bob")
|
|
set_settings(pg_conn, "tool", str(tool["id"]), {"viewers_can_use_in_agents": False}, OWNER)
|
|
repo = UserToolsRepository(pg_conn)
|
|
assert resolve_tool_by_id(str(tool["id"]), "bob", user_tools_repo=repo) is None
|
|
# Editors keep use_in_own when the viewer switch is off.
|
|
_share(pg_conn, tool["id"], "carol", "editor")
|
|
assert resolve_tool_by_id(str(tool["id"]), "carol", user_tools_repo=repo) is not None
|
|
|
|
def test_stranger_does_not_resolve(self, pg_conn):
|
|
tool = _tool(pg_conn)
|
|
assert resolve_tool_by_id(str(tool["id"]), "eve", user_tools_repo=UserToolsRepository(pg_conn)) is None
|
|
|
|
|
|
class TestAgentToolsForOwner:
|
|
def _agent(self, conn, owner, tool_ids):
|
|
key = f"k-{uuid.uuid4().hex[:8]}"
|
|
AgentsRepository(conn).create(owner, "A", "published", key=key, tools=[str(t) for t in tool_ids])
|
|
return key
|
|
|
|
def test_team_tool_on_grantees_agent_resolves_then_drops_on_revoke(self, use_conn):
|
|
tool = _tool(use_conn)
|
|
team = _share(use_conn, tool["id"], "bob")
|
|
key = self._agent(use_conn, "bob", [tool["id"]])
|
|
# Carol chats with Bob's agent: tools resolve under Bob, the agent owner.
|
|
tools = ToolExecutor(user_api_key=key, user="carol")._get_tools_by_api_key(key)
|
|
assert str(tool["id"]) in tools
|
|
assert tools[str(tool["id"])]["user_id"] == OWNER
|
|
_revoke(use_conn, team, tool["id"])
|
|
assert ToolExecutor(user_api_key=key, user="carol")._get_tools_by_api_key(key) == {}
|
|
|
|
def test_explicit_ids_path_checks_the_caller(self, use_conn):
|
|
tool = _tool(use_conn)
|
|
executor = ToolExecutor(user="bob")
|
|
assert executor._get_tools_by_ids([str(tool["id"])]) == {}
|
|
_share(use_conn, tool["id"], "bob")
|
|
assert str(tool["id"]) in executor._get_tools_by_ids([str(tool["id"])])
|
|
|
|
|
|
class TestAgentlessChatTools:
|
|
def test_includes_in_chat_team_tools(self, use_conn):
|
|
own = _tool(use_conn, name="brave")
|
|
shared = _tool(use_conn, name="read_webpage")
|
|
team = _share(use_conn, shared["id"], "bob")
|
|
UserToolsRepository(use_conn).create("bob", "cryptoprice", status=True)
|
|
names = lambda: sorted( # noqa: E731
|
|
t["name"] for t in ToolExecutor(user="bob")._get_user_tools("bob").values() if t.get("user_id")
|
|
)
|
|
assert names() == ["cryptoprice"] # sharing alone never adds to chats
|
|
UserToolPreferencesRepository(use_conn).set_in_chat("bob", str(shared["id"]), True)
|
|
assert names() == ["cryptoprice", "read_webpage"]
|
|
set_settings(use_conn, "tool", str(shared["id"]), {"viewers_can_use_in_agents": False}, OWNER)
|
|
assert names() == ["cryptoprice"]
|
|
set_settings(use_conn, "tool", str(shared["id"]), {"viewers_can_use_in_agents": True}, OWNER)
|
|
_revoke(use_conn, team, shared["id"])
|
|
assert names() == ["cryptoprice"]
|
|
assert own["id"] # the owner's own tool never leaks into bob's chat
|
|
|
|
|
|
class TestRuntimeCredentials:
|
|
def test_mcp_tool_loads_with_owner_identity(self, monkeypatch):
|
|
executor = ToolExecutor(user="bob")
|
|
mock_tm = Mock()
|
|
monkeypatch.setattr("docsgpt.agents.tool_executor.ToolManager", lambda config: mock_tm)
|
|
tool_data = {"id": str(uuid.uuid4()), "name": "mcp_tool", "user_id": OWNER, "config": {}}
|
|
executor._get_or_load_tool(tool_data, "t1", "act")
|
|
assert mock_tm.load_tool.call_args.kwargs["user_id"] == OWNER
|
|
|
|
def test_api_tool_secret_values_decrypt_with_owner(self):
|
|
from docsgpt.agents.tool_executor import api_tool_action_with_secrets
|
|
|
|
blob = encrypt_credentials({"a": {"headers": {"X-Key": "sk"}}}, OWNER)
|
|
tool_data = {
|
|
"user_id": OWNER,
|
|
"config": {
|
|
"encrypted_action_secrets": blob,
|
|
"actions": {"a": {"url": "https://x", "headers": {"properties": {
|
|
"X-Key": {"value": "", "has_value": True, "filled_by_llm": False}}}}},
|
|
},
|
|
}
|
|
action = api_tool_action_with_secrets(tool_data, "a", "bob")
|
|
assert action["headers"]["properties"]["X-Key"]["value"] == "sk"
|
|
# The stored row is not mutated.
|
|
assert tool_data["config"]["actions"]["a"]["headers"]["properties"]["X-Key"]["value"] == ""
|