mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 22:15:10 +00:00
Hold resumed turns and public-link visitors to the right wiki rules
A paused turn is found by the agent owner's id, so anyone holding one of the owner's agent keys could resume the owner's own chat with its saved wiki edit rights. A resume now counts as an API or widget caller when either the saved state or the resuming request is one, cuts the wiki tool to wiki_view unless the wiki allows outside edits, and gives the tool executor the same flags. A request that names an agent, by key or id, may only resume that agent's turn; otherwise the claim is released and the request refused. Public-link visitors run as themselves and reach only wikis they may edit, so the wiki switch no longer applies to them. Instead every wiki write in a public-link run waits for the visitor's approval, so the agent owner's prompt or sources can't steer an edit to the visitor's wiki unasked.
This commit is contained in:
arc53-machine committed
1 parent
91e70d2f0c
commit
051452c438
10 files changed
+405
-64
No files matched your search
@@ -0,0 +1,201 @@
|
||||
"""A resumed turn keeps the wiki and write rules of whoever resumes it.
|
||||
|
||||
Continuation state is looked up by the agent owner, so a request carrying the
|
||||
agent's key (a widget key is public) could resume the owner's own paused chat.
|
||||
The resumed run counts as an outside caller when either the saved state or the
|
||||
resuming request says so, and a request may only resume a turn of the agent it
|
||||
names.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
AGENT = "11111111-1111-1111-1111-111111111111"
|
||||
OTHER_AGENT = "22222222-2222-2222-2222-222222222222"
|
||||
OWNER = "owner"
|
||||
|
||||
|
||||
def _wiki_entry(**config):
|
||||
from docsgpt.agents.tools.wiki import WIKI_TOOL_ID, add_wiki_tool
|
||||
|
||||
tools = {}
|
||||
add_wiki_tool(tools, {"source_id": "wiki-1", "source_owner_id": OWNER, "user": OWNER, **config})
|
||||
return tools[WIKI_TOOL_ID]
|
||||
|
||||
|
||||
def _state(agent_id=AGENT, **flags):
|
||||
return {
|
||||
"messages": [],
|
||||
"pending_tool_calls": [],
|
||||
"tools_dict": {"wiki": _wiki_entry()},
|
||||
"tool_schemas": [],
|
||||
"client_tools": None,
|
||||
"agent_config": {
|
||||
"model_id": "m1", "llm_name": "openai", "api_key": "k", "user_api_key": "agent-key",
|
||||
"agent_type": "ClassicAgent", "agent_id": agent_id, **flags,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def resume(monkeypatch):
|
||||
"""Resume a saved state; returns (tools_dict, executor, continuation service)."""
|
||||
from docsgpt.agents import agent_creator as ac_mod
|
||||
from docsgpt.api.answer.services import continuation_service as cont_mod
|
||||
from docsgpt.api.answer.services import stream_processor as sp_mod
|
||||
from docsgpt.llm import llm_creator as llm_creator_mod
|
||||
from docsgpt.llm.handlers import handler_creator as handler_mod
|
||||
|
||||
@contextmanager
|
||||
def _noop():
|
||||
yield None
|
||||
|
||||
agents = {"agent-key": {"id": AGENT, "user_id": OWNER}, "other-key": {"id": OTHER_AGENT, "user_id": OWNER}}
|
||||
|
||||
class _Agents:
|
||||
def __init__(self, conn):
|
||||
pass
|
||||
|
||||
def find_by_key(self, key):
|
||||
return agents.get(key)
|
||||
|
||||
live = {"allowed": False}
|
||||
|
||||
class _Sources:
|
||||
def __init__(self, conn):
|
||||
pass
|
||||
|
||||
def get_by_id(self, sid):
|
||||
return {"id": sid, "wiki_outside_edits": live["allowed"]}
|
||||
|
||||
def _get_agent_key(self, agent_id, user_id):
|
||||
# Only the owner reaches the agent directly; anyone else by its link.
|
||||
self.public_link_usage = user_id != OWNER
|
||||
return "agent-key", user_id != OWNER, None
|
||||
|
||||
monkeypatch.setattr(sp_mod.StreamProcessor, "_get_agent_key", _get_agent_key)
|
||||
monkeypatch.setattr(sp_mod, "db_readonly", _noop)
|
||||
monkeypatch.setattr(sp_mod, "AgentsRepository", _Agents)
|
||||
monkeypatch.setattr("docsgpt.agents.tools.wiki.db_readonly", _noop)
|
||||
monkeypatch.setattr("docsgpt.agents.tools.wiki.SourcesRepository", _Sources)
|
||||
monkeypatch.setattr(llm_creator_mod.LLMCreator, "create_llm", lambda *a, **kw: MagicMock())
|
||||
monkeypatch.setattr(handler_mod.LLMHandlerCreator, "create_handler", lambda *a, **kw: MagicMock())
|
||||
created = {}
|
||||
monkeypatch.setattr(ac_mod.AgentCreator, "create_agent", lambda *a, **kw: created.update(kw) or MagicMock())
|
||||
|
||||
def _run(state, data, token, *, external_caller=False, allowed=False):
|
||||
live["allowed"] = allowed
|
||||
cont_service = MagicMock()
|
||||
cont_service.claim_state.return_value = copy.deepcopy(state)
|
||||
monkeypatch.setattr(cont_mod, "ContinuationService", lambda: cont_service)
|
||||
processor = sp_mod.StreamProcessor(data, token, external_caller=external_caller)
|
||||
result = processor.resume_from_tool_actions(tool_actions=[], conversation_id=str(uuid.uuid4()))
|
||||
return result[2], created["tool_executor"], cont_service
|
||||
|
||||
return _run
|
||||
|
||||
|
||||
def _actions(tools_dict):
|
||||
return [a["name"] for a in tools_dict["wiki"]["actions"]]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestResumeByOutsideCaller:
|
||||
def test_widget_key_resuming_the_owners_chat_gets_read_only_wiki(self, resume):
|
||||
tools, executor, _ = resume(_state(), {"api_key": "agent-key"}, None)
|
||||
assert executor.external_caller is True
|
||||
assert _actions(tools) == ["wiki_view"]
|
||||
assert tools["wiki"]["config"]["outside_caller"] is True
|
||||
|
||||
def test_signed_in_stranger_with_the_key_is_outside(self, resume):
|
||||
tools, executor, _ = resume(_state(), {"api_key": "agent-key"}, {"sub": "stranger"})
|
||||
assert executor.external_caller is True
|
||||
assert _actions(tools) == ["wiki_view"]
|
||||
|
||||
def test_edits_stay_when_the_wiki_allows_them(self, resume):
|
||||
tools, executor, _ = resume(_state(), {"api_key": "agent-key"}, None, allowed=True)
|
||||
assert executor.external_caller is True
|
||||
assert "wiki_create" in _actions(tools)
|
||||
# The tool still re-checks the live setting on every write.
|
||||
assert tools["wiki"]["config"]["outside_caller"] is True
|
||||
|
||||
def test_v1_key_holder_is_outside(self, resume):
|
||||
tools, executor, _ = resume(_state(), {"api_key": "agent-key"}, {"sub": OWNER}, external_caller=True)
|
||||
assert executor.external_caller is True
|
||||
assert _actions(tools) == ["wiki_view"]
|
||||
|
||||
def test_saved_outside_state_stays_outside_for_the_owner(self, resume):
|
||||
tools, executor, _ = resume(_state(external_api_caller=True), {}, {"sub": OWNER})
|
||||
assert executor.external_caller is True
|
||||
assert _actions(tools) == ["wiki_view"]
|
||||
|
||||
def test_owner_resuming_in_app_is_unaffected(self, resume):
|
||||
tools, executor, _ = resume(_state(), {}, {"sub": OWNER})
|
||||
assert executor.external_caller is False
|
||||
assert executor.public_link_caller is False
|
||||
assert "wiki_create" in _actions(tools)
|
||||
assert tools["wiki"]["config"]["outside_caller"] is False
|
||||
|
||||
def test_owner_previewing_with_the_key_is_unaffected(self, resume):
|
||||
tools, executor, _ = resume(_state(), {"api_key": "agent-key"}, {"sub": OWNER})
|
||||
assert executor.external_caller is False
|
||||
assert "wiki_create" in _actions(tools)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestResumeTargetsTheSameAgent:
|
||||
def test_a_key_for_another_agent_is_refused(self, resume):
|
||||
with pytest.raises(ValueError):
|
||||
resume(_state(), {"api_key": "other-key"}, None)
|
||||
|
||||
def test_release_on_refusal(self, resume, monkeypatch):
|
||||
from docsgpt.api.answer.services import continuation_service as cont_mod
|
||||
from docsgpt.api.answer.services import stream_processor as sp_mod
|
||||
|
||||
cont_service = MagicMock()
|
||||
cont_service.claim_state.return_value = _state()
|
||||
monkeypatch.setattr(cont_mod, "ContinuationService", lambda: cont_service)
|
||||
processor = sp_mod.StreamProcessor({"api_key": "other-key"}, None)
|
||||
conversation_id = str(uuid.uuid4())
|
||||
with pytest.raises(ValueError):
|
||||
processor.resume_from_tool_actions(tool_actions=[], conversation_id=conversation_id)
|
||||
cont_service.release_claim.assert_called_once_with(conversation_id, OWNER)
|
||||
|
||||
def test_a_key_cannot_resume_an_agentless_chat(self, resume):
|
||||
with pytest.raises(ValueError):
|
||||
resume(_state(agent_id=None), {"api_key": "agent-key"}, None)
|
||||
|
||||
def test_a_named_agent_must_match(self, resume):
|
||||
with pytest.raises(ValueError):
|
||||
resume(_state(), {"agent_id": OTHER_AGENT}, {"sub": OWNER})
|
||||
|
||||
def test_same_agent_passes(self, resume):
|
||||
tools, _executor, _ = resume(_state(), {"agent_id": AGENT.upper()}, {"sub": OWNER})
|
||||
assert "wiki_create" in _actions(tools)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestResumePublicLink:
|
||||
def test_saved_public_link_run_keeps_wiki_writes_behind_approval(self, resume):
|
||||
state = _state(public_link_caller=True)
|
||||
tools, executor, _ = resume(state, {}, {"sub": "visitor"})
|
||||
assert executor.public_link_caller is True
|
||||
writes = [a for a in tools["wiki"]["actions"] if a["name"] != "wiki_view"]
|
||||
assert writes and all(a.get("require_approval") for a in writes)
|
||||
view = next(a for a in tools["wiki"]["actions"] if a["name"] == "wiki_view")
|
||||
assert not view.get("require_approval")
|
||||
# Public-link visitors aren't covered by the switch.
|
||||
assert tools["wiki"]["config"]["outside_caller"] is False
|
||||
|
||||
def test_a_request_naming_the_agent_by_its_link_is_public(self, resume):
|
||||
tools, executor, _ = resume(_state(), {"agent_id": AGENT}, {"sub": "visitor"})
|
||||
assert executor.public_link_caller is True
|
||||
assert all(
|
||||
a.get("require_approval") for a in tools["wiki"]["actions"] if a["name"] != "wiki_view"
|
||||
)
|
||||
@@ -119,7 +119,7 @@ class TestWikiConfigAccess:
|
||||
|
||||
|
||||
class TestWikiOutsideEdits:
|
||||
"""API, widget and public-link runs edit a wiki only when its owner allows."""
|
||||
"""API-key runs edit a wiki only when its owner allows; link visitors are asked."""
|
||||
|
||||
def _tools(self, conn, caller, sid, agent_config):
|
||||
from docsgpt.agents.tools.wiki import WIKI_TOOL_ID, add_wiki_tool
|
||||
@@ -130,19 +130,44 @@ class TestWikiOutsideEdits:
|
||||
proc.decoded_token = {"sub": caller}
|
||||
proc.agent_config = agent_config
|
||||
cfg = proc._build_wiki_config()
|
||||
if cfg is None:
|
||||
return None
|
||||
tools = {}
|
||||
add_wiki_tool(tools, cfg)
|
||||
return {a["name"] for a in tools[WIKI_TOOL_ID]["actions"]}
|
||||
return {a["name"]: bool(a.get("require_approval")) for a in tools[WIKI_TOOL_ID]["actions"]}
|
||||
|
||||
@pytest.mark.parametrize("flag", ["external_api_caller", "public_link_caller"])
|
||||
def test_outside_caller_reads_until_the_owner_allows_edits(self, use_conn, flag):
|
||||
sid = str(SourcesRepository(use_conn).create("W", user_id=OWNER, config={"kind": "wiki"})["id"])
|
||||
assert self._tools(use_conn, OWNER, sid, {flag: True}) == {"wiki_view"}
|
||||
def _wiki(self, conn, owner=OWNER):
|
||||
return str(SourcesRepository(conn).create("W", user_id=owner, config={"kind": "wiki"})["id"])
|
||||
|
||||
def test_api_key_caller_reads_until_the_owner_allows_edits(self, use_conn):
|
||||
# A widget or API run acts as the agent owner, who owns the wiki.
|
||||
sid = self._wiki(use_conn)
|
||||
assert self._tools(use_conn, OWNER, sid, {"external_api_caller": True}) == {"wiki_view": False}
|
||||
SourcesRepository(use_conn).set_wiki_outside_edits(sid, OWNER, True)
|
||||
assert "wiki_create" in self._tools(use_conn, OWNER, sid, {flag: True})
|
||||
tools = self._tools(use_conn, OWNER, sid, {"external_api_caller": True})
|
||||
assert tools["wiki_create"] is False
|
||||
|
||||
def test_owner_and_team_editor_keep_every_action(self, use_conn):
|
||||
sid = str(SourcesRepository(use_conn).create("W", user_id=OWNER, config={"kind": "wiki"})["id"])
|
||||
sid = self._wiki(use_conn)
|
||||
_share(use_conn, "source", sid, "ed", "editor")
|
||||
assert "wiki_create" in self._tools(use_conn, OWNER, sid, {})
|
||||
assert "wiki_create" in self._tools(use_conn, "ed", sid, {})
|
||||
assert self._tools(use_conn, OWNER, sid, {})["wiki_create"] is False
|
||||
assert self._tools(use_conn, "ed", sid, {})["wiki_create"] is False
|
||||
|
||||
def test_link_visitor_edits_their_own_wiki_after_approving(self, use_conn):
|
||||
# Visitor owns a wiki that the agent (someone else's) also uses.
|
||||
sid = self._wiki(use_conn, owner="visitor")
|
||||
tools = self._tools(use_conn, "visitor", sid, {"public_link_caller": True})
|
||||
assert tools["wiki_view"] is False
|
||||
assert tools["wiki_create"] is True and tools["wiki_delete"] is True
|
||||
|
||||
def test_link_visitor_who_edits_the_wiki_through_a_team(self, use_conn):
|
||||
sid = self._wiki(use_conn)
|
||||
_share(use_conn, "source", sid, "visitor", "editor")
|
||||
tools = self._tools(use_conn, "visitor", sid, {"public_link_caller": True})
|
||||
assert tools["wiki_create"] is True
|
||||
|
||||
def test_link_visitor_without_edit_rights_gets_no_wiki_tool(self, use_conn):
|
||||
sid = self._wiki(use_conn)
|
||||
assert self._tools(use_conn, "visitor", sid, {"public_link_caller": True}) is None
|
||||
_share(use_conn, "source", sid, "visitor", "viewer")
|
||||
assert self._tools(use_conn, "visitor", sid, {"public_link_caller": True}) is None
|
||||
Reference in new issue
Block a user