mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 22:17:00 +00:00
Connections own their credentials
Migration 0038 moves every stored secret (OAuth tokens, MCP OAuth tokens and client registrations, API keys) into the connection's encrypted envelope, links API-key tools to one connection per distinct credential, allows several accounts per provider, and adds credential_mode to sources and tools. OAuth MCP tools keep resolving each member's own token, as they did before. docsgpt.connectors.service is now the only reader of OAuth tokens: get_valid_token_info refreshes under a row lock and persists rotated refresh tokens, and a revoked grant flags the connection, pauses its sources and notifies the owner. Loaders build from a connection (BaseConnectorLoader.from_connection), so scheduled sync covers Drive, SharePoint and Confluence sources with no browser. S3 and Reddit keys stay on the connection instead of in remote_data. New endpoints: POST /api/connections, /setup, /reconnect, /picker-token, /claim, DELETE /api/connections/<id>, per-action permissions and MCP refresh-tools. Upload, file listing, sync and validate-session take a connection_id; session tokens keep working for this release. The tool executor reads credentials from the resolved connection (owner or member mode) and pauses on a Connect card when a connection needs signing in. docsgpt connectors reencrypt rewrites stored credentials after a key rotation.
This commit is contained in:
1 parent
667d4bbab0
commit
3511004f50
54 files changed
+4007
-754
No files matched your search
@@ -11,12 +11,16 @@ from __future__ import annotations
|
||||
from flask import current_app, jsonify, make_response, request
|
||||
from flask_restx import Namespace, Resource
|
||||
|
||||
import uuid
|
||||
|
||||
from docsgpt.api import api
|
||||
from docsgpt.api.user.authz import ROLE_ADMIN, has_role
|
||||
from docsgpt.connectors import service
|
||||
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
|
||||
|
||||
_FREQUENCIES = ("never", "daily", "weekly", "monthly")
|
||||
|
||||
connections_ns = Namespace("connections", description="Connectors and connections", path="/api")
|
||||
api.add_namespace(connections_ns)
|
||||
|
||||
@@ -52,8 +56,64 @@ class ConnectorCatalog(Resource):
|
||||
return make_response(jsonify({"success": True, "connectors": entries}), 200)
|
||||
|
||||
|
||||
def _json_body() -> dict:
|
||||
body = request.get_json(silent=True)
|
||||
return body if isinstance(body, dict) else {}
|
||||
|
||||
|
||||
def _owned(conn, connection_id: str, user_id: str):
|
||||
return ConnectorSessionsRepository(conn).get_for_user(connection_id, user_id)
|
||||
|
||||
|
||||
def _error(message: str, status: int, **extra):
|
||||
return make_response(jsonify({"success": False, "error": message, **extra}), status)
|
||||
|
||||
|
||||
@connections_ns.route("/connections")
|
||||
class ConnectionsList(Resource):
|
||||
@api.doc(
|
||||
description=(
|
||||
"Create a connection from pasted credentials: "
|
||||
"{connector_key, credentials, label?}. Same credentials reuse the same connection."
|
||||
)
|
||||
)
|
||||
def post(self):
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
body = _json_body()
|
||||
definition = catalog.get_definition(body.get("connector_key"))
|
||||
if definition is None or definition.auth_kind != "api_key":
|
||||
return _error("This connector does not take pasted credentials", 400)
|
||||
if definition.missing_settings:
|
||||
return _error("This connector needs admin setup", 400, code="needs_setup")
|
||||
credentials = body.get("credentials")
|
||||
if not isinstance(credentials, dict):
|
||||
return _error("credentials must be an object", 400)
|
||||
try:
|
||||
with db_session() as conn:
|
||||
row, created = service.create_api_key_connection(
|
||||
conn, user_id, definition, credentials, label=(body.get("label") or None),
|
||||
)
|
||||
except service.EncryptionKeyNotConfigured as err:
|
||||
return _error(str(err), 400, code="encryption_key_default")
|
||||
except ValueError as err:
|
||||
return _error(str(err), 400)
|
||||
except Exception as err:
|
||||
current_app.logger.error(f"Error creating connection: {err}", exc_info=True)
|
||||
return _error("Failed to create connection", 500)
|
||||
return make_response(
|
||||
jsonify(
|
||||
{
|
||||
"success": True,
|
||||
"created": created,
|
||||
"connection": service.serialize_connection(row),
|
||||
"setup": dict(definition.setup),
|
||||
}
|
||||
),
|
||||
201 if created else 200,
|
||||
)
|
||||
|
||||
@api.doc(description="The caller's connections with status and linked resource counts")
|
||||
def get(self):
|
||||
user_id = _user_id()
|
||||
@@ -86,6 +146,35 @@ class ConnectionDetail(Resource):
|
||||
return make_response(jsonify({"success": False, "error": "Failed to load connection"}), 500)
|
||||
return make_response(jsonify({"success": True, "connection": detail}), 200)
|
||||
|
||||
@api.doc(
|
||||
description=(
|
||||
"Remove a connection: {sources: keep | delete, tools: delete | keep}. "
|
||||
"Kept sources keep their content and stop syncing."
|
||||
)
|
||||
)
|
||||
def delete(self, connection_id: str):
|
||||
from docsgpt.api.user.sources.routes import delete_source
|
||||
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
body = _json_body()
|
||||
sources_mode = body.get("sources", "keep")
|
||||
tools_mode = body.get("tools", "delete")
|
||||
if sources_mode not in ("keep", "delete") or tools_mode not in ("keep", "delete"):
|
||||
return _error("sources and tools must be keep or delete", 400)
|
||||
try:
|
||||
with db_session() as conn:
|
||||
row = _owned(conn, connection_id, user_id)
|
||||
if row is None:
|
||||
return _not_found()
|
||||
to_delete = service.remove_connection(conn, row, sources=sources_mode, tools=tools_mode)
|
||||
failed = [str(doc["id"]) for doc in to_delete if not delete_source(user_id, doc)]
|
||||
except Exception as err:
|
||||
current_app.logger.error(f"Error removing connection: {err}", exc_info=True)
|
||||
return _error("Failed to remove connection", 500)
|
||||
return make_response(jsonify({"success": True, "failed_sources": failed}), 200)
|
||||
|
||||
|
||||
@connections_ns.route("/connections/<string:connection_id>/disconnect")
|
||||
class ConnectionDisconnect(Resource):
|
||||
@@ -104,3 +193,258 @@ class ConnectionDisconnect(Resource):
|
||||
current_app.logger.error(f"Error disconnecting connection: {err}", exc_info=True)
|
||||
return make_response(jsonify({"success": False, "error": "Failed to disconnect"}), 500)
|
||||
return make_response(jsonify({"success": True, "connection": connection}), 200)
|
||||
|
||||
|
||||
@connections_ns.route("/connections/<string:connection_id>/setup")
|
||||
class ConnectionSetup(Resource):
|
||||
@api.doc(
|
||||
description=(
|
||||
"Apply the connect wizard's choices: {create_tools, tool_permissions?, "
|
||||
"sync?: {items, frequency, name?}}. Honours an Idempotency-Key header for the sync."
|
||||
)
|
||||
)
|
||||
def post(self, connection_id: str):
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
body = _json_body()
|
||||
try:
|
||||
with db_session() as conn:
|
||||
row = _owned(conn, connection_id, user_id)
|
||||
if row is None:
|
||||
return _not_found()
|
||||
if service.normalize_status(row) != service.STATUS_CONNECTED:
|
||||
return _error("Reconnect before setting up", 409, code="reconnect")
|
||||
tools = []
|
||||
if body.get("create_tools", True):
|
||||
tools = service.ensure_connection_tools(
|
||||
conn, user_id, row, permissions=body.get("tool_permissions") or None,
|
||||
)
|
||||
tool_payload = [service.serialize_tool(tool) for tool in tools]
|
||||
sources = []
|
||||
if body.get("sync"):
|
||||
started = _start_sync(user_id, row, body["sync"])
|
||||
if isinstance(started, tuple):
|
||||
return _error(*started)
|
||||
sources.append(started)
|
||||
except Exception as err:
|
||||
current_app.logger.error(f"Error setting up connection: {err}", exc_info=True)
|
||||
return _error("Failed to set up connection", 500)
|
||||
return make_response(jsonify({"success": True, "tools": tool_payload, "sources": sources}), 200)
|
||||
|
||||
|
||||
def _start_sync(user_id: str, row: dict, sync: dict):
|
||||
"""Queue the first ingest of a source synced from ``row``.
|
||||
|
||||
Returns the source summary, or ``(message, status)`` on a bad request.
|
||||
"""
|
||||
from docsgpt.api.user.sources.upload import (
|
||||
_claim_task_or_get_cached,
|
||||
_derive_source_id,
|
||||
_read_idempotency_key,
|
||||
_scoped_idempotency_key,
|
||||
)
|
||||
from docsgpt.api.user.tasks import ingest_connector_task, ingest_remote
|
||||
|
||||
definition = catalog.get_definition(catalog.connector_key_for_row(row))
|
||||
if definition is None or not definition.sync_ingestor:
|
||||
return ("This connector does not sync content", 400)
|
||||
items = sync.get("items") or {}
|
||||
if not isinstance(items, dict):
|
||||
return ("items must be an object", 400)
|
||||
frequency = sync.get("frequency") or definition.default_sync_frequency
|
||||
if frequency not in _FREQUENCIES:
|
||||
return ("Unknown sync frequency", 400)
|
||||
name = (sync.get("name") or "").strip() or definition.name
|
||||
idempotency_key, _ = _read_idempotency_key()
|
||||
scoped_key = _scoped_idempotency_key(idempotency_key, user_id)
|
||||
task_id = None
|
||||
if scoped_key:
|
||||
task_id, cached = _claim_task_or_get_cached(scoped_key, "connection_setup_sync")
|
||||
if cached is not None:
|
||||
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,
|
||||
)
|
||||
return {"id": source_id, "task_id": task_id or task.id, "name": name, "sync_frequency": frequency}
|
||||
|
||||
|
||||
@connections_ns.route("/connections/<string:connection_id>/reconnect")
|
||||
class ConnectionReconnect(Resource):
|
||||
@api.doc(
|
||||
description=(
|
||||
"OAuth: returns an authorization URL for the same account. "
|
||||
"API key: accepts {credentials} and replaces the stored ones."
|
||||
)
|
||||
)
|
||||
def post(self, connection_id: str):
|
||||
from docsgpt.api.connector.routes import build_authorization
|
||||
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
body = _json_body()
|
||||
try:
|
||||
with db_readonly() as conn:
|
||||
row = _owned(conn, connection_id, user_id)
|
||||
if row is None:
|
||||
return _not_found()
|
||||
key = catalog.connector_key_for_row(row)
|
||||
definition = catalog.get_definition(key)
|
||||
auth_kind = row.get("auth_kind") or (definition.auth_kind if definition else None)
|
||||
if auth_kind == "oauth":
|
||||
started = build_authorization(row["provider"], user_id, connection_id)
|
||||
return make_response(jsonify({"success": True, "kind": "oauth", **started}), 200)
|
||||
if auth_kind == "mcp_oauth":
|
||||
# The MCP client runs the OAuth dance (dynamic registration,
|
||||
# PKCE); the frontend starts it through /api/mcp_server/test.
|
||||
return make_response(
|
||||
jsonify({"success": True, "kind": "mcp_oauth", "server_url": row.get("server_url")}), 200,
|
||||
)
|
||||
credentials = body.get("credentials")
|
||||
if not isinstance(credentials, dict) or not credentials:
|
||||
return _error("credentials are required", 400)
|
||||
service.ensure_can_store_credentials()
|
||||
with db_session() as conn:
|
||||
locked = ConnectorSessionsRepository(conn).get_for_update(connection_id)
|
||||
stored = service.load_secrets(locked) or {}
|
||||
merged = {**(stored.get("credentials") or {}), **{k: v for k, v in credentials.items() if v}}
|
||||
service.write_secrets(
|
||||
conn, locked, {**stored, "credentials": merged},
|
||||
status=service.STATUS_CONNECTED, last_error=None,
|
||||
)
|
||||
service.resume_sources(conn, connection_id)
|
||||
connection = service.serialize_connection(ConnectorSessionsRepository(conn).get(connection_id))
|
||||
except service.EncryptionKeyNotConfigured as err:
|
||||
return _error(str(err), 400, code="encryption_key_default")
|
||||
except Exception as err:
|
||||
current_app.logger.error(f"Error reconnecting: {err}", exc_info=True)
|
||||
return _error("Failed to reconnect", 500)
|
||||
return make_response(jsonify({"success": True, "kind": "api_key", "connection": connection}), 200)
|
||||
|
||||
|
||||
@connections_ns.route("/connections/<string:connection_id>/picker-token")
|
||||
class ConnectionPickerToken(Resource):
|
||||
@api.doc(description="A short-lived access token for a browser-side file picker. Owner only; never a refresh token.")
|
||||
def post(self, connection_id: str):
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
with db_readonly() as conn:
|
||||
row = _owned(conn, connection_id, user_id)
|
||||
if row is None or (row.get("auth_kind") or "oauth") != "oauth":
|
||||
return _not_found()
|
||||
try:
|
||||
token = service.picker_token(connection_id)
|
||||
except service.ConnectionUnavailable:
|
||||
return _error("Reconnect to continue", 409, code="reconnect")
|
||||
except service.TransientConnectionError:
|
||||
return _error("The provider is not responding. Try again.", 503)
|
||||
return make_response(jsonify({"success": True, **token}), 200)
|
||||
|
||||
|
||||
@connections_ns.route("/connections/claim")
|
||||
class ConnectionClaim(Resource):
|
||||
@api.doc(
|
||||
description=(
|
||||
"One-time link of a legacy browser session token ({provider, session_token}) "
|
||||
"to the caller's connection. Removed next release."
|
||||
)
|
||||
)
|
||||
def post(self):
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
body = _json_body()
|
||||
provider, token = body.get("provider"), body.get("session_token")
|
||||
if not provider or not token:
|
||||
return _error("provider and session_token are required", 400)
|
||||
with db_readonly() as conn:
|
||||
row = service.claim_session_token(conn, user_id, str(provider), str(token))
|
||||
if row is None:
|
||||
return _not_found()
|
||||
return make_response(jsonify({"success": True, "connection_id": str(row["id"])}), 200)
|
||||
|
||||
|
||||
@connections_ns.route("/connections/<string:connection_id>/tools/<string:tool_id>/permissions")
|
||||
class ConnectionToolPermissions(Resource):
|
||||
@api.doc(description="Set per-action permissions: {permissions: {action: always | ask | off}}")
|
||||
def put(self, connection_id: str, tool_id: str):
|
||||
from docsgpt.connectors.permissions import PERMISSIONS
|
||||
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
permissions = _json_body().get("permissions")
|
||||
if not isinstance(permissions, dict) or any(p not in PERMISSIONS for p in permissions.values()):
|
||||
return _error("permissions must map action names to always, ask or off", 400)
|
||||
with db_session() as conn:
|
||||
row = _owned(conn, connection_id, user_id)
|
||||
if row is None:
|
||||
return _not_found()
|
||||
tool = service.set_tool_permissions(conn, user_id, tool_id, permissions)
|
||||
if tool is None or str(tool.get("connection_id")) != connection_id:
|
||||
return _not_found()
|
||||
payload = service.serialize_tool(tool)
|
||||
return make_response(jsonify({"success": True, "tool": payload}), 200)
|
||||
|
||||
|
||||
@connections_ns.route("/connections/<string:connection_id>/refresh-tools")
|
||||
class ConnectionRefreshTools(Resource):
|
||||
@api.doc(description="MCP: re-scan the server's actions and return what was added and removed")
|
||||
def post(self, connection_id: str):
|
||||
from docsgpt.connectors.mcp import refresh_mcp_tools
|
||||
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
with db_readonly() as conn:
|
||||
row = _owned(conn, connection_id, user_id)
|
||||
if row is None:
|
||||
return _not_found()
|
||||
try:
|
||||
diff = refresh_mcp_tools(user_id, row)
|
||||
except service.ConnectionUnavailable:
|
||||
return _error("Reconnect to continue", 409, code="reconnect")
|
||||
except Exception as err:
|
||||
current_app.logger.error(f"Error refreshing MCP tools: {err}", exc_info=True)
|
||||
return _error("Failed to refresh tools", 502)
|
||||
return make_response(jsonify({"success": True, **diff}), 200)
|
||||
+108
-109
@@ -1,7 +1,6 @@
|
||||
import base64
|
||||
import html
|
||||
import json
|
||||
import uuid
|
||||
from typing import Optional
|
||||
from urllib.parse import urlencode, urlsplit
|
||||
|
||||
@@ -20,11 +19,11 @@ from docsgpt.api import api
|
||||
from docsgpt.api.user.tasks import (
|
||||
ingest_connector_task,
|
||||
)
|
||||
from docsgpt.connectors import service
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.parser.connectors.connector_creator import ConnectorCreator
|
||||
from docsgpt.storage.db.repositories.connector_sessions import (
|
||||
ConnectorSessionsRepository,
|
||||
owns_connector_session,
|
||||
)
|
||||
from docsgpt.storage.db.repositories.sources import SourcesRepository
|
||||
from docsgpt.storage.db.session import db_readonly, db_session
|
||||
@@ -100,6 +99,7 @@ def _js_literal(value) -> str:
|
||||
|
||||
def _render_callback_page(
|
||||
status: str, message: str, provider_raw: str, session_token: str = "", user_email: str = "",
|
||||
connection_id: str = "",
|
||||
):
|
||||
"""Popup page that reports an OAuth result to the opener on allowed origins only."""
|
||||
status = status if status in ("success", "error", "cancelled") else "error"
|
||||
@@ -110,7 +110,14 @@ def _render_callback_page(
|
||||
)
|
||||
payload = None
|
||||
if provider_key and status == "success" and session_token:
|
||||
payload = {"type": f"{provider_key}_auth_success", "session_token": session_token, "user_email": user_email}
|
||||
payload = {
|
||||
"type": f"{provider_key}_auth_success",
|
||||
# The connection id is what current frontends use; the session
|
||||
# token is kept for one release for frontends from before it.
|
||||
"connection_id": connection_id,
|
||||
"session_token": session_token,
|
||||
"user_email": user_email,
|
||||
}
|
||||
elif provider_key and status == "error":
|
||||
# The frontend shows its own localized failure message; cancellations are
|
||||
# reported when the popup closes.
|
||||
@@ -170,6 +177,27 @@ def _render_callback_page(
|
||||
|
||||
|
||||
|
||||
def build_authorization(provider: str, user_id: str, connection_id: Optional[str] = None) -> dict:
|
||||
"""Start an OAuth sign-in for ``provider`` and return its authorization URL.
|
||||
|
||||
Raises:
|
||||
service.EncryptionKeyNotConfigured: See ``ensure_can_store_credentials``.
|
||||
service.ConnectionUnavailable: ``connection_id`` is not the caller's.
|
||||
"""
|
||||
service.ensure_can_store_credentials()
|
||||
with db_session() as conn:
|
||||
session_row = service.begin_oauth(conn, user_id, provider, connection_id)
|
||||
state = base64.urlsafe_b64encode(
|
||||
json.dumps({"provider": provider, "object_id": str(session_row["id"])}).encode()
|
||||
).decode()
|
||||
auth = ConnectorCreator.create_auth(provider)
|
||||
return {
|
||||
"authorization_url": auth.get_authorization_url(state=state),
|
||||
"state": state,
|
||||
"callback_origin": _origin_of(settings.CONNECTOR_REDIRECT_BASE_URI),
|
||||
}
|
||||
|
||||
|
||||
@connectors_ns.route("/api/connectors/auth")
|
||||
class ConnectorAuth(Resource):
|
||||
@api.doc(description="Get connector OAuth authorization URL", params={"provider": "Connector provider (e.g., google_drive)"})
|
||||
@@ -186,20 +214,14 @@ class ConnectorAuth(Resource):
|
||||
if not decoded_token:
|
||||
return make_response(jsonify({"success": False, "error": "Unauthorized"}), 401)
|
||||
user_id = decoded_token.get('sub')
|
||||
|
||||
with db_session() as conn:
|
||||
session_row = ConnectorSessionsRepository(conn).upsert(
|
||||
user_id, provider, status="pending",
|
||||
try:
|
||||
started = build_authorization(provider, user_id, request.args.get("connection_id") or None)
|
||||
except service.EncryptionKeyNotConfigured as err:
|
||||
return make_response(
|
||||
jsonify({"success": False, "error": str(err), "code": "encryption_key_default"}), 400,
|
||||
)
|
||||
session_pg_id = str(session_row["id"])
|
||||
state_dict = {
|
||||
"provider": provider,
|
||||
"object_id": session_pg_id,
|
||||
}
|
||||
state = base64.urlsafe_b64encode(json.dumps(state_dict).encode()).decode()
|
||||
|
||||
auth = ConnectorCreator.create_auth(provider)
|
||||
authorization_url = auth.get_authorization_url(state=state)
|
||||
except service.ConnectionUnavailable:
|
||||
return make_response(jsonify({"success": False, "error": "Connection not found"}), 404)
|
||||
# The popup drops results for origins outside the allowlist, which the
|
||||
# user only sees as a cancelled sign-in; name the missing origin here.
|
||||
request_origin = _origin_of(request.headers.get("Origin"))
|
||||
@@ -208,12 +230,7 @@ class ConnectorAuth(Resource):
|
||||
f"Connector sign-in requested from {request_origin}, which cannot receive the result; "
|
||||
"add it to CONNECTOR_ALLOWED_ORIGINS"
|
||||
)
|
||||
return make_response(jsonify({
|
||||
"success": True,
|
||||
"authorization_url": authorization_url,
|
||||
"state": state,
|
||||
"callback_origin": _origin_of(settings.CONNECTOR_REDIRECT_BASE_URI),
|
||||
}), 200)
|
||||
return make_response(jsonify({"success": True, **started}), 200)
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Error generating connector auth URL: {e}", exc_info=True)
|
||||
return make_response(jsonify({"success": False, "error": "Failed to generate authorization URL"}), 500)
|
||||
@@ -269,13 +286,11 @@ class ConnectorsCallback(Resource):
|
||||
auth = ConnectorCreator.create_auth(provider)
|
||||
token_info = auth.exchange_code_for_tokens(authorization_code)
|
||||
|
||||
session_token = str(uuid.uuid4())
|
||||
|
||||
try:
|
||||
if provider == "google_drive":
|
||||
credentials = auth.create_credentials_from_token_info(token_info)
|
||||
service = auth.build_drive_service(credentials)
|
||||
user_info = service.about().get(fields="user").execute()
|
||||
drive_service = auth.build_drive_service(credentials)
|
||||
user_info = drive_service.about().get(fields="user").execute()
|
||||
user_email = user_info.get('user', {}).get('emailAddress', 'Connected User')
|
||||
else:
|
||||
user_email = token_info.get('user_info', {}).get('email', 'Connected User')
|
||||
@@ -288,29 +303,23 @@ class ConnectorsCallback(Resource):
|
||||
|
||||
# ``object_id`` in the OAuth state is the PG session row
|
||||
# UUID (new flow) or a legacy Mongo ObjectId (pre-cutover
|
||||
# issued state). Try UUID update first; fall back to
|
||||
# legacy id path.
|
||||
patch = {
|
||||
"session_token": session_token,
|
||||
"token_info": sanitized_token_info,
|
||||
"user_email": user_email,
|
||||
"status": "authorized",
|
||||
}
|
||||
# issued state).
|
||||
with db_session() as conn:
|
||||
repo = ConnectorSessionsRepository(conn)
|
||||
if state_object_id:
|
||||
value = str(state_object_id)
|
||||
updated = False
|
||||
if len(value) == 36 and "-" in value:
|
||||
updated = repo.update(value, patch)
|
||||
if not updated:
|
||||
repo.update_by_legacy_id(value, patch)
|
||||
value = str(state_object_id or "")
|
||||
state_row = repo.get(value) or (repo.get_by_legacy_id(value) if value else None)
|
||||
if state_row is None or state_row.get("provider") != provider:
|
||||
raise ValueError("OAuth state names no pending connection")
|
||||
connection = service.complete_oauth(
|
||||
conn, state_row, provider, sanitized_token_info, user_email,
|
||||
)
|
||||
|
||||
# Render instead of redirecting so the session token never
|
||||
# lands in a URL (browser history, access logs, Referer).
|
||||
return _render_callback_page(
|
||||
"success", "Authentication successful", provider,
|
||||
session_token=session_token, user_email=user_email,
|
||||
session_token=connection.get("session_token") or "", user_email=user_email,
|
||||
connection_id=str(connection["id"]),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
@@ -333,7 +342,8 @@ class ConnectorsCallback(Resource):
|
||||
class ConnectorFiles(Resource):
|
||||
@api.expect(api.model("ConnectorFilesModel", {
|
||||
"provider": fields.String(required=True),
|
||||
"session_token": fields.String(required=True),
|
||||
"connection_id": fields.String(required=False),
|
||||
"session_token": fields.String(required=False, description="Legacy; use connection_id"),
|
||||
"folder_id": fields.String(required=False),
|
||||
"limit": fields.Integer(required=False),
|
||||
"page_token": fields.String(required=False),
|
||||
@@ -344,26 +354,29 @@ class ConnectorFiles(Resource):
|
||||
try:
|
||||
data = request.get_json()
|
||||
provider = data.get('provider')
|
||||
session_token = data.get('session_token')
|
||||
limit = data.get('limit', 10)
|
||||
|
||||
if not provider or not session_token:
|
||||
return make_response(jsonify({"success": False, "error": "provider and session_token are required"}), 400)
|
||||
if not provider or not (data.get('connection_id') or data.get('session_token')):
|
||||
return make_response(
|
||||
jsonify({"success": False, "error": "provider and connection_id are required"}), 400,
|
||||
)
|
||||
|
||||
decoded_token = request.decoded_token
|
||||
if not decoded_token:
|
||||
return make_response(jsonify({"success": False, "error": "Unauthorized"}), 401)
|
||||
user = decoded_token.get('sub')
|
||||
with db_readonly() as conn:
|
||||
session = ConnectorSessionsRepository(conn).get_by_session_token(
|
||||
session_token,
|
||||
)
|
||||
if not owns_connector_session(session, user, provider):
|
||||
session = service.resolve_request_connection(user, provider, data)
|
||||
if session is None:
|
||||
return make_response(jsonify({"success": False, "error": "Invalid or unauthorized session"}), 401)
|
||||
|
||||
loader = ConnectorCreator.create_connector(provider, session_token)
|
||||
try:
|
||||
loader = ConnectorCreator.create_connector(provider, connection_id=str(session["id"]))
|
||||
except service.ConnectionUnavailable:
|
||||
return make_response(
|
||||
jsonify({"success": False, "error": "Reconnect to continue", "reconnect": True}), 401,
|
||||
)
|
||||
|
||||
generic_keys = {'provider', 'session_token'}
|
||||
generic_keys = {'provider', 'session_token', 'connection_id'}
|
||||
input_config = {
|
||||
k: v for k, v in data.items() if k not in generic_keys
|
||||
}
|
||||
@@ -408,65 +421,49 @@ class ConnectorFiles(Resource):
|
||||
|
||||
@connectors_ns.route("/api/connectors/validate-session")
|
||||
class ConnectorValidateSession(Resource):
|
||||
@api.expect(api.model("ConnectorValidateSessionModel", {"provider": fields.String(required=True), "session_token": fields.String(required=True)}))
|
||||
@api.doc(description="Validate connector session token and return user info and access token")
|
||||
@api.expect(api.model("ConnectorValidateSessionModel", {
|
||||
"provider": fields.String(required=True),
|
||||
"connection_id": fields.String(required=False),
|
||||
"session_token": fields.String(required=False, description="Legacy; use connection_id"),
|
||||
}))
|
||||
@api.doc(description="Validate a connection and return the account and a short-lived access token")
|
||||
def post(self):
|
||||
try:
|
||||
data = request.get_json()
|
||||
data = request.get_json() or {}
|
||||
provider = data.get('provider')
|
||||
session_token = data.get('session_token')
|
||||
if not provider or not session_token:
|
||||
return make_response(jsonify({"success": False, "error": "provider and session_token are required"}), 400)
|
||||
if not provider or not (data.get('connection_id') or data.get('session_token')):
|
||||
return make_response(
|
||||
jsonify({"success": False, "error": "provider and connection_id are required"}), 400,
|
||||
)
|
||||
|
||||
decoded_token = request.decoded_token
|
||||
if not decoded_token:
|
||||
return make_response(jsonify({"success": False, "error": "Unauthorized"}), 401)
|
||||
user = decoded_token.get('sub')
|
||||
|
||||
with db_readonly() as conn:
|
||||
session = ConnectorSessionsRepository(conn).get_by_session_token(
|
||||
session_token,
|
||||
)
|
||||
if not owns_connector_session(session, user, provider) or not session.get("token_info"):
|
||||
session = service.resolve_request_connection(user, provider, data)
|
||||
if session is None:
|
||||
return make_response(jsonify({"success": False, "error": "Invalid or expired session"}), 401)
|
||||
|
||||
token_info = session["token_info"]
|
||||
auth = ConnectorCreator.create_auth(provider)
|
||||
is_expired = auth.is_token_expired(token_info)
|
||||
|
||||
if is_expired and token_info.get('refresh_token'):
|
||||
try:
|
||||
refreshed_token_info = auth.refresh_access_token(token_info.get('refresh_token'))
|
||||
sanitized_token_info = auth.sanitize_token_info(refreshed_token_info)
|
||||
with db_session() as conn:
|
||||
repo = ConnectorSessionsRepository(conn)
|
||||
row = repo.get_by_session_token(session_token)
|
||||
if row:
|
||||
repo.update(str(row["id"]), {"token_info": sanitized_token_info})
|
||||
token_info = sanitized_token_info
|
||||
is_expired = False
|
||||
except Exception as refresh_error:
|
||||
current_app.logger.error(f"Failed to refresh token: {refresh_error}")
|
||||
|
||||
if is_expired:
|
||||
try:
|
||||
token = service.picker_token(str(session["id"]))
|
||||
except service.ConnectionUnavailable:
|
||||
return make_response(jsonify({
|
||||
"success": False,
|
||||
"expired": True,
|
||||
"error": "Session token has expired. Please reconnect."
|
||||
}), 401)
|
||||
except service.TransientConnectionError:
|
||||
return make_response(
|
||||
jsonify({"success": False, "error": "The provider is not responding. Try again."}), 503,
|
||||
)
|
||||
|
||||
_base_fields = {"access_token", "refresh_token", "token_uri", "expiry"}
|
||||
provider_extras = {k: v for k, v in token_info.items() if k not in _base_fields}
|
||||
|
||||
response_data = {
|
||||
return make_response(jsonify({
|
||||
"success": True,
|
||||
"expired": False,
|
||||
"user_email": session.get('user_email', 'Connected User'),
|
||||
"access_token": token_info.get('access_token'),
|
||||
**provider_extras,
|
||||
}
|
||||
|
||||
return make_response(jsonify(response_data), 200)
|
||||
"connection_id": str(session["id"]),
|
||||
"user_email": session.get('account_label') or session.get('user_email') or 'Connected User',
|
||||
**token,
|
||||
}), 200)
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Error validating connector session: {e}", exc_info=True)
|
||||
return make_response(jsonify({"success": False, "error": "Failed to validate session"}), 500)
|
||||
@@ -483,15 +480,13 @@ class ConnectorDisconnect(Resource):
|
||||
try:
|
||||
data = request.get_json()
|
||||
provider = data.get('provider')
|
||||
session_token = data.get('session_token')
|
||||
if not provider:
|
||||
return make_response(jsonify({"success": False, "error": "provider is required"}), 400)
|
||||
|
||||
if session_token:
|
||||
session = service.resolve_request_connection(decoded_token.get('sub'), provider, data)
|
||||
if session is not None:
|
||||
with db_session() as conn:
|
||||
ConnectorSessionsRepository(conn).delete_by_session_token(
|
||||
session_token, decoded_token.get('sub'),
|
||||
)
|
||||
service.disconnect(conn, session)
|
||||
|
||||
return make_response(jsonify({"success": True}), 200)
|
||||
except Exception as e:
|
||||
@@ -506,7 +501,10 @@ class ConnectorSync(Resource):
|
||||
"ConnectorSyncModel",
|
||||
{
|
||||
"source_id": fields.String(required=True, description="Source ID to sync"),
|
||||
"session_token": fields.String(required=True, description="Authentication token")
|
||||
"connection_id": fields.String(
|
||||
required=False, description="Connection to sync with; defaults to the source's own",
|
||||
),
|
||||
"session_token": fields.String(required=False, description="Legacy; use connection_id")
|
||||
},
|
||||
)
|
||||
)
|
||||
@@ -517,16 +515,15 @@ class ConnectorSync(Resource):
|
||||
return make_response(jsonify({"success": False}), 401)
|
||||
|
||||
try:
|
||||
data = request.get_json()
|
||||
data = request.get_json() or {}
|
||||
source_id = data.get('source_id')
|
||||
session_token = data.get('session_token')
|
||||
|
||||
if not all([source_id, session_token]):
|
||||
if not source_id:
|
||||
return make_response(
|
||||
jsonify({
|
||||
"success": False,
|
||||
"error": "source_id and session_token are required"
|
||||
}),
|
||||
"error": "source_id is required"
|
||||
}),
|
||||
400
|
||||
)
|
||||
user_id = decoded_token.get('sub')
|
||||
@@ -562,9 +559,11 @@ class ConnectorSync(Resource):
|
||||
400
|
||||
)
|
||||
|
||||
with db_readonly() as conn:
|
||||
session = ConnectorSessionsRepository(conn).get_by_session_token(session_token)
|
||||
if not owns_connector_session(session, user_id, source_type):
|
||||
lookup = dict(data)
|
||||
if not (lookup.get('connection_id') or lookup.get('session_token')) and source.get('connection_id'):
|
||||
lookup['connection_id'] = str(source['connection_id'])
|
||||
session = service.resolve_request_connection(user_id, source_type, lookup)
|
||||
if session is None:
|
||||
return make_response(
|
||||
jsonify({"success": False, "error": "Invalid or unauthorized session"}),
|
||||
401,
|
||||
@@ -580,7 +579,7 @@ class ConnectorSync(Resource):
|
||||
job_name=source.get('name'),
|
||||
user=decoded_token.get('sub'),
|
||||
source_type=source_type,
|
||||
session_token=session_token,
|
||||
connection_id=str(session["id"]),
|
||||
file_ids=file_ids,
|
||||
folder_ids=folder_ids,
|
||||
recursive=recursive,
|
||||
|
||||
@@ -231,6 +231,63 @@ class PaginatedSources(Resource):
|
||||
return make_response(jsonify({"success": False}), 400)
|
||||
|
||||
|
||||
def delete_source(user: str, doc: dict) -> bool:
|
||||
"""Delete a source's index, stored files and row. Returns whether it worked.
|
||||
|
||||
Args:
|
||||
user: The owner deleting it.
|
||||
doc: The source row, already authorised for ``user``.
|
||||
"""
|
||||
storage = StorageCreator.get_storage()
|
||||
resolved_id = str(doc["id"])
|
||||
source_id = resolved_id
|
||||
|
||||
try:
|
||||
if settings.VECTOR_STORE == "faiss":
|
||||
index_path = f"indexes/{resolved_id}"
|
||||
# index.pkl is the legacy sidecar; index.json the current one.
|
||||
# Older sources have only the former, so clear whichever exist.
|
||||
for index_file in ("index.faiss", "index.json", "index.pkl"):
|
||||
if storage.file_exists(f"{index_path}/{index_file}"):
|
||||
storage.delete_file(f"{index_path}/{index_file}")
|
||||
else:
|
||||
vectorstore = VectorCreator.create_vectorstore(
|
||||
settings.VECTOR_STORE, source_id=source_id
|
||||
)
|
||||
vectorstore.delete_index()
|
||||
if "file_path" in doc and doc["file_path"]:
|
||||
file_path = doc["file_path"]
|
||||
if storage.is_directory(file_path):
|
||||
files = storage.list_files(file_path)
|
||||
for f in files:
|
||||
storage.delete_file(f)
|
||||
else:
|
||||
storage.delete_file(file_path)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except Exception as err:
|
||||
current_app.logger.error(
|
||||
f"Error deleting files and indexes: {err}", exc_info=True
|
||||
)
|
||||
return False
|
||||
try:
|
||||
with db_session() as conn:
|
||||
SourcesRepository(conn).delete(resolved_id, user)
|
||||
record_event(
|
||||
conn,
|
||||
"source.deleted",
|
||||
actor=user,
|
||||
source_id=resolved_id,
|
||||
name=doc.get("name"),
|
||||
)
|
||||
except Exception as err:
|
||||
current_app.logger.error(
|
||||
f"Error deleting source row: {err}", exc_info=True
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@sources_ns.route("/delete_old")
|
||||
class DeleteOldIndexes(Resource):
|
||||
@api.doc(
|
||||
@@ -255,51 +312,7 @@ class DeleteOldIndexes(Resource):
|
||||
return make_response(jsonify({"success": False}), 400)
|
||||
if not doc:
|
||||
return make_response(jsonify({"status": "not found"}), 404)
|
||||
storage = StorageCreator.get_storage()
|
||||
resolved_id = str(doc["id"])
|
||||
|
||||
try:
|
||||
if settings.VECTOR_STORE == "faiss":
|
||||
index_path = f"indexes/{resolved_id}"
|
||||
# index.pkl is the legacy sidecar; index.json the current one.
|
||||
# Older sources have only the former, so clear whichever exist.
|
||||
for index_file in ("index.faiss", "index.json", "index.pkl"):
|
||||
if storage.file_exists(f"{index_path}/{index_file}"):
|
||||
storage.delete_file(f"{index_path}/{index_file}")
|
||||
else:
|
||||
vectorstore = VectorCreator.create_vectorstore(
|
||||
settings.VECTOR_STORE, source_id=resolved_id
|
||||
)
|
||||
vectorstore.delete_index()
|
||||
if "file_path" in doc and doc["file_path"]:
|
||||
file_path = doc["file_path"]
|
||||
if storage.is_directory(file_path):
|
||||
files = storage.list_files(file_path)
|
||||
for f in files:
|
||||
storage.delete_file(f)
|
||||
else:
|
||||
storage.delete_file(file_path)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except Exception as err:
|
||||
current_app.logger.error(
|
||||
f"Error deleting files and indexes: {err}", exc_info=True
|
||||
)
|
||||
return make_response(jsonify({"success": False}), 400)
|
||||
try:
|
||||
with db_session() as conn:
|
||||
SourcesRepository(conn).delete(resolved_id, user)
|
||||
record_event(
|
||||
conn,
|
||||
"source.deleted",
|
||||
actor=user,
|
||||
source_id=resolved_id,
|
||||
name=doc.get("name"),
|
||||
)
|
||||
except Exception as err:
|
||||
current_app.logger.error(
|
||||
f"Error deleting source row: {err}", exc_info=True
|
||||
)
|
||||
if not delete_source(user, doc):
|
||||
return make_response(jsonify({"success": False}), 400)
|
||||
return make_response(jsonify({"success": True}), 200)
|
||||
|
||||
|
||||
@@ -28,7 +28,6 @@ from docsgpt.security.zip_archive import (
|
||||
)
|
||||
from docsgpt.storage.db.repositories.connector_sessions import (
|
||||
ConnectorSessionsRepository,
|
||||
owns_connector_session,
|
||||
)
|
||||
from docsgpt.storage.db.repositories.idempotency import IdempotencyRepository
|
||||
from docsgpt.storage.db.repositories.sources import SourcesRepository
|
||||
@@ -451,6 +450,43 @@ class UploadFile(Resource):
|
||||
return make_response(jsonify(response_payload), 200)
|
||||
|
||||
|
||||
def _remote_credentials(user, source, config):
|
||||
"""Split an S3 / Reddit request into loader config and the connection holding its keys.
|
||||
|
||||
A request naming a ``connection_id`` uses that connection's stored keys.
|
||||
A request carrying keys (the form before connections) stores them on a
|
||||
connection, so they are entered once and never land in
|
||||
``sources.remote_data``. When a multi-user install still runs on the
|
||||
public default encryption key, the keys stay with the source as before.
|
||||
|
||||
Returns:
|
||||
``(source_data, connection_id, error_response)``.
|
||||
"""
|
||||
from docsgpt.connectors import catalog, service
|
||||
|
||||
definition = catalog.get_definition(source)
|
||||
credential_keys = {f.key for f in definition.credential_fields}
|
||||
public = {k: v for k, v in config.items() if k not in credential_keys and k != "connection_id"}
|
||||
connection_id = config.get("connection_id")
|
||||
if connection_id:
|
||||
with db_readonly() as conn:
|
||||
row = ConnectorSessionsRepository(conn).get_for_user(str(connection_id), user)
|
||||
if row is None or catalog.connector_key_for_row(row) != source:
|
||||
return None, None, make_response(
|
||||
jsonify({"success": False, "error": "Invalid or unauthorized connection"}), 401,
|
||||
)
|
||||
return public, str(row["id"]), None
|
||||
provided = {k: config[k] for k in credential_keys if config.get(k) not in (None, "")}
|
||||
if not provided:
|
||||
return config, None, None
|
||||
try:
|
||||
with db_session() as conn:
|
||||
row, _ = service.create_api_key_connection(conn, user, definition, provided)
|
||||
except (service.EncryptionKeyNotConfigured, ValueError):
|
||||
return config, None, None
|
||||
return public, str(row["id"]), None
|
||||
|
||||
|
||||
@sources_upload_ns.route("/remote")
|
||||
class UploadRemote(Resource):
|
||||
@api.expect(
|
||||
@@ -520,32 +556,35 @@ class UploadRemote(Resource):
|
||||
try:
|
||||
config = json.loads(data["data"])
|
||||
source_data = None
|
||||
connection_id = None
|
||||
|
||||
if data["source"] == "github":
|
||||
source_data = config.get("repo_url")
|
||||
elif data["source"] in ["crawler", "url", "sitemap"]:
|
||||
source_data = config.get("url")
|
||||
elif data["source"] == "reddit":
|
||||
source_data = config
|
||||
elif data["source"] == "s3":
|
||||
source_data = config
|
||||
elif data["source"] in ("reddit", "s3"):
|
||||
source_data, connection_id, error = _remote_credentials(user, data["source"], config)
|
||||
if error is not None:
|
||||
if scoped_key:
|
||||
_release_claim(scoped_key)
|
||||
return error
|
||||
elif data["source"] in ConnectorCreator.get_supported_connectors():
|
||||
session_token = config.get("session_token")
|
||||
if not session_token:
|
||||
if not (config.get("connection_id") or config.get("session_token")):
|
||||
if scoped_key:
|
||||
_release_claim(scoped_key)
|
||||
return make_response(
|
||||
jsonify(
|
||||
{
|
||||
"success": False,
|
||||
"error": f"Missing session_token in {data['source']} configuration",
|
||||
"error": f"Missing connection_id in {data['source']} configuration",
|
||||
}
|
||||
),
|
||||
400,
|
||||
)
|
||||
with db_readonly() as conn:
|
||||
connector_session = ConnectorSessionsRepository(conn).get_by_session_token(session_token)
|
||||
if not owns_connector_session(connector_session, user, data["source"]):
|
||||
from docsgpt.connectors import service as connection_service
|
||||
|
||||
connector_session = connection_service.resolve_request_connection(user, data["source"], config)
|
||||
if connector_session is None:
|
||||
if scoped_key:
|
||||
_release_claim(scoped_key)
|
||||
return make_response(
|
||||
@@ -576,7 +615,7 @@ class UploadRemote(Resource):
|
||||
"job_name": data["name"],
|
||||
"user": user,
|
||||
"source_type": data["source"],
|
||||
"session_token": session_token,
|
||||
"connection_id": str(connector_session["id"]),
|
||||
"file_ids": file_ids,
|
||||
"folder_ids": folder_ids,
|
||||
"recursive": config.get("recursive", False),
|
||||
@@ -611,6 +650,7 @@ class UploadRemote(Resource):
|
||||
remote_kwargs = {
|
||||
"kwargs": {
|
||||
"source_data": source_data,
|
||||
"connection_id": connection_id,
|
||||
"job_name": data["name"],
|
||||
"user": user,
|
||||
"loader": data["source"],
|
||||
|
||||
@@ -79,6 +79,12 @@ DURABLE_TASK = dict(
|
||||
)
|
||||
|
||||
|
||||
def _connection_unavailable():
|
||||
from docsgpt.connectors.service import ConnectionUnavailable
|
||||
|
||||
return ConnectionUnavailable
|
||||
|
||||
|
||||
def durable_task(**overrides) -> Dict:
|
||||
"""Return ``DURABLE_TASK`` with per-task overrides applied.
|
||||
|
||||
@@ -167,13 +173,16 @@ def ingest(
|
||||
@with_idempotency(task_name="ingest_remote", on_poison=_emit_ingest_poison_event)
|
||||
def ingest_remote(
|
||||
self, source_data, job_name, user, loader,
|
||||
config=None, idempotency_key=None, source_id=None,
|
||||
config=None, idempotency_key=None, source_id=None, connection_id=None,
|
||||
sync_frequency="never",
|
||||
):
|
||||
resp = remote_worker(
|
||||
self, source_data, job_name, user, loader,
|
||||
sync_frequency=sync_frequency,
|
||||
config=config,
|
||||
idempotency_key=idempotency_key,
|
||||
source_id=source_id,
|
||||
connection_id=connection_id,
|
||||
)
|
||||
return resp
|
||||
|
||||
@@ -249,6 +258,15 @@ def schedule_syncs(self, frequency):
|
||||
return resp
|
||||
|
||||
|
||||
@celery.task(bind=True, acks_late=True, autoretry_for=(Exception,), max_retries=3, retry_backoff=60,
|
||||
dont_autoretry_for=(_connection_unavailable(),))
|
||||
def sync_connector_source(self, source_id):
|
||||
"""Re-sync one connector source from its connection, with no browser involved."""
|
||||
from docsgpt.worker import sync_connector_source as run
|
||||
|
||||
return run(self, source_id)
|
||||
|
||||
|
||||
@celery.task(bind=True)
|
||||
def sync_source(
|
||||
self,
|
||||
@@ -408,7 +426,9 @@ except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@celery.task(**DURABLE_TASK)
|
||||
# A revoked or disconnected connection will not heal by retrying; the
|
||||
# service has already paused the source and told its owner to reconnect.
|
||||
@celery.task(**durable_task(dont_autoretry_for=(DocumentParseError, _connection_unavailable())))
|
||||
@with_idempotency(
|
||||
task_name="ingest_connector_task", on_poison=_emit_ingest_poison_event,
|
||||
)
|
||||
@@ -428,6 +448,7 @@ def ingest_connector_task(
|
||||
config=None,
|
||||
idempotency_key=None,
|
||||
source_id=None,
|
||||
connection_id=None,
|
||||
):
|
||||
from docsgpt.worker import ingest_connector
|
||||
|
||||
@@ -437,6 +458,7 @@ def ingest_connector_task(
|
||||
user,
|
||||
source_type,
|
||||
session_token=session_token,
|
||||
connection_id=connection_id,
|
||||
file_ids=file_ids,
|
||||
folder_ids=folder_ids,
|
||||
recursive=recursive,
|
||||
|
||||
@@ -75,6 +75,43 @@ def _validate_mcp_server_url(config: dict) -> None:
|
||||
raise ValueError(f"Invalid server URL: {exc}") from exc
|
||||
|
||||
|
||||
def _mcp_connection(user, config, auth_type, auth_credentials, display_name):
|
||||
"""The connection an MCP tool runs with; created on first save.
|
||||
|
||||
OAuth servers already have one (the sign-in stored its tokens there).
|
||||
Key, bearer and basic auth store their secret on a connection; servers
|
||||
with no auth get a credential-less connection so they still appear on
|
||||
the Connectors page. Returns None when a multi-user install runs on the
|
||||
default encryption key, which keeps the legacy per-tool secret.
|
||||
"""
|
||||
from docsgpt.connectors import catalog, service
|
||||
|
||||
base = catalog.base_url(config.get("server_url"))
|
||||
if not base:
|
||||
return None
|
||||
if auth_type == "oauth":
|
||||
with db_readonly() as conn:
|
||||
row = service._mcp_row(conn, user, base, None)
|
||||
return str(row["id"]) if row else None
|
||||
definition = catalog.get_definition("custom_mcp")
|
||||
host = base.split("://")[-1]
|
||||
try:
|
||||
with db_session() as conn:
|
||||
if auth_credentials:
|
||||
row, _ = service.create_api_key_connection(
|
||||
conn, user, definition, auth_credentials, server_url=base, display_name=display_name,
|
||||
)
|
||||
else:
|
||||
repo = ConnectorSessionsRepository(conn)
|
||||
row = repo.find_account(user, "custom_mcp", server_url=base, account_label=host) or repo.create(
|
||||
user, "custom_mcp", connector_key="custom_mcp", auth_kind="none",
|
||||
display_name=display_name, account_label=host, server_url=base,
|
||||
)
|
||||
except service.EncryptionKeyNotConfigured:
|
||||
return None
|
||||
return str(row["id"]) if row else None
|
||||
|
||||
|
||||
@tools_mcp_ns.route("/mcp_server/test")
|
||||
class TestMCPServerConfig(Resource):
|
||||
@api.expect(
|
||||
@@ -283,9 +320,19 @@ class MCPServerSave(Resource):
|
||||
"redirect_uri",
|
||||
]:
|
||||
storage_config.pop(field, None)
|
||||
transformed_actions = transform_actions(actions_metadata)
|
||||
from docsgpt.connectors.permissions import apply_default_permissions
|
||||
|
||||
transformed_actions = apply_default_permissions(
|
||||
"mcp_tool", transform_actions(actions_metadata),
|
||||
)
|
||||
|
||||
display_name = data["displayName"]
|
||||
connection_id = _mcp_connection(
|
||||
user, storage_config, auth_type, auth_credentials, display_name,
|
||||
) or (str(existing_doc["connection_id"]) if existing_doc and existing_doc.get("connection_id") else None)
|
||||
if connection_id and auth_type != "oauth":
|
||||
# The secret lives on the connection only.
|
||||
storage_config.pop("encrypted_credentials", None)
|
||||
description = f"MCP Server: {storage_config.get('server_url', 'Unknown')}"
|
||||
status_bool = bool(data.get("status", True))
|
||||
|
||||
@@ -301,6 +348,7 @@ class MCPServerSave(Resource):
|
||||
"config": storage_config,
|
||||
"actions": transformed_actions,
|
||||
"status": status_bool,
|
||||
"connection_id": connection_id,
|
||||
},
|
||||
)
|
||||
saved_id = str(existing_doc["id"])
|
||||
@@ -328,6 +376,7 @@ class MCPServerSave(Resource):
|
||||
"config": storage_config,
|
||||
"actions": transformed_actions,
|
||||
"status": status_bool,
|
||||
"connection_id": connection_id,
|
||||
},
|
||||
)
|
||||
saved_id = str(existing_by_name["id"])
|
||||
@@ -347,6 +396,7 @@ class MCPServerSave(Resource):
|
||||
config_requirements={},
|
||||
actions=transformed_actions,
|
||||
status=status_bool,
|
||||
connection_id=connection_id,
|
||||
)
|
||||
saved_id = str(created["id"])
|
||||
response_data = {
|
||||
@@ -461,50 +511,34 @@ class MCPAuthStatus(Resource):
|
||||
jsonify({"success": True, "statuses": {}}), 200
|
||||
)
|
||||
|
||||
oauth_server_urls: dict = {}
|
||||
from docsgpt.connectors import service
|
||||
|
||||
# Read from connection status alone: status checks never
|
||||
# decrypt credentials.
|
||||
statuses: dict = {}
|
||||
for tool in mcp_tools:
|
||||
tool_id = str(tool["id"])
|
||||
config = tool.get("config") or {}
|
||||
auth_type = config.get("auth_type", "none")
|
||||
if auth_type == "oauth":
|
||||
server_url = config.get("server_url", "")
|
||||
if server_url:
|
||||
parsed = urlparse(server_url)
|
||||
base_url = f"{parsed.scheme}://{parsed.netloc}"
|
||||
oauth_server_urls[tool_id] = base_url
|
||||
row = None
|
||||
if tool.get("connection_id"):
|
||||
row = sessions_repo.get(str(tool["connection_id"]))
|
||||
elif auth_type == "oauth" and config.get("server_url"):
|
||||
parsed = urlparse(config["server_url"])
|
||||
row = sessions_repo.get_by_user_provider(
|
||||
user, service.mcp_provider(f"{parsed.scheme}://{parsed.netloc}"),
|
||||
)
|
||||
if row is not None:
|
||||
connected = service.normalize_status(row) == service.STATUS_CONNECTED
|
||||
if auth_type == "oauth" or not connected:
|
||||
statuses[tool_id] = "connected" if connected else "needs_auth"
|
||||
else:
|
||||
statuses[tool_id] = "needs_auth"
|
||||
statuses[tool_id] = "configured"
|
||||
elif auth_type == "oauth":
|
||||
statuses[tool_id] = "needs_auth"
|
||||
else:
|
||||
statuses[tool_id] = "configured"
|
||||
|
||||
if oauth_server_urls:
|
||||
# Look up a session per distinct base URL. MCP sessions
|
||||
# are stored with ``provider = "mcp:<server_url>"``
|
||||
# and the URL in ``server_url``; reuse the repo's
|
||||
# per-URL accessor rather than an ad-hoc $in query.
|
||||
url_has_tokens: dict = {}
|
||||
for base_url in set(oauth_server_urls.values()):
|
||||
session = sessions_repo.get_by_user_and_server_url(
|
||||
user, base_url,
|
||||
)
|
||||
tokens = (
|
||||
(session or {}).get("session_data", {}) or {}
|
||||
).get("tokens", {}) or {}
|
||||
# MCP code also stashes tokens into token_info on
|
||||
# the row; consider either present as "connected".
|
||||
token_info = (session or {}).get("token_info") or {}
|
||||
url_has_tokens[base_url] = bool(
|
||||
tokens.get("access_token")
|
||||
or token_info.get("access_token")
|
||||
)
|
||||
|
||||
for tool_id, base_url in oauth_server_urls.items():
|
||||
if url_has_tokens.get(base_url):
|
||||
statuses[tool_id] = "connected"
|
||||
else:
|
||||
statuses[tool_id] = "needs_auth"
|
||||
|
||||
return make_response(jsonify({"success": True, "statuses": statuses}), 200)
|
||||
except Exception as e:
|
||||
current_app.logger.error(
|
||||
|
||||
@@ -271,6 +271,10 @@ class GetTools(Resource):
|
||||
):
|
||||
tool_copy["config"]["has_encrypted_credentials"] = True
|
||||
tool_copy["config"].pop("encrypted_credentials", None)
|
||||
if tool_copy.get("connection_id"):
|
||||
# The secret lives on the connection; the form must not
|
||||
# ask for it again.
|
||||
tool_copy.setdefault("config", {})["has_encrypted_credentials"] = True
|
||||
tool_copy["ownership"] = ownership
|
||||
return tool_copy
|
||||
|
||||
@@ -383,6 +387,11 @@ class CreateTool(Resource):
|
||||
f"Error getting tool actions: {err}", exc_info=True
|
||||
)
|
||||
return make_response(jsonify({"success": False}), 400)
|
||||
definition = definition_for_tool(data["name"])
|
||||
if definition is not None:
|
||||
connected = _create_connected_tool(user, data, definition, tool_instance)
|
||||
if connected is not None:
|
||||
return connected
|
||||
try:
|
||||
config_requirements = tool_instance.get_config_requirements()
|
||||
if config_requirements:
|
||||
@@ -422,6 +431,74 @@ class CreateTool(Resource):
|
||||
return make_response(jsonify({"id": new_id}), 200)
|
||||
|
||||
|
||||
def _create_connected_tool(user, data, definition, tool_instance):
|
||||
"""Create a service tool whose secret lives on a connection, not the tool.
|
||||
|
||||
Uses ``connection_id`` when given, otherwise stores the pasted secret on
|
||||
a connection (reusing an identical one). Returns None to fall back to the
|
||||
legacy path when a multi-user install still runs on the default key.
|
||||
"""
|
||||
from docsgpt.connectors import catalog as connector_catalog
|
||||
from docsgpt.connectors import service as connection_service
|
||||
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
|
||||
|
||||
config_requirements = tool_instance.get_config_requirements()
|
||||
public, secrets = connection_service.split_secrets(data.get("config") or {}, config_requirements)
|
||||
connection_id = data.get("connection_id")
|
||||
if not connection_id:
|
||||
validation_errors = _validate_config(data.get("config") or {}, config_requirements)
|
||||
if validation_errors:
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Validation failed", "errors": validation_errors}), 400,
|
||||
)
|
||||
try:
|
||||
with db_session() as conn:
|
||||
if connection_id:
|
||||
connection = ConnectorSessionsRepository(conn).get_for_user(str(connection_id), user)
|
||||
if connection is None or connector_catalog.connector_key_for_row(connection) != definition.key:
|
||||
return make_response(jsonify({"success": False, "message": "Connection not found"}), 404)
|
||||
else:
|
||||
connection, _ = connection_service.create_api_key_connection(conn, user, definition, secrets)
|
||||
created = connection_service.create_tool_for_connection(
|
||||
conn,
|
||||
user,
|
||||
connection,
|
||||
template=data["name"],
|
||||
display_name=data.get("customName") or data.get("displayName") or definition.name,
|
||||
config=public,
|
||||
status=bool(data.get("status", True)),
|
||||
)
|
||||
except connection_service.EncryptionKeyNotConfigured:
|
||||
return None
|
||||
except ValueError as err:
|
||||
return make_response(jsonify({"success": False, "message": str(err)}), 400)
|
||||
return make_response(jsonify({"id": str(created["id"]), "connection_id": str(connection["id"])}), 200)
|
||||
|
||||
|
||||
def _update_connection_secrets(conn, user, tool_doc, config, config_requirements):
|
||||
"""Write changed secrets onto the tool's connection; return the tool's public config.
|
||||
|
||||
Returns None when the caller does not own the connection (an editor on a
|
||||
team share may change actions, never the owner's credentials).
|
||||
"""
|
||||
from docsgpt.connectors import service as connection_service
|
||||
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
|
||||
|
||||
public, secrets = connection_service.split_secrets(config or {}, config_requirements)
|
||||
if secrets:
|
||||
connection = ConnectorSessionsRepository(conn).get_for_user(str(tool_doc["connection_id"]), user)
|
||||
if connection is None:
|
||||
return None
|
||||
stored = connection_service.read_secrets(connection)
|
||||
credentials = {**(stored.get("credentials") or {}), **secrets}
|
||||
connection_service.write_secrets(
|
||||
conn, connection, {**stored, "credentials": credentials},
|
||||
status=connection_service.STATUS_CONNECTED, last_error=None,
|
||||
)
|
||||
connection_service.resume_sources(conn, str(connection["id"]))
|
||||
return public
|
||||
|
||||
|
||||
@tools_ns.route("/update_tool")
|
||||
class UpdateTool(Resource):
|
||||
@api.expect(
|
||||
@@ -528,6 +605,17 @@ class UpdateTool(Resource):
|
||||
)
|
||||
existing_config = tool_doc.get("config", {}) or {}
|
||||
has_existing_secrets = "encrypted_credentials" in existing_config
|
||||
if tool_doc.get("connection_id"):
|
||||
has_existing_secrets = True
|
||||
new_config = _update_connection_secrets(
|
||||
conn, user, tool_doc, data["config"], config_requirements,
|
||||
)
|
||||
if new_config is None:
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Only the owner can change the credentials"}),
|
||||
403,
|
||||
)
|
||||
data = {**data, "config": new_config}
|
||||
|
||||
if config_requirements:
|
||||
validation_errors = _validate_config(
|
||||
|
||||
Reference in new issue
Block a user