Address review: MCP token scope, API-key identity, reconnect lock, setup retries

- An MCP connection's tokens only go to its own server, and a
  connection_id can no longer come from a client: the MCP test/save
  routes drop it and the tool executor uses only the resolved one.
- API-key connections are reused only when they hold the same
  credentials. Two keys that share a hint or a label get separate
  connections ("…abcd (2)"), at runtime and in the 0038 backfill.
- Reconnecting a connection whose secrets cannot be decrypted no longer
  flags it from a second transaction while holding its row lock.
- Connection setup validates the sync request before claiming its
  Idempotency-Key, so a corrected retry is queued.
- A request that names no connection resolves to none without opening
  a database connection; fixes two tests that reached a real database.

docsgpt/agents/tool_executor.py

@@ -1560,6 +1560,9 @@ class ToolExecutor:
                 tool_config["body_encoding_rules"] = action_config.get("body_encoding_rules", {})
         else:
             tool_config = tool_data["config"].copy() if tool_data["config"] else {}
+            # Whose MCP tokens a tool uses is decided by resolving its
+            # connection below, never by a value stored in its config.
+            tool_config.pop("connection_id", None)
             # Credentials are PBKDF2-bound to the tool OWNER's sub, not the
             # invoker's. Decrypt with the tool row's user_id so a team member
             # running an owner's shared tool authenticates with the owner's

docsgpt/alembic/versions/0038_connections.py

@@ -292,7 +292,12 @@ def _link_api_key_tools(bind) -> None:
 
     from sqlalchemy import text
 
-    from docsgpt.security.encryption import decrypt_credentials, encrypt_json
+    from docsgpt.security.encryption import (
+        CredentialDecryptionError,
+        decrypt_credentials,
+        decrypt_json,
+        encrypt_json,
+    )
 
     rows = bind.execute(
         text(
@@ -323,18 +328,32 @@ def _link_api_key_tools(bind) -> None:
             server_url = None
             connector_key = row.name
             display_name = _TOOL_CONNECTORS[row.name]
-        label = _credential_hint(credentials)
-        existing = bind.execute(
-            text(
-                """
-                SELECT id, encrypted_credentials FROM connector_sessions
-                WHERE user_id = :user_id AND provider = :provider
-                  AND COALESCE(server_url, '') = COALESCE(:server_url, '')
-                  AND COALESCE(account_label, '') = :label
-                """
-            ),
-            {"user_id": row.user_id, "provider": connector_key, "server_url": server_url, "label": label},
-        ).fetchone()
+        # The hint is not an identity: reuse a connection only when it holds
+        # the same credentials, and give a different key its own label.
+        hint = _credential_hint(credentials)
+        label, suffix = hint, 1
+        while True:
+            existing = bind.execute(
+                text(
+                    """
+                    SELECT id, encrypted_credentials FROM connector_sessions
+                    WHERE user_id = :user_id AND provider = :provider
+                      AND COALESCE(server_url, '') = COALESCE(:server_url, '')
+                      AND COALESCE(account_label, '') = :label
+                    """
+                ),
+                {"user_id": row.user_id, "provider": connector_key, "server_url": server_url, "label": label},
+            ).fetchone()
+            if existing is None:
+                break
+            try:
+                stored = decrypt_json(existing.encrypted_credentials or "", row.user_id).get("credentials")
+            except CredentialDecryptionError:
+                stored = None
+            if stored == credentials:
+                break
+            suffix += 1
+            label = f"{hint} ({suffix})"
         if existing is None:
             connection_id = bind.execute(
                 text(

docsgpt/api/connector/connections.py

@@ -18,6 +18,7 @@ from docsgpt.api.user.authz import ROLE_ADMIN, has_role
 from docsgpt.connectors import catalog, service
 from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
 from docsgpt.storage.db.session import db_readonly, db_session
+from docsgpt.security.encryption import CredentialDecryptionError
 
 _FREQUENCIES = ("never", "daily", "weekly", "monthly")
 
@@ -258,6 +259,39 @@ def _start_sync(user_id: str, row: dict, sync: dict):
     if frequency not in _FREQUENCIES:
         return ("Unknown sync frequency", 400)
     name = (sync.get("name") or "").strip() or definition.name
+    # Validate before claiming the idempotency key: a rejected request must
+    # leave the key free for the corrected retry.
+    if definition.auth_kind == "oauth":
+        file_ids = [str(i) for i in items.get("file_ids") or [] if i]
+        folder_ids = [str(i) for i in items.get("folder_ids") or [] if i]
+        if not file_ids and not folder_ids:
+            return ("Pick at least one file or folder", 400)
+        task_fn = ingest_connector_task
+        kwargs = {
+            "job_name": name,
+            "user": user_id,
+            "source_type": definition.sync_ingestor,
+            "connection_id": str(row["id"]),
+            "file_ids": file_ids,
+            "folder_ids": folder_ids,
+            "recursive": bool(items.get("recursive", True)),
+            "sync_frequency": frequency,
+        }
+    else:
+        fields = {f.key for f in definition.setup_fields}
+        source_data = {k: v for k, v in items.items() if k in fields and v not in (None, "")}
+        missing = [f.label for f in definition.setup_fields if f.required and f.key not in source_data]
+        if missing:
+            return (f"Missing: {', '.join(missing)}", 400)
+        task_fn = ingest_remote
+        kwargs = {
+            "source_data": source_data,
+            "job_name": name,
+            "user": user_id,
+            "loader": definition.sync_ingestor,
+            "connection_id": str(row["id"]),
+            "sync_frequency": frequency,
+        }
     idempotency_key, _ = _read_idempotency_key()
     scoped_key = _scoped_idempotency_key(idempotency_key, user_id)
     task_id = None
@@ -267,45 +301,9 @@ def _start_sync(user_id: str, row: dict, sync: dict):
             return {"id": cached.get("source_id"), "task_id": cached.get("task_id"), "name": name}
     source_id = str(_derive_source_id(scoped_key)) if scoped_key else str(uuid.uuid4())
     options = {"task_id": task_id} if task_id else {}
-    if definition.auth_kind == "oauth":
-        file_ids = [str(i) for i in items.get("file_ids") or [] if i]
-        folder_ids = [str(i) for i in items.get("folder_ids") or [] if i]
-        if not file_ids and not folder_ids:
-            return ("Pick at least one file or folder", 400)
-        task = ingest_connector_task.apply_async(
-            kwargs={
-                "job_name": name,
-                "user": user_id,
-                "source_type": definition.sync_ingestor,
-                "connection_id": str(row["id"]),
-                "file_ids": file_ids,
-                "folder_ids": folder_ids,
-                "recursive": bool(items.get("recursive", True)),
-                "sync_frequency": frequency,
-                "idempotency_key": scoped_key,
-                "source_id": source_id,
-            },
-            **options,
-        )
-    else:
-        fields = {f.key for f in definition.setup_fields}
-        source_data = {k: v for k, v in items.items() if k in fields and v not in (None, "")}
-        missing = [f.label for f in definition.setup_fields if f.required and f.key not in source_data]
-        if missing:
-            return (f"Missing: {', '.join(missing)}", 400)
-        task = ingest_remote.apply_async(
-            kwargs={
-                "source_data": source_data,
-                "job_name": name,
-                "user": user_id,
-                "loader": definition.sync_ingestor,
-                "connection_id": str(row["id"]),
-                "sync_frequency": frequency,
-                "idempotency_key": scoped_key,
-                "source_id": source_id,
-            },
-            **options,
-        )
+    task = task_fn.apply_async(
+        kwargs={**kwargs, "idempotency_key": scoped_key, "source_id": source_id}, **options,
+    )
     return {"id": source_id, "task_id": task_id or task.id, "name": name, "sync_frequency": frequency}
 
 
@@ -347,7 +345,13 @@ class ConnectionReconnect(Resource):
             service.ensure_can_store_credentials()
             with db_session() as conn:
                 locked = ConnectorSessionsRepository(conn).get_for_update(connection_id)
-                stored = service.load_secrets(locked) or {}
+                # read_secrets, not load_secrets: flagging an unreadable row
+                # would write it from a second transaction while this one
+                # holds its lock. The new credentials replace it anyway.
+                try:
+                    stored = service.read_secrets(locked)
+                except CredentialDecryptionError:
+                    stored = {}
                 merged = {**(stored.get("credentials") or {}), **{k: v for k, v in credentials.items() if v}}
                 service.write_secrets(
                     conn, locked, {**stored, "credentials": merged},

docsgpt/api/user/tools/mcp.py

@@ -27,13 +27,16 @@ def _sanitize_mcp_transport(config):
     """Normalise and validate the transport_type field.
 
     Strips ``command`` / ``args`` keys that are only valid for local STDIO
-    transports and returns the cleaned transport type string.
+    transports, and ``connection_id``, which only the tool executor sets
+    (it picks whose MCP tokens the tool uses). Returns the cleaned
+    transport type string.
     """
     transport_type = (config.get("transport_type") or "auto").lower()
     if transport_type not in _ALLOWED_TRANSPORTS:
         raise ValueError(f"Unsupported transport_type: {transport_type}")
     config.pop("command", None)
     config.pop("args", None)
+    config.pop("connection_id", None)
     config["transport_type"] = transport_type
     return transport_type
 

docsgpt/connectors/service.py

@@ -647,6 +647,40 @@ def get_credentials(row: dict) -> dict:
     return dict(secrets.get("credentials") or {})
 
 
+def _api_key_account(
+    repo: ConnectorSessionsRepository,
+    user_id: str,
+    connector_key: str,
+    server_url: Optional[str],
+    label: str,
+    credentials: dict,
+) -> tuple[Optional[dict], str]:
+    """The connection holding exactly ``credentials``, or a free label for a new one.
+
+    A label (a hint of the key, or what the user typed) is not an identity:
+    two different keys can share it. A row under the label is reused only
+    when it holds the same credentials; otherwise the label gets a
+    ``(2)``, ``(3)`` suffix until it names a matching row or no row.
+
+    Returns:
+        ``(row, label)``: the row to reuse (None to create one) and its label.
+    """
+    from docsgpt.security.encryption import CredentialDecryptionError
+
+    candidate, suffix = label, 1
+    while True:
+        existing = repo.find_account(user_id, connector_key, server_url=server_url, account_label=candidate)
+        if existing is None:
+            return None, candidate
+        try:
+            if (read_secrets(existing).get("credentials") or {}) == credentials:
+                return existing, candidate
+        except CredentialDecryptionError:
+            pass
+        suffix += 1
+        candidate = f"{label} ({suffix})"
+
+
 def create_api_key_connection(
     conn,
     user_id: str,
@@ -689,9 +723,10 @@ def create_api_key_connection(
     ensure_can_store_credentials()
     ensure_connector_allowed(conn, definition.key)
     secret_values = {k: v for k, v in credentials.items() if (fields.get(k).secret if fields.get(k) else True)}
-    account_label = label or credential_hint(secret_values or credentials)
     repo = ConnectorSessionsRepository(conn)
-    existing = repo.find_account(user_id, definition.key, server_url=server_url, account_label=account_label)
+    existing, account_label = _api_key_account(
+        repo, user_id, definition.key, server_url, label or credential_hint(secret_values or credentials), credentials,
+    )
     if existing is not None:
         write_secrets(conn, existing, {"credentials": credentials}, status=STATUS_CONNECTED, last_error=None)
         resume_sources(conn, str(existing["id"]))
@@ -833,16 +868,16 @@ def resolve_request_connection(user_id: str, provider: Optional[str], data: dict
     """
     connection_id = data.get("connection_id")
     session_token = data.get("session_token")
+    if not connection_id and not session_token:
+        return None
     with db_readonly() as conn:
         repo = ConnectorSessionsRepository(conn)
         if connection_id:
             row = repo.get_for_user(str(connection_id), user_id)
-        elif session_token:
+        else:
             row = repo.get_by_session_token(session_token)
             if not owns_connector_session(row, user_id, provider):
                 return None
-        else:
-            return None
     if row is None:
         return None
     if provider and (row.get("provider") or "").lower() != provider.lower():
@@ -867,7 +902,13 @@ def _mcp_row(conn, user_id: str, base_url: str, connection_id: Optional[str], *,
 
     repo = ConnectorSessionsRepository(conn)
     if connection_id:
-        return repo.get_for_update(connection_id) if lock else repo.get(connection_id)
+        row = repo.get_for_update(connection_id) if lock else repo.get(connection_id)
+        # A connection's tokens only ever go to its own server. Ownership is
+        # checked where the id is chosen (the tool executor); clients cannot
+        # supply one (see ``_sanitize_mcp_transport``).
+        if row is None or row.get("provider") != mcp_provider(base_url):
+            return None
+        return row
     result = conn.execute(
         text(
             "SELECT * FROM connector_sessions WHERE user_id = :user_id AND provider = :provider "

tests/connectors/test_connections_api.py

@@ -141,6 +141,26 @@ class TestSetupSync:
                          body={"sync": {"items": {}}}, args=[cid])
         assert resp.status_code == 400
 
+    def test_rejected_request_does_not_claim_its_idempotency_key(self, app, pg_conn):
+        from docsgpt.api.connector.connections import ConnectionSetup
+
+        @contextmanager
+        def _yield():
+            yield pg_conn
+
+        cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={"token_info": {}})
+        headers = {"Idempotency-Key": "setup-retry-1"}
+        with _db(pg_conn), patch("docsgpt.api.user.sources.upload.db_session", _yield), patch(
+            "docsgpt.api.user.tasks.ingest_connector_task.apply_async", return_value=MagicMock(id="t"),
+        ) as apply:
+            bad = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup",
+                        body={"sync": {"items": {}}}, headers=headers, args=[cid])
+            fixed = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup",
+                          body={"sync": {"items": {"folder_ids": ["f1"]}}}, headers=headers, args=[cid])
+        assert bad.status_code == 400
+        assert fixed.status_code == 200
+        apply.assert_called_once()
+
     def test_s3_keys_stay_on_the_connection(self, app, pg_conn):
         from docsgpt.api.connector.connections import ConnectionSetup
 
@@ -186,6 +206,23 @@ class TestReconnect:
         meta = pg_conn.execute(text("SELECT metadata FROM sources WHERE user_id = 'alice'")).scalar()
         assert "sync_state" not in meta
 
+    def test_reconnect_recovers_undecryptable_credentials_without_a_second_write(self, app, pg_conn):
+        """After a lost key the reconnect replaces the blob; it must not flag the row it holds locked."""
+        from docsgpt.api.connector.connections import ConnectionReconnect
+
+        cid = _connection(pg_conn, status="reconnect_needed")
+        pg_conn.execute(text(
+            "UPDATE connector_sessions SET encrypted_credentials = :blob WHERE id = CAST(:c AS uuid)"
+        ), {"c": cid, "blob": encrypt_json({"credentials": {"token": "old"}}, "someone-else")})
+        with _db(pg_conn), patch.object(service, "mark_reconnect_needed") as flag:
+            resp = _call(app, ConnectionReconnect, "post", f"/api/connections/{cid}/reconnect",
+                         body={"credentials": {"token": "new"}}, args=[cid])
+        assert resp.status_code == 200
+        flag.assert_not_called()
+        row = pg_conn.execute(text("SELECT * FROM connector_sessions WHERE id = CAST(:c AS uuid)"),
+                              {"c": cid}).one()._mapping
+        assert service.read_secrets(dict(row))["credentials"] == {"token": "new"}
+
     def test_oauth_reconnect_returns_authorization_url(self, app, pg_conn):
         from docsgpt.api.connector.connections import ConnectionReconnect
 

tests/connectors/test_runtime.py

@@ -173,6 +173,17 @@ class TestExecutor:
         assert config["connection_id"] == cid
         assert "secret" not in str(config)
 
+    def test_stored_connection_id_in_config_is_ignored(self, pg_conn):
+        """Only a resolved connection reaches the tool; a config value never does."""
+        victim = _connection(pg_conn, user="victim", provider="mcp:https://m.example.com", auth_kind="mcp_oauth",
+                             server_url="https://m.example.com", secrets={"tokens": {"access_token": "v"}})
+        tool = {**_tool(None, name="mcp_tool"), "config": {"server_url": "https://m.example.com/mcp",
+                                                          "auth_type": "oauth", "connection_id": victim}}
+        with _service_db(pg_conn), patch("docsgpt.agents.tool_executor.ToolManager") as manager:
+            _executor()._get_or_load_tool(tool, "t1", "search")
+        config = manager.return_value.load_tool.call_args.kwargs["tool_config"]
+        assert "connection_id" not in config
+
 
 class TestScheduledSync:
     def test_connector_sources_with_a_connection_are_dispatched(self, pg_conn):

tests/connectors/test_service_credentials.py

@@ -249,6 +249,36 @@ class TestApiKeyConnections:
         assert created is False and again["id"] == row["id"]
         assert service.get_credentials(again) == {"token": "123456:ABCDEFG"}
 
+    def test_different_keys_with_the_same_hint_stay_apart(self, pg_conn):
+        """Two keys ending in the same four characters are two accounts, not one."""
+        from docsgpt.connectors import catalog
+
+        postgres = catalog.get_definition("postgres")
+        first, _ = service.create_api_key_connection(
+            pg_conn, "alice", postgres, {"token": "postgresql://ro@db-a/app"},
+        )
+        second, created = service.create_api_key_connection(
+            pg_conn, "alice", postgres, {"token": "postgresql://rw@db-b/app"},
+        )
+        assert created is True and second["id"] != first["id"]
+        assert second["account_label"] != first["account_label"]
+        assert service.get_credentials(_row(pg_conn, str(first["id"]))) == {"token": "postgresql://ro@db-a/app"}
+        again, created = service.create_api_key_connection(
+            pg_conn, "alice", postgres, {"token": "postgresql://rw@db-b/app"},
+        )
+        assert created is False and again["id"] == second["id"]
+
+    def test_same_label_different_key_does_not_overwrite(self, pg_conn):
+        from docsgpt.connectors import catalog
+
+        brave = catalog.get_definition("brave")
+        first, _ = service.create_api_key_connection(pg_conn, "alice", brave, {"token": "key-one-111111"}, label="Team")
+        second, created = service.create_api_key_connection(
+            pg_conn, "alice", brave, {"token": "key-two-222222"}, label="Team",
+        )
+        assert created is True and second["id"] != first["id"]
+        assert service.get_credentials(_row(pg_conn, str(first["id"]))) == {"token": "key-one-111111"}
+
     def test_missing_field(self, pg_conn):
         from docsgpt.connectors import catalog
 
@@ -418,3 +448,42 @@ class TestRevoke:
 
         with patch("requests.post", side_effect=requests.exceptions.Timeout()):
             assert not service.revoke_at_provider({"provider": "google_drive"}, {"token_info": {"access_token": "a"}})
+
+
+class TestMcpConnectionScope:
+    def _mcp(self, conn, user="alice", base="https://mcp.notion.com"):
+        return _connection(
+            conn, user=user, provider=f"mcp:{base}", auth_kind="mcp_oauth", server_url=base,
+            secrets={"tokens": {"access_token": f"{user}-mcp-token"}},
+        )
+
+    def test_reads_tokens_of_the_named_connection(self, pg_conn):
+        cid = self._mcp(pg_conn)
+        with _patch_service_db(pg_conn):
+            data = service.read_mcp_secrets("alice", "https://mcp.notion.com", cid)
+        assert data["tokens"]["access_token"] == "alice-mcp-token"
+
+    def test_connection_for_another_server_yields_nothing(self, pg_conn):
+        """A connection id never sends its tokens to a different server."""
+        cid = self._mcp(pg_conn)
+        with _patch_service_db(pg_conn):
+            assert service.read_mcp_secrets("alice", "https://attacker.example", cid) == {}
+
+    def test_writes_never_land_on_another_servers_connection(self, pg_conn):
+        cid = self._mcp(pg_conn)
+        with _patch_service_db(pg_conn):
+            row = service.update_mcp_secrets(
+                "alice", "https://attacker.example", {"tokens": {"access_token": "planted"}}, connection_id=cid,
+            )
+            assert str(row["id"]) != cid
+            assert service.read_mcp_secrets("alice", "https://mcp.notion.com", cid)["tokens"]["access_token"] == (
+                "alice-mcp-token"
+            )
+
+    def test_mcp_routes_drop_client_supplied_connection_id(self):
+        """Only the tool executor may pick the connection whose tokens a tool uses."""
+        from docsgpt.api.user.tools.mcp import _sanitize_mcp_transport
+
+        config = {"transport_type": "http", "connection_id": "someone-elses-connection"}
+        _sanitize_mcp_transport(config)
+        assert "connection_id" not in config

tests/parser/connectors/test_auth_token_leakage.py

@@ -55,7 +55,7 @@ def _patches(session_return: Optional[Dict[str, Any]]):
             fake_repo_cls,
         ),
         patch(
-            "docsgpt.storage.db.session.db_readonly",
+            "docsgpt.connectors.service.db_readonly",
             lambda: _FakeReadonlyCtx(),
         ),
     )

tests/storage/db/test_migration_0038.py

@@ -272,6 +272,32 @@ class TestMigration0038Credentials:
             ).scalar()
             assert "encrypted_credentials" in config
 
+    def test_different_keys_with_the_same_hint_get_their_own_connections(self, pg_engine):
+        from docsgpt.connectors.service import read_secrets
+        from docsgpt.security.encryption import encrypt_credentials
+
+        def seed(conn):
+            ids = {}
+            for name, token in (("a", "111111:SAMEEND1"), ("b", "222222:SAMEEND1")):
+                ids[name] = conn.execute(
+                    text("INSERT INTO user_tools (user_id, name, config) VALUES ('frank', 'telegram', "
+                         "CAST(:c AS jsonb)) RETURNING id"),
+                    {"c": json.dumps({"encrypted_credentials": encrypt_credentials({"token": token}, "frank")})},
+                ).scalar()
+            return ids
+
+        _, ids = self._upgrade_with(pg_engine, seed)
+        with pg_engine.connect() as conn:
+            links = dict(conn.execute(
+                text("SELECT id, connection_id FROM user_tools WHERE user_id = 'frank'")
+            ).fetchall())
+            assert links[ids["a"]] != links[ids["b"]]
+            for name, token in (("a", "111111:SAMEEND1"), ("b", "222222:SAMEEND1")):
+                row = dict(conn.execute(
+                    text("SELECT * FROM connector_sessions WHERE id = :i"), {"i": links[ids[name]]}
+                ).one()._mapping)
+                assert read_secrets(row) == {"credentials": {"token": token}}
+
     def test_oauth_mcp_tools_keep_member_credentials(self, pg_engine):
         _, ids = self._upgrade_with(pg_engine, _seed_secrets)
         with pg_engine.connect() as conn: