MCP and OAuth

octomate.managers.mcp

Per-user MCP installations: install, enable and authorize them, and pool their clients.

logger module-attribute

logger = logging.getLogger(__name__)

McpUnavailable

Bases: LookupError

The MCP instance or tentacle named is unknown, not the caller's, or not serving.

Source code in octomate/managers/mcp.py

class McpUnavailable(LookupError):
    """The MCP instance or tentacle named is unknown, not the caller's, or not
    serving."""

    def __init__(self) -> None:
        super().__init__("MCP instance unavailable")

McpClientKey

Bases: NamedTuple

The owner and installation a pooled client is cached under.

Source code in octomate/managers/mcp.py

class McpClientKey(NamedTuple):
    """The owner and installation a pooled client is cached under."""

    user_id: UUID7
    mcp_id: UUID7

user_id instance-attribute

user_id: UUID7

mcp_id instance-attribute

mcp_id: UUID7

McpManager

Bases: Manager, Locks[UUID7]

A user's MCP installations: install, enable, authorize and remove them, and pool their clients.

All connections advertise MCP Apps support. Client acquisition and capability negotiation are independent of channel rendering and sandbox configuration.

Source code in octomate/managers/mcp.py

class McpManager(Manager, Locks[UUID7]):
    """A user's MCP installations: install, enable, authorize and remove them, and
    pool their clients.

    All connections advertise MCP Apps support. Client acquisition and capability
    negotiation are independent of channel rendering and sandbox configuration.
    """

    def __init__(
        self,
        users: UserManager,
        cipher: OAuthCipher | None,
        *,
        idle_timeout: float | None = None,
        httpx_client_factory: McpHttpClientFactory = mcp_http_client,
        oauth: OAuthManager | None = None,
    ) -> None:
        self.users: UserManager = users
        self.cipher: OAuthCipher | None = cipher
        self.idle_timeout: float = (
            idle_timeout if idle_timeout is not None else McpPoolConfig().idle_timeout
        )
        self.httpx_client_factory: McpHttpClientFactory = httpx_client_factory
        self.tentacles: dict[str, McpTentacle] = {}
        self.oauth = oauth
        self.clients: dict[McpClientKey, Client] = {}
        # Latest discovery per pooled connection, bounded by least-recent use.
        self.tool_catalogs: LRUCache[McpClientKey, McpToolCatalog] = LRUCache(
            maxsize=256
        )
        self.sweaps: dict[McpClientKey, asyncio.Task[None]] = {}

    def available(self) -> list[McpTentacleInfo]:
        return [
            tentacle.info for tentacle in self.tentacles.values() if tentacle.serving
        ]

    async def install(self, user_id: UUID7, request: McpInstallRequest) -> Mcp:
        mcp_id = uuid7()
        namespace = f"personal/{request.namespace}"
        instance: Mcp
        tentacle_id = request.tentacle_id
        auth = request.auth or NoAuth()
        instructions = ""
        if tentacle_id is not None:
            tentacle = self.tentacles.get(tentacle_id)
            if tentacle is None or not tentacle.serving:
                raise McpUnavailable
            if tentacle.upstream != str(request.url):
                raise ValueError("MCP URL must match the tentacle")
            instructions = tentacle.instructions
            if tentacle.auth_kind == "bearer":
                if tentacle.token is None:
                    raise ValueError("Bearer MCPs require a token")
                auth = BearerAuth(token=tentacle.token)
            elif tentacle.auth_kind == "oauth":
                auth = OAuth()
        match auth:
            case BearerAuth():
                if self.cipher is None:
                    raise ValueError(
                        "Configure oauth.encryption_key before installing bearer MCPs"
                    )
                instance = BearerMcp(
                    id=mcp_id,
                    user_id=user_id,
                    name=request.name,
                    namespace=namespace,
                    url=str(request.url),
                    instructions=instructions,
                    tentacle_id=tentacle_id,
                    encrypted_token=self.cipher.encrypt(
                        auth.token.get_secret_value(),
                        context=f"mcp:{user_id}:{mcp_id}",
                    ),
                )
            case OAuth():
                if self.cipher is None or self.oauth is None:
                    raise ValueError(
                        "Configure oauth.encryption_key before installing OAuth MCPs"
                    )
                if tentacle_id is not None:
                    connector = self.oauth.connector(tentacle_id)
                    if connector.mcp_url is None or str(connector.mcp_url) != str(
                        request.url
                    ):
                        raise ValueError(
                            "MCP URL must match the configured OAuth tentacle"
                        )
                elif self.oauth.callback_base_uri is None:
                    raise ValueError(
                        "Configure oauth.callback_base_uri before installing dynamic OAuth MCPs"
                    )
                instance = OAuthMcp(
                    id=mcp_id,
                    user_id=user_id,
                    name=request.name,
                    namespace=namespace,
                    url=str(request.url),
                    instructions=instructions,
                    tentacle_id=tentacle_id,
                )
            case NoAuth():
                instance = NoAuthMcp(
                    id=mcp_id,
                    user_id=user_id,
                    name=request.name,
                    namespace=namespace,
                    url=str(request.url),
                    instructions=instructions,
                    tentacle_id=tentacle_id,
                )
        async with async_session() as session:
            session.add(instance)
            try:
                await session.commit()
            except IntegrityError as error:
                await session.rollback()
                duplicate = await session.one_or_none(
                    Mcp,
                    expressions=[
                        Mcp["user_id"] == user_id,
                        Mcp["namespace"] == namespace,
                    ],
                )
                if duplicate is None:
                    raise
                raise ValueError("MCP namespace is already installed") from error
        return instance

    async def list[McpT: Mcp](
        self,
        user_id: UUID7,
        *,
        enabled: bool = False,
        mcp_type: type[McpT] = Mcp,
    ) -> list[McpT]:
        expressions = [Mcp["user_id"] == user_id]
        if enabled:
            expressions.append(Mcp["enabled"].is_(True))
        async with async_session() as session:
            return list(
                await session.list(
                    mcp_type,
                    expressions=expressions,
                    limit=None,
                )
            )

    async def summary(self, user: User, mcp: Mcp) -> McpServerSummary:
        if mcp.user_id != user.id:
            raise McpUnavailable
        summary = mcp.summary
        if isinstance(mcp, OAuthMcp):
            authorization: McpAuthorizationResult = McpAuthorizationStatus(status=None)
            flows: list[OAuthFlowKind] = []
            if self.oauth is not None:
                authorization = await self.authorization_status(user, mcp)
                if mcp.tentacle_id is None:
                    if self.oauth.callback_base_uri is not None:
                        flows = ["authorization_code"]
                else:
                    connector = self.oauth.connectors.get(mcp.tentacle_id)
                    if connector is not None:
                        flows = [flow.kind for flow in connector.flows]
            summary.oauth = McpOAuthSummary(status=authorization.status, flows=flows)
        return summary

    async def enable(self, user_id: UUID7, mcp_id: UUID7) -> Mcp:
        async with self.lock(mcp_id):
            async with async_session() as session:
                instance = await session.one_or_none(
                    Mcp,
                    expressions=[
                        Mcp["id"] == mcp_id,
                        Mcp["user_id"] == user_id,
                    ],
                )
                if instance is None:
                    raise McpUnavailable
                instance.enabled = True
                instance.updated_at = datetime.now(UTC)
                await session.commit()
            return instance

    async def disable(self, user_id: UUID7, mcp_id: UUID7) -> Mcp:
        async with self.lock(mcp_id):
            async with async_session() as session:
                instance = await session.one_or_none(
                    Mcp,
                    expressions=[
                        Mcp["id"] == mcp_id,
                        Mcp["user_id"] == user_id,
                    ],
                )
                if instance is None:
                    raise McpUnavailable
                async with (
                    self.oauth.lock(
                        OAuthLockKey(
                            user_id=user_id,
                            mcp_id=mcp_id,
                            connector_id=instance.tentacle_id or "mcp",
                        )
                    )
                    if isinstance(instance, OAuthMcp) and self.oauth is not None
                    else nullcontext()
                ):
                    instance.enabled = False
                    instance.updated_at = datetime.now(UTC)
                    await session.commit()
            await self.evict(McpClientKey(user_id=user_id, mcp_id=mcp_id))
            return instance

    async def uninstall(self, user_id: UUID7, mcp_id: UUID7) -> None:
        async with self.lock(mcp_id):
            async with async_session() as session:
                instance = await session.one_or_none(
                    Mcp,
                    expressions=[
                        Mcp["id"] == mcp_id,
                        Mcp["user_id"] == user_id,
                    ],
                )
                if instance is None:
                    raise McpUnavailable
                async with (
                    self.oauth.lock(
                        OAuthLockKey(
                            user_id=user_id,
                            mcp_id=mcp_id,
                            connector_id=instance.tentacle_id or "mcp",
                        )
                    )
                    if isinstance(instance, OAuthMcp) and self.oauth is not None
                    else nullcontext()
                ):
                    await session.delete(instance)
                    await session.commit()
            await self.evict(McpClientKey(user_id=user_id, mcp_id=mcp_id))

    async def authorizable(
        self,
        user_id: UUID7,
        *,
        mcp_id: UUID7 | None = None,
        namespace: str | None = None,
    ) -> OAuthMcp:
        if mcp_id is None and namespace is None:
            raise ValueError("MCP ID or namespace is required")
        expressions = [
            OAuthMcp["user_id"] == user_id,
            OAuthMcp["enabled"].is_(True),
        ]
        if mcp_id is not None:
            expressions.append(OAuthMcp["id"] == mcp_id)
        if namespace is not None:
            expressions.append(OAuthMcp["namespace"] == namespace)
        async with async_session() as session:
            instance = await session.one_or_none(OAuthMcp, expressions=expressions)
        if instance is None:
            raise McpUnavailable
        return instance

    async def connect(
        self,
        user: User,
        mcp_id: UUID7,
        *,
        profile: UserProfile | None = None,
        flow: OAuthFlowKind | None = None,
    ) -> OAuthStartResult:
        if self.oauth is None:
            raise McpUnavailable
        instance = await self.authorizable(user_id=user.id, mcp_id=mcp_id)
        async with self.oauth.lock(
            OAuthLockKey(
                user_id=user.id,
                mcp_id=mcp_id,
                connector_id=instance.tentacle_id or "mcp",
            )
        ):
            return await self.oauth.start(
                user,
                instance.tentacle_id or "mcp",
                mcp_id=mcp_id,
                profile=profile,
                flow=flow,
            )

    async def pending_authorization(
        self, user: User, mcp_id: UUID7
    ) -> OAuthStartResult | None:
        if self.oauth is None:
            raise McpUnavailable
        mcp = await self.authorizable(user_id=user.id, mcp_id=mcp_id)
        return await self.oauth.pending_authorization(
            user, mcp.tentacle_id or "mcp", mcp_id=mcp_id
        )

    async def cancel_authorization(self, user: User, mcp_id: UUID7) -> None:
        if self.oauth is None:
            raise McpUnavailable
        mcp = await self.authorizable(user_id=user.id, mcp_id=mcp_id)
        connector_id = mcp.tentacle_id or "mcp"
        async with (
            self.oauth.lock(
                OAuthLockKey(user_id=user.id, mcp_id=mcp_id, connector_id=connector_id)
            ),
            async_session() as session,
        ):
            operations = await session.list(
                OAuthOperation,
                limit=None,
                expressions=[
                    OAuthOperation["user_id"] == user.id,
                    OAuthOperation["mcp_id"] == mcp_id,
                    OAuthOperation["connector_id"] == connector_id,
                    OAuthOperation["consumed_at"].is_(None),
                ],
            )
            now = datetime.now(UTC)
            for operation in operations:
                operation.consumed_at = now
            await session.commit()

    async def confirm(
        self,
        user: User,
        mcp_id: UUID7,
        *,
        profile: UserProfile | None = None,
    ) -> McpAuthorizationResult:
        if self.oauth is None:
            raise McpUnavailable
        instance = await self.authorizable(user_id=user.id, mcp_id=mcp_id)
        connector_id = instance.tentacle_id or "mcp"
        connector = await self.oauth.resolve_connector(
            connector_id, user_id=user.id, mcp_id=mcp_id
        )
        if any(isinstance(flow, DeviceAuthorizationFlow) for flow in connector.flows):
            try:
                result = await self.oauth.complete_latest(
                    user, connector_id, mcp_id=mcp_id, profile=profile
                )
            except NoPendingAuthorization:
                pass
            else:
                if isinstance(result, OAuthPending):
                    return McpDeviceAuthorizationPending(
                        retry_after_seconds=result.retry_after_seconds
                    )
        return await self.authorization_status(user, instance)

    async def authorization_status(
        self, user: User, mcp: OAuthMcp
    ) -> McpAuthorizationResult:
        if self.oauth is None or mcp.user_id != user.id:
            raise McpUnavailable
        connector_id = mcp.tentacle_id or "mcp"
        status = await self.oauth.connection_status(user, connector_id, mcp_id=mcp.id)
        if status != "active":
            async with async_session() as session:
                pending = await session.first(
                    OAuthOperation,
                    expressions=[
                        OAuthOperation["user_id"] == user.id,
                        OAuthOperation["mcp_id"] == mcp.id,
                        OAuthOperation["connector_id"] == connector_id,
                        OAuthOperation["consumed_at"].is_(None),
                        OAuthOperation["expires_at"] > datetime.now(UTC),
                    ],
                    order_bys=[OAuthOperation["id"].desc()],
                )
            if pending is not None:
                if pending.interval_seconds is not None:
                    return McpDeviceAuthorizationPending(
                        retry_after_seconds=pending.interval_seconds
                    )
                return McpBrowserAuthorizationPending()
        return McpAuthorizationStatus(status=status)

    async def resolve(
        self,
        scope: OctomateSession | User,
        namespace: str | None = None,
        *,
        mcp_id: UUID7 | None = None,
    ) -> tuple[User, Mcp]:
        """Resolve one enabled installation owned by the caller, by namespace or UUID.

        Historical invocations supply only their saved UUID. A reused namespace
        therefore cannot substitute a different installation.
        """
        if (namespace is None) == (mcp_id is None):
            raise ValueError("Supply exactly one MCP namespace or installation ID")
        if isinstance(scope, User):
            user = scope
        else:
            user = (
                scope.user_profile.user.value
                if scope.user_profile is not None
                and scope.user_profile.user_id is not None
                else None
            )
        if user is None:
            raise McpUnavailable
        async with async_session() as session:
            instance = await session.one_or_none(
                Mcp,
                expressions=[
                    Mcp["user_id"] == user.id,
                    Mcp["id"] == mcp_id
                    if mcp_id is not None
                    else Mcp["namespace"] == namespace,
                    Mcp["enabled"].is_(True),
                ],
            )
        if instance is None:
            raise McpUnavailable
        return user, instance

    async def catalog(
        self,
        scope: OctomateSession | User,
        namespace: str | None = None,
        *,
        mcp_id: UUID7 | None = None,
        refresh: bool = False,
    ) -> McpToolCatalog:
        """Read cached discovery, loading it when absent or explicitly refreshed.

        Ownership, enabled state, and authorization are checked even on cache hits.
        The bounded catalog follows its pooled connection's lifetime.
        """
        user, instance = await self.resolve(scope, namespace, mcp_id=mcp_id)
        async with (
            self.acquire(user, mcp_id=instance.id) as client,
            self.lock(instance.id),
        ):
            key = McpClientKey(user.id, instance.id)
            catalog = self.tool_catalogs.get(key)
            if catalog is not None and not refresh:
                return catalog
            instructions = "\n".join(
                part for part in (instance.instructions, client.instructions) if part
            )
            catalog = McpToolCatalog(
                instructions=instructions, tools=await client.list_tools()
            )
            self.tool_catalogs[key] = catalog
            return catalog

    async def call_tool(
        self,
        scope: OctomateSession,
        namespace: str,
        name: str,
        arguments: dict[str, JsonValue],
    ) -> CallToolResult:
        """Execute one upstream call and retain its declared app before publishing.

        This boundary sees each call independently, including calls a runtime
        nests inside a code wrapper. Sessions without thread attribution return
        the upstream result without producing an artifact or ledger reference.
        App declarations reuse the connection's latest discovery, loading once
        when absent. Explicit discovery refreshes them; client eviction clears them.
        """
        user, instance = await self.resolve(scope, namespace)
        address = scope.conversation_address
        channel = scope.channels.get(address.channel_tentacle_id) if address else None
        async with self.acquire(user, mcp_id=instance.id) as client:
            if (
                scope.native
                or scope.thread_id is None
                or scope.threads is None
                or channel is None
                or address is None
            ):
                return await client.call_tool_mcp(name, arguments)
            catalog = await self.catalog(user, mcp_id=instance.id)
            tool = catalog.tools_by_name.get(name)
            ui = McpAppToolMeta.model_validate(tool.meta or {}) if tool else None
            result = await client.call_tool_mcp(name, arguments)
        if tool is None or ui is None or ui.resource_uri is None:
            return result
        message = await scope.threads.record_artifact(
            scope.thread_id,
            McpAppArtifact(
                owner_id=user.id,
                title=tool.title or tool.name,
                invocation=McpAppInvocation(
                    installation_id=instance.id,
                    namespace=instance.namespace,
                    provider_name=instance.name,
                    tool=tool,
                    resource_uri=ui.resource_uri,
                    arguments=arguments,
                    result=result,
                ),
            ),
            sender=channel.self_profile,
            agent_tentacle_id=scope.current_agent_id,
        )
        await scope.publish(
            address,
            ArtifactAvailableEvent(
                message_id=message.id, artifact_id=message.artifact_id
            ),
        )
        return result

    async def app_resource(
        self, user: User, invocation: McpAppInvocation
    ) -> McpAppResource:
        """Read only the saved UI URI from its still-authorized installation UUID.

        This performs one resource read, never discovery or tool execution. The
        response must contain one matching MCP Apps HTML resource, as text or
        base64-encoded UTF-8. UI metadata is validated before reaching a browser.
        """
        async with self.acquire(user, mcp_id=invocation.installation_id) as client:
            contents = await client.read_resource(invocation.resource_uri)
        if len(contents) != 1:
            raise ValueError("An MCP app must return exactly one UI resource")
        resource = contents[0]
        if (
            resource.uri != str(invocation.resource_uri)
            or resource.mime_type != UI_MIME_TYPE
        ):
            raise ValueError("MCP app resource URI or MIME type does not match")
        html = (
            resource.text
            if isinstance(resource, TextResourceContents)
            else b64decode(resource.blob, validate=True).decode("utf-8")
        )
        return McpAppResource(
            uri=invocation.resource_uri,
            html=html,
            ui=(resource.meta or {}).get("ui", {}),
        )

    @asynccontextmanager
    async def acquire(
        self,
        scope: OctomateSession | User,
        namespace: str | None = None,
        *,
        mcp_id: UUID7 | None = None,
    ) -> AsyncGenerator[Client]:
        user, instance = await self.resolve(scope, namespace, mcp_id=mcp_id)
        user_id = user.id
        async with self.lock(instance.id):
            # The first lookup finds the lock; reload under it so disabling or
            # deleting an instance wins before another call is dispatched.
            async with async_session() as session:
                current = await session.one_or_none(
                    Mcp,
                    expressions=[Mcp["id"] == instance.id],
                )
            if current is None or not current.enabled or current.user_id != user_id:
                raise McpUnavailable
            if isinstance(current, OAuthMcp) and (
                self.oauth is None
                or await self.oauth.connection_status(
                    user,
                    current.tentacle_id or "mcp",
                    mcp_id=current.id,
                )
                != "active"
            ):
                raise McpUnavailable
            key = McpClientKey(user_id=user_id, mcp_id=current.id)
            client = self.clients.get(key)
            if client is not None and not client.is_connected():
                await self.evict(key)
                client = None
            if client is None:
                headers: dict[str, str] = {}
                tentacle = self.tentacles.get(current.tentacle_id or "")
                if isinstance(current, BearerMcp):
                    if self.cipher is None:
                        raise ValueError("MCP credential encryption is not configured")
                    token = self.cipher.decrypt(
                        current.encrypted_token, context=f"mcp:{user_id}:{current.id}"
                    )
                    headers["Authorization"] = f"Bearer {token}"
                auth: McpBearerAuth | None = None
                if isinstance(current, OAuthMcp) and self.oauth is not None:
                    if isinstance(tentacle, LarkTentacle):
                        auth = LarkMcpBearerAuth(
                            manager=self.oauth,
                            user=user,
                            mcp_id=current.id,
                            url=current.url,
                            tentacle=tentacle,
                        )
                    else:
                        auth = McpBearerAuth(
                            manager=self.oauth,
                            user=user,
                            mcp_id=current.id,
                            connector_id=current.tentacle_id or "mcp",
                            url=current.url,
                        )
                client = Client(
                    transport=StreamableHttpTransport(
                        url=current.url,
                        headers=headers,
                        auth=auth,
                        httpx_client_factory=self.httpx_client_factory,
                    ),
                    mode="auto",
                    cache=False,
                    extensions=[
                        advertise(
                            UI_EXTENSION_ID,
                            {"mimeTypes": [UI_MIME_TYPE]},
                        )
                    ],
                )
                await client.__aenter__()
                self.clients[key] = client
            # Acquire the call's reference before releasing the lock, so removal
            # cannot close the client between checking access and starting a call.
            try:
                await client.__aenter__()
            except Exception:
                if not client.is_connected():
                    await self.evict(key)
                raise
            if cleanup := self.sweaps.pop(key, None):
                cleanup.cancel()
                with suppress(asyncio.CancelledError):
                    await cleanup
            self.sweaps[key] = asyncio.create_task(self.expire(key))
        try:
            yield client
        finally:
            try:
                await client.__aexit__(None, None, None)
            finally:
                if not client.is_connected():
                    async with self.lock(key.mcp_id):
                        if self.clients.get(key) is client:
                            await self.evict(key)

    async def expire(self, key: McpClientKey) -> None:
        await asyncio.sleep(self.idle_timeout)
        async with self.lock(key.mcp_id):
            await self.evict(key)

    async def evict(self, key: McpClientKey) -> None:
        """Close a cached client while the caller holds its MCP lock."""
        self.tool_catalogs.pop(key, None)
        if cleanup := self.sweaps.pop(key, None):
            if cleanup is not asyncio.current_task():
                cleanup.cancel()
                with suppress(asyncio.CancelledError):
                    await cleanup
        if client := self.clients.get(key):
            try:
                await client.close()
            except Exception:
                logger.exception("Failed to close MCP client %s", key.mcp_id)
            finally:
                del self.clients[key]

    @asynccontextmanager
    async def lifespan(self) -> AsyncGenerator[None]:
        try:
            yield
        finally:
            while self.clients:
                key = next(iter(self.clients))
                async with self.lock(key.mcp_id):
                    await self.evict(key)

users instance-attribute

users: UserManager = users

cipher instance-attribute

cipher: OAuthCipher | None = cipher

idle_timeout instance-attribute

idle_timeout: float = (
    idle_timeout
    if idle_timeout is not None
    else McpPoolConfig().idle_timeout
)

httpx_client_factory instance-attribute

httpx_client_factory: McpHttpClientFactory = (
    httpx_client_factory
)

tentacles instance-attribute

tentacles: dict[str, McpTentacle] = {}

oauth instance-attribute

oauth = oauth

clients instance-attribute

clients: dict[McpClientKey, Client] = {}

tool_catalogs instance-attribute

tool_catalogs: LRUCache[McpClientKey, McpToolCatalog] = (
    LRUCache(maxsize=256)
)

sweaps instance-attribute

sweaps: dict[McpClientKey, Task[None]] = {}

available

available() -> list[McpTentacleInfo]

Source code in octomate/managers/mcp.py

def available(self) -> list[McpTentacleInfo]:
    return [
        tentacle.info for tentacle in self.tentacles.values() if tentacle.serving
    ]

install async

install(user_id: UUID7, request: McpInstallRequest) -> Mcp

Source code in octomate/managers/mcp.py

async def install(self, user_id: UUID7, request: McpInstallRequest) -> Mcp:
    mcp_id = uuid7()
    namespace = f"personal/{request.namespace}"
    instance: Mcp
    tentacle_id = request.tentacle_id
    auth = request.auth or NoAuth()
    instructions = ""
    if tentacle_id is not None:
        tentacle = self.tentacles.get(tentacle_id)
        if tentacle is None or not tentacle.serving:
            raise McpUnavailable
        if tentacle.upstream != str(request.url):
            raise ValueError("MCP URL must match the tentacle")
        instructions = tentacle.instructions
        if tentacle.auth_kind == "bearer":
            if tentacle.token is None:
                raise ValueError("Bearer MCPs require a token")
            auth = BearerAuth(token=tentacle.token)
        elif tentacle.auth_kind == "oauth":
            auth = OAuth()
    match auth:
        case BearerAuth():
            if self.cipher is None:
                raise ValueError(
                    "Configure oauth.encryption_key before installing bearer MCPs"
                )
            instance = BearerMcp(
                id=mcp_id,
                user_id=user_id,
                name=request.name,
                namespace=namespace,
                url=str(request.url),
                instructions=instructions,
                tentacle_id=tentacle_id,
                encrypted_token=self.cipher.encrypt(
                    auth.token.get_secret_value(),
                    context=f"mcp:{user_id}:{mcp_id}",
                ),
            )
        case OAuth():
            if self.cipher is None or self.oauth is None:
                raise ValueError(
                    "Configure oauth.encryption_key before installing OAuth MCPs"
                )
            if tentacle_id is not None:
                connector = self.oauth.connector(tentacle_id)
                if connector.mcp_url is None or str(connector.mcp_url) != str(
                    request.url
                ):
                    raise ValueError(
                        "MCP URL must match the configured OAuth tentacle"
                    )
            elif self.oauth.callback_base_uri is None:
                raise ValueError(
                    "Configure oauth.callback_base_uri before installing dynamic OAuth MCPs"
                )
            instance = OAuthMcp(
                id=mcp_id,
                user_id=user_id,
                name=request.name,
                namespace=namespace,
                url=str(request.url),
                instructions=instructions,
                tentacle_id=tentacle_id,
            )
        case NoAuth():
            instance = NoAuthMcp(
                id=mcp_id,
                user_id=user_id,
                name=request.name,
                namespace=namespace,
                url=str(request.url),
                instructions=instructions,
                tentacle_id=tentacle_id,
            )
    async with async_session() as session:
        session.add(instance)
        try:
            await session.commit()
        except IntegrityError as error:
            await session.rollback()
            duplicate = await session.one_or_none(
                Mcp,
                expressions=[
                    Mcp["user_id"] == user_id,
                    Mcp["namespace"] == namespace,
                ],
            )
            if duplicate is None:
                raise
            raise ValueError("MCP namespace is already installed") from error
    return instance

list async

list(
    user_id: UUID7,
    *,
    enabled: bool = False,
    mcp_type: type[McpT] = Mcp,
) -> list[McpT]

Source code in octomate/managers/mcp.py

async def list[McpT: Mcp](
    self,
    user_id: UUID7,
    *,
    enabled: bool = False,
    mcp_type: type[McpT] = Mcp,
) -> list[McpT]:
    expressions = [Mcp["user_id"] == user_id]
    if enabled:
        expressions.append(Mcp["enabled"].is_(True))
    async with async_session() as session:
        return list(
            await session.list(
                mcp_type,
                expressions=expressions,
                limit=None,
            )
        )

summary async

summary(user: User, mcp: Mcp) -> McpServerSummary

Source code in octomate/managers/mcp.py

async def summary(self, user: User, mcp: Mcp) -> McpServerSummary:
    if mcp.user_id != user.id:
        raise McpUnavailable
    summary = mcp.summary
    if isinstance(mcp, OAuthMcp):
        authorization: McpAuthorizationResult = McpAuthorizationStatus(status=None)
        flows: list[OAuthFlowKind] = []
        if self.oauth is not None:
            authorization = await self.authorization_status(user, mcp)
            if mcp.tentacle_id is None:
                if self.oauth.callback_base_uri is not None:
                    flows = ["authorization_code"]
            else:
                connector = self.oauth.connectors.get(mcp.tentacle_id)
                if connector is not None:
                    flows = [flow.kind for flow in connector.flows]
        summary.oauth = McpOAuthSummary(status=authorization.status, flows=flows)
    return summary

enable async

enable(user_id: UUID7, mcp_id: UUID7) -> Mcp

Source code in octomate/managers/mcp.py

async def enable(self, user_id: UUID7, mcp_id: UUID7) -> Mcp:
    async with self.lock(mcp_id):
        async with async_session() as session:
            instance = await session.one_or_none(
                Mcp,
                expressions=[
                    Mcp["id"] == mcp_id,
                    Mcp["user_id"] == user_id,
                ],
            )
            if instance is None:
                raise McpUnavailable
            instance.enabled = True
            instance.updated_at = datetime.now(UTC)
            await session.commit()
        return instance

disable async

disable(user_id: UUID7, mcp_id: UUID7) -> Mcp

Source code in octomate/managers/mcp.py

async def disable(self, user_id: UUID7, mcp_id: UUID7) -> Mcp:
    async with self.lock(mcp_id):
        async with async_session() as session:
            instance = await session.one_or_none(
                Mcp,
                expressions=[
                    Mcp["id"] == mcp_id,
                    Mcp["user_id"] == user_id,
                ],
            )
            if instance is None:
                raise McpUnavailable
            async with (
                self.oauth.lock(
                    OAuthLockKey(
                        user_id=user_id,
                        mcp_id=mcp_id,
                        connector_id=instance.tentacle_id or "mcp",
                    )
                )
                if isinstance(instance, OAuthMcp) and self.oauth is not None
                else nullcontext()
            ):
                instance.enabled = False
                instance.updated_at = datetime.now(UTC)
                await session.commit()
        await self.evict(McpClientKey(user_id=user_id, mcp_id=mcp_id))
        return instance

uninstall async

uninstall(user_id: UUID7, mcp_id: UUID7) -> None

Source code in octomate/managers/mcp.py

async def uninstall(self, user_id: UUID7, mcp_id: UUID7) -> None:
    async with self.lock(mcp_id):
        async with async_session() as session:
            instance = await session.one_or_none(
                Mcp,
                expressions=[
                    Mcp["id"] == mcp_id,
                    Mcp["user_id"] == user_id,
                ],
            )
            if instance is None:
                raise McpUnavailable
            async with (
                self.oauth.lock(
                    OAuthLockKey(
                        user_id=user_id,
                        mcp_id=mcp_id,
                        connector_id=instance.tentacle_id or "mcp",
                    )
                )
                if isinstance(instance, OAuthMcp) and self.oauth is not None
                else nullcontext()
            ):
                await session.delete(instance)
                await session.commit()
        await self.evict(McpClientKey(user_id=user_id, mcp_id=mcp_id))

authorizable async

authorizable(
    user_id: UUID7,
    *,
    mcp_id: UUID7 | None = None,
    namespace: str | None = None,
) -> OAuthMcp

Source code in octomate/managers/mcp.py

async def authorizable(
    self,
    user_id: UUID7,
    *,
    mcp_id: UUID7 | None = None,
    namespace: str | None = None,
) -> OAuthMcp:
    if mcp_id is None and namespace is None:
        raise ValueError("MCP ID or namespace is required")
    expressions = [
        OAuthMcp["user_id"] == user_id,
        OAuthMcp["enabled"].is_(True),
    ]
    if mcp_id is not None:
        expressions.append(OAuthMcp["id"] == mcp_id)
    if namespace is not None:
        expressions.append(OAuthMcp["namespace"] == namespace)
    async with async_session() as session:
        instance = await session.one_or_none(OAuthMcp, expressions=expressions)
    if instance is None:
        raise McpUnavailable
    return instance

connect async

connect(
    user: User,
    mcp_id: UUID7,
    *,
    profile: UserProfile | None = None,
    flow: OAuthFlowKind | None = None,
) -> OAuthStartResult

Source code in octomate/managers/mcp.py

async def connect(
    self,
    user: User,
    mcp_id: UUID7,
    *,
    profile: UserProfile | None = None,
    flow: OAuthFlowKind | None = None,
) -> OAuthStartResult:
    if self.oauth is None:
        raise McpUnavailable
    instance = await self.authorizable(user_id=user.id, mcp_id=mcp_id)
    async with self.oauth.lock(
        OAuthLockKey(
            user_id=user.id,
            mcp_id=mcp_id,
            connector_id=instance.tentacle_id or "mcp",
        )
    ):
        return await self.oauth.start(
            user,
            instance.tentacle_id or "mcp",
            mcp_id=mcp_id,
            profile=profile,
            flow=flow,
        )

pending_authorization async

pending_authorization(
    user: User, mcp_id: UUID7
) -> OAuthStartResult | None

Source code in octomate/managers/mcp.py

async def pending_authorization(
    self, user: User, mcp_id: UUID7
) -> OAuthStartResult | None:
    if self.oauth is None:
        raise McpUnavailable
    mcp = await self.authorizable(user_id=user.id, mcp_id=mcp_id)
    return await self.oauth.pending_authorization(
        user, mcp.tentacle_id or "mcp", mcp_id=mcp_id
    )

cancel_authorization async

cancel_authorization(user: User, mcp_id: UUID7) -> None

Source code in octomate/managers/mcp.py

async def cancel_authorization(self, user: User, mcp_id: UUID7) -> None:
    if self.oauth is None:
        raise McpUnavailable
    mcp = await self.authorizable(user_id=user.id, mcp_id=mcp_id)
    connector_id = mcp.tentacle_id or "mcp"
    async with (
        self.oauth.lock(
            OAuthLockKey(user_id=user.id, mcp_id=mcp_id, connector_id=connector_id)
        ),
        async_session() as session,
    ):
        operations = await session.list(
            OAuthOperation,
            limit=None,
            expressions=[
                OAuthOperation["user_id"] == user.id,
                OAuthOperation["mcp_id"] == mcp_id,
                OAuthOperation["connector_id"] == connector_id,
                OAuthOperation["consumed_at"].is_(None),
            ],
        )
        now = datetime.now(UTC)
        for operation in operations:
            operation.consumed_at = now
        await session.commit()

confirm async

confirm(
    user: User,
    mcp_id: UUID7,
    *,
    profile: UserProfile | None = None,
) -> McpAuthorizationResult

Source code in octomate/managers/mcp.py

async def confirm(
    self,
    user: User,
    mcp_id: UUID7,
    *,
    profile: UserProfile | None = None,
) -> McpAuthorizationResult:
    if self.oauth is None:
        raise McpUnavailable
    instance = await self.authorizable(user_id=user.id, mcp_id=mcp_id)
    connector_id = instance.tentacle_id or "mcp"
    connector = await self.oauth.resolve_connector(
        connector_id, user_id=user.id, mcp_id=mcp_id
    )
    if any(isinstance(flow, DeviceAuthorizationFlow) for flow in connector.flows):
        try:
            result = await self.oauth.complete_latest(
                user, connector_id, mcp_id=mcp_id, profile=profile
            )
        except NoPendingAuthorization:
            pass
        else:
            if isinstance(result, OAuthPending):
                return McpDeviceAuthorizationPending(
                    retry_after_seconds=result.retry_after_seconds
                )
    return await self.authorization_status(user, instance)

authorization_status async

authorization_status(
    user: User, mcp: OAuthMcp
) -> McpAuthorizationResult

Source code in octomate/managers/mcp.py

async def authorization_status(
    self, user: User, mcp: OAuthMcp
) -> McpAuthorizationResult:
    if self.oauth is None or mcp.user_id != user.id:
        raise McpUnavailable
    connector_id = mcp.tentacle_id or "mcp"
    status = await self.oauth.connection_status(user, connector_id, mcp_id=mcp.id)
    if status != "active":
        async with async_session() as session:
            pending = await session.first(
                OAuthOperation,
                expressions=[
                    OAuthOperation["user_id"] == user.id,
                    OAuthOperation["mcp_id"] == mcp.id,
                    OAuthOperation["connector_id"] == connector_id,
                    OAuthOperation["consumed_at"].is_(None),
                    OAuthOperation["expires_at"] > datetime.now(UTC),
                ],
                order_bys=[OAuthOperation["id"].desc()],
            )
        if pending is not None:
            if pending.interval_seconds is not None:
                return McpDeviceAuthorizationPending(
                    retry_after_seconds=pending.interval_seconds
                )
            return McpBrowserAuthorizationPending()
    return McpAuthorizationStatus(status=status)

resolve async

resolve(
    scope: OctomateSession | User,
    namespace: str | None = None,
    *,
    mcp_id: UUID7 | None = None,
) -> tuple[User, Mcp]

Resolve one enabled installation owned by the caller, by namespace or UUID.

Historical invocations supply only their saved UUID. A reused namespace therefore cannot substitute a different installation.

Source code in octomate/managers/mcp.py

async def resolve(
    self,
    scope: OctomateSession | User,
    namespace: str | None = None,
    *,
    mcp_id: UUID7 | None = None,
) -> tuple[User, Mcp]:
    """Resolve one enabled installation owned by the caller, by namespace or UUID.

    Historical invocations supply only their saved UUID. A reused namespace
    therefore cannot substitute a different installation.
    """
    if (namespace is None) == (mcp_id is None):
        raise ValueError("Supply exactly one MCP namespace or installation ID")
    if isinstance(scope, User):
        user = scope
    else:
        user = (
            scope.user_profile.user.value
            if scope.user_profile is not None
            and scope.user_profile.user_id is not None
            else None
        )
    if user is None:
        raise McpUnavailable
    async with async_session() as session:
        instance = await session.one_or_none(
            Mcp,
            expressions=[
                Mcp["user_id"] == user.id,
                Mcp["id"] == mcp_id
                if mcp_id is not None
                else Mcp["namespace"] == namespace,
                Mcp["enabled"].is_(True),
            ],
        )
    if instance is None:
        raise McpUnavailable
    return user, instance

catalog async

catalog(
    scope: OctomateSession | User,
    namespace: str | None = None,
    *,
    mcp_id: UUID7 | None = None,
    refresh: bool = False,
) -> McpToolCatalog

Read cached discovery, loading it when absent or explicitly refreshed.

Ownership, enabled state, and authorization are checked even on cache hits. The bounded catalog follows its pooled connection's lifetime.

Source code in octomate/managers/mcp.py

async def catalog(
    self,
    scope: OctomateSession | User,
    namespace: str | None = None,
    *,
    mcp_id: UUID7 | None = None,
    refresh: bool = False,
) -> McpToolCatalog:
    """Read cached discovery, loading it when absent or explicitly refreshed.

    Ownership, enabled state, and authorization are checked even on cache hits.
    The bounded catalog follows its pooled connection's lifetime.
    """
    user, instance = await self.resolve(scope, namespace, mcp_id=mcp_id)
    async with (
        self.acquire(user, mcp_id=instance.id) as client,
        self.lock(instance.id),
    ):
        key = McpClientKey(user.id, instance.id)
        catalog = self.tool_catalogs.get(key)
        if catalog is not None and not refresh:
            return catalog
        instructions = "\n".join(
            part for part in (instance.instructions, client.instructions) if part
        )
        catalog = McpToolCatalog(
            instructions=instructions, tools=await client.list_tools()
        )
        self.tool_catalogs[key] = catalog
        return catalog

call_tool async

call_tool(
    scope: OctomateSession,
    namespace: str,
    name: str,
    arguments: dict[str, JsonValue],
) -> CallToolResult

Execute one upstream call and retain its declared app before publishing.

This boundary sees each call independently, including calls a runtime nests inside a code wrapper. Sessions without thread attribution return the upstream result without producing an artifact or ledger reference. App declarations reuse the connection's latest discovery, loading once when absent. Explicit discovery refreshes them; client eviction clears them.

Source code in octomate/managers/mcp.py

async def call_tool(
    self,
    scope: OctomateSession,
    namespace: str,
    name: str,
    arguments: dict[str, JsonValue],
) -> CallToolResult:
    """Execute one upstream call and retain its declared app before publishing.

    This boundary sees each call independently, including calls a runtime
    nests inside a code wrapper. Sessions without thread attribution return
    the upstream result without producing an artifact or ledger reference.
    App declarations reuse the connection's latest discovery, loading once
    when absent. Explicit discovery refreshes them; client eviction clears them.
    """
    user, instance = await self.resolve(scope, namespace)
    address = scope.conversation_address
    channel = scope.channels.get(address.channel_tentacle_id) if address else None
    async with self.acquire(user, mcp_id=instance.id) as client:
        if (
            scope.native
            or scope.thread_id is None
            or scope.threads is None
            or channel is None
            or address is None
        ):
            return await client.call_tool_mcp(name, arguments)
        catalog = await self.catalog(user, mcp_id=instance.id)
        tool = catalog.tools_by_name.get(name)
        ui = McpAppToolMeta.model_validate(tool.meta or {}) if tool else None
        result = await client.call_tool_mcp(name, arguments)
    if tool is None or ui is None or ui.resource_uri is None:
        return result
    message = await scope.threads.record_artifact(
        scope.thread_id,
        McpAppArtifact(
            owner_id=user.id,
            title=tool.title or tool.name,
            invocation=McpAppInvocation(
                installation_id=instance.id,
                namespace=instance.namespace,
                provider_name=instance.name,
                tool=tool,
                resource_uri=ui.resource_uri,
                arguments=arguments,
                result=result,
            ),
        ),
        sender=channel.self_profile,
        agent_tentacle_id=scope.current_agent_id,
    )
    await scope.publish(
        address,
        ArtifactAvailableEvent(
            message_id=message.id, artifact_id=message.artifact_id
        ),
    )
    return result

app_resource async

app_resource(
    user: User, invocation: McpAppInvocation
) -> McpAppResource

Read only the saved UI URI from its still-authorized installation UUID.

This performs one resource read, never discovery or tool execution. The response must contain one matching MCP Apps HTML resource, as text or base64-encoded UTF-8. UI metadata is validated before reaching a browser.

Source code in octomate/managers/mcp.py

async def app_resource(
    self, user: User, invocation: McpAppInvocation
) -> McpAppResource:
    """Read only the saved UI URI from its still-authorized installation UUID.

    This performs one resource read, never discovery or tool execution. The
    response must contain one matching MCP Apps HTML resource, as text or
    base64-encoded UTF-8. UI metadata is validated before reaching a browser.
    """
    async with self.acquire(user, mcp_id=invocation.installation_id) as client:
        contents = await client.read_resource(invocation.resource_uri)
    if len(contents) != 1:
        raise ValueError("An MCP app must return exactly one UI resource")
    resource = contents[0]
    if (
        resource.uri != str(invocation.resource_uri)
        or resource.mime_type != UI_MIME_TYPE
    ):
        raise ValueError("MCP app resource URI or MIME type does not match")
    html = (
        resource.text
        if isinstance(resource, TextResourceContents)
        else b64decode(resource.blob, validate=True).decode("utf-8")
    )
    return McpAppResource(
        uri=invocation.resource_uri,
        html=html,
        ui=(resource.meta or {}).get("ui", {}),
    )

acquire async

acquire(
    scope: OctomateSession | User,
    namespace: str | None = None,
    *,
    mcp_id: UUID7 | None = None,
) -> AsyncGenerator[Client]

Source code in octomate/managers/mcp.py

@asynccontextmanager
async def acquire(
    self,
    scope: OctomateSession | User,
    namespace: str | None = None,
    *,
    mcp_id: UUID7 | None = None,
) -> AsyncGenerator[Client]:
    user, instance = await self.resolve(scope, namespace, mcp_id=mcp_id)
    user_id = user.id
    async with self.lock(instance.id):
        # The first lookup finds the lock; reload under it so disabling or
        # deleting an instance wins before another call is dispatched.
        async with async_session() as session:
            current = await session.one_or_none(
                Mcp,
                expressions=[Mcp["id"] == instance.id],
            )
        if current is None or not current.enabled or current.user_id != user_id:
            raise McpUnavailable
        if isinstance(current, OAuthMcp) and (
            self.oauth is None
            or await self.oauth.connection_status(
                user,
                current.tentacle_id or "mcp",
                mcp_id=current.id,
            )
            != "active"
        ):
            raise McpUnavailable
        key = McpClientKey(user_id=user_id, mcp_id=current.id)
        client = self.clients.get(key)
        if client is not None and not client.is_connected():
            await self.evict(key)
            client = None
        if client is None:
            headers: dict[str, str] = {}
            tentacle = self.tentacles.get(current.tentacle_id or "")
            if isinstance(current, BearerMcp):
                if self.cipher is None:
                    raise ValueError("MCP credential encryption is not configured")
                token = self.cipher.decrypt(
                    current.encrypted_token, context=f"mcp:{user_id}:{current.id}"
                )
                headers["Authorization"] = f"Bearer {token}"
            auth: McpBearerAuth | None = None
            if isinstance(current, OAuthMcp) and self.oauth is not None:
                if isinstance(tentacle, LarkTentacle):
                    auth = LarkMcpBearerAuth(
                        manager=self.oauth,
                        user=user,
                        mcp_id=current.id,
                        url=current.url,
                        tentacle=tentacle,
                    )
                else:
                    auth = McpBearerAuth(
                        manager=self.oauth,
                        user=user,
                        mcp_id=current.id,
                        connector_id=current.tentacle_id or "mcp",
                        url=current.url,
                    )
            client = Client(
                transport=StreamableHttpTransport(
                    url=current.url,
                    headers=headers,
                    auth=auth,
                    httpx_client_factory=self.httpx_client_factory,
                ),
                mode="auto",
                cache=False,
                extensions=[
                    advertise(
                        UI_EXTENSION_ID,
                        {"mimeTypes": [UI_MIME_TYPE]},
                    )
                ],
            )
            await client.__aenter__()
            self.clients[key] = client
        # Acquire the call's reference before releasing the lock, so removal
        # cannot close the client between checking access and starting a call.
        try:
            await client.__aenter__()
        except Exception:
            if not client.is_connected():
                await self.evict(key)
            raise
        if cleanup := self.sweaps.pop(key, None):
            cleanup.cancel()
            with suppress(asyncio.CancelledError):
                await cleanup
        self.sweaps[key] = asyncio.create_task(self.expire(key))
    try:
        yield client
    finally:
        try:
            await client.__aexit__(None, None, None)
        finally:
            if not client.is_connected():
                async with self.lock(key.mcp_id):
                    if self.clients.get(key) is client:
                        await self.evict(key)

expire async

expire(key: McpClientKey) -> None

Source code in octomate/managers/mcp.py

async def expire(self, key: McpClientKey) -> None:
    await asyncio.sleep(self.idle_timeout)
    async with self.lock(key.mcp_id):
        await self.evict(key)

evict async

evict(key: McpClientKey) -> None

Close a cached client while the caller holds its MCP lock.

Source code in octomate/managers/mcp.py

async def evict(self, key: McpClientKey) -> None:
    """Close a cached client while the caller holds its MCP lock."""
    self.tool_catalogs.pop(key, None)
    if cleanup := self.sweaps.pop(key, None):
        if cleanup is not asyncio.current_task():
            cleanup.cancel()
            with suppress(asyncio.CancelledError):
                await cleanup
    if client := self.clients.get(key):
        try:
            await client.close()
        except Exception:
            logger.exception("Failed to close MCP client %s", key.mcp_id)
        finally:
            del self.clients[key]

lifespan async

lifespan() -> AsyncGenerator[None]

Source code in octomate/managers/mcp.py

@asynccontextmanager
async def lifespan(self) -> AsyncGenerator[None]:
    try:
        yield
    finally:
        while self.clients:
            key = next(iter(self.clients))
            async with self.lock(key.mcp_id):
                await self.evict(key)

octomate.managers.oauth

OAuth connectors and the manager that runs their authorizations for a user, stores the resulting grants, and refreshes them.

NoPendingAuthorization

Bases: ValueError

Nothing of this connector's is waiting on this user to authorize it.

A ValueError still, so the channel handlers reporting a failed connection keep catching it unchanged; the narrower type is what lets a tool tell the model to start an authorization rather than reporting the turn as broken.

Source code in octomate/managers/oauth.py

class NoPendingAuthorization(ValueError):
    """Nothing of this connector's is waiting on this user to authorize it.

    A `ValueError` still, so the channel handlers reporting a failed connection
    keep catching it unchanged; the narrower type is what lets a tool tell the
    model to start an authorization rather than reporting the turn as broken.
    """

UnusableOAuthOperation

Bases: ValueError

The authorization a browser came back to is gone, spent, or not its own.

One type for every way the round trip can fail to find its operation, because the browser is told the same thing by all of them — this link is finished, ask again — and only the log distinguishes them. A ValueError for the same reason NoPendingAuthorization is one.

Source code in octomate/managers/oauth.py

class UnusableOAuthOperation(ValueError):
    """The authorization a browser came back to is gone, spent, or not its own.

    One type for every way the round trip can fail to find its operation, because
    the browser is told the same thing by all of them — this link is finished, ask
    again — and only the log distinguishes them. A `ValueError` for the same reason
    `NoPendingAuthorization` is one.
    """

OAuthConnector pydantic-model

Bases: BaseModel

One registered integration's injected OAuth composition.

It selects the upstream flow and, for authorization code, the deployment callback transport. It carries no user identity; OAuthManager supplies the current channel owner when an authorization starts.

Config:

• frozen: True
• arbitrary_types_allowed: True

Fields:

• id (str)
• flows (list[DeviceAuthorizationFlow | AuthorizationCodeFlow])
• callback_transport (OAuthCallbackTransport | None)
• mcp_url (HttpsUrl | None)

Validators:

• callback_matches_flow

Source code in octomate/managers/oauth.py

class OAuthConnector(BaseModel):
    """One registered integration's injected OAuth composition.

    It selects the upstream flow and, for authorization code, the deployment callback
    transport. It carries no user identity; ``OAuthManager`` supplies the current
    channel owner when an authorization starts.
    """

    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)

    id: str = Field(min_length=1)
    flows: list[DeviceAuthorizationFlow | AuthorizationCodeFlow] = Field(min_length=1)
    callback_transport: OAuthCallbackTransport | None = None
    mcp_url: HttpsUrl | None = None

    @model_validator(mode="after")
    def callback_matches_flow(self) -> Self:
        if len({flow.kind for flow in self.flows}) != len(self.flows):
            raise ValueError("OAuth flow types must be unique")
        if all(isinstance(flow, DeviceAuthorizationFlow) for flow in self.flows):
            if self.callback_transport is not None:
                raise ValueError("device OAuth does not use a callback transport")
            return self
        if self.callback_transport is None:
            raise ValueError("authorization-code OAuth requires a callback transport")
        return self

    def select_flow(
        self, kind: OAuthFlowKind | None = None
    ) -> DeviceAuthorizationFlow | AuthorizationCodeFlow:
        if kind is None:
            return self.flows[0]
        for flow in self.flows:
            if flow.kind == kind:
                return flow
        raise ValueError(f"OAuth connector {self.id!r} does not support {kind!r}")

    async def resolve_profile(self, grant: OAuthGrant) -> UserProfile | None:
        """An ownerless snapshot verified by this channel's OAuth provider, if supported."""
        return None

model_config class-attribute instance-attribute

model_config = ConfigDict(
    frozen=True, arbitrary_types_allowed=True
)

id pydantic-field

id: str

flows pydantic-field

flows: list[DeviceAuthorizationFlow | AuthorizationCodeFlow]

callback_transport pydantic-field

callback_transport: OAuthCallbackTransport | None = None

mcp_url pydantic-field

mcp_url: HttpsUrl | None = None

callback_matches_flow pydantic-validator

callback_matches_flow() -> Self

Source code in octomate/managers/oauth.py

@model_validator(mode="after")
def callback_matches_flow(self) -> Self:
    if len({flow.kind for flow in self.flows}) != len(self.flows):
        raise ValueError("OAuth flow types must be unique")
    if all(isinstance(flow, DeviceAuthorizationFlow) for flow in self.flows):
        if self.callback_transport is not None:
            raise ValueError("device OAuth does not use a callback transport")
        return self
    if self.callback_transport is None:
        raise ValueError("authorization-code OAuth requires a callback transport")
    return self

select_flow

select_flow(
    kind: OAuthFlowKind | None = None,
) -> DeviceAuthorizationFlow | AuthorizationCodeFlow

Source code in octomate/managers/oauth.py

def select_flow(
    self, kind: OAuthFlowKind | None = None
) -> DeviceAuthorizationFlow | AuthorizationCodeFlow:
    if kind is None:
        return self.flows[0]
    for flow in self.flows:
        if flow.kind == kind:
            return flow
    raise ValueError(f"OAuth connector {self.id!r} does not support {kind!r}")

resolve_profile async

resolve_profile(grant: OAuthGrant) -> UserProfile | None

An ownerless snapshot verified by this channel's OAuth provider, if supported.

Source code in octomate/managers/oauth.py

async def resolve_profile(self, grant: OAuthGrant) -> UserProfile | None:
    """An ownerless snapshot verified by this channel's OAuth provider, if supported."""
    return None

OAuthCallback

Bases: NamedTuple

What a completed authorization-code callback yields: the grant and the account bound to the operation.

Source code in octomate/managers/oauth.py

class OAuthCallback(NamedTuple):
    """What a completed authorization-code callback yields: the grant and the
    account bound to the operation."""

    grant: OAuthGrant
    user: User  # The authenticated account bound to the completed operation.

grant instance-attribute

grant: OAuthGrant

user instance-attribute

user: User

OAuthLockKey

Bases: NamedTuple

The (user, mcp, connector) triple an authorization is serialized under.

Source code in octomate/managers/oauth.py

class OAuthLockKey(NamedTuple):
    """The (user, mcp, connector) triple an authorization is serialized under."""

    user_id: UUID7
    mcp_id: UUID7 | None
    connector_id: str

user_id instance-attribute

user_id: UUID7

mcp_id instance-attribute

mcp_id: UUID7 | None

connector_id instance-attribute

connector_id: str

OAuthManager

Bases: Manager, Locks[OAuthLockKey]

Registered OAuth connectors bound to the current channel user.

Source code in octomate/managers/oauth.py

class OAuthManager(Manager, Locks[OAuthLockKey]):
    """Registered OAuth connectors bound to the current channel user."""

    def __init__(
        self,
        *,
        users: UserManager,
        authorization_lifetime: timedelta,
        token_refresh_leeway: timedelta,
        encryption_key: SecretStr | None = None,
        connectors: Iterable[OAuthConnector] = (),
        callback_base_uri: AnyHttpUrl | None = None,
        client_metadata_url: HttpsUrl | None = None,
        httpx_client_factory: McpHttpClientFactory = mcp_http_client,
    ) -> None:
        self.users = users
        self.cipher = (
            OAuthCipher(encryption_key) if encryption_key is not None else None
        )
        self.connectors: dict[str, OAuthConnector] = {}
        self.callback_base_uri = callback_base_uri
        self.authorization_lifetime: timedelta = authorization_lifetime
        self.client_metadata_url = client_metadata_url
        self.token_refresh_leeway: timedelta = token_refresh_leeway
        self.httpx_client_factory = httpx_client_factory
        for connector in connectors:
            self.register(connector)

    def register(self, connector: OAuthConnector) -> OAuthConnector:
        if connector.id in self.connectors:
            raise ValueError(f"OAuth connector {connector.id!r} is already registered")
        self.connectors[connector.id] = connector
        return connector

    def connector(self, connector_id: str) -> OAuthConnector:
        connector = self.connectors.get(connector_id)
        if connector is None:
            raise ValueError(f"unknown OAuth connector {connector_id!r}")
        return connector

    async def resolve_connector(
        self,
        connector_id: str,
        *,
        user_id: UUID7,
        mcp_id: UUID7 | None,
        discovery_state: OAuthDiscoveryState | None = None,
    ) -> OAuthConnector:
        if mcp_id is None:
            return self.connector(connector_id)
        async with async_session() as session:
            mcp = await session.get(OAuthMcp, mcp_id)
        if mcp is None or mcp.user_id != user_id or not mcp.enabled:
            raise UnusableOAuthOperation("MCP authorization is unavailable")
        if connector_id != (mcp.tentacle_id or "mcp"):
            raise UnusableOAuthOperation("MCP authorization names another connector")
        url = HTTPS_URL.validate_python(mcp.url)
        if mcp.tentacle_id is not None:
            connector = self.connector(mcp.tentacle_id)
            if connector.mcp_url is None or str(connector.mcp_url) != str(url):
                raise UnusableOAuthOperation(
                    "MCP URL does not match the configured tentacle"
                )
            selected = connector.select_flow()
            if not isinstance(selected, AuthorizationCodeFlow) or isinstance(
                selected.tokens, OAuthTokenExchange
            ):
                return connector
        if self.callback_base_uri is None:
            raise ValueError(
                "Configure oauth.callback_base_uri before connecting dynamic MCPs"
            )
        return OAuthConnector(
            id=connector_id,
            flows=[
                AuthorizationCodeFlow(
                    tokens=McpOAuthDiscovery(
                        url=url,
                        httpx_client_factory=self.httpx_client_factory,
                        client_metadata_url=self.client_metadata_url,
                        discovery_state=discovery_state,
                    ),
                    authorization_lifetime=self.authorization_lifetime,
                )
            ],
            callback_transport=DirectHttpOAuthCallbackTransport(self.callback_base_uri),
        )

    async def start(
        self,
        user: User,
        connector_id: str,
        *,
        mcp_id: UUID7 | None = None,
        profile: UserProfile | None = None,
        flow: OAuthFlowKind | None = None,
    ) -> OAuthStartResult:
        if profile is not None and profile.user_id != user.id:
            raise ValueError("OAuth profile does not belong to this user")

        connector = await self.resolve_connector(
            connector_id, user_id=user.id, mcp_id=mcp_id
        )
        selected = connector.select_flow(flow)
        context = OAuthFlowContext(
            operation_id=uuid7(),
            connector_id=connector.id,
            mcp_id=mcp_id,
            user=user,
            profile=profile,
        )
        if isinstance(selected, DeviceAuthorizationFlow):
            cipher = self.cipher
            if cipher is None:
                raise ValueError("OAuth persistence requires an encryption key")
            resumed = await self.pending_authorization(
                user,
                connector.id,
                profile=profile,
                mcp_id=mcp_id,
                flow="device",
            )
            if resumed is not None:
                return resumed
            authorization = await selected.start(context)
            payload = DeviceOperationPayload(
                device_code=authorization.device_code,
                user_code=authorization.user_code,
                verification_uri=authorization.verification_uri,
                verification_uri_complete=authorization.verification_uri_complete,
            )
            operation = OAuthOperation(
                id=context.operation_id,
                user_id=user.id,
                profile_id=profile.id if profile is not None else None,
                connector_id=connector.id,
                mcp_id=mcp_id,
                encrypted_data=cipher.encrypt(
                    payload.model_dump_json(),
                    context=f"operation:{context.operation_id}",
                ),
                expires_at=authorization.expires_at,
                interval_seconds=authorization.interval_seconds,
            )
            async with async_session() as session:
                session.add(operation)
                await session.commit()
            return DeviceAuthorization(
                operation_id=context.operation_id,
                verification_uri=authorization.verification_uri,
                verification_uri_complete=authorization.verification_uri_complete,
                user_code=authorization.user_code,
                expires_at=authorization.expires_at,
                interval_seconds=authorization.interval_seconds,
            )

        if not isinstance(selected, AuthorizationCodeFlow):
            raise TypeError(f"unsupported OAuth flow {type(selected).__name__}")
        transport = connector.callback_transport
        if transport is None:
            raise ValueError("authorization-code OAuth requires a callback transport")
        cipher = self.cipher
        if cipher is None:
            raise ValueError("OAuth persistence requires an encryption key")

        callback_uri = await transport.callback_uri(connector.id)
        # The state names its own operation because the callback cannot: a provider
        # redirects to one URI per connector, with no room to say which authorization
        # came back. The random half is what makes it unforgeable, and it is why the
        # operation id alone will not do — that id travels in the public start link,
        # where anyone the link reaches can read it.
        state = SecretStr(f"{context.operation_id}.{secrets.token_urlsafe(32)}")
        try:
            authorization = await selected.start(context, callback_uri, state)
        except ValidationError:
            raise ValueError("OAuth authorization returned invalid metadata") from None
        payload = AuthorizationCodeOperationPayload(
            state=state,
            code_verifier=authorization.code_verifier,
            callback_uri=callback_uri,
            authorization_uri=authorization.authorization_uri,
            discovery_state=authorization.discovery_state,
        )
        operation = OAuthOperation(
            id=context.operation_id,
            user_id=user.id,
            profile_id=profile.id if profile is not None else None,
            connector_id=connector.id,
            mcp_id=mcp_id,
            encrypted_data=cipher.encrypt(
                payload.model_dump_json(),
                context=f"operation:{context.operation_id}",
            ),
            expires_at=authorization.expires_at,
        )
        async with async_session() as session:
            session.add(operation)
            await session.commit()
        # Staged before the link exists, because the direct-HTTP start route serves
        # the link by reading exactly this row back.
        public_uri = await transport.prepare_authorization(
            context,
            authorization.authorization_uri,
        )
        return AuthorizationLink(
            operation_id=context.operation_id,
            authorization_uri=public_uri,
            expires_at=authorization.expires_at,
        )

    async def pending_authorization(
        self,
        user: User,
        connector_id: str,
        *,
        profile: UserProfile | None = None,
        mcp_id: UUID7 | None = None,
        flow: OAuthFlowKind | None = None,
    ) -> OAuthStartResult | None:
        """Restore a live authorization without starting another provider flow."""
        if profile is not None and profile.user_id != user.id:
            raise ValueError("OAuth profile does not belong to this user")
        cipher = self.cipher
        if cipher is None:
            raise ValueError("OAuth persistence requires an encryption key")
        expressions = [
            OAuthOperation["user_id"] == user.id,
            OAuthOperation["connector_id"] == connector_id,
            OAuthOperation["mcp_id"] == mcp_id,
            OAuthOperation["consumed_at"].is_(None),
            OAuthOperation["expires_at"] > datetime.now(UTC),
        ]
        if profile is not None or mcp_id is None:
            expressions.append(
                OAuthOperation["profile_id"] == (profile.id if profile else None)
            )
        if flow is not None:
            expressions.append(
                OAuthOperation["interval_seconds"].is_not(None)
                if flow == "device"
                else OAuthOperation["interval_seconds"].is_(None)
            )
        async with async_session() as session:
            operation = await session.first(
                OAuthOperation,
                order_bys=[OAuthOperation["id"].desc()],
                expressions=expressions,
            )
            if operation is not None and operation.profile_id is not None:
                profile = await session.get(UserProfile, operation.profile_id)
        if operation is None:
            return None
        if operation.interval_seconds is None:
            payload = self.authorization_payload(operation)
            connector = await self.resolve_connector(
                connector_id,
                user_id=user.id,
                mcp_id=mcp_id,
                discovery_state=payload.discovery_state,
            )
            transport = connector.callback_transport
            if transport is None:
                raise ValueError(
                    "authorization-code OAuth requires a callback transport"
                )
            return AuthorizationLink(
                operation_id=operation.id,
                authorization_uri=await transport.prepare_authorization(
                    OAuthFlowContext(
                        operation_id=operation.id,
                        connector_id=connector_id,
                        user=user,
                        profile=profile,
                        mcp_id=mcp_id,
                    ),
                    payload.authorization_uri,
                ),
                expires_at=operation.expires_at,
            )
        device = DeviceOperationPayload.model_validate_json(
            cipher.decrypt(
                operation.encrypted_data,
                context=f"operation:{operation.id}",
            )
        )
        return DeviceAuthorization(
            operation_id=operation.id,
            verification_uri=device.verification_uri,
            verification_uri_complete=device.verification_uri_complete,
            user_code=device.user_code,
            expires_at=operation.expires_at,
            interval_seconds=operation.interval_seconds,
        )

    async def complete_latest(
        self,
        user: User,
        connector_id: str,
        *,
        mcp_id: UUID7 | None = None,
        profile: UserProfile | None = None,
    ) -> OAuthGrant | OAuthPending:
        """Poll this user's newest device operation from the originating context."""
        if profile is not None and profile.user_id != user.id:
            raise ValueError("OAuth profile does not belong to this user")
        expressions = [
            OAuthOperation["user_id"] == user.id,
            OAuthOperation["connector_id"] == connector_id,
            OAuthOperation["mcp_id"] == mcp_id,
            OAuthOperation["consumed_at"].is_(None),
            OAuthOperation["interval_seconds"].is_not(None),
        ]
        if profile is not None or mcp_id is None:
            expressions.append(
                OAuthOperation["profile_id"] == (profile.id if profile else None)
            )
        async with async_session() as session:
            operations = await session.list(
                OAuthOperation,
                limit=1,
                order_bys=[OAuthOperation["id"].desc()],
                expressions=expressions,
            )
        if not operations:
            raise NoPendingAuthorization(f"no pending {connector_id} authorization")
        return await self.complete(user, operations[0].id, profile=profile)

    async def complete(
        self,
        user: User,
        operation_id: UUID7,
        *,
        profile: UserProfile | None = None,
    ) -> OAuthGrant | OAuthPending:
        """Complete a device operation owned by this user and originating profile."""
        cipher = self.cipher
        if cipher is None:
            raise ValueError("OAuth persistence requires an encryption key")
        if profile is not None and profile.user_id != user.id:
            raise ValueError("OAuth profile does not belong to this user")

        async with async_session() as session:
            operation = await session.get(OAuthOperation, operation_id)
            if operation is None:
                raise ValueError("unknown OAuth operation")
            key = OAuthLockKey(
                user_id=user.id,
                mcp_id=operation.mcp_id,
                connector_id=operation.connector_id,
            )

        async with self.lock(key), async_session() as session:
            operation = await session.get(OAuthOperation, operation_id)
            if (
                operation is None
                or operation.user_id != user.id
                or (profile is not None and operation.profile_id != profile.id)
            ):
                raise ValueError("unknown OAuth operation")
            if operation.consumed_at is not None:
                raise ValueError("OAuth operation has already been consumed")
            expires_at = operation.expires_at
            if expires_at <= datetime.now(UTC):
                operation.consumed_at = datetime.now(UTC)
                await session.commit()
                raise ValueError("OAuth operation has expired")

            connector = await self.resolve_connector(
                operation.connector_id, user_id=user.id, mcp_id=operation.mcp_id
            )
            flow = connector.select_flow("device")
            if (
                not isinstance(flow, DeviceAuthorizationFlow)
                or operation.interval_seconds is None
            ):
                raise ValueError("OAuth operation is not a device authorization")
            context = OAuthFlowContext(
                operation_id=operation.id,
                connector_id=connector.id,
                user=user,
                profile=profile,
                mcp_id=operation.mcp_id,
                interval_seconds=operation.interval_seconds,
            )
            payload = DeviceOperationPayload.model_validate_json(
                cipher.decrypt(
                    operation.encrypted_data,
                    context=f"operation:{operation.id}",
                )
            )
            result = await flow.complete(context, payload.device_code)
            if isinstance(result, OAuthPending):
                operation.interval_seconds = result.retry_after_seconds
                await session.commit()
                return result

            # Reauthorization replaces this grant, including when it is invalid.
            existing = await session.one_or_none(
                OAuthConnection,
                expressions=[
                    OAuthConnection["user_id"] == user.id,
                    OAuthConnection["connector_id"] == connector.id,
                    OAuthConnection["mcp_id"] == operation.mcp_id,
                ],
            )
            session.add(
                self.granted_connection(
                    result,
                    user_id=user.id,
                    connector_id=connector.id,
                    existing=existing,
                    flow=flow.kind,
                    mcp_id=operation.mcp_id,
                )
            )
            operation.consumed_at = datetime.now(UTC)
            await session.commit()
        return result

    def granted_connection(
        self,
        grant: OAuthGrant,
        *,
        user_id: UUID7,
        connector_id: str,
        existing: OAuthConnection | None,
        flow: OAuthFlowKind,
        mcp_id: UUID7 | None = None,
    ) -> OAuthConnection:
        """This user's connection to a connector, carrying a grant, encrypted.

        The one place credentials become storable, shared by every way of obtaining
        them — a device poll, a browser callback, a refresh — so replacement is one
        behaviour rather than three that drift. The whole token envelope is rewritten
        together: a refresh that kept half of the old one would leave a rotated
        refresh token pointing at a grant that no longer exists.

        Touches no session. Persisting is the caller's, because only the caller knows
        what else has to land in the same transaction — consuming the operation that
        produced the grant has to commit with it or not at all.
        """
        cipher = self.cipher
        if cipher is None:
            raise ValueError("OAuth persistence requires an encryption key")
        connection = existing or OAuthConnection(
            user_id=user_id,
            connector_id=connector_id,
            encrypted_tokens=b"",
            mcp_id=mcp_id,
            subject=grant.subject,
            account_label=grant.account_label,
        )
        connection.status = "active"
        connection.encrypted_tokens = cipher.encrypt(
            OAuthTokenPayload(
                flow=flow,
                access_token=grant.access_token,
                refresh_token=grant.refresh_token,
                token_type=grant.token_type,
                discovery_state=grant.discovery_state,
            ).model_dump_json(),
            context=f"connection:{connection.id}",
        )
        connection.subject = grant.subject
        connection.account_label = grant.account_label
        connection.scopes = grant.scopes
        connection.expires_at = grant.expires_at
        connection.updated_at = datetime.now(UTC)
        return connection

    async def staged_authorization(
        self,
        connector_id: str,
        operation_id: UUID7,
    ) -> AuthorizationCodeOperationPayload:
        """The provider request a UUID-only start link stands for.

        Reading it is what keeps the provider's URI — client id, scopes, state, PKCE
        challenge — out of the message the user was sent, which carries nothing but
        this operation's id.
        """
        async with async_session() as session:
            operation = self.usable_operation(
                await session.get(OAuthOperation, operation_id), connector_id
            )
            return self.authorization_payload(operation)

    async def complete_callback(
        self,
        connector_id: str,
        *,
        state: str,
        code: str,
        issuer: str | None = None,
    ) -> OAuthCallback:
        """Finish an authorization-code operation from the provider's callback.

        Nothing about this request proves who sent it — it is an ordinary browser GET
        arriving out of band — so everything binding it to a person was decided when
        the authorization started and is checked again here. The state matches an
        operation that has not been spent or expired; that operation names the profile
        that asked; and the profile still has to resolve to the same registered user,
        because a YAML declaration removed mid-flow must not finish as a connection.
        """
        async with async_session() as session:
            operation, _ = await self.operation_for_state(session, connector_id, state)
            key = OAuthLockKey(
                user_id=operation.user_id,
                mcp_id=operation.mcp_id,
                connector_id=operation.connector_id,
            )

        async with self.lock(key), async_session() as session:
            operation, payload = await self.operation_for_state(
                session, connector_id, state
            )
            profile = (
                await session.get(UserProfile, operation.profile_id)
                if operation.profile_id is not None
                else None
            )
            if operation.profile_id is None:
                user = await session.get(User, operation.user_id)
            else:
                user = await self.users.owner(profile) if profile is not None else None
            if user is None or user.id != operation.user_id:
                raise UnusableOAuthOperation(
                    "the profile that started this authorization is no longer linked "
                    "to the user it was started for"
                )
            connector = await self.resolve_connector(
                operation.connector_id,
                user_id=user.id,
                mcp_id=operation.mcp_id,
                discovery_state=payload.discovery_state,
            )
            if payload.discovery_state is not None:
                validate_authorization_response_iss(
                    issuer, payload.discovery_state.metadata
                )
            flow = connector.select_flow("authorization_code")
            if not isinstance(flow, AuthorizationCodeFlow):
                raise UnusableOAuthOperation(
                    "OAuth operation is not an authorization-code authorization"
                )
            # Consume before contacting the provider, including when the exchange fails.
            operation.consumed_at = datetime.now(UTC)
            await session.commit()
            grant = await flow.exchange(
                OAuthFlowContext(
                    operation_id=operation.id,
                    connector_id=connector.id,
                    user=user,
                    profile=profile,
                    mcp_id=operation.mcp_id,
                ),
                code=code,
                code_verifier=payload.code_verifier,
                callback_uri=payload.callback_uri,
            )
            # Reauthorization replaces this grant, including when it is invalid.
            existing = await session.one_or_none(
                OAuthConnection,
                expressions=[
                    OAuthConnection["user_id"] == user.id,
                    OAuthConnection["connector_id"] == connector.id,
                    OAuthConnection["mcp_id"] == operation.mcp_id,
                ],
            )
            session.add(
                self.granted_connection(
                    grant,
                    user_id=user.id,
                    connector_id=connector.id,
                    existing=existing,
                    flow=flow.kind,
                    mcp_id=operation.mcp_id,
                )
            )
            await session.commit()
        return OAuthCallback(grant=grant, user=user)

    async def link_profile(
        self, connector_id: str, grant: OAuthGrant, user: User
    ) -> bool:
        """Link the verified channel account to the user who started OAuth."""
        connector = self.connectors.get(connector_id)
        if connector is None or self.users.authorization_base_uri is None:
            return False
        observed = await connector.resolve_profile(grant)
        if observed is None:
            return False
        if observed.user_id is not None:
            raise ValueError(
                "OAuth profile resolution must not assign an Octomate owner"
            )
        profile = await self.users.ensure_profile(connector.id, observed)
        await self.users.link_verified_profile(profile, user)
        return True

    async def abandon_callback(
        self, connector_id: str, *, state: str, issuer: str | None = None
    ) -> None:
        """Close an operation the provider says the user turned down.

        A denial is an answer, so the operation is spent rather than left to expire —
        otherwise the link the user declined stays live for its full lifetime.
        """
        async with async_session() as session:
            operation, _ = await self.operation_for_state(session, connector_id, state)
            key = OAuthLockKey(
                user_id=operation.user_id,
                mcp_id=operation.mcp_id,
                connector_id=operation.connector_id,
            )

        async with self.lock(key), async_session() as session:
            operation, payload = await self.operation_for_state(
                session, connector_id, state
            )
            if payload.discovery_state is not None:
                validate_authorization_response_iss(
                    issuer, payload.discovery_state.metadata
                )
            operation.consumed_at = datetime.now(UTC)
            await session.commit()

    async def operation_for_state(
        self,
        session: AsyncSession,
        connector_id: str,
        state: str,
    ) -> tuple[OAuthOperation, AuthorizationCodeOperationPayload]:
        """The live operation a callback's state belongs to, and its sealed payload.

        The state's own prefix names the operation, since the row cannot be searched
        by a value stored encrypted; the comparison that follows is against the whole
        state, so naming an operation is not the same as being able to finish it.
        """
        operation_id, _, _ = state.partition(".")
        try:
            found = await session.get(OAuthOperation, uuid.UUID(operation_id))
        except ValueError:
            raise UnusableOAuthOperation("malformed OAuth state") from None
        operation = self.usable_operation(found, connector_id)
        payload = self.authorization_payload(operation)
        if not secrets.compare_digest(payload.state.get_secret_value(), state):
            raise UnusableOAuthOperation("OAuth state does not match its operation")
        return operation, payload

    def usable_operation(
        self,
        operation: OAuthOperation | None,
        connector_id: str,
    ) -> OAuthOperation:
        """Refuse an operation that is missing, another connector's, spent, or stale."""
        if operation is None or operation.connector_id != connector_id:
            raise UnusableOAuthOperation("unknown OAuth operation")
        if operation.consumed_at is not None:
            raise UnusableOAuthOperation("OAuth operation has already been consumed")
        expires_at = operation.expires_at
        if expires_at <= datetime.now(UTC):
            raise UnusableOAuthOperation("OAuth operation has expired")
        return operation

    def authorization_payload(
        self,
        operation: OAuthOperation,
    ) -> AuthorizationCodeOperationPayload:
        cipher = self.cipher
        if cipher is None:
            raise ValueError("OAuth persistence requires an encryption key")
        return AuthorizationCodeOperationPayload.model_validate_json(
            cipher.decrypt(
                operation.encrypted_data,
                context=f"operation:{operation.id}",
            )
        )

    async def invalidate(
        self,
        user: User,
        connector_id: str,
        *,
        mcp_id: UUID7 | None = None,
        expected_token: SecretStr | None = None,
    ) -> None:
        """Record that this user's credentials for a connector no longer work.

        Only the provider can say so — a revoked token answers 401 while nothing
        about the stored row changed — so something that spoke to it has to bring
        the news back. Marking it is what stops `access_token` handing the dead
        credential out again, which is what turns a capability that can only fail
        into an offer to authorize afresh.

        Nothing to mark is a normal outcome: a second 401 from the same dead
        session finds it already recorded.
        """
        async with (
            self.lock(
                OAuthLockKey(user_id=user.id, mcp_id=mcp_id, connector_id=connector_id)
            ),
            async_session() as session,
        ):
            connection = await session.one_or_none(
                OAuthConnection,
                expressions=[
                    OAuthConnection["user_id"] == user.id,
                    OAuthConnection["connector_id"] == connector_id,
                    OAuthConnection["mcp_id"] == mcp_id,
                    OAuthConnection["status"] == "active",
                ],
            )
            if connection is None:
                return
            if expected_token is not None:
                if self.cipher is None:
                    raise ValueError("OAuth persistence requires an encryption key")
                payload = OAuthTokenPayload.model_validate_json(
                    self.cipher.decrypt(
                        connection.encrypted_tokens,
                        context=f"connection:{connection.id}",
                    )
                )
                if payload.access_token != expected_token:
                    return
            connection.status = "invalid"
            connection.updated_at = datetime.now(UTC)
            await session.commit()

    async def connection_status(
        self,
        user: User,
        connector_id: str,
        *,
        mcp_id: UUID7 | None = None,
    ) -> OAuthConnectionStatus | None:
        """Whether this user has a connection to a connector, and whether it works.

        `None` is never having connected; `"invalid"` is having connected and lost
        it. `access_token` collapses both to no token, which is right for using one
        and wrong for explaining its absence — a user who was connected a moment ago
        cannot see that they no longer are, and only this tells them apart.
        """
        async with async_session() as session:
            connection = await session.one_or_none(
                OAuthConnection,
                expressions=[
                    OAuthConnection["user_id"] == user.id,
                    OAuthConnection["connector_id"] == connector_id,
                    OAuthConnection["mcp_id"] == mcp_id,
                ],
            )
        return connection.status if connection is not None else None

    async def access_token(
        self,
        user: User,
        connector_id: str,
        *,
        mcp_id: UUID7 | None = None,
    ) -> SecretStr | None:
        """Return this user's active connector token."""
        cipher = self.cipher
        if cipher is None:
            return None
        async with async_session() as session:
            connection = await session.one_or_none(
                OAuthConnection,
                expressions=[
                    OAuthConnection["user_id"] == user.id,
                    OAuthConnection["connector_id"] == connector_id,
                    OAuthConnection["mcp_id"] == mcp_id,
                    OAuthConnection["status"] == "active",
                ],
            )
        if connection is None:
            return None
        payload = OAuthTokenPayload.model_validate_json(
            cipher.decrypt(
                connection.encrypted_tokens,
                context=f"connection:{connection.id}",
            )
        )
        if connection.expires_at is None or not self.expiring(connection.expires_at):
            return payload.access_token
        if payload.refresh_token is None:
            # Expiry is the one death the row can announce by itself; record it
            # so the next run offers a fresh authorization rather than re-reading
            # the same dead token.
            await self.invalidate(
                user,
                connector_id,
                mcp_id=mcp_id,
                expected_token=payload.access_token,
            )
            return None
        return await self.refresh(user, connector_id, mcp_id=mcp_id)

    async def refresh(
        self,
        user: User,
        connector_id: str,
        *,
        mcp_id: UUID7 | None = None,
    ) -> SecretStr | None:
        """Spend this user's refresh token for a credential worth starting a run on.

        Serialized, and re-read once the lock is held, because a rotating refresh
        token is spendable exactly once: two runs racing on the same near-expired
        connection would send the loser to the provider with a token the winner had
        already rotated, and the connection would die of being used correctly.

        A provider that refuses the refresh has ended the connection, so it retires
        the same way a 401 does. A provider that merely failed to answer raises,
        because nothing about the credential was learned.
        """
        cipher = self.cipher
        if cipher is None:
            return None
        async with (
            self.lock(
                OAuthLockKey(user_id=user.id, mcp_id=mcp_id, connector_id=connector_id)
            ),
            async_session() as session,
        ):
            connection = await session.one_or_none(
                OAuthConnection,
                expressions=[
                    OAuthConnection["user_id"] == user.id,
                    OAuthConnection["connector_id"] == connector_id,
                    OAuthConnection["mcp_id"] == mcp_id,
                    OAuthConnection["status"] == "active",
                ],
            )
            if connection is None:
                return None
            payload = OAuthTokenPayload.model_validate_json(
                cipher.decrypt(
                    connection.encrypted_tokens,
                    context=f"connection:{connection.id}",
                )
            )
            # Re-read, so a run that lost the race finds the winner's credential
            # instead of spending a refresh token that has already been rotated.
            if connection.expires_at is None or not self.expiring(
                connection.expires_at
            ):
                return payload.access_token
            if payload.refresh_token is None:
                return None
            connector = await self.resolve_connector(
                connector_id,
                user_id=user.id,
                mcp_id=mcp_id,
                discovery_state=payload.discovery_state,
            )
            flow = connector.select_flow(payload.flow)
            try:
                grant = await flow.refresh(payload.refresh_token)
            except ValueError as error:
                if isinstance(error, ValidationError):
                    raise ValueError(
                        "OAuth refresh returned invalid credentials"
                    ) from None
                if not isinstance(error, OAuthRefreshRejected):
                    raise
                connection.status = "invalid"
                connection.updated_at = datetime.now(UTC)
                await session.commit()
                return None
            # Already loaded and already this user's, so it is handed straight
            # back rather than looked up again.
            self.granted_connection(
                grant,
                user_id=user.id,
                connector_id=connector_id,
                existing=connection,
                flow=flow.kind,
                mcp_id=mcp_id,
            )
            await session.commit()
        return grant.access_token

    def expiring(self, expires_at: datetime) -> bool:
        """Whether a credential has too little life left to start a run on."""
        return expires_at <= datetime.now(UTC) + self.token_refresh_leeway

users instance-attribute

users = users

cipher instance-attribute

cipher = (
    OAuthCipher(encryption_key)
    if encryption_key is not None
    else None
)

connectors instance-attribute

connectors: dict[str, OAuthConnector] = {}

callback_base_uri instance-attribute

callback_base_uri = callback_base_uri

authorization_lifetime instance-attribute

authorization_lifetime: timedelta = authorization_lifetime

client_metadata_url instance-attribute

client_metadata_url = client_metadata_url

token_refresh_leeway instance-attribute

token_refresh_leeway: timedelta = token_refresh_leeway

httpx_client_factory instance-attribute

httpx_client_factory = httpx_client_factory

register

register(connector: OAuthConnector) -> OAuthConnector

Source code in octomate/managers/oauth.py

def register(self, connector: OAuthConnector) -> OAuthConnector:
    if connector.id in self.connectors:
        raise ValueError(f"OAuth connector {connector.id!r} is already registered")
    self.connectors[connector.id] = connector
    return connector

connector

connector(connector_id: str) -> OAuthConnector

Source code in octomate/managers/oauth.py

def connector(self, connector_id: str) -> OAuthConnector:
    connector = self.connectors.get(connector_id)
    if connector is None:
        raise ValueError(f"unknown OAuth connector {connector_id!r}")
    return connector

resolve_connector async

resolve_connector(
    connector_id: str,
    *,
    user_id: UUID7,
    mcp_id: UUID7 | None,
    discovery_state: OAuthDiscoveryState | None = None,
) -> OAuthConnector

Source code in octomate/managers/oauth.py

async def resolve_connector(
    self,
    connector_id: str,
    *,
    user_id: UUID7,
    mcp_id: UUID7 | None,
    discovery_state: OAuthDiscoveryState | None = None,
) -> OAuthConnector:
    if mcp_id is None:
        return self.connector(connector_id)
    async with async_session() as session:
        mcp = await session.get(OAuthMcp, mcp_id)
    if mcp is None or mcp.user_id != user_id or not mcp.enabled:
        raise UnusableOAuthOperation("MCP authorization is unavailable")
    if connector_id != (mcp.tentacle_id or "mcp"):
        raise UnusableOAuthOperation("MCP authorization names another connector")
    url = HTTPS_URL.validate_python(mcp.url)
    if mcp.tentacle_id is not None:
        connector = self.connector(mcp.tentacle_id)
        if connector.mcp_url is None or str(connector.mcp_url) != str(url):
            raise UnusableOAuthOperation(
                "MCP URL does not match the configured tentacle"
            )
        selected = connector.select_flow()
        if not isinstance(selected, AuthorizationCodeFlow) or isinstance(
            selected.tokens, OAuthTokenExchange
        ):
            return connector
    if self.callback_base_uri is None:
        raise ValueError(
            "Configure oauth.callback_base_uri before connecting dynamic MCPs"
        )
    return OAuthConnector(
        id=connector_id,
        flows=[
            AuthorizationCodeFlow(
                tokens=McpOAuthDiscovery(
                    url=url,
                    httpx_client_factory=self.httpx_client_factory,
                    client_metadata_url=self.client_metadata_url,
                    discovery_state=discovery_state,
                ),
                authorization_lifetime=self.authorization_lifetime,
            )
        ],
        callback_transport=DirectHttpOAuthCallbackTransport(self.callback_base_uri),
    )

start async

start(
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
    profile: UserProfile | None = None,
    flow: OAuthFlowKind | None = None,
) -> OAuthStartResult

Source code in octomate/managers/oauth.py

async def start(
    self,
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
    profile: UserProfile | None = None,
    flow: OAuthFlowKind | None = None,
) -> OAuthStartResult:
    if profile is not None and profile.user_id != user.id:
        raise ValueError("OAuth profile does not belong to this user")

    connector = await self.resolve_connector(
        connector_id, user_id=user.id, mcp_id=mcp_id
    )
    selected = connector.select_flow(flow)
    context = OAuthFlowContext(
        operation_id=uuid7(),
        connector_id=connector.id,
        mcp_id=mcp_id,
        user=user,
        profile=profile,
    )
    if isinstance(selected, DeviceAuthorizationFlow):
        cipher = self.cipher
        if cipher is None:
            raise ValueError("OAuth persistence requires an encryption key")
        resumed = await self.pending_authorization(
            user,
            connector.id,
            profile=profile,
            mcp_id=mcp_id,
            flow="device",
        )
        if resumed is not None:
            return resumed
        authorization = await selected.start(context)
        payload = DeviceOperationPayload(
            device_code=authorization.device_code,
            user_code=authorization.user_code,
            verification_uri=authorization.verification_uri,
            verification_uri_complete=authorization.verification_uri_complete,
        )
        operation = OAuthOperation(
            id=context.operation_id,
            user_id=user.id,
            profile_id=profile.id if profile is not None else None,
            connector_id=connector.id,
            mcp_id=mcp_id,
            encrypted_data=cipher.encrypt(
                payload.model_dump_json(),
                context=f"operation:{context.operation_id}",
            ),
            expires_at=authorization.expires_at,
            interval_seconds=authorization.interval_seconds,
        )
        async with async_session() as session:
            session.add(operation)
            await session.commit()
        return DeviceAuthorization(
            operation_id=context.operation_id,
            verification_uri=authorization.verification_uri,
            verification_uri_complete=authorization.verification_uri_complete,
            user_code=authorization.user_code,
            expires_at=authorization.expires_at,
            interval_seconds=authorization.interval_seconds,
        )

    if not isinstance(selected, AuthorizationCodeFlow):
        raise TypeError(f"unsupported OAuth flow {type(selected).__name__}")
    transport = connector.callback_transport
    if transport is None:
        raise ValueError("authorization-code OAuth requires a callback transport")
    cipher = self.cipher
    if cipher is None:
        raise ValueError("OAuth persistence requires an encryption key")

    callback_uri = await transport.callback_uri(connector.id)
    # The state names its own operation because the callback cannot: a provider
    # redirects to one URI per connector, with no room to say which authorization
    # came back. The random half is what makes it unforgeable, and it is why the
    # operation id alone will not do — that id travels in the public start link,
    # where anyone the link reaches can read it.
    state = SecretStr(f"{context.operation_id}.{secrets.token_urlsafe(32)}")
    try:
        authorization = await selected.start(context, callback_uri, state)
    except ValidationError:
        raise ValueError("OAuth authorization returned invalid metadata") from None
    payload = AuthorizationCodeOperationPayload(
        state=state,
        code_verifier=authorization.code_verifier,
        callback_uri=callback_uri,
        authorization_uri=authorization.authorization_uri,
        discovery_state=authorization.discovery_state,
    )
    operation = OAuthOperation(
        id=context.operation_id,
        user_id=user.id,
        profile_id=profile.id if profile is not None else None,
        connector_id=connector.id,
        mcp_id=mcp_id,
        encrypted_data=cipher.encrypt(
            payload.model_dump_json(),
            context=f"operation:{context.operation_id}",
        ),
        expires_at=authorization.expires_at,
    )
    async with async_session() as session:
        session.add(operation)
        await session.commit()
    # Staged before the link exists, because the direct-HTTP start route serves
    # the link by reading exactly this row back.
    public_uri = await transport.prepare_authorization(
        context,
        authorization.authorization_uri,
    )
    return AuthorizationLink(
        operation_id=context.operation_id,
        authorization_uri=public_uri,
        expires_at=authorization.expires_at,
    )

pending_authorization async

pending_authorization(
    user: User,
    connector_id: str,
    *,
    profile: UserProfile | None = None,
    mcp_id: UUID7 | None = None,
    flow: OAuthFlowKind | None = None,
) -> OAuthStartResult | None

Restore a live authorization without starting another provider flow.

Source code in octomate/managers/oauth.py

async def pending_authorization(
    self,
    user: User,
    connector_id: str,
    *,
    profile: UserProfile | None = None,
    mcp_id: UUID7 | None = None,
    flow: OAuthFlowKind | None = None,
) -> OAuthStartResult | None:
    """Restore a live authorization without starting another provider flow."""
    if profile is not None and profile.user_id != user.id:
        raise ValueError("OAuth profile does not belong to this user")
    cipher = self.cipher
    if cipher is None:
        raise ValueError("OAuth persistence requires an encryption key")
    expressions = [
        OAuthOperation["user_id"] == user.id,
        OAuthOperation["connector_id"] == connector_id,
        OAuthOperation["mcp_id"] == mcp_id,
        OAuthOperation["consumed_at"].is_(None),
        OAuthOperation["expires_at"] > datetime.now(UTC),
    ]
    if profile is not None or mcp_id is None:
        expressions.append(
            OAuthOperation["profile_id"] == (profile.id if profile else None)
        )
    if flow is not None:
        expressions.append(
            OAuthOperation["interval_seconds"].is_not(None)
            if flow == "device"
            else OAuthOperation["interval_seconds"].is_(None)
        )
    async with async_session() as session:
        operation = await session.first(
            OAuthOperation,
            order_bys=[OAuthOperation["id"].desc()],
            expressions=expressions,
        )
        if operation is not None and operation.profile_id is not None:
            profile = await session.get(UserProfile, operation.profile_id)
    if operation is None:
        return None
    if operation.interval_seconds is None:
        payload = self.authorization_payload(operation)
        connector = await self.resolve_connector(
            connector_id,
            user_id=user.id,
            mcp_id=mcp_id,
            discovery_state=payload.discovery_state,
        )
        transport = connector.callback_transport
        if transport is None:
            raise ValueError(
                "authorization-code OAuth requires a callback transport"
            )
        return AuthorizationLink(
            operation_id=operation.id,
            authorization_uri=await transport.prepare_authorization(
                OAuthFlowContext(
                    operation_id=operation.id,
                    connector_id=connector_id,
                    user=user,
                    profile=profile,
                    mcp_id=mcp_id,
                ),
                payload.authorization_uri,
            ),
            expires_at=operation.expires_at,
        )
    device = DeviceOperationPayload.model_validate_json(
        cipher.decrypt(
            operation.encrypted_data,
            context=f"operation:{operation.id}",
        )
    )
    return DeviceAuthorization(
        operation_id=operation.id,
        verification_uri=device.verification_uri,
        verification_uri_complete=device.verification_uri_complete,
        user_code=device.user_code,
        expires_at=operation.expires_at,
        interval_seconds=operation.interval_seconds,
    )

complete_latest async

complete_latest(
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
    profile: UserProfile | None = None,
) -> OAuthGrant | OAuthPending

Poll this user's newest device operation from the originating context.

Source code in octomate/managers/oauth.py

async def complete_latest(
    self,
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
    profile: UserProfile | None = None,
) -> OAuthGrant | OAuthPending:
    """Poll this user's newest device operation from the originating context."""
    if profile is not None and profile.user_id != user.id:
        raise ValueError("OAuth profile does not belong to this user")
    expressions = [
        OAuthOperation["user_id"] == user.id,
        OAuthOperation["connector_id"] == connector_id,
        OAuthOperation["mcp_id"] == mcp_id,
        OAuthOperation["consumed_at"].is_(None),
        OAuthOperation["interval_seconds"].is_not(None),
    ]
    if profile is not None or mcp_id is None:
        expressions.append(
            OAuthOperation["profile_id"] == (profile.id if profile else None)
        )
    async with async_session() as session:
        operations = await session.list(
            OAuthOperation,
            limit=1,
            order_bys=[OAuthOperation["id"].desc()],
            expressions=expressions,
        )
    if not operations:
        raise NoPendingAuthorization(f"no pending {connector_id} authorization")
    return await self.complete(user, operations[0].id, profile=profile)

complete async

complete(
    user: User,
    operation_id: UUID7,
    *,
    profile: UserProfile | None = None,
) -> OAuthGrant | OAuthPending

Complete a device operation owned by this user and originating profile.

Source code in octomate/managers/oauth.py

async def complete(
    self,
    user: User,
    operation_id: UUID7,
    *,
    profile: UserProfile | None = None,
) -> OAuthGrant | OAuthPending:
    """Complete a device operation owned by this user and originating profile."""
    cipher = self.cipher
    if cipher is None:
        raise ValueError("OAuth persistence requires an encryption key")
    if profile is not None and profile.user_id != user.id:
        raise ValueError("OAuth profile does not belong to this user")

    async with async_session() as session:
        operation = await session.get(OAuthOperation, operation_id)
        if operation is None:
            raise ValueError("unknown OAuth operation")
        key = OAuthLockKey(
            user_id=user.id,
            mcp_id=operation.mcp_id,
            connector_id=operation.connector_id,
        )

    async with self.lock(key), async_session() as session:
        operation = await session.get(OAuthOperation, operation_id)
        if (
            operation is None
            or operation.user_id != user.id
            or (profile is not None and operation.profile_id != profile.id)
        ):
            raise ValueError("unknown OAuth operation")
        if operation.consumed_at is not None:
            raise ValueError("OAuth operation has already been consumed")
        expires_at = operation.expires_at
        if expires_at <= datetime.now(UTC):
            operation.consumed_at = datetime.now(UTC)
            await session.commit()
            raise ValueError("OAuth operation has expired")

        connector = await self.resolve_connector(
            operation.connector_id, user_id=user.id, mcp_id=operation.mcp_id
        )
        flow = connector.select_flow("device")
        if (
            not isinstance(flow, DeviceAuthorizationFlow)
            or operation.interval_seconds is None
        ):
            raise ValueError("OAuth operation is not a device authorization")
        context = OAuthFlowContext(
            operation_id=operation.id,
            connector_id=connector.id,
            user=user,
            profile=profile,
            mcp_id=operation.mcp_id,
            interval_seconds=operation.interval_seconds,
        )
        payload = DeviceOperationPayload.model_validate_json(
            cipher.decrypt(
                operation.encrypted_data,
                context=f"operation:{operation.id}",
            )
        )
        result = await flow.complete(context, payload.device_code)
        if isinstance(result, OAuthPending):
            operation.interval_seconds = result.retry_after_seconds
            await session.commit()
            return result

        # Reauthorization replaces this grant, including when it is invalid.
        existing = await session.one_or_none(
            OAuthConnection,
            expressions=[
                OAuthConnection["user_id"] == user.id,
                OAuthConnection["connector_id"] == connector.id,
                OAuthConnection["mcp_id"] == operation.mcp_id,
            ],
        )
        session.add(
            self.granted_connection(
                result,
                user_id=user.id,
                connector_id=connector.id,
                existing=existing,
                flow=flow.kind,
                mcp_id=operation.mcp_id,
            )
        )
        operation.consumed_at = datetime.now(UTC)
        await session.commit()
    return result

granted_connection

granted_connection(
    grant: OAuthGrant,
    *,
    user_id: UUID7,
    connector_id: str,
    existing: OAuthConnection | None,
    flow: OAuthFlowKind,
    mcp_id: UUID7 | None = None,
) -> OAuthConnection

This user's connection to a connector, carrying a grant, encrypted.

The one place credentials become storable, shared by every way of obtaining them — a device poll, a browser callback, a refresh — so replacement is one behaviour rather than three that drift. The whole token envelope is rewritten together: a refresh that kept half of the old one would leave a rotated refresh token pointing at a grant that no longer exists.

Touches no session. Persisting is the caller's, because only the caller knows what else has to land in the same transaction — consuming the operation that produced the grant has to commit with it or not at all.

Source code in octomate/managers/oauth.py

def granted_connection(
    self,
    grant: OAuthGrant,
    *,
    user_id: UUID7,
    connector_id: str,
    existing: OAuthConnection | None,
    flow: OAuthFlowKind,
    mcp_id: UUID7 | None = None,
) -> OAuthConnection:
    """This user's connection to a connector, carrying a grant, encrypted.

    The one place credentials become storable, shared by every way of obtaining
    them — a device poll, a browser callback, a refresh — so replacement is one
    behaviour rather than three that drift. The whole token envelope is rewritten
    together: a refresh that kept half of the old one would leave a rotated
    refresh token pointing at a grant that no longer exists.

    Touches no session. Persisting is the caller's, because only the caller knows
    what else has to land in the same transaction — consuming the operation that
    produced the grant has to commit with it or not at all.
    """
    cipher = self.cipher
    if cipher is None:
        raise ValueError("OAuth persistence requires an encryption key")
    connection = existing or OAuthConnection(
        user_id=user_id,
        connector_id=connector_id,
        encrypted_tokens=b"",
        mcp_id=mcp_id,
        subject=grant.subject,
        account_label=grant.account_label,
    )
    connection.status = "active"
    connection.encrypted_tokens = cipher.encrypt(
        OAuthTokenPayload(
            flow=flow,
            access_token=grant.access_token,
            refresh_token=grant.refresh_token,
            token_type=grant.token_type,
            discovery_state=grant.discovery_state,
        ).model_dump_json(),
        context=f"connection:{connection.id}",
    )
    connection.subject = grant.subject
    connection.account_label = grant.account_label
    connection.scopes = grant.scopes
    connection.expires_at = grant.expires_at
    connection.updated_at = datetime.now(UTC)
    return connection

staged_authorization async

staged_authorization(
    connector_id: str, operation_id: UUID7
) -> AuthorizationCodeOperationPayload

The provider request a UUID-only start link stands for.

Reading it is what keeps the provider's URI — client id, scopes, state, PKCE challenge — out of the message the user was sent, which carries nothing but this operation's id.

Source code in octomate/managers/oauth.py

async def staged_authorization(
    self,
    connector_id: str,
    operation_id: UUID7,
) -> AuthorizationCodeOperationPayload:
    """The provider request a UUID-only start link stands for.

    Reading it is what keeps the provider's URI — client id, scopes, state, PKCE
    challenge — out of the message the user was sent, which carries nothing but
    this operation's id.
    """
    async with async_session() as session:
        operation = self.usable_operation(
            await session.get(OAuthOperation, operation_id), connector_id
        )
        return self.authorization_payload(operation)

complete_callback async

complete_callback(
    connector_id: str,
    *,
    state: str,
    code: str,
    issuer: str | None = None,
) -> OAuthCallback

Finish an authorization-code operation from the provider's callback.

Nothing about this request proves who sent it — it is an ordinary browser GET arriving out of band — so everything binding it to a person was decided when the authorization started and is checked again here. The state matches an operation that has not been spent or expired; that operation names the profile that asked; and the profile still has to resolve to the same registered user, because a YAML declaration removed mid-flow must not finish as a connection.

Source code in octomate/managers/oauth.py

async def complete_callback(
    self,
    connector_id: str,
    *,
    state: str,
    code: str,
    issuer: str | None = None,
) -> OAuthCallback:
    """Finish an authorization-code operation from the provider's callback.

    Nothing about this request proves who sent it — it is an ordinary browser GET
    arriving out of band — so everything binding it to a person was decided when
    the authorization started and is checked again here. The state matches an
    operation that has not been spent or expired; that operation names the profile
    that asked; and the profile still has to resolve to the same registered user,
    because a YAML declaration removed mid-flow must not finish as a connection.
    """
    async with async_session() as session:
        operation, _ = await self.operation_for_state(session, connector_id, state)
        key = OAuthLockKey(
            user_id=operation.user_id,
            mcp_id=operation.mcp_id,
            connector_id=operation.connector_id,
        )

    async with self.lock(key), async_session() as session:
        operation, payload = await self.operation_for_state(
            session, connector_id, state
        )
        profile = (
            await session.get(UserProfile, operation.profile_id)
            if operation.profile_id is not None
            else None
        )
        if operation.profile_id is None:
            user = await session.get(User, operation.user_id)
        else:
            user = await self.users.owner(profile) if profile is not None else None
        if user is None or user.id != operation.user_id:
            raise UnusableOAuthOperation(
                "the profile that started this authorization is no longer linked "
                "to the user it was started for"
            )
        connector = await self.resolve_connector(
            operation.connector_id,
            user_id=user.id,
            mcp_id=operation.mcp_id,
            discovery_state=payload.discovery_state,
        )
        if payload.discovery_state is not None:
            validate_authorization_response_iss(
                issuer, payload.discovery_state.metadata
            )
        flow = connector.select_flow("authorization_code")
        if not isinstance(flow, AuthorizationCodeFlow):
            raise UnusableOAuthOperation(
                "OAuth operation is not an authorization-code authorization"
            )
        # Consume before contacting the provider, including when the exchange fails.
        operation.consumed_at = datetime.now(UTC)
        await session.commit()
        grant = await flow.exchange(
            OAuthFlowContext(
                operation_id=operation.id,
                connector_id=connector.id,
                user=user,
                profile=profile,
                mcp_id=operation.mcp_id,
            ),
            code=code,
            code_verifier=payload.code_verifier,
            callback_uri=payload.callback_uri,
        )
        # Reauthorization replaces this grant, including when it is invalid.
        existing = await session.one_or_none(
            OAuthConnection,
            expressions=[
                OAuthConnection["user_id"] == user.id,
                OAuthConnection["connector_id"] == connector.id,
                OAuthConnection["mcp_id"] == operation.mcp_id,
            ],
        )
        session.add(
            self.granted_connection(
                grant,
                user_id=user.id,
                connector_id=connector.id,
                existing=existing,
                flow=flow.kind,
                mcp_id=operation.mcp_id,
            )
        )
        await session.commit()
    return OAuthCallback(grant=grant, user=user)

link_profile async

link_profile(
    connector_id: str, grant: OAuthGrant, user: User
) -> bool

Link the verified channel account to the user who started OAuth.

Source code in octomate/managers/oauth.py

async def link_profile(
    self, connector_id: str, grant: OAuthGrant, user: User
) -> bool:
    """Link the verified channel account to the user who started OAuth."""
    connector = self.connectors.get(connector_id)
    if connector is None or self.users.authorization_base_uri is None:
        return False
    observed = await connector.resolve_profile(grant)
    if observed is None:
        return False
    if observed.user_id is not None:
        raise ValueError(
            "OAuth profile resolution must not assign an Octomate owner"
        )
    profile = await self.users.ensure_profile(connector.id, observed)
    await self.users.link_verified_profile(profile, user)
    return True

abandon_callback async

abandon_callback(
    connector_id: str,
    *,
    state: str,
    issuer: str | None = None,
) -> None

Close an operation the provider says the user turned down.

A denial is an answer, so the operation is spent rather than left to expire — otherwise the link the user declined stays live for its full lifetime.

Source code in octomate/managers/oauth.py

async def abandon_callback(
    self, connector_id: str, *, state: str, issuer: str | None = None
) -> None:
    """Close an operation the provider says the user turned down.

    A denial is an answer, so the operation is spent rather than left to expire —
    otherwise the link the user declined stays live for its full lifetime.
    """
    async with async_session() as session:
        operation, _ = await self.operation_for_state(session, connector_id, state)
        key = OAuthLockKey(
            user_id=operation.user_id,
            mcp_id=operation.mcp_id,
            connector_id=operation.connector_id,
        )

    async with self.lock(key), async_session() as session:
        operation, payload = await self.operation_for_state(
            session, connector_id, state
        )
        if payload.discovery_state is not None:
            validate_authorization_response_iss(
                issuer, payload.discovery_state.metadata
            )
        operation.consumed_at = datetime.now(UTC)
        await session.commit()

operation_for_state async

operation_for_state(
    session: AsyncSession, connector_id: str, state: str
) -> tuple[
    OAuthOperation, AuthorizationCodeOperationPayload
]

The live operation a callback's state belongs to, and its sealed payload.

The state's own prefix names the operation, since the row cannot be searched by a value stored encrypted; the comparison that follows is against the whole state, so naming an operation is not the same as being able to finish it.

Source code in octomate/managers/oauth.py

async def operation_for_state(
    self,
    session: AsyncSession,
    connector_id: str,
    state: str,
) -> tuple[OAuthOperation, AuthorizationCodeOperationPayload]:
    """The live operation a callback's state belongs to, and its sealed payload.

    The state's own prefix names the operation, since the row cannot be searched
    by a value stored encrypted; the comparison that follows is against the whole
    state, so naming an operation is not the same as being able to finish it.
    """
    operation_id, _, _ = state.partition(".")
    try:
        found = await session.get(OAuthOperation, uuid.UUID(operation_id))
    except ValueError:
        raise UnusableOAuthOperation("malformed OAuth state") from None
    operation = self.usable_operation(found, connector_id)
    payload = self.authorization_payload(operation)
    if not secrets.compare_digest(payload.state.get_secret_value(), state):
        raise UnusableOAuthOperation("OAuth state does not match its operation")
    return operation, payload

usable_operation

usable_operation(
    operation: OAuthOperation | None, connector_id: str
) -> OAuthOperation

Refuse an operation that is missing, another connector's, spent, or stale.

Source code in octomate/managers/oauth.py

def usable_operation(
    self,
    operation: OAuthOperation | None,
    connector_id: str,
) -> OAuthOperation:
    """Refuse an operation that is missing, another connector's, spent, or stale."""
    if operation is None or operation.connector_id != connector_id:
        raise UnusableOAuthOperation("unknown OAuth operation")
    if operation.consumed_at is not None:
        raise UnusableOAuthOperation("OAuth operation has already been consumed")
    expires_at = operation.expires_at
    if expires_at <= datetime.now(UTC):
        raise UnusableOAuthOperation("OAuth operation has expired")
    return operation

authorization_payload

authorization_payload(
    operation: OAuthOperation,
) -> AuthorizationCodeOperationPayload

Source code in octomate/managers/oauth.py

def authorization_payload(
    self,
    operation: OAuthOperation,
) -> AuthorizationCodeOperationPayload:
    cipher = self.cipher
    if cipher is None:
        raise ValueError("OAuth persistence requires an encryption key")
    return AuthorizationCodeOperationPayload.model_validate_json(
        cipher.decrypt(
            operation.encrypted_data,
            context=f"operation:{operation.id}",
        )
    )

invalidate async

invalidate(
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
    expected_token: SecretStr | None = None,
) -> None

Record that this user's credentials for a connector no longer work.

Only the provider can say so — a revoked token answers 401 while nothing about the stored row changed — so something that spoke to it has to bring the news back. Marking it is what stops access_token handing the dead credential out again, which is what turns a capability that can only fail into an offer to authorize afresh.

Nothing to mark is a normal outcome: a second 401 from the same dead session finds it already recorded.

Source code in octomate/managers/oauth.py

async def invalidate(
    self,
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
    expected_token: SecretStr | None = None,
) -> None:
    """Record that this user's credentials for a connector no longer work.

    Only the provider can say so — a revoked token answers 401 while nothing
    about the stored row changed — so something that spoke to it has to bring
    the news back. Marking it is what stops `access_token` handing the dead
    credential out again, which is what turns a capability that can only fail
    into an offer to authorize afresh.

    Nothing to mark is a normal outcome: a second 401 from the same dead
    session finds it already recorded.
    """
    async with (
        self.lock(
            OAuthLockKey(user_id=user.id, mcp_id=mcp_id, connector_id=connector_id)
        ),
        async_session() as session,
    ):
        connection = await session.one_or_none(
            OAuthConnection,
            expressions=[
                OAuthConnection["user_id"] == user.id,
                OAuthConnection["connector_id"] == connector_id,
                OAuthConnection["mcp_id"] == mcp_id,
                OAuthConnection["status"] == "active",
            ],
        )
        if connection is None:
            return
        if expected_token is not None:
            if self.cipher is None:
                raise ValueError("OAuth persistence requires an encryption key")
            payload = OAuthTokenPayload.model_validate_json(
                self.cipher.decrypt(
                    connection.encrypted_tokens,
                    context=f"connection:{connection.id}",
                )
            )
            if payload.access_token != expected_token:
                return
        connection.status = "invalid"
        connection.updated_at = datetime.now(UTC)
        await session.commit()

connection_status async

connection_status(
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
) -> OAuthConnectionStatus | None

Whether this user has a connection to a connector, and whether it works.

None is never having connected; "invalid" is having connected and lost it. access_token collapses both to no token, which is right for using one and wrong for explaining its absence — a user who was connected a moment ago cannot see that they no longer are, and only this tells them apart.

Source code in octomate/managers/oauth.py

async def connection_status(
    self,
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
) -> OAuthConnectionStatus | None:
    """Whether this user has a connection to a connector, and whether it works.

    `None` is never having connected; `"invalid"` is having connected and lost
    it. `access_token` collapses both to no token, which is right for using one
    and wrong for explaining its absence — a user who was connected a moment ago
    cannot see that they no longer are, and only this tells them apart.
    """
    async with async_session() as session:
        connection = await session.one_or_none(
            OAuthConnection,
            expressions=[
                OAuthConnection["user_id"] == user.id,
                OAuthConnection["connector_id"] == connector_id,
                OAuthConnection["mcp_id"] == mcp_id,
            ],
        )
    return connection.status if connection is not None else None

access_token async

access_token(
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
) -> SecretStr | None

Return this user's active connector token.

Source code in octomate/managers/oauth.py

async def access_token(
    self,
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
) -> SecretStr | None:
    """Return this user's active connector token."""
    cipher = self.cipher
    if cipher is None:
        return None
    async with async_session() as session:
        connection = await session.one_or_none(
            OAuthConnection,
            expressions=[
                OAuthConnection["user_id"] == user.id,
                OAuthConnection["connector_id"] == connector_id,
                OAuthConnection["mcp_id"] == mcp_id,
                OAuthConnection["status"] == "active",
            ],
        )
    if connection is None:
        return None
    payload = OAuthTokenPayload.model_validate_json(
        cipher.decrypt(
            connection.encrypted_tokens,
            context=f"connection:{connection.id}",
        )
    )
    if connection.expires_at is None or not self.expiring(connection.expires_at):
        return payload.access_token
    if payload.refresh_token is None:
        # Expiry is the one death the row can announce by itself; record it
        # so the next run offers a fresh authorization rather than re-reading
        # the same dead token.
        await self.invalidate(
            user,
            connector_id,
            mcp_id=mcp_id,
            expected_token=payload.access_token,
        )
        return None
    return await self.refresh(user, connector_id, mcp_id=mcp_id)

refresh async

refresh(
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
) -> SecretStr | None

Spend this user's refresh token for a credential worth starting a run on.

Serialized, and re-read once the lock is held, because a rotating refresh token is spendable exactly once: two runs racing on the same near-expired connection would send the loser to the provider with a token the winner had already rotated, and the connection would die of being used correctly.

A provider that refuses the refresh has ended the connection, so it retires the same way a 401 does. A provider that merely failed to answer raises, because nothing about the credential was learned.

Source code in octomate/managers/oauth.py

async def refresh(
    self,
    user: User,
    connector_id: str,
    *,
    mcp_id: UUID7 | None = None,
) -> SecretStr | None:
    """Spend this user's refresh token for a credential worth starting a run on.

    Serialized, and re-read once the lock is held, because a rotating refresh
    token is spendable exactly once: two runs racing on the same near-expired
    connection would send the loser to the provider with a token the winner had
    already rotated, and the connection would die of being used correctly.

    A provider that refuses the refresh has ended the connection, so it retires
    the same way a 401 does. A provider that merely failed to answer raises,
    because nothing about the credential was learned.
    """
    cipher = self.cipher
    if cipher is None:
        return None
    async with (
        self.lock(
            OAuthLockKey(user_id=user.id, mcp_id=mcp_id, connector_id=connector_id)
        ),
        async_session() as session,
    ):
        connection = await session.one_or_none(
            OAuthConnection,
            expressions=[
                OAuthConnection["user_id"] == user.id,
                OAuthConnection["connector_id"] == connector_id,
                OAuthConnection["mcp_id"] == mcp_id,
                OAuthConnection["status"] == "active",
            ],
        )
        if connection is None:
            return None
        payload = OAuthTokenPayload.model_validate_json(
            cipher.decrypt(
                connection.encrypted_tokens,
                context=f"connection:{connection.id}",
            )
        )
        # Re-read, so a run that lost the race finds the winner's credential
        # instead of spending a refresh token that has already been rotated.
        if connection.expires_at is None or not self.expiring(
            connection.expires_at
        ):
            return payload.access_token
        if payload.refresh_token is None:
            return None
        connector = await self.resolve_connector(
            connector_id,
            user_id=user.id,
            mcp_id=mcp_id,
            discovery_state=payload.discovery_state,
        )
        flow = connector.select_flow(payload.flow)
        try:
            grant = await flow.refresh(payload.refresh_token)
        except ValueError as error:
            if isinstance(error, ValidationError):
                raise ValueError(
                    "OAuth refresh returned invalid credentials"
                ) from None
            if not isinstance(error, OAuthRefreshRejected):
                raise
            connection.status = "invalid"
            connection.updated_at = datetime.now(UTC)
            await session.commit()
            return None
        # Already loaded and already this user's, so it is handed straight
        # back rather than looked up again.
        self.granted_connection(
            grant,
            user_id=user.id,
            connector_id=connector_id,
            existing=connection,
            flow=flow.kind,
            mcp_id=mcp_id,
        )
        await session.commit()
    return grant.access_token

expiring

expiring(expires_at: datetime) -> bool

Whether a credential has too little life left to start a run on.

Source code in octomate/managers/oauth.py

def expiring(self, expires_at: datetime) -> bool:
    """Whether a credential has too little life left to start a run on."""
    return expires_at <= datetime.now(UTC) + self.token_refresh_leeway

octomate.managers.spills

The database-backed spill store for oversized tool output.

COMPRESSION_LEVEL module-attribute

COMPRESSION_LEVEL = 6

SpillStore dataclass

Bases: Manager

The tool-output-limits OverflowStore, backed by the Octomate database.

The harness ships only a local-file store, whose handles resolve against one host's filesystem. A spill has to outlive the request that produced it — the handle is what the model reads back through read_tool_result, possibly turns later — and Octomate answers a conversation from whichever process picks the message up. Sharing the database is what makes those the same payload.

Payloads are zlib-compressed. A spill is by definition the oversized end of what tools return, and that content is the compressible kind, so this is most of what keeps a blob column reasonable on SQLite.

Source code in octomate/managers/spills.py

@dataclass
class SpillStore(Manager):
    """The tool-output-limits `OverflowStore`, backed by the Octomate database.

    The harness ships only a local-file store, whose handles resolve against one
    host's filesystem. A spill has to outlive the request that produced it — the
    handle is what the model reads back through `read_tool_result`, possibly turns
    later — and Octomate answers a conversation from whichever process picks the
    message up. Sharing the database is what makes those the same payload.

    Payloads are zlib-compressed. A spill is by definition the oversized end of what
    tools return, and that content is the compressible kind, so this is most of what
    keeps a blob column reasonable on SQLite.
    """

    # None keeps payloads forever, matching the harness's own store. A window is the
    # honest default here instead: a spill nobody read back within it is dead weight,
    # since the model has moved on and the tool can simply be called again.
    retention: timedelta | None = None

    async def write(self, key: str, data: bytes) -> str:
        """Persist `data` under `key`, returning it unchanged as the handle.

        The key already encodes the run, the tool call, and the retry, so it is
        self-contained — nothing else is needed to find the row again.
        """
        # Off the event loop: the work is proportional to a payload this store puts
        # no ceiling on, and a multi-megabyte return costs tens of milliseconds.
        packed = await to_thread.run_sync(zlib.compress, data, COMPRESSION_LEVEL)
        async with async_session() as session:
            session.add(ToolOutputSpill(handle=key, payload=packed))
            if self.retention is not None:
                # One statement rather than loading the expired rows to delete them
                # one by one: a transmuter always materializes whole, so listing them
                # would pull every stale payload through memory on the way to a DELETE.
                await session.execute(
                    delete(ToolOutputSpill).where(
                        ToolOutputSpill["created_at"]
                        < datetime.now(UTC) - self.retention
                    )
                )
            await session.commit()
        return key

    async def read(self, handle: str) -> bytes:
        async with async_session() as session:
            spill = await session.one_or_none(
                ToolOutputSpill,
                expressions=[ToolOutputSpill["handle"] == handle],
            )
        if spill is None:
            # The protocol's own signal for an unknown handle: `read_tool_result`
            # catches OSError and answers the model with a correction rather than
            # spending a retry on it.
            raise FileNotFoundError(handle)
        return await to_thread.run_sync(zlib.decompress, spill.payload)

retention class-attribute instance-attribute

retention: timedelta | None = None

write async

write(key: str, data: bytes) -> str

Persist data under key, returning it unchanged as the handle.

The key already encodes the run, the tool call, and the retry, so it is self-contained — nothing else is needed to find the row again.

Source code in octomate/managers/spills.py

async def write(self, key: str, data: bytes) -> str:
    """Persist `data` under `key`, returning it unchanged as the handle.

    The key already encodes the run, the tool call, and the retry, so it is
    self-contained — nothing else is needed to find the row again.
    """
    # Off the event loop: the work is proportional to a payload this store puts
    # no ceiling on, and a multi-megabyte return costs tens of milliseconds.
    packed = await to_thread.run_sync(zlib.compress, data, COMPRESSION_LEVEL)
    async with async_session() as session:
        session.add(ToolOutputSpill(handle=key, payload=packed))
        if self.retention is not None:
            # One statement rather than loading the expired rows to delete them
            # one by one: a transmuter always materializes whole, so listing them
            # would pull every stale payload through memory on the way to a DELETE.
            await session.execute(
                delete(ToolOutputSpill).where(
                    ToolOutputSpill["created_at"]
                    < datetime.now(UTC) - self.retention
                )
            )
        await session.commit()
    return key

read async

read(handle: str) -> bytes

Source code in octomate/managers/spills.py

async def read(self, handle: str) -> bytes:
    async with async_session() as session:
        spill = await session.one_or_none(
            ToolOutputSpill,
            expressions=[ToolOutputSpill["handle"] == handle],
        )
    if spill is None:
        # The protocol's own signal for an unknown handle: `read_tool_result`
        # catches OSError and answers the model with a correction rather than
        # spending a retry on it.
        raise FileNotFoundError(handle)
    return await to_thread.run_sync(zlib.decompress, spill.payload)