MCP rollback for connectors

tests/api/user/agents/test_roles_access.py

@@ -246,14 +246,22 @@ class TestUpdateAgent:
                          {"name": "n", "folder_id": str(folder["id"])})
         assert _status(resp) == 200
 
-    def test_workflow_validated_against_owner(self, app, pg_conn):
+    def test_workflow_validated_against_owner_and_owner_only_to_change(self, app, pg_conn):
         from docsgpt.storage.db.repositories.workflows import WorkflowsRepository
 
         wf = WorkflowsRepository(pg_conn).create(OWNER, "wf")
         agent_id = _agent(pg_conn, agent_type="workflow")
+        # An editor can't point the agent at another of the owner's workflows.
         resp = self._put(app, pg_conn, agent_id, EDITOR, {"workflow": str(wf["id"])})
+        assert _status(resp) == 403
+        assert _row(pg_conn, agent_id)["workflow_id"] is None
+
+        resp = self._put(app, pg_conn, agent_id, OWNER, {"workflow": str(wf["id"])})
         assert _status(resp) == 200
         assert str(_row(pg_conn, agent_id)["workflow_id"]) == str(wf["id"])
+        # Re-sending the current one is a plain save.
+        resp = self._put(app, pg_conn, agent_id, EDITOR, {"workflow": str(wf["id"])})
+        assert _status(resp) == 200
 
         mine = WorkflowsRepository(pg_conn).create(EDITOR, "editor-wf")
         resp = self._put(app, pg_conn, agent_id, EDITOR, {"workflow": str(mine["id"])})
@@ -269,10 +277,14 @@ class TestUpdateAgent:
         _team_share(pg_conn, "tool", shared_tool, owner=STRANGER)
         agent_id = _agent(pg_conn)
 
-        ok = self._put(app, pg_conn, agent_id, EDITOR, {"tools": [owner_tool, shared_tool]})
-        assert _status(ok) == 200
+        # The owner owning a tool isn't enough: it must reach the editor.
+        denied = self._put(app, pg_conn, agent_id, EDITOR, {"tools": [owner_tool]})
+        assert _status(denied) == 403
         denied = self._put(app, pg_conn, agent_id, EDITOR, {"tools": [foreign_tool]})
         assert _status(denied) == 403
+        assert _status(self._put(app, pg_conn, agent_id, OWNER, {"tools": [owner_tool]})) == 200
+        ok = self._put(app, pg_conn, agent_id, EDITOR, {"tools": [owner_tool, shared_tool]})
+        assert _status(ok) == 200
 
         set_settings(pg_conn, "tool", shared_tool, {"viewers_can_use_in_agents": False}, STRANGER)
         # Already attached: keeping it is fine.

tests/api/user/test_attach_owner_refs.py

@@ -0,0 +1,212 @@
+"""What a team editor may newly reference from the owner's agent or workflow.
+
+An agent (and its workflow) runs as its owner, so owning a resource is not
+enough for an editor to wire it in: the editor must be able to use it
+themselves. Otherwise an editor of one shared agent could attach the owner's
+private tool (run with the owner's credentials), source or prompt, or the
+workflow of another of the owner's agents and then edit that graph.
+
+Also covers id casing: an uppercase UUID must resolve the same switches as
+the canonical lowercase one. Uses real repositories on ``pg_conn``.
+"""
+
+from __future__ import annotations
+
+import uuid
+
+import pytest
+
+from docsgpt.api.user.resource_access import resolve, set_settings
+from docsgpt.storage.db.repositories.agents import AgentsRepository
+from docsgpt.storage.db.repositories.prompts import PromptsRepository
+from docsgpt.storage.db.repositories.sources import SourcesRepository
+from docsgpt.storage.db.repositories.team_resource_grants import (
+    TeamResourceGrantsRepository,
+)
+from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
+from docsgpt.storage.db.repositories.workflows import WorkflowsRepository
+
+from tests.api.user.test_resource_sponsors import (
+    EDITOR,
+    OWNER,
+    _agent,
+    _call,
+    _put,
+    _row,
+    _status,
+    _wf_body,
+)
+
+
+@pytest.fixture
+def app():
+    from flask import Flask
+
+    return Flask(__name__)
+
+
+def _owner_private(conn):
+    """A tool, prompt and source the owner never shared."""
+    tool = str(UserToolsRepository(conn).create(OWNER, "api_tool")["id"])
+    prompt = str(PromptsRepository(conn).create(OWNER, "private", "Owner prompt")["id"])
+    source = str(SourcesRepository(conn).create("owner-src", user_id=OWNER)["id"])
+    return tool, prompt, source
+
+
+def _share_tool_with_editor(conn, team_id, tool, level="viewer"):
+    TeamResourceGrantsRepository(conn).grant(
+        team_id, "tool", tool, OWNER, OWNER, access_level=level, target_user_id=EDITOR
+    )
+
+
+class TestAgentAttach:
+    def test_editor_cannot_attach_owner_private_tool(self, app, pg_conn):
+        agent_id, _ = _agent(pg_conn)
+        tool, _, _ = _owner_private(pg_conn)
+        resp = _put(app, pg_conn, agent_id, EDITOR, {"tools": [tool]})
+        assert _status(resp) == 403
+        assert _row(pg_conn, agent_id)["tools"] in (None, [])
+
+    def test_editor_cannot_attach_owner_private_source(self, app, pg_conn):
+        agent_id, _ = _agent(pg_conn)
+        _, _, source = _owner_private(pg_conn)
+        resp = _put(app, pg_conn, agent_id, EDITOR, {"sources": [source]})
+        assert _status(resp) == 403
+
+    def test_editor_cannot_attach_owner_private_prompt(self, app, pg_conn):
+        agent_id, _ = _agent(pg_conn)
+        _, prompt, _ = _owner_private(pg_conn)
+        resp = _put(app, pg_conn, agent_id, EDITOR, {"prompt_id": prompt})
+        assert _status(resp) == 403
+
+    def test_editor_may_attach_owner_tool_shared_with_them(self, app, pg_conn):
+        agent_id, team_id = _agent(pg_conn)
+        tool, _, _ = _owner_private(pg_conn)
+        _share_tool_with_editor(pg_conn, team_id, tool)
+        resp = _put(app, pg_conn, agent_id, EDITOR, {"tools": [tool]})
+        assert _status(resp) == 200, resp.get_json()
+        row = _row(pg_conn, agent_id)
+        assert [str(t) for t in row["tools"]] == [tool]
+        # The owner owns it, so it runs as the owner: no sponsor.
+        assert not row.get("resource_sponsors")
+
+    def test_editor_keeps_owner_refs_already_on_agent(self, app, pg_conn):
+        tool, prompt, source = _owner_private(pg_conn)
+        agent_id, _ = _agent(pg_conn, tools=[tool], prompt_id=prompt, source_id=source)
+        resp = _put(
+            app, pg_conn, agent_id, EDITOR,
+            {"tools": [tool], "prompt_id": prompt, "source": source},
+        )
+        assert _status(resp) == 200, resp.get_json()
+
+    def test_owner_may_attach_own_private_tool(self, app, pg_conn):
+        agent_id, _ = _agent(pg_conn)
+        tool, _, _ = _owner_private(pg_conn)
+        assert _status(_put(app, pg_conn, agent_id, OWNER, {"tools": [tool]})) == 200
+
+
+class TestAgentWorkflowSwap:
+    def _two_workflow_agents(self, pg_conn):
+        """Agent A (shared with EDITOR) and the owner's private agent B, each with a workflow."""
+        wf_a = str(WorkflowsRepository(pg_conn).create(OWNER, "A")["id"])
+        wf_b = str(WorkflowsRepository(pg_conn).create(OWNER, "B")["id"])
+        agent_a, _ = _agent(pg_conn, agent_type="workflow", workflow_id=wf_a)
+        AgentsRepository(pg_conn).create(
+            OWNER, "Private B", "published", description="d", key=f"k-{uuid.uuid4().hex}",
+            agent_type="workflow", workflow_id=wf_b,
+        )
+        return agent_a, wf_a, wf_b
+
+    def test_editor_cannot_swap_in_another_owner_workflow(self, app, pg_conn):
+        agent_a, wf_a, wf_b = self._two_workflow_agents(pg_conn)
+        resp = _put(app, pg_conn, agent_a, EDITOR, {"workflow": wf_b})
+        assert _status(resp) == 403
+        assert str(_row(pg_conn, agent_a)["workflow_id"]) == wf_a
+
+    def test_editor_may_resend_current_workflow(self, app, pg_conn):
+        agent_a, wf_a, _ = self._two_workflow_agents(pg_conn)
+        resp = _put(app, pg_conn, agent_a, EDITOR, {"workflow": wf_a})
+        assert _status(resp) == 200, resp.get_json()
+
+    def test_owner_may_swap_workflow(self, app, pg_conn):
+        agent_a, _, wf_b = self._two_workflow_agents(pg_conn)
+        assert _status(_put(app, pg_conn, agent_a, OWNER, {"workflow": wf_b})) == 200
+        assert str(_row(pg_conn, agent_a)["workflow_id"]) == wf_b
+
+
+class TestWorkflowNodeAttach:
+    def _setup(self, pg_conn):
+        wf = WorkflowsRepository(pg_conn).create(OWNER, "wf")
+        _, team_id = _agent(pg_conn, agent_type="workflow", workflow_id=str(wf["id"]))
+        return str(wf["id"]), team_id
+
+    def _put_wf(self, app, pg_conn, wid, user, body):
+        from docsgpt.api.user.workflows.routes import WorkflowDetail
+
+        return _call(app, pg_conn, WorkflowDetail, "put", f"/api/workflows/{wid}", user,
+                     json=body, args=(wid,))
+
+    def test_editor_cannot_add_owner_private_tool_to_node(self, app, pg_conn):
+        wid, _ = self._setup(pg_conn)
+        tool, _, _ = _owner_private(pg_conn)
+        assert _status(self._put_wf(app, pg_conn, wid, EDITOR, _wf_body(tool=tool))) == 403
+
+    def test_editor_cannot_add_owner_private_source_to_node(self, app, pg_conn):
+        wid, _ = self._setup(pg_conn)
+        _, _, source = _owner_private(pg_conn)
+        assert _status(self._put_wf(app, pg_conn, wid, EDITOR, _wf_body(source=source))) == 403
+
+    def test_editor_keeps_owner_node_refs_already_in_graph(self, app, pg_conn):
+        wid, _ = self._setup(pg_conn)
+        tool, _, source = _owner_private(pg_conn)
+        body = _wf_body(tool=tool, source=source)
+        assert _status(self._put_wf(app, pg_conn, wid, OWNER, body)) == 200
+        resp = self._put_wf(app, pg_conn, wid, EDITOR, body)
+        assert _status(resp) == 200, resp.get_json()
+
+    def test_editor_may_add_tool_shared_with_them(self, app, pg_conn):
+        wid, team_id = self._setup(pg_conn)
+        tool, _, _ = _owner_private(pg_conn)
+        _share_tool_with_editor(pg_conn, team_id, tool)
+        resp = self._put_wf(app, pg_conn, wid, EDITOR, _wf_body(tool=tool))
+        assert _status(resp) == 200, resp.get_json()
+
+
+class TestUppercaseIds:
+    def _tool_not_usable_in_own(self, pg_conn, team_id):
+        """An owner tool the EDITOR only views, with ``viewers_can_use_in_agents`` off."""
+        tool, _, _ = _owner_private(pg_conn)
+        _share_tool_with_editor(pg_conn, team_id, tool)
+        set_settings(pg_conn, "tool", tool, {"viewers_can_use_in_agents": False}, OWNER)
+        return tool
+
+    def test_resolve_applies_switches_to_uppercase_id(self, pg_conn):
+        _, team_id = _agent(pg_conn)
+        tool = self._tool_not_usable_in_own(pg_conn, team_id)
+        lower = resolve(pg_conn, "tool", tool, EDITOR)
+        upper = resolve(pg_conn, "tool", tool.upper(), EDITOR)
+        assert lower is not None and upper is not None
+        assert not upper.can("use_in_own")
+        assert upper.settings == lower.settings
+        assert upper.resource_id == tool
+
+    def test_uppercase_tool_id_does_not_bypass_switch_on_attach(self, app, pg_conn):
+        agent_id, team_id = _agent(pg_conn)
+        tool = self._tool_not_usable_in_own(pg_conn, team_id)
+        resp = _put(app, pg_conn, agent_id, EDITOR, {"tools": [tool.upper()]})
+        assert _status(resp) == 403
+
+    def test_attached_ids_are_stored_canonical(self, app, pg_conn):
+        agent_id, _ = _agent(pg_conn)
+        tool, prompt, source = _owner_private(pg_conn)
+        resp = _put(
+            app, pg_conn, agent_id, OWNER,
+            {"tools": [tool.upper()], "prompt_id": prompt.upper(), "sources": [source.upper()]},
+        )
+        assert _status(resp) == 200, resp.get_json()
+        row = _row(pg_conn, agent_id)
+        assert [str(t) for t in row["tools"]] == [tool]
+        assert str(row["prompt_id"]) == prompt
+        assert [str(s) for s in row["extra_source_ids"] or []] + (
+            [str(row["source_id"])] if row.get("source_id") else []
+        ) == [source]

tests/api/user/test_tools_access.py

@@ -503,17 +503,44 @@ class TestMCPSaveAccess:
         resp, _ = self._save(app, pg_conn, "ed", body)
         assert resp.status_code == 403
 
-    def test_editor_oauth_edit_without_reconnect_keeps_actions(self, app, pg_conn):
+    def test_editor_oauth_save_without_reconnect_is_owner_only(self, app, pg_conn):
+        # A shared OAuth server's connection is the owner's: an editor can't
+        # save it, even with the same URL (the path, scopes or client could
+        # still change while reusing the owner's tokens).
         tool = _mcp_tool(pg_conn, auth_type="oauth")
         _share(pg_conn, tool["id"], "ed", "editor")
         body = {"id": str(tool["id"]), "displayName": "Renamed",
                 "config": {"server_url": "https://mcp.example.com/mcp", "auth_type": "oauth",
                            "transport_type": "http"}}
         resp, _ = self._save(app, pg_conn, "ed", body)
-        assert resp.status_code == 200, resp.json
-        row = _row(pg_conn, tool["id"])
-        assert row["display_name"] == "Renamed"
-        assert [a["name"] for a in row["actions"]] == ["old"]
+        assert resp.status_code == 403
+        assert _row(pg_conn, tool["id"])["display_name"] == "M"
+
+    def test_editor_cannot_switch_oauth_server_to_other_auth(self, app, pg_conn):
+        tool = _mcp_tool(pg_conn, auth_type="oauth")
+        _share(pg_conn, tool["id"], "ed", "editor")
+        body = {"id": str(tool["id"]), "displayName": "M",
+                "config": {"server_url": "https://mcp.example.com/mcp", "auth_type": "bearer",
+                           "transport_type": "http", "bearer_token": "mine"}}
+        assert self._save(app, pg_conn, "ed", body)[0].status_code == 403
+        assert _row(pg_conn, tool["id"])["config"]["auth_type"] == "oauth"
+
+    def test_editor_cannot_switch_server_to_oauth(self, app, pg_conn):
+        tool = _mcp_tool(pg_conn, secrets={"bearer_token": "tok"})
+        _share(pg_conn, tool["id"], "ed", "editor")
+        body = {"id": str(tool["id"]), "displayName": "M",
+                "config": {"server_url": "https://mcp.example.com/mcp", "auth_type": "oauth",
+                           "transport_type": "http"}}
+        assert self._save(app, pg_conn, "ed", body)[0].status_code == 403
+
+    def test_owner_oauth_save_needs_a_completed_sign_in(self, app, pg_conn):
+        tool = _mcp_tool(pg_conn, auth_type="oauth")
+        body = {"id": str(tool["id"]), "displayName": "Renamed",
+                "config": {"server_url": "https://mcp.example.com/mcp", "auth_type": "oauth",
+                           "transport_type": "http"}}
+        resp, _ = self._save(app, pg_conn, OWNER, body)
+        assert resp.status_code == 400
+        assert _row(pg_conn, tool["id"])["display_name"] == "M"
 
     def test_auth_status_includes_team_mcp_tools(self, app, pg_conn):
         from docsgpt.api.user.tools.mcp import MCPAuthStatus
@@ -562,6 +589,12 @@ class TestMCPTestEndpointAccess:
         assert self._test(app, pg_conn, "vi", self._body(tool))[0].status_code == 403
         assert self._test(app, pg_conn, "eve", self._body(tool))[0].status_code == 404
 
+    def test_editor_test_of_stored_oauth_server_is_owner_only(self, app, pg_conn):
+        tool = _mcp_tool(pg_conn, auth_type="oauth")
+        _share(pg_conn, tool["id"], "ed", "editor")
+        resp, cls = self._test(app, pg_conn, "ed", self._body(tool, bearer_token="mine"))
+        assert resp.status_code == 403 and not cls.called
+
     def test_editor_oauth_test_is_owner_only(self, app, pg_conn):
         tool = _mcp_tool(pg_conn, auth_type="oauth")
         _share(pg_conn, tool["id"], "ed", "editor")
@@ -569,3 +602,62 @@ class TestMCPTestEndpointAccess:
                                                   "auth_type": "oauth", "transport_type": "http"}}
         resp, cls = self._test(app, pg_conn, "ed", body)
         assert resp.status_code == 403 and not cls.called
+
+
+# ---------------------------------------------------------------------------
+# 6. Shared OAuth MCP servers: the connection is the owner's
+# ---------------------------------------------------------------------------
+class TestSharedOAuthMCPConfig:
+    """``/api/update_tool`` and ``/api/update_tool_config`` can't move a shared
+    OAuth server or switch a shared server to OAuth: OAuth tokens are looked up
+    by owner + server URL, so either would run on the owner's sign-in
+    elsewhere. Connection changes on OAuth servers are the owner's."""
+
+    OTHER = {"server_url": "https://other-mcp.example.org/mcp", "auth_type": "oauth"}
+
+    def _oauth_tool(self, conn):
+        tool = _mcp_tool(conn, auth_type="oauth")
+        _share(conn, tool["id"], "ed", "editor")
+        return tool
+
+    def test_update_tool_config_cannot_move_oauth_server(self, app, pg_conn):
+        from docsgpt.api.user.tools.routes import UpdateToolConfig
+
+        tool = self._oauth_tool(pg_conn)
+        body = {"id": str(tool["id"]), "config": dict(self.OTHER)}
+        assert _call(app, pg_conn, UpdateToolConfig, "ed", json=body).status_code == 403
+        assert _row(pg_conn, tool["id"])["config"]["server_url"] == "https://mcp.example.com/mcp"
+
+    def test_update_tool_cannot_move_oauth_server(self, app, pg_conn):
+        from docsgpt.api.user.tools.routes import UpdateTool
+
+        tool = self._oauth_tool(pg_conn)
+        body = {"id": str(tool["id"]), "config": dict(self.OTHER)}
+        assert _call(app, pg_conn, UpdateTool, "ed", json=body).status_code == 403
+        assert _row(pg_conn, tool["id"])["config"]["server_url"] == "https://mcp.example.com/mcp"
+
+    def test_update_tool_config_cannot_switch_server_to_oauth(self, app, pg_conn):
+        from docsgpt.api.user.tools.routes import UpdateToolConfig
+
+        tool = _mcp_tool(pg_conn, secrets={"bearer_token": "tok"})
+        _share(pg_conn, tool["id"], "ed", "editor")
+        body = {"id": str(tool["id"]),
+                "config": {"server_url": "https://mcp.example.com/mcp", "auth_type": "oauth"}}
+        assert _call(app, pg_conn, UpdateToolConfig, "ed", json=body).status_code == 403
+        assert _row(pg_conn, tool["id"])["config"]["auth_type"] == "bearer"
+
+    def test_editor_still_renames_oauth_server_without_config(self, app, pg_conn):
+        from docsgpt.api.user.tools.routes import UpdateTool
+
+        tool = self._oauth_tool(pg_conn)
+        body = {"id": str(tool["id"]), "customName": "Renamed"}
+        assert _call(app, pg_conn, UpdateTool, "ed", json=body).status_code == 200
+        assert _row(pg_conn, tool["id"])["custom_name"] == "Renamed"
+
+    def test_owner_may_change_oauth_server_config(self, app, pg_conn):
+        from docsgpt.api.user.tools.routes import UpdateToolConfig
+
+        tool = _mcp_tool(pg_conn, auth_type="oauth")
+        body = {"id": str(tool["id"]), "config": dict(self.OTHER)}
+        assert _call(app, pg_conn, UpdateToolConfig, OWNER, json=body).status_code == 200
+        assert _row(pg_conn, tool["id"])["config"]["server_url"] == self.OTHER["server_url"]