mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-11 03:12:55 +00:00
Keep pre-fetch off live-approval tools and outside callers' owner writes
Pre-fetch judged someone else's tool by its stored approval flags, which a remote device or the code executor decides per call, and it took a widget or API run for the owner because the run carries the owner's id. Those tools no longer pre-fetch for anyone but their owner, an API-key or public-link run treats every tool as someone else's, and writes with the owner's credentials are never pre-fetched for them.
This commit is contained in:
1 parent
5c8b97b609
commit
bfa67d3138
2 files changed
+48
-16
No files matched your search
@@ -128,6 +128,10 @@ def get_prompt(prompt_id: str, prompts_collection=None) -> str:
|
||||
|
||||
_PROMPT_PRESETS_WITHOUT_ROW = ("reduce",)
|
||||
|
||||
# Tools whose approval is decided per call from live state, not the stored
|
||||
# ``require_approval`` flags (see ``ToolExecutor.check_pause``).
|
||||
_LIVE_APPROVAL_TOOLS = frozenset({"remote_device", "code_executor"})
|
||||
|
||||
|
||||
def authorized_prompt_id(prompt_id: Any, principal: Optional[str], agent: Optional[dict] = None) -> Any:
|
||||
"""``prompt_id`` if ``principal`` (or the agent's sponsor) may use it, else ``"default"``.
|
||||
@@ -1407,6 +1411,9 @@ class StreamProcessor:
|
||||
|
||||
try:
|
||||
user_id = self.initial_user_id or "local"
|
||||
outside_caller = bool(
|
||||
self.agent_config.get("external_api_caller") or self.agent_config.get("public_link_caller")
|
||||
)
|
||||
# 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.
|
||||
@@ -1443,9 +1450,10 @@ class StreamProcessor:
|
||||
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.
|
||||
if owner and (owner != user_id or outside_caller):
|
||||
# Someone else's tool (a widget or API run carries the
|
||||
# owner's id but isn't the owner): 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
|
||||
@@ -1473,8 +1481,11 @@ class StreamProcessor:
|
||||
"""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.
|
||||
needs the caller's own connection, a tool that decides approval per
|
||||
call (a remote device, the code executor) can't be judged from its
|
||||
stored flags, an approval-gated action waits for a person, and a
|
||||
write with the owner's credentials needs the owner's say-so;
|
||||
pre-fetch has none of these, so all are left out.
|
||||
|
||||
Args:
|
||||
tool_doc: The tool row, owned by someone other than the caller.
|
||||
@@ -1484,18 +1495,15 @@ class StreamProcessor:
|
||||
Returns:
|
||||
The action names to run, empty when there are none.
|
||||
"""
|
||||
if tool_doc.get("connection_id"):
|
||||
from docsgpt.connectors.permissions import owner_credential_writes, tool_actions
|
||||
|
||||
if tool_doc.get("connection_id") or tool_doc.get("name") in _LIVE_APPROVAL_TOOLS:
|
||||
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 []
|
||||
owner_writes = set(owner_credential_writes(tool_doc))
|
||||
unasked = {
|
||||
action.get("name") for action in actions
|
||||
action.get("name") for action in tool_actions(tool_doc)
|
||||
if action.get("name") and action.get("active", True) and not action.get("require_approval")
|
||||
and action.get("name") not in owner_writes
|
||||
}
|
||||
if required_actions is None or None in required_actions:
|
||||
return unasked
|
||||
|
||||
@@ -424,14 +424,14 @@ class TestRunTime:
|
||||
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):
|
||||
def _prefetched(self, pg_conn, agent_id, caller, required, **flags):
|
||||
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.agent_config = {"user_api_key": agent["key"], **flags}
|
||||
processor._required_tool_actions = required
|
||||
fetched = {}
|
||||
|
||||
@@ -479,6 +479,30 @@ class TestToolPrefetch:
|
||||
assert self._prefetched(pg_conn, agent_id, VIEWER, required) == {api: {"fetch"}}
|
||||
|
||||
|
||||
def test_tools_that_decide_approval_live_are_not_prefetched_for_others(self, pg_conn):
|
||||
agent_id, _ = _agent(pg_conn)
|
||||
device = self._tool(pg_conn, OWNER, "remote_device", [{"name": "run_command", "active": True}])
|
||||
code = self._tool(pg_conn, OWNER, "code_executor", [{"name": "execute_code", "active": True}])
|
||||
AgentsRepository(pg_conn).update_by_id(agent_id, {"tools": [device, code]})
|
||||
required = {"remote_device": {None}, "code_executor": {None}}
|
||||
assert self._prefetched(pg_conn, agent_id, OWNER, required) == {device: {None}, code: {None}}
|
||||
assert self._prefetched(pg_conn, agent_id, VIEWER, required) == {}
|
||||
|
||||
def test_api_key_callers_prefetch_like_someone_else(self, pg_conn):
|
||||
"""A widget or API run carries the owner's id, but the caller is not the owner."""
|
||||
agent_id, _ = _agent(pg_conn)
|
||||
api = str(UserToolsRepository(pg_conn).create(OWNER, "api_tool", config={"actions": {
|
||||
"status": {"url": "https://x.test/s", "method": "GET", "active": True},
|
||||
"notify": {"url": "https://x.test/n", "method": "POST", "active": True},
|
||||
}})["id"])
|
||||
AgentsRepository(pg_conn).update_by_id(agent_id, {"tools": [api]})
|
||||
required = {"api_tool": {None}}
|
||||
assert self._prefetched(pg_conn, agent_id, OWNER, required) == {api: {None}}
|
||||
# Writes on the owner's credentials never pre-fetch for them.
|
||||
assert self._prefetched(pg_conn, agent_id, OWNER, required, external_api_caller=True) == {api: {"status"}}
|
||||
assert self._prefetched(pg_conn, agent_id, VIEWER, required, public_link_caller=True) == {api: {"status"}}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# workflows
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in new issue
Block a user