diff --git a/docsgpt/api/answer/services/stream_processor.py b/docsgpt/api/answer/services/stream_processor.py index a96e6036..891bcca8 100644 --- a/docsgpt/api/answer/services/stream_processor.py +++ b/docsgpt/api/answer/services/stream_processor.py @@ -10,7 +10,6 @@ from flask import after_this_request from docsgpt import tracing from docsgpt.agents.agent_creator import AgentCreator -from docsgpt.agents.default_tools import synthesized_default_tools from docsgpt.api.answer.services.compression import CompressionOrchestrator from docsgpt.api.answer.services.compression.token_counter import TokenCounter from docsgpt.api.answer.services.compression.types import is_compression_summary_row @@ -38,8 +37,6 @@ from docsgpt.storage.db.repositories.attachments import AttachmentsRepository from docsgpt.storage.db.repositories.prompts import PromptsRepository from docsgpt.storage.db.repositories.sources import SourcesRepository from docsgpt.storage.db.repositories.team_scope import TeamScopeRepository -from docsgpt.storage.db.repositories.user_tools import UserToolsRepository -from docsgpt.storage.db.repositories.users import UsersRepository from docsgpt.api.user.team_sharing import can_access from docsgpt.storage.db.session import db_readonly, db_session from docsgpt.storage.db.source_config import SourceConfig @@ -1376,7 +1373,16 @@ class StreamProcessor: return None, None def pre_fetch_tools(self) -> Optional[Dict[str, Any]]: - """Pre-fetch tool data for template rendering before agent creation""" + """Pre-fetch tool data for template rendering before agent creation. + + Runs the actions the prompt template names on the toolset the agent + run gets, so a teammate or public-link user renders the owner's + prompt with the owner's tools, never their own. + + Returns: + Action results keyed by tool name and tool id, or None when + nothing was fetched. + """ if not settings.ENABLE_TOOL_PREFETCH: logger.info( "Tool pre-fetching disabled globally via ENABLE_TOOL_PREFETCH setting" @@ -1392,17 +1398,14 @@ class StreamProcessor: try: user_id = self.initial_user_id or "local" - agentless = self.agent_id is None - with db_readonly() as conn: - user_tools = UserToolsRepository(conn).list_active_for_user(user_id) - user_doc = ( - UsersRepository(conn).get(user_id) if agentless else None - ) - - default_docs = ( - synthesized_default_tools(user_doc) if agentless else [] - ) - tool_docs = list(user_tools) + default_docs + # The same toolset the run gets: an agent's own tools (resolved as + # its owner, or the editor who attached them), else the caller's + # tools plus defaults. Explicit rows first, so they claim names. + run_tools = [ + tool for tool in self._run_tool_executor().get_tools().values() + if isinstance(tool, dict) and not tool.get("client_side") + ] + tool_docs = sorted(run_tools, key=lambda tool: bool(tool.get("default"))) if not tool_docs: return None @@ -1430,6 +1433,14 @@ class StreamProcessor: continue required_actions = None + owner = tool_doc.get("user_id") + if owner and owner != user_id: + # Someone else's tool: pre-fetch asks nobody, so only + # what the run would do without asking. + required_actions = self._unasked_actions(tool_doc, required_actions) + if not required_actions: + continue + tool_data = self._fetch_tool_data(tool_doc, required_actions) if tool_data: # Explicit rows claim the name key; a default tool takes @@ -1446,6 +1457,57 @@ class StreamProcessor: logger.warning(f"Failed to pre-fetch tools: {type(e).__name__}") return None + @staticmethod + def _unasked_actions( + tool_doc: Dict[str, Any], required_actions: Optional[Set[Optional[str]]] + ) -> Set[Optional[str]]: + """The required actions of someone else's tool that run without asking. + + A tool on someone else's connected account runs on their account or + needs the caller's own connection, and an approval-gated action waits + for a person; pre-fetch has neither, so both are left out. + + Args: + tool_doc: The tool row, owned by someone other than the caller. + required_actions: Action names the template needs; None, or a set + holding None, means all of them. + + Returns: + The action names to run, empty when there are none. + """ + if tool_doc.get("connection_id"): + return set() + if tool_doc.get("name") == "api_tool": + actions = [ + {"name": name, **(action or {})} + for name, action in ((tool_doc.get("config") or {}).get("actions") or {}).items() + ] + else: + actions = tool_doc.get("actions") or [] + unasked = { + action.get("name") for action in actions + if action.get("name") and action.get("active", True) and not action.get("require_approval") + } + if required_actions is None or None in required_actions: + return unasked + return {name for name in required_actions if name in unasked} + + def _run_tool_executor(self): + """A ``ToolExecutor`` resolving the toolset this turn's agent run gets. + + Returns: + ToolExecutor: Built with the run's key, user and agent. + """ + from docsgpt.agents.tool_executor import ToolExecutor + + user = self.decoded_token.get("sub") if self.decoded_token else None + return ToolExecutor( + user_api_key=self.agent_config.get("user_api_key"), + user=user, + decoded_token=self.decoded_token, + agent_id=self.agent_id, + ) + def _enabled_tool_names(self) -> Optional[set]: """Resolve the tool names enabled for this turn, for ``tools.enabled`` gating. @@ -1455,15 +1517,7 @@ class StreamProcessor: (keeps the section) rather than hiding guidance when resolution breaks. """ try: - from docsgpt.agents.tool_executor import ToolExecutor - - user = self.decoded_token.get("sub") if self.decoded_token else None - tool_executor = ToolExecutor( - user_api_key=self.agent_config.get("user_api_key"), - user=user, - decoded_token=self.decoded_token, - agent_id=self.agent_id, - ) + tool_executor = self._run_tool_executor() client_tools = self.data.get("client_tools") if client_tools: tool_executor.client_tools = client_tools diff --git a/tests/api/answer/services/test_stream_processor_utils.py b/tests/api/answer/services/test_stream_processor_utils.py index e6751a20..411ec014 100644 --- a/tests/api/answer/services/test_stream_processor_utils.py +++ b/tests/api/answer/services/test_stream_processor_utils.py @@ -22,7 +22,7 @@ def _patch_db(conn): "docsgpt.api.answer.services.stream_processor.db_readonly", _yield ), patch( "docsgpt.api.answer.services.stream_processor.db_session", _yield - ): + ), patch("docsgpt.agents.tool_executor.db_readonly", _yield): yield diff --git a/tests/api/user/test_resource_sponsors.py b/tests/api/user/test_resource_sponsors.py index 5ce2ec45..4762e852 100644 --- a/tests/api/user/test_resource_sponsors.py +++ b/tests/api/user/test_resource_sponsors.py @@ -386,6 +386,64 @@ class TestRunTime: assert _authorized_source_ids(pg_conn, agent, [source]) == [] +class TestToolPrefetch: + """Prompt tool pre-fetch runs the agent's tools, as the agent run resolves them.""" + + def _prefetched(self, pg_conn, agent_id, caller, required): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + from docsgpt.core.settings import settings + + agent = _row(pg_conn, agent_id) + processor = StreamProcessor({"agent_id": agent_id}, {"sub": caller}) + processor.agent_id = agent_id + processor.agent_config = {"user_api_key": agent["key"]} + processor._required_tool_actions = required + fetched = {} + + def _fake_fetch(tool_doc, required_actions): + fetched[str(tool_doc["id"])] = required_actions + return {"ok": True} + + with _patch_db(pg_conn), patch.object(settings, "ENABLE_TOOL_PREFETCH", True), \ + patch.object(processor, "_fetch_tool_data", _fake_fetch): + processor.pre_fetch_tools() + return fetched + + def _tool(self, pg_conn, user, name, actions, **extra): + return str(UserToolsRepository(pg_conn).create(user, name, actions=actions, **extra)["id"]) + + def test_prefetch_runs_the_agents_tools_not_the_callers(self, pg_conn): + agent_id, _ = _agent(pg_conn) + owner_tool = self._tool(pg_conn, OWNER, "read_webpage", [{"name": "fetch", "active": True}]) + self._tool(pg_conn, VIEWER, "cryptoprice", [{"name": "price", "active": True}]) + AgentsRepository(pg_conn).update_by_id(agent_id, {"tools": [owner_tool]}) + required = {"read_webpage": {None}, "cryptoprice": {None}} + assert self._prefetched(pg_conn, agent_id, OWNER, required) == {owner_tool: {None}} + assert list(self._prefetched(pg_conn, agent_id, VIEWER, required)) == [owner_tool] + + def test_someone_elses_tool_prefetches_only_what_runs_unasked(self, pg_conn): + from sqlalchemy import text + + agent_id, _ = _agent(pg_conn) + api = self._tool(pg_conn, OWNER, "read_webpage", [ + {"name": "fetch", "active": True}, + {"name": "post", "active": True, "require_approval": True}, + ]) + connection = str(pg_conn.execute(text( + "INSERT INTO connector_sessions (user_id, provider, connector_key, auth_kind, status) " + "VALUES (:u, 'telegram', 'telegram', 'api_key', 'connected') RETURNING id" + ), {"u": OWNER}).scalar()) + connected = self._tool(pg_conn, OWNER, "telegram", [{"name": "telegram_get_updates", "active": True}], + connection_id=connection) + AgentsRepository(pg_conn).update_by_id(agent_id, {"tools": [api, connected]}) + required = {"read_webpage": {None}, "telegram": {None}} + + assert self._prefetched(pg_conn, agent_id, OWNER, required) == {api: {None}, connected: {None}} + # Nobody is asked during pre-fetch: a teammate gets no approval-gated + # action and nothing on the owner's connected account. + assert self._prefetched(pg_conn, agent_id, VIEWER, required) == {api: {"fetch"}} + + # --------------------------------------------------------------------------- # workflows # ---------------------------------------------------------------------------