Complete OAuth callbacks host-side by flow state

Co-Authored-By: GPT-5 Codex <codex@openai.com>

Author: dcramer

packages/ash-sandbox-cli/src/ash_sandbox_cli/commands/capability.py

@@ -258,6 +258,52 @@ def auth_complete(
     )
 
 
+@auth_app.command("complete-callback")
+def auth_complete_callback(
+    callback_url: Annotated[
+        str | None,
+        typer.Option("--callback-url", help="OAuth callback URL"),
+    ] = None,
+    code: Annotated[
+        str | None,
+        typer.Option("--code", help="Authorization code"),
+    ] = None,
+    capability: Annotated[
+        str | None,
+        typer.Option("--capability", "-c", help="Optional namespaced capability id"),
+    ] = None,
+    account_hint: Annotated[
+        str | None,
+        typer.Option("--account", help="Optional account reference hint"),
+    ] = None,
+) -> None:
+    """Complete capability auth by callback/code with host-side flow resolution."""
+    if not callback_url and not code:
+        typer.echo("Error: Must specify either --callback-url or --code", err=True)
+        raise typer.Exit(1)
+
+    params: dict[str, Any] = {}
+    if callback_url:
+        params["callback_url"] = callback_url
+    if code:
+        params["code"] = code
+    if capability:
+        params["capability"] = capability
+    if account_hint:
+        params["account_hint"] = account_hint
+
+    result = _call("capability.auth.complete_callback", params)
+    if not result.get("ok"):
+        typer.echo("Error: capability auth completion failed", err=True)
+        raise typer.Exit(1)
+    typer.echo(
+        "Capability auth completed "
+        f"(flow_id={result.get('flow_id', '')}, "
+        f"capability={result.get('capability', '')}, "
+        f"account_ref={result.get('account_ref', '')})"
+    )
+
+
 @auth_app.command("poll")
 def auth_poll(
     flow_id: Annotated[

src/ash/capabilities/manager.py

@@ -454,6 +454,125 @@ async def auth_complete(
 
         return {"ok": True, "account_ref": account_ref}
 
+    async def auth_complete_callback(
+        self,
+        *,
+        user_id: str,
+        callback_url: str | None = None,
+        code: str | None = None,
+        capability_id: str | None = None,
+        account_hint: str | None = None,
+        chat_id: str | None = None,
+        chat_type: str | None = None,
+        provider: str | None = None,
+        thread_id: str | None = None,
+        session_key: str | None = None,
+        source_username: str | None = None,
+        source_display_name: str | None = None,
+    ) -> dict[str, Any]:
+        """Complete pending auth by callback URL/code with host-side flow resolution.
+
+        Prefers deterministic state matching when callback state is present.
+        """
+        normalized_user_id = _required_text(
+            value=user_id,
+            code="capability_invalid_input",
+            message="user_id is required",
+        )
+        normalized_capability_id = (
+            _required_capability_id(capability_id)
+            if capability_id is not None
+            else None
+        )
+        normalized_account_hint = _optional_text(account_hint)
+        normalized_chat_id = _optional_text(chat_id)
+        normalized_chat_type = _optional_text(chat_type)
+        normalized_provider = _optional_text(provider)
+        normalized_thread_id = _optional_text(thread_id)
+        normalized_session_key = _optional_text(session_key)
+        normalized_source_username = _optional_text(source_username)
+        normalized_source_display_name = _optional_text(source_display_name)
+
+        try:
+            normalized_completion = normalize_auth_completion(
+                callback_url=_optional_text(callback_url),
+                code=_optional_text(code),
+                expected_state=None,
+            )
+        except AuthNormalizationError as e:
+            raise CapabilityError(e.code, str(e)) from e
+
+        callback_state = _optional_text(normalized_completion.state)
+
+        async with self._lock:
+            self._prune_expired_flows_locked()
+            eligible = [
+                flow
+                for flow in self._auth_flows.values()
+                if flow.user_id == normalized_user_id
+                and (
+                    normalized_capability_id is None
+                    or flow.capability_id == normalized_capability_id
+                )
+                and (
+                    normalized_account_hint is None
+                    or flow.account_hint == normalized_account_hint
+                )
+            ]
+
+            if not eligible:
+                raise CapabilityError(
+                    "capability_auth_flow_invalid",
+                    "no pending auth flows match caller scope",
+                )
+
+            selected: CapabilityAuthFlow | None = None
+            if callback_state is not None:
+                state_matches = [
+                    flow
+                    for flow in eligible
+                    if _optional_text(flow.expected_callback_state) == callback_state
+                ]
+                if not state_matches:
+                    raise CapabilityError(
+                        "capability_auth_state_mismatch",
+                        "callback_url state does not match auth flow",
+                    )
+                if len(state_matches) > 1:
+                    raise CapabilityError(
+                        "capability_auth_flow_ambiguous",
+                        "multiple pending auth flows matched callback state",
+                    )
+                selected = state_matches[0]
+            elif len(eligible) == 1:
+                selected = eligible[0]
+            else:
+                raise CapabilityError(
+                    "capability_auth_flow_ambiguous",
+                    "multiple pending auth flows; callback state is required",
+                )
+
+        result = await self.auth_complete(
+            flow_id=selected.flow_id,
+            user_id=normalized_user_id,
+            chat_id=normalized_chat_id,
+            chat_type=normalized_chat_type,
+            provider=normalized_provider,
+            thread_id=normalized_thread_id,
+            session_key=normalized_session_key,
+            source_username=normalized_source_username,
+            source_display_name=normalized_source_display_name,
+            callback_url=normalized_completion.raw_callback_url,
+            code=normalized_completion.authorization_code,
+        )
+        return {
+            "ok": bool(result.get("ok")),
+            "account_ref": result.get("account_ref"),
+            "flow_id": selected.flow_id,
+            "capability": selected.capability_id,
+            "account_hint": selected.account_hint,
+        }
+
     async def auth_poll(
         self,
         *,

src/ash/providers/telegram/handlers/message_handler.py

@@ -5,6 +5,7 @@
 import asyncio
 import inspect
 import logging
+import re
 from datetime import UTC, datetime, timedelta
 from typing import TYPE_CHECKING, Any
 
@@ -43,6 +44,7 @@
     from ash.tools.registry import ToolRegistry
 
 logger = logging.getLogger("telegram")
+_LOCALHOST_CALLBACK_URL_PATTERN = re.compile(r"https?://localhost[^\s]*[?&]code=[^\s]*")
 
 
 def _extract_tool_calls_from_session(session: SessionState) -> list[dict[str, Any]]:
@@ -216,6 +218,126 @@ def _create_tool_tracker(self, message: IncomingMessage) -> ToolTracker:
             skill_registry=self._skill_registry,
         )
 
+    def _get_capability_manager(self) -> Any | None:
+        """Best-effort lookup of the capability manager via use_skill tool wiring."""
+        if self._tool_registry is None:
+            return None
+        if not hasattr(self._tool_registry, "has") or not self._tool_registry.has(
+            "use_skill"
+        ):
+            return None
+        try:
+            skill_tool = self._tool_registry.get("use_skill")
+        except Exception:
+            return None
+        return getattr(skill_tool, "_capability_manager", None)
+
+    @staticmethod
+    def _extract_localhost_callback_url(message_text: str) -> str | None:
+        """Extract a localhost OAuth callback URL from message text."""
+        match = _LOCALHOST_CALLBACK_URL_PATTERN.search(message_text)
+        if not match:
+            return None
+        callback_url = match.group(0).strip().rstrip(".,;")
+        return callback_url or None
+
+    async def _try_handle_capability_oauth_callback(
+        self, message: IncomingMessage
+    ) -> bool:
+        """Complete capability auth from pasted callback URL without LLM orchestration."""
+        callback_url = self._extract_localhost_callback_url(message.text or "")
+        if not callback_url:
+            return False
+
+        manager = self._get_capability_manager()
+        if manager is None:
+            return False
+
+        try:
+            completion = await manager.auth_complete_callback(
+                user_id=message.user_id,
+                callback_url=callback_url,
+                chat_id=message.chat_id,
+                chat_type=message.metadata.get("chat_type"),
+                provider=self._provider.name,
+                thread_id=message.metadata.get("thread_id"),
+                source_username=message.username,
+                source_display_name=message.display_name,
+            )
+            pending_flows = await manager.list_auth_flows(user_id=message.user_id)
+        except Exception as e:
+            code = getattr(e, "code", "")
+            if isinstance(code, str) and code.startswith("capability_auth_"):
+                reply = (
+                    "I could not apply that OAuth callback yet "
+                    f"({code}). Please continue with the latest auth URL."
+                )
+                await self._provider.send(
+                    OutgoingMessage(
+                        chat_id=message.chat_id,
+                        text=reply,
+                        reply_to_message_id=message.id,
+                    )
+                )
+                return True
+            return False
+
+        capability = str(completion.get("capability", "")).strip()
+        account_hint = str(completion.get("account_hint", "")).strip() or "default"
+        capability_label = {
+            "gog.email": "Gmail",
+            "gog.calendar": "Google Calendar",
+        }.get(capability, capability or "Google capability")
+
+        pending_caps = sorted(
+            {
+                str(flow.get("capability", "")).strip()
+                for flow in pending_flows
+                if isinstance(flow, dict) and flow.get("capability")
+            }
+        )
+        if pending_caps:
+            pending_names = ", ".join(
+                sorted(
+                    {
+                        {
+                            "gog.email": "Gmail",
+                            "gog.calendar": "Google Calendar",
+                        }.get(cap, cap)
+                        for cap in pending_caps
+                    }
+                )
+            )
+            reply = (
+                f"{capability_label} connected for account '{account_hint}'. "
+                f"Still pending: {pending_names}. Paste the next callback URL when ready."
+            )
+        else:
+            reply = (
+                f"{capability_label} connected for account '{account_hint}'. "
+                "Google setup is complete."
+            )
+
+        response_external_id = await self._provider.send(
+            OutgoingMessage(
+                chat_id=message.chat_id,
+                text=reply,
+                reply_to_message_id=message.id,
+            )
+        )
+        await self._session_handler.persist_messages(
+            chat_id=message.chat_id,
+            user_id=message.user_id,
+            user_message=message.text,
+            assistant_message=reply,
+            external_id=message.id,
+            response_external_id=response_external_id,
+            thread_id=message.metadata.get("thread_id"),
+            username=message.username,
+            display_name=message.display_name,
+        )
+        return True
+
     def _log_response(self, text: str | None) -> None:
         bot_name = self._provider.bot_username or "bot"
         logger.info(
@@ -370,6 +492,9 @@ async def _process_single_message_inner(
             if isinstance(candidate, IncomingMessage):
                 message = candidate
 
+        if await self._try_handle_capability_oauth_callback(message):
+            return
+
         # Check if there's an active interactive subagent stack for this session
         thread_id = message.metadata.get("thread_id")
         session_key = make_session_key(

src/ash/rpc/methods/capability.py

@@ -127,6 +127,39 @@ async def capability_auth_complete(params: dict[str, Any]) -> dict[str, Any]:
             raise ValueError(f"{e.code}: {e}") from e
         return {"ok": bool(result.get("ok")), "account_ref": result["account_ref"]}
 
+    async def capability_auth_complete_callback(
+        params: dict[str, Any],
+    ) -> dict[str, Any]:
+        user_id = _required_text(params, "user_id")
+        callback_url = _optional_text(params, "callback_url")
+        code = _optional_text(params, "code")
+        capability_id = _optional_text(params, "capability")
+        account_hint = _optional_text(params, "account_hint")
+        try:
+            result = await manager.auth_complete_callback(
+                user_id=user_id,
+                callback_url=callback_url,
+                code=code,
+                capability_id=capability_id,
+                account_hint=account_hint,
+                chat_id=_optional_text(params, "chat_id"),
+                chat_type=_optional_text(params, "chat_type"),
+                provider=_optional_text(params, "provider"),
+                thread_id=_optional_text(params, "thread_id"),
+                session_key=_optional_text(params, "session_key"),
+                source_username=_optional_text(params, "source_username"),
+                source_display_name=_optional_text(params, "source_display_name"),
+            )
+        except CapabilityError as e:
+            raise ValueError(f"{e.code}: {e}") from e
+        return {
+            "ok": bool(result.get("ok")),
+            "account_ref": result["account_ref"],
+            "flow_id": result["flow_id"],
+            "capability": result["capability"],
+            "account_hint": result.get("account_hint"),
+        }
+
     async def capability_auth_poll(params: dict[str, Any]) -> dict[str, Any]:
         flow_id = _required_text(params, "flow_id")
         user_id = _required_text(params, "user_id")
@@ -150,6 +183,9 @@ async def capability_auth_poll(params: dict[str, Any]) -> dict[str, Any]:
     server.register("capability.auth.begin", capability_auth_begin)
     server.register("capability.auth.list", capability_auth_list)
     server.register("capability.auth.complete", capability_auth_complete)
+    server.register(
+        "capability.auth.complete_callback", capability_auth_complete_callback
+    )
     server.register("capability.auth.poll", capability_auth_poll)
 
     logger.debug("Registered capability RPC methods")