mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 12:14:00 +00:00
Stop stale sponsor records from vouching for re-added resources
A resource attached in a save now ignores any sponsor recorded for it before it was removed: the caller must be able to sponsor it and confirm. YAML import prunes the sponsors of resources it drops, for agents and workflow graphs. Sponsor details, with the resources' names, go only to people who may edit the agent. The workflow read returns the names of every node tool and source, so editors can remove the owner's private ones. An agent image is stored only once the save is known to go ahead.
This commit is contained in:
1 parent
2f6f4edac2
commit
4d001f6a20
5 files changed
+292
-26
No files matched your search
@@ -1491,18 +1491,24 @@ def _apply_workflow(
|
||||
if existing_id:
|
||||
existing = wf_repo.get(str(existing_id), user)
|
||||
if existing is not None:
|
||||
from docsgpt.api.user.resource_access import prune_sponsors
|
||||
from docsgpt.api.user.workflows.routes import _node_refs
|
||||
|
||||
pg_workflow_id = str(existing["id"])
|
||||
next_version = get_workflow_graph_version(existing) + 1
|
||||
_write_graph(conn, pg_workflow_id, next_version, nodes_data, edges_data)
|
||||
wf_repo.update(
|
||||
pg_workflow_id,
|
||||
user,
|
||||
{
|
||||
"name": name,
|
||||
"description": description,
|
||||
"current_graph_version": next_version,
|
||||
},
|
||||
)
|
||||
workflow_fields = {
|
||||
"name": name,
|
||||
"description": description,
|
||||
"current_graph_version": next_version,
|
||||
}
|
||||
# Sponsors of node resources the file dropped go too, so a stale
|
||||
# record can't vouch for the resource if someone adds it back.
|
||||
sponsors = existing.get("resource_sponsors") or {}
|
||||
pruned = prune_sponsors(sponsors, _node_refs(nodes_data))
|
||||
if pruned != sponsors:
|
||||
workflow_fields["resource_sponsors"] = pruned
|
||||
wf_repo.update(pg_workflow_id, user, workflow_fields)
|
||||
WorkflowNodesRepository(conn).delete_other_versions(pg_workflow_id, next_version)
|
||||
WorkflowEdgesRepository(conn).delete_other_versions(pg_workflow_id, next_version)
|
||||
return pg_workflow_id
|
||||
@@ -1512,6 +1518,25 @@ def _apply_workflow(
|
||||
return str(created["id"])
|
||||
|
||||
|
||||
def _prune_agent_sponsors(agents_repo: AgentsRepository, agent_id: str, user: str) -> None:
|
||||
"""Drop sponsor records of resources the imported agent no longer references.
|
||||
|
||||
Args:
|
||||
agents_repo: Repository on the import's connection.
|
||||
agent_id: The updated agent.
|
||||
user: Its owner.
|
||||
"""
|
||||
from docsgpt.api.user.resource_access import agent_refs, prune_sponsors
|
||||
|
||||
row = agents_repo.get(agent_id, user)
|
||||
if not row:
|
||||
return
|
||||
sponsors = row.get("resource_sponsors") or {}
|
||||
pruned = prune_sponsors(sponsors, agent_refs(row))
|
||||
if pruned != sponsors:
|
||||
agents_repo.update(agent_id, user, {"resource_sponsors": pruned})
|
||||
|
||||
|
||||
def apply_import(conn, user: str, doc: dict, resolution: Optional[dict] = None) -> dict:
|
||||
"""Create or update an agent from a parsed YAML doc.
|
||||
|
||||
@@ -1648,6 +1673,7 @@ def apply_import(conn, user: str, doc: dict, resolution: Optional[dict] = None)
|
||||
# Only a brand-new agent (the create path below) starts as a draft.
|
||||
fields = {**authoritative, **optional, "name": spec.get("name")}
|
||||
if agents_repo.update(str(target["agent_id"]), user, fields):
|
||||
_prune_agent_sponsors(agents_repo, str(target["agent_id"]), user)
|
||||
if orphaned_workflow_id and not agents_repo.count_by_workflow(
|
||||
orphaned_workflow_id, user
|
||||
):
|
||||
|
||||
@@ -566,8 +566,10 @@ class GetAgent(Resource):
|
||||
ra = resolve(conn, "agent", agent_id, user)
|
||||
if ra is not None:
|
||||
agent = AgentsRepository(conn).get_by_id(ra.resource_id)
|
||||
# Edit-page detail: who vouches for resources the owner can't use.
|
||||
if agent and ra.can("view"):
|
||||
# Edit-page detail: who vouches for resources the owner can't
|
||||
# use, with their names. Only for people who may edit the
|
||||
# agent: the names can be an editor's private resources.
|
||||
if agent and ra.can("edit"):
|
||||
sponsored = sponsor_details(conn, "agent", agent, viewer=user)
|
||||
if not agent:
|
||||
return {"status": "Not found"}, 404
|
||||
@@ -1058,14 +1060,11 @@ class UpdateAgent(Resource):
|
||||
)
|
||||
pg_agent_id = str(existing_agent["id"])
|
||||
existing_image = existing_agent.get("image", "") or ""
|
||||
image_url, image_error = handle_image_upload(
|
||||
request,
|
||||
existing_image,
|
||||
existing_agent.get("user_id") or user,
|
||||
storage,
|
||||
)
|
||||
if image_error:
|
||||
return image_error
|
||||
# The image is stored only once the save is known to go
|
||||
# ahead, so a refused save (a sponsor confirmation round
|
||||
# trip, a validation error) leaves no orphaned file.
|
||||
image_file = request.files.get("image")
|
||||
has_new_image = bool(image_file and image_file.filename)
|
||||
|
||||
update_fields: dict = {}
|
||||
allowed_fields = [
|
||||
@@ -1313,9 +1312,7 @@ class UpdateAgent(Resource):
|
||||
f"Field '{field}' cannot be empty", user, field
|
||||
)
|
||||
update_fields[field] = value
|
||||
if image_url and image_url != existing_image:
|
||||
update_fields["image"] = image_url
|
||||
if not update_fields:
|
||||
if not update_fields and not has_new_image:
|
||||
return _reject("No valid update data provided", user)
|
||||
|
||||
newly_generated_key = None
|
||||
@@ -1467,6 +1464,18 @@ class UpdateAgent(Resource):
|
||||
if plan.sponsors != (existing_agent.get("resource_sponsors") or {}):
|
||||
update_fields["resource_sponsors"] = plan.sponsors
|
||||
|
||||
if has_new_image:
|
||||
image_url, image_error = handle_image_upload(
|
||||
request,
|
||||
existing_image,
|
||||
existing_agent.get("user_id") or user,
|
||||
storage,
|
||||
)
|
||||
if image_error:
|
||||
return image_error
|
||||
if image_url and image_url != existing_image:
|
||||
update_fields["image"] = image_url
|
||||
|
||||
# Apply update. Owner writes use the dual-key guard; team-editor
|
||||
# writes go by-id (already authorized) with an optimistic-lock
|
||||
# check when the client supplies the row's expected updated_at.
|
||||
|
||||
@@ -615,7 +615,9 @@ def plan_sponsors(
|
||||
* a newly attached one the caller may sponsor but didn't confirm goes to
|
||||
``needs_confirmation``; one they may not sponsor goes to ``not_allowed``;
|
||||
* one already attached keeps its old record even when that sponsor lost
|
||||
access: it stays stopped until someone confirms taking it over.
|
||||
access: it stays stopped until someone confirms taking it over;
|
||||
* a newly attached one ignores any record left from before it was
|
||||
removed, so a stale sponsor never vouches for it again.
|
||||
|
||||
Resources the owner can use, presets and builtin tools need no sponsor.
|
||||
Removed refs drop out. A confirmed key that names none of the refs the
|
||||
@@ -652,19 +654,38 @@ def plan_sponsors(
|
||||
caller_may = caller != owner_id and can_sponsor_ref(conn, resource_type, resource_id, caller)
|
||||
if caller_may:
|
||||
eligible.add(key)
|
||||
live = active_sponsor(conn, holder_type, holder, resource_type, resource_id)
|
||||
# A record only vouches for a resource that stayed attached: one left
|
||||
# behind by a path that dropped the resource never covers it again.
|
||||
live = None if is_new else active_sponsor(conn, holder_type, holder, resource_type, resource_id)
|
||||
if live:
|
||||
plan.sponsors[key] = live
|
||||
elif caller_may and key in confirmed:
|
||||
plan.sponsors[key] = caller
|
||||
elif is_new and caller != owner_id:
|
||||
(plan.needs_confirmation if caller_may else plan.not_allowed).append((resource_type, resource_id))
|
||||
elif previous.get(key):
|
||||
elif not is_new and previous.get(key):
|
||||
plan.sponsors[key] = previous[key]
|
||||
plan.unexpected = sorted(confirmed - eligible)
|
||||
return plan
|
||||
|
||||
|
||||
def prune_sponsors(sponsors: Optional[dict], refs: Iterable[tuple[str, str]]) -> dict:
|
||||
"""``sponsors`` without the keys of resources no longer referenced.
|
||||
|
||||
For paths that rewrite a holder's references without going through
|
||||
:func:`plan_sponsors` (YAML import, a workflow graph written by import).
|
||||
|
||||
Args:
|
||||
sponsors: The stored ``resource_sponsors`` map.
|
||||
refs: Every ``(type, id)`` the holder references now.
|
||||
|
||||
Returns:
|
||||
dict: The map to store.
|
||||
"""
|
||||
keep = {sponsor_key(t, str(i).lower()) for t, i in refs}
|
||||
return {k: v for k, v in (sponsors or {}).items() if k.lower() in keep}
|
||||
|
||||
|
||||
def ref_names(conn: Connection, refs: Iterable[tuple[str, str]]) -> dict[str, str]:
|
||||
"""Display names of referenced resources, looked up by id (owner-agnostic).
|
||||
|
||||
|
||||
@@ -127,6 +127,33 @@ def _node_refs(nodes: List[Dict]) -> List[Tuple[str, str]]:
|
||||
return refs
|
||||
|
||||
|
||||
def _node_ref_details(nodes: List[Dict]) -> Dict[str, List[Dict]]:
|
||||
"""Names of every tool and source the graph's agent nodes reference.
|
||||
|
||||
Looked up by id whoever owns them, so an editor's node pickers can show
|
||||
(and remove) the owner's private tools and sources.
|
||||
|
||||
Args:
|
||||
nodes: Nodes in builder shape.
|
||||
|
||||
Returns:
|
||||
dict: ``tools`` as ``[{id, name, display_name}]`` and ``sources`` as
|
||||
``[{id, name}]``, each id once.
|
||||
"""
|
||||
from docsgpt.api.user.base import resolve_source_details, resolve_tool_details
|
||||
|
||||
tool_ids: List[str] = []
|
||||
source_ids: List[str] = []
|
||||
for resource_type, resource_id in _node_refs(nodes):
|
||||
bucket = tool_ids if resource_type == "tool" else source_ids
|
||||
if resource_id not in bucket:
|
||||
bucket.append(resource_id)
|
||||
return {
|
||||
"tools": resolve_tool_details(tool_ids),
|
||||
"sources": resolve_source_details(source_ids),
|
||||
}
|
||||
|
||||
|
||||
def _new_node_ref_denied(
|
||||
conn, previous_nodes: List[Dict], new_nodes: List[Dict], caller: str
|
||||
) -> Optional[AccessDenied]:
|
||||
@@ -634,15 +661,18 @@ class WorkflowDetail(Resource):
|
||||
pg_workflow_id, graph_version,
|
||||
)
|
||||
sponsored = sponsor_details(conn, "workflow", workflow, viewer=user_id)
|
||||
serialized_nodes = [serialize_node(n) for n in nodes]
|
||||
ref_details = _node_ref_details(serialized_nodes)
|
||||
except Exception as err:
|
||||
return _workflow_error_response("Failed to fetch workflow", err)
|
||||
|
||||
return success_response(
|
||||
{
|
||||
"workflow": serialize_workflow(workflow),
|
||||
"nodes": [serialize_node(n) for n in nodes],
|
||||
"nodes": serialized_nodes,
|
||||
"edges": [serialize_edge(e) for e in edges],
|
||||
"resource_sponsors": sponsored,
|
||||
"ref_details": ref_details,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -431,3 +431,183 @@ class TestWorkflowConsent:
|
||||
resp = _call(app, pg_conn, WorkflowDetail, "get", f"/api/workflows/{wid}", OWNER, args=(wid,))
|
||||
[detail] = _body(resp)["data"]["resource_sponsors"]
|
||||
assert detail["state"] == "active" and detail["reason"] is None and detail["name"] == "api_tool"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Review follow-ups
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _viewer_share(conn, team_id, resource_type, resource_id, owner, user):
|
||||
TeamResourceGrantsRepository(conn).grant(team_id, resource_type, resource_id, owner, owner,
|
||||
access_level="viewer", target_user_id=user)
|
||||
|
||||
|
||||
class TestNoResurrectedSponsor:
|
||||
"""A sponsor record left behind never vouches for a newly attached resource."""
|
||||
|
||||
def _setup(self, app, pg_conn):
|
||||
"""EDITOR sponsors OTHER's tool T (edit share); STRANGER edits the agent, views T."""
|
||||
agent_id, team_id = _agent(pg_conn)
|
||||
tool = str(UserToolsRepository(pg_conn).create(OTHER, "api_tool")["id"])
|
||||
TeamMembersRepository(pg_conn).add_member(team_id, OTHER)
|
||||
TeamResourceGrantsRepository(pg_conn).grant(team_id, "tool", tool, OTHER, OTHER, access_level="editor",
|
||||
target_user_id=EDITOR)
|
||||
_add_agent_editor(pg_conn, team_id, agent_id, STRANGER)
|
||||
_viewer_share(pg_conn, team_id, "tool", tool, OTHER, STRANGER)
|
||||
assert _status(_put(app, pg_conn, agent_id, EDITOR,
|
||||
{"tools": [tool], "confirm_sponsor": [f"tool:{tool}"]})) == 200
|
||||
return agent_id, team_id, tool
|
||||
|
||||
def test_stale_record_does_not_cover_a_new_attachment(self, app, pg_conn):
|
||||
agent_id, _, tool = self._setup(app, pg_conn)
|
||||
# Some path drops the tool but leaves EDITOR's record behind.
|
||||
AgentsRepository(pg_conn).update_by_id(agent_id, {"tools": []})
|
||||
assert _row(pg_conn, agent_id)["resource_sponsors"] == {sponsor_key("tool", tool): EDITOR}
|
||||
resp = _put(app, pg_conn, agent_id, STRANGER, {"tools": [tool]})
|
||||
assert _status(resp) == 403
|
||||
assert _body(resp)["code"] == CODE_NOT_ALLOWED
|
||||
assert not _row(pg_conn, agent_id)["tools"]
|
||||
|
||||
def test_new_attachment_by_a_sponsor_capable_editor_asks_them(self, app, pg_conn):
|
||||
agent_id, team_id, tool = self._setup(app, pg_conn)
|
||||
AgentsRepository(pg_conn).update_by_id(agent_id, {"tools": []})
|
||||
_add_agent_editor(pg_conn, team_id, agent_id, OTHER)
|
||||
resp = _put(app, pg_conn, agent_id, OTHER, {"tools": [tool]})
|
||||
assert _status(resp) == 409
|
||||
resp = _put(app, pg_conn, agent_id, OTHER, {"tools": [tool], "confirm_sponsor": [f"tool:{tool}"]})
|
||||
assert _status(resp) == 200
|
||||
assert _row(pg_conn, agent_id)["resource_sponsors"] == {sponsor_key("tool", tool): OTHER}
|
||||
|
||||
def test_reimport_prunes_dropped_sponsors(self, app, pg_conn):
|
||||
from docsgpt.api.user.agents.portability import API_VERSION, apply_import, ensure_agent_slug
|
||||
|
||||
agent_id, _, tool = self._setup(app, pg_conn)
|
||||
agent = AgentsRepository(pg_conn).get(agent_id, OWNER)
|
||||
slug = ensure_agent_slug(pg_conn, agent, OWNER)
|
||||
doc = {"apiVersion": API_VERSION, "kind": "Agent", "metadata": {"id": agent_id, "slug": slug},
|
||||
"spec": {"name": "Shared", "description": "d", "retriever": "classic"}}
|
||||
assert apply_import(pg_conn, OWNER, doc)["action"] == "updated"
|
||||
row = _row(pg_conn, agent_id)
|
||||
assert not row["tools"]
|
||||
assert row["resource_sponsors"] == {}
|
||||
# The repro: STRANGER only views T, re-adds it, and must not run it as EDITOR.
|
||||
resp = _put(app, pg_conn, agent_id, STRANGER, {"tools": [tool]})
|
||||
assert _status(resp) == 403
|
||||
assert _row(pg_conn, agent_id)["resource_sponsors"] == {}
|
||||
|
||||
def test_workflow_reimport_prunes_dropped_sponsors(self, app, pg_conn):
|
||||
from docsgpt.api.user.agents.portability import (
|
||||
agent_to_yaml,
|
||||
apply_import,
|
||||
parse_agent_yaml,
|
||||
serialize_agent,
|
||||
)
|
||||
from docsgpt.api.user.workflows.routes import WorkflowDetail
|
||||
|
||||
wf = WorkflowsRepository(pg_conn).create(OWNER, "wf")
|
||||
wid = str(wf["id"])
|
||||
agent_id, team_id = _agent(pg_conn, agent_type="workflow", workflow_id=wid)
|
||||
tool, _, _ = _editor_resources(pg_conn)
|
||||
resp = _call(app, pg_conn, WorkflowDetail, "put", f"/api/workflows/{wid}", EDITOR,
|
||||
json={**_wf_body(tool), "confirm_sponsor": [f"tool:{tool}"]}, args=(wid,))
|
||||
assert _status(resp) == 200, _body(resp)
|
||||
assert WorkflowsRepository(pg_conn).get_by_id(wid)["resource_sponsors"]
|
||||
|
||||
agent = AgentsRepository(pg_conn).get(agent_id, OWNER)
|
||||
doc = parse_agent_yaml(agent_to_yaml(serialize_agent(pg_conn, agent, OWNER)))
|
||||
for node in doc["spec"]["workflow"]["nodes"]:
|
||||
if node.get("type") == "agent":
|
||||
node["config"]["tools"] = []
|
||||
apply_import(pg_conn, OWNER, doc)
|
||||
assert WorkflowsRepository(pg_conn).get_by_id(wid)["resource_sponsors"] == {}
|
||||
|
||||
|
||||
class TestTakeOver:
|
||||
def test_confirming_a_stopped_attached_item_takes_it_over(self, app, pg_conn):
|
||||
from docsgpt.api.user.agents.routes import GetAgent
|
||||
|
||||
agent_id, team_id = _agent(pg_conn)
|
||||
tool = str(UserToolsRepository(pg_conn).create(OTHER, "api_tool")["id"])
|
||||
TeamMembersRepository(pg_conn).add_member(team_id, OTHER)
|
||||
TeamResourceGrantsRepository(pg_conn).grant(team_id, "tool", tool, OTHER, OTHER, access_level="editor")
|
||||
_add_agent_editor(pg_conn, team_id, agent_id, OTHER)
|
||||
assert _status(_put(app, pg_conn, agent_id, EDITOR,
|
||||
{"tools": [tool], "confirm_sponsor": [f"tool:{tool}"]})) == 200
|
||||
TeamResourceGrantsRepository(pg_conn).revoke(team_id, "agent", agent_id, target_user_id=EDITOR)
|
||||
|
||||
path = f"/api/get_agent?id={agent_id}"
|
||||
[stopped] = _call(app, pg_conn, GetAgent, "get", path, OTHER).get_json()["resource_sponsors"]
|
||||
assert stopped["can_confirm"] is True and stopped["reason"] == REASON_CANNOT_EDIT_HOLDER
|
||||
# A plain save (no change to the tools) that confirms the stopped key.
|
||||
resp = _put(app, pg_conn, agent_id, OTHER, {"name": "Kept", "confirm_sponsor": [stopped["key"]]})
|
||||
assert _status(resp) == 200, _body(resp)
|
||||
[running] = _call(app, pg_conn, GetAgent, "get", path, OTHER).get_json()["resource_sponsors"]
|
||||
assert running["user_id"] == OTHER and running["state"] == "active"
|
||||
|
||||
|
||||
class TestDetailsVisibility:
|
||||
def test_viewers_get_no_sponsor_details(self, app, pg_conn):
|
||||
from docsgpt.api.user.agents.routes import GetAgent
|
||||
from tests.api.user.test_resource_sponsors import VIEWER
|
||||
|
||||
agent_id, _ = _agent(pg_conn)
|
||||
tool, _, _ = _editor_resources(pg_conn)
|
||||
assert _status(_put(app, pg_conn, agent_id, EDITOR,
|
||||
{"tools": [tool], "confirm_sponsor": [f"tool:{tool}"]})) == 200
|
||||
path = f"/api/get_agent?id={agent_id}"
|
||||
assert _call(app, pg_conn, GetAgent, "get", path, VIEWER).get_json()["resource_sponsors"] == []
|
||||
assert _call(app, pg_conn, GetAgent, "get", path, OWNER).get_json()["resource_sponsors"]
|
||||
|
||||
def test_workflow_viewers_are_refused(self, app, pg_conn):
|
||||
from docsgpt.api.user.workflows.routes import WorkflowDetail
|
||||
from tests.api.user.test_resource_sponsors import VIEWER
|
||||
|
||||
wf = WorkflowsRepository(pg_conn).create(OWNER, "wf")
|
||||
wid = str(wf["id"])
|
||||
_agent(pg_conn, agent_type="workflow", workflow_id=wid)
|
||||
resp = _call(app, pg_conn, WorkflowDetail, "get", f"/api/workflows/{wid}", VIEWER, args=(wid,))
|
||||
assert _status(resp) in (403, 404)
|
||||
|
||||
|
||||
class TestWorkflowNodeRefDetails:
|
||||
def test_get_names_every_node_tool_and_source(self, app, pg_conn):
|
||||
from docsgpt.api.user.workflows.routes import WorkflowDetail
|
||||
|
||||
wf = WorkflowsRepository(pg_conn).create(OWNER, "wf")
|
||||
wid = str(wf["id"])
|
||||
_agent(pg_conn, agent_type="workflow", workflow_id=wid)
|
||||
tool = str(UserToolsRepository(pg_conn).create(OWNER, "api_tool", display_name="Owner API")["id"])
|
||||
source = str(SourcesRepository(pg_conn).create("owner-src", user_id=OWNER)["id"])
|
||||
assert _status(_call(app, pg_conn, WorkflowDetail, "put", f"/api/workflows/{wid}", OWNER,
|
||||
json=_wf_body(tool, source), args=(wid,))) == 200
|
||||
data = _body(_call(app, pg_conn, WorkflowDetail, "get", f"/api/workflows/{wid}", EDITOR,
|
||||
args=(wid,)))["data"]
|
||||
assert data["ref_details"] == {
|
||||
"tools": [{"id": tool, "name": "api_tool", "display_name": "Owner API"}],
|
||||
"sources": [{"id": source, "name": "owner-src"}],
|
||||
}
|
||||
|
||||
|
||||
class TestImageAfterSponsorCheck:
|
||||
def test_refused_save_stores_no_image(self, app, pg_conn):
|
||||
import io
|
||||
from unittest.mock import patch as _patch
|
||||
|
||||
from docsgpt.api.user.agents.routes import UpdateAgent
|
||||
from tests.api.user.test_resource_sponsors import _patch_db
|
||||
|
||||
agent_id, _ = _agent(pg_conn)
|
||||
tool, _, _ = _editor_resources(pg_conn)
|
||||
data = {"tools": f'["{tool}"]', "image": (io.BytesIO(b"png"), "a.png")}
|
||||
with _patch_db(pg_conn), \
|
||||
_patch("docsgpt.api.user.agents.routes.handle_image_upload",
|
||||
return_value=("stored.png", None)) as upload, \
|
||||
app.test_request_context(f"/api/update_agent/{agent_id}", method="PUT", data=data,
|
||||
content_type="multipart/form-data"):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": EDITOR}
|
||||
resp = UpdateAgent().put(agent_id)
|
||||
assert _status(resp) == 409
|
||||
upload.assert_not_called()
|
||||
Reference in new issue
Block a user