diff --git a/docsgpt/api/user/agents/portability.py b/docsgpt/api/user/agents/portability.py index 289f2512..005b6473 100644 --- a/docsgpt/api/user/agents/portability.py +++ b/docsgpt/api/user/agents/portability.py @@ -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 ): diff --git a/docsgpt/api/user/agents/routes.py b/docsgpt/api/user/agents/routes.py index 65089a9a..4680915d 100644 --- a/docsgpt/api/user/agents/routes.py +++ b/docsgpt/api/user/agents/routes.py @@ -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. diff --git a/docsgpt/api/user/resource_access.py b/docsgpt/api/user/resource_access.py index d9a35d48..4fcf3bb1 100644 --- a/docsgpt/api/user/resource_access.py +++ b/docsgpt/api/user/resource_access.py @@ -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). diff --git a/docsgpt/api/user/workflows/routes.py b/docsgpt/api/user/workflows/routes.py index 65625f80..5c4729a9 100644 --- a/docsgpt/api/user/workflows/routes.py +++ b/docsgpt/api/user/workflows/routes.py @@ -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, } ) diff --git a/tests/api/user/test_sponsor_consent.py b/tests/api/user/test_sponsor_consent.py index d0c01db3..25471cb0 100644 --- a/tests/api/user/test_sponsor_consent.py +++ b/tests/api/user/test_sponsor_consent.py @@ -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()