klaude-code 1.2.6__py3-none-any.whl → 1.2.7__py3-none-any.whl

klaude_code/auth/__init__.py

@@ -0,0 +1,24 @@
+"""Authentication module.
+
+Currently includes Codex OAuth helpers in ``klaude_code.auth.codex``.
+"""
+
+from klaude_code.auth.codex import (
+    CodexAuthError,
+    CodexAuthState,
+    CodexNotLoggedInError,
+    CodexOAuth,
+    CodexOAuthError,
+    CodexTokenExpiredError,
+    CodexTokenManager,
+)
+
+__all__ = [
+    "CodexAuthError",
+    "CodexAuthState",
+    "CodexNotLoggedInError",
+    "CodexOAuth",
+    "CodexOAuthError",
+    "CodexTokenExpiredError",
+    "CodexTokenManager",
+]

klaude_code/auth/codex/__init__.py

@@ -0,0 +1,20 @@
+"""Codex authentication helpers."""
+
+from klaude_code.auth.codex.exceptions import (
+    CodexAuthError,
+    CodexNotLoggedInError,
+    CodexOAuthError,
+    CodexTokenExpiredError,
+)
+from klaude_code.auth.codex.oauth import CodexOAuth
+from klaude_code.auth.codex.token_manager import CodexAuthState, CodexTokenManager
+
+__all__ = [
+    "CodexAuthError",
+    "CodexAuthState",
+    "CodexNotLoggedInError",
+    "CodexOAuth",
+    "CodexOAuthError",
+    "CodexTokenExpiredError",
+    "CodexTokenManager",
+]

klaude_code/auth/codex/exceptions.py

@@ -0,0 +1,17 @@
+"""Exceptions for Codex authentication."""
+
+
+class CodexAuthError(Exception):
+    """Base exception for Codex authentication errors."""
+
+
+class CodexNotLoggedInError(CodexAuthError):
+    """User has not logged in to Codex."""
+
+
+class CodexTokenExpiredError(CodexAuthError):
+    """Token expired and refresh failed."""
+
+
+class CodexOAuthError(CodexAuthError):
+    """OAuth flow failed."""

klaude_code/auth/codex/jwt_utils.py

@@ -0,0 +1,45 @@
+"""JWT parsing utilities for Codex authentication."""
+
+import base64
+import json
+from typing import Any
+
+
+def decode_jwt_payload(token: str) -> dict[str, Any]:
+    """Decode JWT payload without verification.
+
+    Args:
+        token: JWT token string
+
+    Returns:
+        Decoded payload as a dictionary
+    """
+    parts = token.split(".")
+    if len(parts) != 3:
+        raise ValueError("Invalid JWT format")
+
+    payload = parts[1]
+    # Add padding if needed
+    padding = 4 - len(payload) % 4
+    if padding != 4:
+        payload += "=" * padding
+
+    decoded = base64.urlsafe_b64decode(payload)
+    return json.loads(decoded)
+
+
+def extract_account_id(token: str) -> str:
+    """Extract ChatGPT account ID from JWT token.
+
+    Args:
+        token: JWT access token from OpenAI OAuth
+
+    Returns:
+        The chatgpt_account_id from the token claims
+    """
+    payload = decode_jwt_payload(token)
+    auth_claim = payload.get("https://api.openai.com/auth", {})
+    account_id = auth_claim.get("chatgpt_account_id")
+    if not account_id:
+        raise ValueError("chatgpt_account_id not found in token")
+    return account_id

klaude_code/auth/codex/oauth.py

@@ -0,0 +1,229 @@
+"""OAuth PKCE flow for Codex authentication."""
+
+import base64
+import hashlib
+import secrets
+import time
+import webbrowser
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from threading import Thread
+from typing import Any
+from urllib.parse import parse_qs, urlencode, urlparse
+
+import httpx
+
+from klaude_code.auth.codex.exceptions import CodexOAuthError
+from klaude_code.auth.codex.jwt_utils import extract_account_id
+from klaude_code.auth.codex.token_manager import CodexAuthState, CodexTokenManager
+
+# OAuth configuration
+CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
+AUTHORIZE_URL = "https://auth.openai.com/oauth/authorize"
+TOKEN_URL = "https://auth.openai.com/oauth/token"
+REDIRECT_URI = "http://localhost:1455/auth/callback"
+REDIRECT_PORT = 1455
+SCOPE = "openid profile email offline_access"
+
+
+def generate_code_verifier() -> str:
+    """Generate a random code verifier for PKCE."""
+    return secrets.token_urlsafe(64)[:128]
+
+
+def generate_code_challenge(verifier: str) -> str:
+    """Generate code challenge from verifier using S256 method."""
+    digest = hashlib.sha256(verifier.encode()).digest()
+    return base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
+
+
+def build_authorize_url(code_challenge: str, state: str) -> str:
+    """Build the authorization URL with all required parameters."""
+    params = {
+        "response_type": "code",
+        "client_id": CLIENT_ID,
+        "redirect_uri": REDIRECT_URI,
+        "scope": SCOPE,
+        "code_challenge": code_challenge,
+        "code_challenge_method": "S256",
+        "state": state,
+        "id_token_add_organizations": "true",
+        "codex_cli_simplified_flow": "true",
+        "originator": "codex_cli_rs",
+    }
+    return f"{AUTHORIZE_URL}?{urlencode(params)}"
+
+
+class OAuthCallbackHandler(BaseHTTPRequestHandler):
+    """HTTP request handler for OAuth callback."""
+
+    code: str | None = None
+    state: str | None = None
+    error: str | None = None
+
+    def log_message(self, format: str, *args: Any) -> None:
+        """Suppress HTTP server logs."""
+        pass
+
+    def do_GET(self) -> None:
+        """Handle GET request from OAuth callback."""
+        parsed = urlparse(self.path)
+        params = parse_qs(parsed.query)
+
+        OAuthCallbackHandler.code = params.get("code", [None])[0]
+        OAuthCallbackHandler.state = params.get("state", [None])[0]
+        OAuthCallbackHandler.error = params.get("error", [None])[0]
+
+        self.send_response(200)
+        self.send_header("Content-Type", "text/html")
+        self.end_headers()
+
+        if OAuthCallbackHandler.error:
+            html = """
+            <html><body style="font-family: sans-serif; text-align: center; padding: 50px;">
+            <h1>Authentication Failed</h1>
+            <p>Error: {}</p>
+            <p>Please close this window and try again.</p>
+            </body></html>
+            """.format(OAuthCallbackHandler.error)
+        else:
+            html = """
+            <html><body style="font-family: sans-serif; text-align: center; padding: 50px;">
+            <h1>Authentication Successful!</h1>
+            <p>You can close this window now.</p>
+            <script>setTimeout(function() { window.close(); }, 2000);</script>
+            </body></html>
+            """
+        self.wfile.write(html.encode())
+
+
+class CodexOAuth:
+    """Handle OAuth PKCE flow for Codex authentication."""
+
+    def __init__(self, token_manager: CodexTokenManager | None = None):
+        self.token_manager = token_manager or CodexTokenManager()
+
+    def login(self) -> CodexAuthState:
+        """Run the complete OAuth login flow."""
+        # Generate PKCE parameters
+        verifier = generate_code_verifier()
+        challenge = generate_code_challenge(verifier)
+        state = secrets.token_urlsafe(32)
+
+        # Build authorization URL
+        auth_url = build_authorize_url(challenge, state)
+
+        # Start callback server
+        OAuthCallbackHandler.code = None
+        OAuthCallbackHandler.state = None
+        OAuthCallbackHandler.error = None
+
+        server = HTTPServer(("localhost", REDIRECT_PORT), OAuthCallbackHandler)
+        server_thread = Thread(target=server.handle_request)
+        server_thread.start()
+
+        # Open browser for user to authenticate
+        webbrowser.open(auth_url)
+
+        # Wait for callback
+        server_thread.join(timeout=300)  # 5 minute timeout
+        server.server_close()
+
+        # Check for errors
+        if OAuthCallbackHandler.error:
+            raise CodexOAuthError(f"OAuth error: {OAuthCallbackHandler.error}")
+
+        if not OAuthCallbackHandler.code:
+            raise CodexOAuthError("No authorization code received")
+
+        if OAuthCallbackHandler.state is None or OAuthCallbackHandler.state != state:
+            raise CodexOAuthError("OAuth state mismatch")
+
+        # Exchange code for tokens
+        auth_state = self._exchange_code(OAuthCallbackHandler.code, verifier)
+
+        # Save tokens
+        self.token_manager.save(auth_state)
+
+        return auth_state
+
+    def _exchange_code(self, code: str, verifier: str) -> CodexAuthState:
+        """Exchange authorization code for tokens."""
+        data = {
+            "grant_type": "authorization_code",
+            "client_id": CLIENT_ID,
+            "code": code,
+            "redirect_uri": REDIRECT_URI,
+            "code_verifier": verifier,
+        }
+
+        with httpx.Client() as client:
+            response = client.post(TOKEN_URL, data=data)
+
+        if response.status_code != 200:
+            raise CodexOAuthError(f"Token exchange failed: {response.text}")
+
+        tokens = response.json()
+        access_token = tokens["access_token"]
+        refresh_token = tokens["refresh_token"]
+        expires_in = tokens.get("expires_in", 3600)
+
+        account_id = extract_account_id(access_token)
+
+        return CodexAuthState(
+            access_token=access_token,
+            refresh_token=refresh_token,
+            expires_at=int(time.time()) + expires_in,
+            account_id=account_id,
+        )
+
+    def refresh(self) -> CodexAuthState:
+        """Refresh the access token using refresh token."""
+        state = self.token_manager.get_state()
+        if state is None:
+            from klaude_code.auth.codex.exceptions import CodexNotLoggedInError
+
+            raise CodexNotLoggedInError("Not logged in to Codex. Run 'klaude login codex' first.")
+
+        data = {
+            "grant_type": "refresh_token",
+            "client_id": CLIENT_ID,
+            "refresh_token": state.refresh_token,
+        }
+
+        with httpx.Client() as client:
+            response = client.post(TOKEN_URL, data=data)
+
+        if response.status_code != 200:
+            from klaude_code.auth.codex.exceptions import CodexTokenExpiredError
+
+            raise CodexTokenExpiredError(f"Token refresh failed: {response.text}")
+
+        tokens = response.json()
+        access_token = tokens["access_token"]
+        refresh_token = tokens.get("refresh_token", state.refresh_token)
+        expires_in = tokens.get("expires_in", 3600)
+
+        account_id = extract_account_id(access_token)
+
+        new_state = CodexAuthState(
+            access_token=access_token,
+            refresh_token=refresh_token,
+            expires_at=int(time.time()) + expires_in,
+            account_id=account_id,
+        )
+
+        self.token_manager.save(new_state)
+        return new_state
+
+    def ensure_valid_token(self) -> str:
+        """Ensure we have a valid access token, refreshing if needed."""
+        state = self.token_manager.get_state()
+        if state is None:
+            from klaude_code.auth.codex.exceptions import CodexNotLoggedInError
+
+            raise CodexNotLoggedInError("Not logged in to Codex. Run 'klaude login codex' first.")
+
+        if state.is_expired():
+            state = self.refresh()
+
+        return state.access_token

klaude_code/auth/codex/token_manager.py

@@ -0,0 +1,84 @@
+"""Token storage and management for Codex authentication."""
+
+import json
+import time
+from pathlib import Path
+
+from pydantic import BaseModel
+
+
+class CodexAuthState(BaseModel):
+    """Stored authentication state for Codex."""
+
+    access_token: str
+    refresh_token: str
+    expires_at: int  # Unix timestamp
+    account_id: str
+
+    def is_expired(self, buffer_seconds: int = 300) -> bool:
+        """Check if token is expired or will expire soon."""
+        return time.time() + buffer_seconds >= self.expires_at
+
+
+CODEX_AUTH_FILE = Path.home() / ".klaude" / "codex-auth.json"
+
+
+class CodexTokenManager:
+    """Manage Codex OAuth tokens."""
+
+    def __init__(self, auth_file: Path | None = None):
+        self.auth_file = auth_file or CODEX_AUTH_FILE
+        self._state: CodexAuthState | None = None
+
+    def load(self) -> CodexAuthState | None:
+        """Load authentication state from file."""
+        if not self.auth_file.exists():
+            return None
+
+        try:
+            data = json.loads(self.auth_file.read_text())
+            self._state = CodexAuthState.model_validate(data)
+            return self._state
+        except (json.JSONDecodeError, ValueError):
+            return None
+
+    def save(self, state: CodexAuthState) -> None:
+        """Save authentication state to file."""
+        self.auth_file.parent.mkdir(parents=True, exist_ok=True)
+        self.auth_file.write_text(state.model_dump_json(indent=2))
+        self._state = state
+
+    def delete(self) -> None:
+        """Delete stored tokens."""
+        if self.auth_file.exists():
+            self.auth_file.unlink()
+        self._state = None
+
+    def is_logged_in(self) -> bool:
+        """Check if user is logged in."""
+        state = self._state or self.load()
+        return state is not None
+
+    def get_state(self) -> CodexAuthState | None:
+        """Get current authentication state."""
+        if self._state is None:
+            self._state = self.load()
+        return self._state
+
+    def get_access_token(self) -> str:
+        """Get access token, raising if not logged in."""
+        state = self.get_state()
+        if state is None:
+            from klaude_code.auth.codex.exceptions import CodexNotLoggedInError
+
+            raise CodexNotLoggedInError("Not logged in to Codex. Run 'klaude login codex' first.")
+        return state.access_token
+
+    def get_account_id(self) -> str:
+        """Get account ID, raising if not logged in."""
+        state = self.get_state()
+        if state is None:
+            from klaude_code.auth.codex.exceptions import CodexNotLoggedInError
+
+            raise CodexNotLoggedInError("Not logged in to Codex. Run 'klaude login codex' first.")
+        return state.account_id

klaude_code/cli/main.py

@@ -1,4 +1,5 @@
 import asyncio
+import datetime
 import os
 import subprocess
 import sys
@@ -43,6 +44,68 @@ register_session_commands(session_app)
 app.add_typer(session_app, name="session")
 
 
+@app.command("login")
+def login_command(
+    provider: str = typer.Argument("codex", help="Provider to login (codex)"),
+) -> None:
+    """Login to a provider using OAuth."""
+    if provider.lower() != "codex":
+        log((f"Error: Unknown provider '{provider}'. Currently only 'codex' is supported.", "red"))
+        raise typer.Exit(1)
+
+    from klaude_code.auth.codex.oauth import CodexOAuth
+    from klaude_code.auth.codex.token_manager import CodexTokenManager
+
+    token_manager = CodexTokenManager()
+
+    # Check if already logged in
+    if token_manager.is_logged_in():
+        state = token_manager.get_state()
+        if state and not state.is_expired():
+            log(("You are already logged in to Codex.", "green"))
+            log(f"  Account ID: {state.account_id[:8]}...")
+            expires_dt = datetime.datetime.fromtimestamp(state.expires_at, tz=datetime.timezone.utc)
+            log(f"  Expires: {expires_dt.strftime('%Y-%m-%d %H:%M:%S UTC')}")
+            if not typer.confirm("Do you want to re-login?"):
+                return
+
+    log("Starting Codex OAuth login flow...")
+    log("A browser window will open for authentication.")
+
+    try:
+        oauth = CodexOAuth(token_manager)
+        state = oauth.login()
+        log(("Login successful!", "green"))
+        log(f"  Account ID: {state.account_id[:8]}...")
+        expires_dt = datetime.datetime.fromtimestamp(state.expires_at, tz=datetime.timezone.utc)
+        log(f"  Expires: {expires_dt.strftime('%Y-%m-%d %H:%M:%S UTC')}")
+    except Exception as e:
+        log((f"Login failed: {e}", "red"))
+        raise typer.Exit(1)
+
+
+@app.command("logout")
+def logout_command(
+    provider: str = typer.Argument("codex", help="Provider to logout (codex)"),
+) -> None:
+    """Logout from a provider."""
+    if provider.lower() != "codex":
+        log((f"Error: Unknown provider '{provider}'. Currently only 'codex' is supported.", "red"))
+        raise typer.Exit(1)
+
+    from klaude_code.auth.codex.token_manager import CodexTokenManager
+
+    token_manager = CodexTokenManager()
+
+    if not token_manager.is_logged_in():
+        log("You are not logged in to Codex.")
+        return
+
+    if typer.confirm("Are you sure you want to logout from Codex?"):
+        token_manager.delete()
+        log(("Logged out from Codex.", "green"))
+
+
 @app.command("list")
 def list_models() -> None:
     """List all models and providers configuration"""

klaude_code/command/status_cmd.py

@@ -13,11 +13,18 @@ def accumulate_session_usage(session: Session) -> tuple[model.Usage, int]:
     """
     total = model.Usage()
     task_count = 0
+    first_currency_set = False
 
     for item in session.conversation_history:
         if isinstance(item, model.ResponseMetadataItem) and item.usage:
             task_count += 1
             usage = item.usage
+            
+            # Set currency from first usage item
+            if not first_currency_set and usage.currency:
+                total.currency = usage.currency
+                first_currency_set = True
+            
             total.input_tokens += usage.input_tokens
             total.cached_tokens += usage.cached_tokens
             total.reasoning_tokens += usage.reasoning_tokens
@@ -50,13 +57,14 @@ def _format_tokens(tokens: int) -> str:
     return str(tokens)
 
 
-def _format_cost(cost: float | None) -> str:
-    """Format cost in USD."""
+def _format_cost(cost: float | None, currency: str = "USD") -> str:
+    """Format cost with currency symbol."""
     if cost is None:
         return "-"
+    symbol = "¥" if currency == "CNY" else "$"
     if cost < 0.01:
-        return f"${cost:.4f}"
-    return f"${cost:.2f}"
+        return f"{symbol}{cost:.4f}"
+    return f"{symbol}{cost:.2f}"
 
 
 def format_status_content(usage: model.Usage) -> str:
@@ -70,7 +78,7 @@ def format_status_content(usage: model.Usage) -> str:
     parts.append(f"Total: {_format_tokens(usage.total_tokens)}")
 
     if usage.total_cost is not None:
-        parts.append(f"Cost: {_format_cost(usage.total_cost)}")
+        parts.append(f"Cost: {_format_cost(usage.total_cost, usage.currency)}")
 
     return ", ".join(parts)
 

klaude_code/config/list_model.py

@@ -1,3 +1,5 @@
+import datetime
+
 from rich.console import Console, Group
 from rich.panel import Panel
 from rich.table import Table
@@ -8,6 +10,51 @@ from klaude_code.protocol.sub_agent import iter_sub_agent_profiles
 from klaude_code.ui.rich.theme import ThemeKey, get_theme
 
 
+def _display_codex_status(console: Console) -> None:
+    """Display Codex OAuth login status."""
+    from klaude_code.auth.codex.token_manager import CodexTokenManager
+
+    token_manager = CodexTokenManager()
+    state = token_manager.get_state()
+
+    if state is None:
+        console.print(
+            Text.assemble(
+                ("Codex Status: ", "bold"),
+                ("Not logged in", ThemeKey.CONFIG_STATUS_ERROR),
+                (" (run 'klaude login codex' to authenticate)", "dim"),
+            )
+        )
+    elif state.is_expired():
+        console.print(
+            Text.assemble(
+                ("Codex Status: ", "bold"),
+                ("Token expired", ThemeKey.CONFIG_STATUS_ERROR),
+                (" (run 'klaude login codex' to re-authenticate)", "dim"),
+            )
+        )
+    else:
+        expires_dt = datetime.datetime.fromtimestamp(state.expires_at, tz=datetime.timezone.utc)
+        console.print(
+            Text.assemble(
+                ("Codex Status: ", "bold"),
+                ("Logged in", ThemeKey.CONFIG_STATUS_OK),
+                (f" (account: {state.account_id[:8]}..., expires: {expires_dt.strftime('%Y-%m-%d %H:%M UTC')})", "dim"),
+            )
+        )
+
+    console.print(
+        Text.assemble(
+            ("Visit ", "dim"),
+            (
+                "https://chatgpt.com/codex/settings/usage",
+                "blue underline link https://chatgpt.com/codex/settings/usage",
+            ),
+            (" for up-to-date information on rate limits and credits", "dim"),
+        )
+    )
+
+
 def mask_api_key(api_key: str | None) -> str:
     """Mask API key to show only first 6 and last 6 characters with *** in between"""
     if not api_key or api_key == "N/A":
@@ -160,3 +207,9 @@ def display_models_and_providers(config: Config):
                 (sub_model_name, ThemeKey.CONFIG_STATUS_PRIMARY),
             )
         )
+
+    # Display Codex login status if any codex provider is configured
+    has_codex_provider = any(p.protocol.value == "codex" for p in config.provider_list)
+    if has_codex_provider:
+        console.print()
+        _display_codex_status(console)

klaude_code/core/prompt.py

@@ -13,7 +13,8 @@ COMMAND_DESCRIPTIONS: dict[str, str] = {
 
 # Mapping from logical prompt keys to resource file paths under the core/prompt directory.
 PROMPT_FILES: dict[str, str] = {
-    "main_codex": "prompts/prompt-codex.md",
+    "main_gpt_5_1": "prompts/prompt-codex-gpt-5-1.md",
+    "main_gpt_5_1_codex_max": "prompts/prompt-codex-gpt-5-1-codex-max.md",
     "main_claude": "prompts/prompt-claude-code.md",
     "main_gemini": "prompts/prompt-gemini.md",  # https://ai.google.dev/gemini-api/docs/prompting-strategies?hl=zh-cn#agentic-si-template
     # Sub-agent prompts keyed by their name
@@ -39,8 +40,10 @@ def get_system_prompt(model_name: str, sub_agent_type: str | None = None) -> str
 
     if sub_agent_type is None:
         match model_name:
+            case "gpt-5.1-codex-max":
+                file_key = "main_gpt_5_1_codex_max"
             case name if "gpt-5" in name:
-                file_key = "main_codex"
+                file_key = "main_gpt_5_1"
             case name if "gemini" in name:
                 file_key = "main_gemini"
             case _:
@@ -53,18 +56,11 @@ def get_system_prompt(model_name: str, sub_agent_type: str | None = None) -> str
     except KeyError as exc:
         raise ValueError(f"Unknown prompt key: {file_key}") from exc
 
-    base_prompt = (
-        files(__package__)
-        .joinpath(prompt_path)
-        .read_text(encoding="utf-8")
-        .format(
-            working_directory=str(cwd),
-            date=today,
-            is_git_repo=is_git_repo,
-            model_name=model_name,
-        )
-        .strip()
-    )
+    base_prompt = files(__package__).joinpath(prompt_path).read_text(encoding="utf-8").strip()
+
+    if model_name == "gpt-5.1-codex-max":
+        # Do not add env info for gpt-5.1-codex-max
+        return base_prompt
 
     env_lines: list[str] = [
         "",