codex-as-api 0.2.1__py3-none-any.whl

codex_as_api/__init__.py

@@ -0,0 +1,3 @@
+"""codex-as-api: Use ChatGPT/Codex OAuth as a local OpenAI-compatible API server."""
+
+__version__ = "0.1.0"

codex_as_api/anthropic_adapter.py

@@ -0,0 +1,447 @@
+from __future__ import annotations
+
+import json
+import uuid
+from typing import Any, Iterator, Sequence
+
+from .messages import Message, MessageRole, ToolCall, ToolSchema, Usage, AssistantResponse
+
+
+# ---------------------------------------------------------------------------
+# Request conversion: Anthropic → internal
+# ---------------------------------------------------------------------------
+
+def anthropic_request_to_internal(
+    *,
+    model: str,
+    messages: list[dict[str, Any]],
+    system: str | list[dict[str, Any]] | None = None,
+    max_tokens: int = 4096,
+    tools: list[dict[str, Any]] | None = None,
+    tool_choice: dict[str, Any] | None = None,
+    stop_sequences: list[str] | None = None,
+    thinking: dict[str, Any] | None = None,
+) -> tuple[list[Message], list[ToolSchema] | None, str | dict | None, list[str] | None, str | None]:
+    """Convert Anthropic Messages request fields to internal types.
+
+    Returns (messages, tools, tool_choice, stop, reasoning_effort).
+    """
+    internal_messages: list[Message] = []
+
+    # System prompt → SYSTEM message
+    if system is not None:
+        sys_text = _extract_system_text(system)
+        if sys_text:
+            internal_messages.append(Message(role=MessageRole.SYSTEM, content=sys_text))
+
+    # Convert messages
+    for msg in messages:
+        role = msg.get("role", "user")
+        content = msg.get("content", "")
+        if role == "user":
+            _convert_user_message(content, internal_messages)
+        elif role == "assistant":
+            _convert_assistant_message(content, internal_messages)
+
+    # Convert tools
+    internal_tools = _convert_tools(tools) if tools else None
+
+    # Convert tool_choice
+    internal_tool_choice = _convert_tool_choice(tool_choice)
+
+    # Convert thinking → reasoning_effort
+    reasoning_effort = _convert_thinking(thinking)
+
+    return internal_messages, internal_tools, internal_tool_choice, stop_sequences, reasoning_effort
+
+
+def _extract_system_text(system: str | list[dict[str, Any]]) -> str:
+    if isinstance(system, str):
+        return system
+    parts: list[str] = []
+    for block in system:
+        if isinstance(block, dict) and block.get("type") == "text":
+            text = block.get("text")
+            if isinstance(text, str) and text:
+                parts.append(text)
+    return "\n\n".join(parts)
+
+
+def _convert_user_message(content: str | list[dict[str, Any]], out: list[Message]) -> None:
+    if isinstance(content, str):
+        out.append(Message(role=MessageRole.USER, content=content))
+        return
+    text_parts: list[str] = []
+    image_urls: list[str] = []
+    for block in content:
+        if not isinstance(block, dict):
+            continue
+        block_type = block.get("type")
+        if block_type == "text":
+            text = block.get("text")
+            if isinstance(text, str):
+                text_parts.append(text)
+        elif block_type == "tool_result":
+            if text_parts or image_urls:
+                out.append(Message(
+                    role=MessageRole.USER,
+                    content="".join(text_parts),
+                    images=tuple(image_urls),
+                ))
+                text_parts = []
+                image_urls = []
+            tool_use_id = block.get("tool_use_id") or "tool-call"
+            result_content = block.get("content", "")
+            tool_result_images: list[str] = []
+            if isinstance(result_content, list):
+                text_pieces: list[str] = []
+                for p in result_content:
+                    if not isinstance(p, dict):
+                        continue
+                    if p.get("type") == "text":
+                        text_pieces.append(p.get("text", ""))
+                    elif p.get("type") == "image":
+                        source = p.get("source", {})
+                        if isinstance(source, dict) and source.get("type") == "base64":
+                            media_type = source.get("media_type", "image/png")
+                            data = source.get("data", "")
+                            tool_result_images.append(f"data:{media_type};base64,{data}")
+                result_content = "".join(text_pieces)
+            elif not isinstance(result_content, str):
+                result_content = str(result_content) if result_content else ""
+            out.append(Message(
+                role=MessageRole.TOOL,
+                content=result_content,
+                tool_call_id=tool_use_id,
+                name=tool_use_id,
+            ))
+            if tool_result_images:
+                out.append(Message(
+                    role=MessageRole.USER,
+                    content="",
+                    images=tuple(tool_result_images),
+                ))
+        elif block_type == "image":
+            source = block.get("source", {})
+            if isinstance(source, dict) and source.get("type") == "base64":
+                media_type = source.get("media_type", "image/png")
+                data = source.get("data", "")
+                image_urls.append(f"data:{media_type};base64,{data}")
+    if text_parts or image_urls:
+        out.append(Message(
+            role=MessageRole.USER,
+            content="".join(text_parts),
+            images=tuple(image_urls),
+        ))
+
+
+def _convert_assistant_message(content: str | list[dict[str, Any]], out: list[Message]) -> None:
+    if isinstance(content, str):
+        out.append(Message(role=MessageRole.ASSISTANT, content=content))
+        return
+    text_parts: list[str] = []
+    tool_calls: list[ToolCall] = []
+    reasoning_content: str | None = None
+    for block in content:
+        if not isinstance(block, dict):
+            continue
+        block_type = block.get("type")
+        if block_type == "text":
+            text = block.get("text")
+            if isinstance(text, str):
+                text_parts.append(text)
+        elif block_type == "tool_use":
+            tool_calls.append(ToolCall(
+                id=block.get("id") or uuid.uuid4().hex,
+                name=block.get("name") or "",
+                arguments=block.get("input") or {},
+            ))
+        elif block_type == "thinking":
+            thinking_text = block.get("thinking")
+            if isinstance(thinking_text, str) and thinking_text:
+                reasoning_content = thinking_text
+    out.append(Message(
+        role=MessageRole.ASSISTANT,
+        content="".join(text_parts),
+        tool_calls=tuple(tool_calls) if tool_calls else (),
+        reasoning_content=reasoning_content,
+    ))
+
+
+def _convert_tools(tools: list[dict[str, Any]]) -> list[ToolSchema]:
+    result: list[ToolSchema] = []
+    for tool in tools:
+        if not isinstance(tool, dict):
+            continue
+        name = tool.get("name")
+        if not name:
+            continue
+        result.append(ToolSchema(
+            name=str(name),
+            description=str(tool.get("description") or ""),
+            parameters=tool.get("input_schema") or {},
+        ))
+    return result
+
+
+def _convert_tool_choice(tc: dict[str, Any] | None) -> str | dict | None:
+    if tc is None:
+        return None
+    tc_type = tc.get("type")
+    if tc_type == "auto":
+        return "auto"
+    if tc_type == "any":
+        return "required"
+    if tc_type == "tool":
+        return {"type": "function", "function": {"name": tc.get("name")}}
+    if tc_type == "none":
+        return "none"
+    return "auto"
+
+
+def _convert_thinking(thinking: dict[str, Any] | None) -> str | None:
+    if thinking is None:
+        return None
+    if thinking.get("type") == "enabled":
+        return "high"
+    if thinking.get("type") == "adaptive":
+        return "medium"
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Non-streaming response: internal → Anthropic
+# ---------------------------------------------------------------------------
+
+def internal_response_to_anthropic(
+    response: AssistantResponse,
+    model: str,
+    request_id: str,
+) -> dict[str, Any]:
+    content: list[dict[str, Any]] = []
+
+    if response.reasoning_content:
+        content.append({
+            "type": "thinking",
+            "thinking": response.reasoning_content,
+            "signature": "sig-placeholder",
+        })
+
+    if response.content:
+        content.append({"type": "text", "text": response.content})
+
+    for tc in response.tool_calls:
+        content.append({
+            "type": "tool_use",
+            "id": tc.id,
+            "name": tc.name,
+            "input": tc.arguments,
+        })
+
+    stop_reason = _map_stop_reason(response.finish_reason, bool(response.tool_calls))
+
+    usage_dict: dict[str, Any] = {"input_tokens": 0, "output_tokens": 0}
+    if response.usage:
+        usage_dict = {
+            "input_tokens": response.usage.prompt_tokens,
+            "output_tokens": response.usage.completion_tokens,
+            "cache_creation_input_tokens": 0,
+            "cache_read_input_tokens": response.usage.cached_tokens,
+        }
+
+    if not content:
+        content.append({"type": "text", "text": ""})
+
+    return {
+        "id": request_id,
+        "type": "message",
+        "role": "assistant",
+        "model": model,
+        "content": content,
+        "stop_reason": stop_reason,
+        "stop_sequence": None,
+        "usage": usage_dict,
+    }
+
+
+def _map_stop_reason(finish_reason: str, has_tool_calls: bool) -> str:
+    if has_tool_calls:
+        return "tool_use"
+    mapping = {
+        "stop": "end_turn",
+        "length": "max_tokens",
+        "max_tokens": "max_tokens",
+        "tool_calls": "tool_use",
+        "tool_use": "tool_use",
+        "stop_sequence": "stop_sequence",
+    }
+    return mapping.get(finish_reason, "end_turn")
+
+
+# ---------------------------------------------------------------------------
+# Streaming adapter: provider events → Anthropic SSE
+# ---------------------------------------------------------------------------
+
+def anthropic_stream_adapter(
+    event_stream: Iterator[dict[str, Any]],
+    model: str,
+    request_id: str,
+) -> Iterator[str]:
+    """Convert provider chat_stream events into Anthropic SSE strings."""
+    # Emit message_start
+    yield _sse("message_start", {
+        "type": "message_start",
+        "message": {
+            "id": request_id,
+            "type": "message",
+            "role": "assistant",
+            "model": model,
+            "content": [],
+            "stop_reason": None,
+            "stop_sequence": None,
+            "usage": {"input_tokens": 0, "output_tokens": 0},
+        },
+    })
+
+    block_index = 0
+    current_block: str | None = None  # "thinking", "text", "tool_use"
+    has_any_content = False
+    output_tokens = 0
+
+    for event in event_stream:
+        typ = event.get("type")
+
+        if typ in ("reasoning_delta", "reasoning_raw_delta"):
+            has_any_content = True
+            text = str(event.get("text", ""))
+            if current_block != "thinking":
+                if current_block is not None:
+                    yield _sse("content_block_stop", {"type": "content_block_stop", "index": block_index})
+                    block_index += 1
+                yield _sse("content_block_start", {
+                    "type": "content_block_start",
+                    "index": block_index,
+                    "content_block": {"type": "thinking", "thinking": "", "signature": ""},
+                })
+                current_block = "thinking"
+            yield _sse("content_block_delta", {
+                "type": "content_block_delta",
+                "index": block_index,
+                "delta": {"type": "thinking_delta", "thinking": text},
+            })
+
+        elif typ == "content":
+            has_any_content = True
+            text = str(event.get("text", ""))
+            if current_block == "thinking":
+                # Close thinking block, emit signature
+                yield _sse("content_block_delta", {
+                    "type": "content_block_delta",
+                    "index": block_index,
+                    "delta": {"type": "signature_delta", "signature": "sig-placeholder"},
+                })
+                yield _sse("content_block_stop", {"type": "content_block_stop", "index": block_index})
+                block_index += 1
+                current_block = None
+            if current_block != "text":
+                if current_block is not None:
+                    yield _sse("content_block_stop", {"type": "content_block_stop", "index": block_index})
+                    block_index += 1
+                yield _sse("content_block_start", {
+                    "type": "content_block_start",
+                    "index": block_index,
+                    "content_block": {"type": "text", "text": ""},
+                })
+                current_block = "text"
+            yield _sse("content_block_delta", {
+                "type": "content_block_delta",
+                "index": block_index,
+                "delta": {"type": "text_delta", "text": text},
+            })
+
+        elif typ == "tool_call":
+            has_any_content = True
+            if current_block is not None:
+                if current_block == "thinking":
+                    yield _sse("content_block_delta", {
+                        "type": "content_block_delta",
+                        "index": block_index,
+                        "delta": {"type": "signature_delta", "signature": "sig-placeholder"},
+                    })
+                yield _sse("content_block_stop", {"type": "content_block_stop", "index": block_index})
+                block_index += 1
+            tool_id = str(event.get("id", ""))
+            tool_name = str(event.get("name", ""))
+            tool_args = event.get("arguments") or {}
+            yield _sse("content_block_start", {
+                "type": "content_block_start",
+                "index": block_index,
+                "content_block": {"type": "tool_use", "id": tool_id, "name": tool_name, "input": {}},
+            })
+            yield _sse("content_block_delta", {
+                "type": "content_block_delta",
+                "index": block_index,
+                "delta": {"type": "input_json_delta", "partial_json": json.dumps(tool_args, ensure_ascii=False)},
+            })
+            yield _sse("content_block_stop", {"type": "content_block_stop", "index": block_index})
+            block_index += 1
+            current_block = None
+
+        elif typ == "finish":
+            if current_block is not None:
+                if current_block == "thinking":
+                    yield _sse("content_block_delta", {
+                        "type": "content_block_delta",
+                        "index": block_index,
+                        "delta": {"type": "signature_delta", "signature": "sig-placeholder"},
+                    })
+                yield _sse("content_block_stop", {"type": "content_block_stop", "index": block_index})
+                current_block = None
+
+            if not has_any_content:
+                yield _sse("content_block_start", {
+                    "type": "content_block_start",
+                    "index": block_index,
+                    "content_block": {"type": "text", "text": ""},
+                })
+                yield _sse("content_block_stop", {"type": "content_block_stop", "index": block_index})
+
+            finish_reason = str(event.get("finish_reason") or "stop")
+            stop_reason = _map_stop_reason(finish_reason, False)
+            usage = event.get("usage")
+            if isinstance(usage, dict):
+                output_tokens = usage.get("output_tokens", usage.get("completion_tokens", 0))
+
+            yield _sse("message_delta", {
+                "type": "message_delta",
+                "delta": {"stop_reason": stop_reason, "stop_sequence": None},
+                "usage": {"output_tokens": output_tokens},
+            })
+            yield _sse("message_stop", {"type": "message_stop"})
+
+
+def _sse(event_type: str, data: dict[str, Any]) -> str:
+    return f"event: {event_type}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
+
+
+# ---------------------------------------------------------------------------
+# Error formatting
+# ---------------------------------------------------------------------------
+
+def format_anthropic_error(status: int, message: str) -> dict[str, Any]:
+    type_map = {
+        400: "invalid_request_error",
+        401: "authentication_error",
+        403: "permission_error",
+        404: "not_found_error",
+        429: "rate_limit_error",
+        500: "api_error",
+        529: "overloaded_error",
+    }
+    return {
+        "type": "error",
+        "error": {
+            "type": type_map.get(status, "api_error"),
+            "message": message,
+        },
+    }

codex_as_api/auth.py

@@ -0,0 +1,239 @@
+from __future__ import annotations
+
+import base64
+import dataclasses
+import datetime as _dt
+import json
+import os
+import pathlib
+import threading
+import urllib.error
+import urllib.request
+from typing import Any
+
+CHATGPT_OAUTH_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
+DEFAULT_AUTH_PATH = "~/.codex/auth.json"
+DEFAULT_REFRESH_URL = "https://auth.openai.com/oauth/token"
+REFRESH_URL_OVERRIDE_ENV = "CODEX_REFRESH_TOKEN_URL_OVERRIDE"
+
+_SECRET_KEYS = (
+    "access_token",
+    "refresh_token",
+    "id_token",
+    "Authorization",
+    "authorization",
+    "ChatGPT-Account-Id",
+    "chatgpt-account-id",
+)
+
+_REFRESH_LOCKS: dict[pathlib.Path, threading.Lock] = {}
+_REFRESH_LOCKS_GUARD = threading.Lock()
+
+
+class ChatGPTOAuthError(RuntimeError):
+    pass
+
+
+class ChatGPTOAuthMissingError(ChatGPTOAuthError):
+    pass
+
+
+class ChatGPTOAuthRefreshError(ChatGPTOAuthError):
+    pass
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class ChatGPTTokenData:
+    auth_path: pathlib.Path
+    access_token: str
+    refresh_token: str
+    id_token: str
+    account_id: str
+    plan_type: str | None
+    user_id: str | None
+    fedramp: bool
+    access_expires_at: _dt.datetime | None
+
+    @property
+    def expired(self) -> bool:
+        return self.access_expires_at is not None and self.access_expires_at <= _dt.datetime.now(_dt.UTC)
+
+
+def resolve_auth_path(raw: str | None = None) -> pathlib.Path:
+    value = raw or os.getenv("CODEX_HOME")
+    if value and raw is None:
+        return pathlib.Path(value).expanduser() / "auth.json"
+    return pathlib.Path(raw or DEFAULT_AUTH_PATH).expanduser()
+
+
+def _jwt_claims(jwt: str) -> dict[str, Any]:
+    parts = jwt.split(".")
+    if len(parts) < 2 or not parts[1]:
+        return {}
+    payload = parts[1] + "=" * ((4 - len(parts[1]) % 4) % 4)
+    try:
+        decoded = base64.urlsafe_b64decode(payload.encode())
+        value = json.loads(decoded)
+    except Exception as exc:  # noqa: BLE001 - invalid JWT is auth-data invalid
+        raise ChatGPTOAuthError("invalid ChatGPT OAuth JWT payload") from exc
+    if not isinstance(value, dict):
+        raise ChatGPTOAuthError("invalid ChatGPT OAuth JWT claims")
+    return value
+
+
+def _expiration(jwt: str) -> _dt.datetime | None:
+    claims = _jwt_claims(jwt)
+    exp = claims.get("exp")
+    if not isinstance(exp, int):
+        return None
+    return _dt.datetime.fromtimestamp(exp, _dt.UTC)
+
+
+def _auth_claims(jwt: str) -> dict[str, Any]:
+    claims = _jwt_claims(jwt)
+    value = claims.get("https://api.openai.com/auth")
+    return value if isinstance(value, dict) else {}
+
+
+def register_token_secrets(*values: str | None) -> None:
+    """No-op in standalone mode; secrets are redacted inline via redact_text()."""
+    pass
+
+
+def redact_text(text: str, *values: str | None) -> str:
+    redacted = str(text)
+    for value in sorted([v for v in values if v], key=len, reverse=True):
+        redacted = redacted.replace(value, "***")
+    return redacted
+
+
+def load_token_data(auth_json_path: str | pathlib.Path | None = None) -> ChatGPTTokenData:
+    path = resolve_auth_path(str(auth_json_path) if auth_json_path is not None else None)
+    try:
+        raw = path.read_text(encoding="utf-8")
+    except FileNotFoundError as exc:
+        raise ChatGPTOAuthMissingError(f"ChatGPT OAuth auth file not found: {path}") from exc
+    try:
+        data = json.loads(raw)
+    except json.JSONDecodeError as exc:
+        raise ChatGPTOAuthError(f"ChatGPT OAuth auth file is invalid JSON: {path}") from exc
+    if not isinstance(data, dict):
+        raise ChatGPTOAuthError("ChatGPT OAuth auth file root must be an object")
+    mode = data.get("auth_mode")
+    if mode not in {"chatgpt", "Chatgpt", "chatgpt_auth_tokens", "ChatgptAuthTokens", None}:
+        raise ChatGPTOAuthError(f"ChatGPT OAuth auth_mode required, got {mode!r}")
+    tokens = data.get("tokens")
+    if not isinstance(tokens, dict):
+        raise ChatGPTOAuthError("ChatGPT OAuth token data is not available")
+    access_token = tokens.get("access_token")
+    refresh_token = tokens.get("refresh_token")
+    id_token = tokens.get("id_token")
+    for name, value in (("access_token", access_token), ("refresh_token", refresh_token), ("id_token", id_token)):
+        if not isinstance(value, str) or value == "":
+            raise ChatGPTOAuthError(f"ChatGPT OAuth {name} is missing")
+    register_token_secrets(access_token, refresh_token, id_token)
+    id_auth = _auth_claims(id_token)
+    access_auth = _auth_claims(access_token)
+    account_id = tokens.get("account_id") or id_auth.get("chatgpt_account_id") or access_auth.get("chatgpt_account_id")
+    if not isinstance(account_id, str) or account_id == "":
+        raise ChatGPTOAuthError("ChatGPT OAuth account id not available; rerun codex login")
+    register_token_secrets(account_id)
+    plan = id_auth.get("chatgpt_plan_type") or access_auth.get("chatgpt_plan_type")
+    user = id_auth.get("chatgpt_user_id") or id_auth.get("user_id") or access_auth.get("chatgpt_user_id") or access_auth.get("user_id")
+    fedramp = bool(id_auth.get("chatgpt_account_is_fedramp") or access_auth.get("chatgpt_account_is_fedramp"))
+    return ChatGPTTokenData(
+        auth_path=path,
+        access_token=access_token,
+        refresh_token=refresh_token,
+        id_token=id_token,
+        account_id=account_id,
+        plan_type=plan if isinstance(plan, str) else None,
+        user_id=user if isinstance(user, str) else None,
+        fedramp=fedramp,
+        access_expires_at=_expiration(access_token),
+    )
+
+
+def is_auth_locally_available(auth_json_path: str | pathlib.Path | None = None) -> bool:
+    try:
+        data = load_token_data(auth_json_path)
+    except ChatGPTOAuthError:
+        return False
+    return bool(data.access_token and data.account_id)
+
+
+def _refresh_lock(path: pathlib.Path) -> threading.Lock:
+    resolved = path.expanduser()
+    with _REFRESH_LOCKS_GUARD:
+        lock = _REFRESH_LOCKS.get(resolved)
+        if lock is None:
+            lock = threading.Lock()
+            _REFRESH_LOCKS[resolved] = lock
+        return lock
+
+
+def _write_auth_json(path: pathlib.Path, data: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_name(f".{path.name}.tmp-{os.getpid()}-{threading.get_ident()}")
+    payload = json.dumps(data, ensure_ascii=False, indent=2) + "\n"
+    flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
+    fd = os.open(tmp, flags, 0o600)
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8") as f:
+            f.write(payload)
+            f.flush()
+            os.fsync(f.fileno())
+        tmp.replace(path)
+        dir_fd = os.open(path.parent, os.O_RDONLY)
+        try:
+            os.fsync(dir_fd)
+        finally:
+            os.close(dir_fd)
+    finally:
+        try:
+            tmp.unlink(missing_ok=True)
+        except OSError:
+            pass
+
+
+def refresh_token(auth_json_path: str | pathlib.Path | None = None) -> ChatGPTTokenData:
+    current = load_token_data(auth_json_path)
+    lock = _refresh_lock(current.auth_path)
+    with lock:
+        current = load_token_data(auth_json_path)
+        endpoint = os.getenv(REFRESH_URL_OVERRIDE_ENV, DEFAULT_REFRESH_URL)
+        body = json.dumps({
+            "client_id": CHATGPT_OAUTH_CLIENT_ID,
+            "grant_type": "refresh_token",
+            "refresh_token": current.refresh_token,
+        }).encode()
+        request = urllib.request.Request(
+            endpoint,
+            data=body,
+            headers={"Content-Type": "application/json"},
+            method="POST",
+        )
+        try:
+            with urllib.request.urlopen(request, timeout=30) as response:
+                payload = json.loads(response.read().decode("utf-8"))
+        except urllib.error.HTTPError as exc:
+            text = exc.read().decode("utf-8", "replace")
+            redacted = redact_text(text, current.access_token, current.refresh_token, current.id_token)
+            if exc.code == 401:
+                raise ChatGPTOAuthRefreshError(f"ChatGPT OAuth refresh token is invalid; rerun codex login: {redacted}") from exc
+            raise ChatGPTOAuthRefreshError(f"ChatGPT OAuth token refresh failed: HTTP {exc.code}: {redacted}") from exc
+        except Exception as exc:  # noqa: BLE001
+            raise ChatGPTOAuthRefreshError(f"ChatGPT OAuth token refresh failed: {exc}") from exc
+        if not isinstance(payload, dict):
+            raise ChatGPTOAuthRefreshError("ChatGPT OAuth token refresh returned invalid JSON")
+        data = json.loads(current.auth_path.read_text(encoding="utf-8"))
+        tokens = data.setdefault("tokens", {})
+        if payload.get("id_token"):
+            tokens["id_token"] = payload["id_token"]
+        if payload.get("access_token"):
+            tokens["access_token"] = payload["access_token"]
+        if payload.get("refresh_token"):
+            tokens["refresh_token"] = payload["refresh_token"]
+        data["last_refresh"] = _dt.datetime.now(_dt.UTC).isoformat().replace("+00:00", "Z")
+        _write_auth_json(current.auth_path, data)
+        return load_token_data(auth_json_path)