mirror of
https://github.com/langgenius/dify.git
synced 2026-11-11 22:02:23 +08:00
123 lines
4.6 KiB
Python
123 lines
4.6 KiB
Python
"""Resolve tool credentials through scoped snapshots and serialized OAuth refreshes."""
|
|
|
|
import time
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass
|
|
from typing import Protocol
|
|
|
|
from services.credentials.refresh import CredentialRefreshLock, credential_needs_refresh
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ToolRuntimeCredential:
|
|
"""Detached provider state; absent ciphertext represents empty credentials."""
|
|
|
|
id: str
|
|
workspace_id: str
|
|
owner_id: str
|
|
provider: str
|
|
credential_type: str
|
|
encrypted_credentials: str | None
|
|
expires_at: int
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class RefreshedToolCredential:
|
|
credentials: Mapping[str, object]
|
|
expires_at: int
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ResolvedToolCredential:
|
|
credentials: Mapping[str, object]
|
|
credential_type: str
|
|
|
|
|
|
class ToolCredentialNotFoundError(ValueError):
|
|
pass
|
|
|
|
|
|
class ToolCredentialRefreshError(ValueError):
|
|
pass
|
|
|
|
|
|
class ToolCredentialStore(Protocol):
|
|
def get_tool_credential(
|
|
self, *, workspace_id: str, providers: tuple[str, ...], credential_id: str | None
|
|
) -> ToolRuntimeCredential | None:
|
|
"""Return a detached snapshot, scoped to workspace and provider, preferring the default when ID is absent."""
|
|
...
|
|
|
|
def update_tool_credential_if_unchanged(
|
|
self, *, record: ToolRuntimeCredential, encrypted_credentials: str, expires_at: int
|
|
) -> bool:
|
|
"""Commit in a short transaction, guarding the owner, auth type, expiry and previous ciphertext."""
|
|
...
|
|
|
|
|
|
class ToolCredentialRuntime(Protocol):
|
|
def check_policy(self, record: ToolRuntimeCredential) -> None: ...
|
|
|
|
def decrypt(self, record: ToolRuntimeCredential) -> dict[str, object]: ...
|
|
|
|
def encrypt(self, record: ToolRuntimeCredential, credentials: Mapping[str, object]) -> str: ...
|
|
|
|
def refresh(self, record: ToolRuntimeCredential, credentials: Mapping[str, object]) -> RefreshedToolCredential: ...
|
|
|
|
def invalidate(self, record: ToolRuntimeCredential) -> None: ...
|
|
|
|
|
|
class ToolCredentialService:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
credentials: ToolCredentialStore,
|
|
runtime: ToolCredentialRuntime,
|
|
refresh_lock: CredentialRefreshLock,
|
|
providers: tuple[str, ...],
|
|
) -> None:
|
|
self._credentials = credentials
|
|
self._runtime = runtime
|
|
self._refresh_lock = refresh_lock
|
|
self._providers = providers
|
|
|
|
def resolve(self, *, workspace_id: str, credential_id: str | None) -> ResolvedToolCredential:
|
|
record = self._load(workspace_id=workspace_id, credential_id=credential_id)
|
|
self._runtime.check_policy(record)
|
|
if not credential_needs_refresh(record.expires_at, now=time.time()):
|
|
return self._resolved(record)
|
|
|
|
with self._refresh_lock.acquire(
|
|
credential_type="tool", workspace_id=workspace_id, credential_id=record.id
|
|
) as lease:
|
|
record = self._load(workspace_id=workspace_id, credential_id=record.id)
|
|
self._runtime.check_policy(record)
|
|
if not credential_needs_refresh(record.expires_at, now=time.time()):
|
|
return self._resolved(record)
|
|
decrypted = self._runtime.decrypt(record)
|
|
lease.ensure_owned()
|
|
refreshed = self._runtime.refresh(record, decrypted)
|
|
encrypted = self._runtime.encrypt(record, refreshed.credentials)
|
|
lease.ensure_owned()
|
|
if not self._credentials.update_tool_credential_if_unchanged(
|
|
record=record, encrypted_credentials=encrypted, expires_at=refreshed.expires_at
|
|
):
|
|
latest = self._load(workspace_id=workspace_id, credential_id=record.id)
|
|
self._runtime.check_policy(latest)
|
|
if credential_needs_refresh(latest.expires_at, now=time.time()):
|
|
raise ToolCredentialRefreshError("Credential changed while it was being refreshed. Please retry.")
|
|
return self._resolved(latest)
|
|
self._runtime.invalidate(record)
|
|
return ResolvedToolCredential(refreshed.credentials, record.credential_type)
|
|
|
|
def _load(self, *, workspace_id: str, credential_id: str | None) -> ToolRuntimeCredential:
|
|
record = self._credentials.get_tool_credential(
|
|
workspace_id=workspace_id, providers=self._providers, credential_id=credential_id
|
|
)
|
|
if record is None:
|
|
raise ToolCredentialNotFoundError("Tool credential not found")
|
|
return record
|
|
|
|
def _resolved(self, record: ToolRuntimeCredential) -> ResolvedToolCredential:
|
|
return ResolvedToolCredential(self._runtime.decrypt(record), record.credential_type)
|