klaude-code 2.5.2__py3-none-any.whl → 2.6.0__py3-none-any.whl

klaude_code/auth/__init__.py

@@ -12,6 +12,12 @@ from klaude_code.auth.codex import (
     CodexTokenExpiredError,
     CodexTokenManager,
 )
+from klaude_code.auth.env import (
+    delete_auth_env,
+    get_auth_env,
+    list_auth_env,
+    set_auth_env,
+)
 
 __all__ = [
     "CodexAuthError",
@@ -21,4 +27,8 @@ __all__ = [
     "CodexOAuthError",
     "CodexTokenExpiredError",
     "CodexTokenManager",
+    "delete_auth_env",
+    "get_auth_env",
+    "list_auth_env",
+    "set_auth_env",
 ]

klaude_code/auth/env.py

@@ -0,0 +1,77 @@
+"""Environment variable configuration stored in klaude-auth.json."""
+
+import json
+from typing import Any
+
+from klaude_code.auth.base import KLAUDE_AUTH_FILE
+
+
+def _load_store() -> dict[str, Any]:
+    """Load the auth store from file."""
+    if not KLAUDE_AUTH_FILE.exists():
+        return {}
+    try:
+        data: Any = json.loads(KLAUDE_AUTH_FILE.read_text())
+        if isinstance(data, dict):
+            return dict(data)
+        return {}
+    except (json.JSONDecodeError, ValueError):
+        return {}
+
+
+def _save_store(data: dict[str, Any]) -> None:
+    """Save the auth store to file."""
+    KLAUDE_AUTH_FILE.parent.mkdir(parents=True, exist_ok=True)
+    KLAUDE_AUTH_FILE.write_text(json.dumps(data, indent=2))
+
+
+def get_auth_env(env_var: str) -> str | None:
+    """Get environment variable value from klaude-auth.json 'env' section.
+
+    This provides a fallback for API keys when real environment variables are not set.
+    Priority: os.environ > klaude-auth.json env
+    """
+    store = _load_store()
+    env_section: Any = store.get("env")
+    if not isinstance(env_section, dict):
+        return None
+    value: Any = env_section.get(env_var)
+    return str(value) if value is not None else None
+
+
+def set_auth_env(env_var: str, value: str) -> None:
+    """Set environment variable value in klaude-auth.json 'env' section."""
+    store = _load_store()
+    env_section: Any = store.get("env")
+    if not isinstance(env_section, dict):
+        env_section = {}
+    env_section[env_var] = value
+    store["env"] = env_section
+    _save_store(store)
+
+
+def delete_auth_env(env_var: str) -> None:
+    """Delete environment variable from klaude-auth.json 'env' section."""
+    store = _load_store()
+    env_section: Any = store.get("env")
+    if not isinstance(env_section, dict):
+        return
+    env_section.pop(env_var, None)
+    if len(env_section) == 0:
+        store.pop("env", None)
+    else:
+        store["env"] = env_section
+    if len(store) == 0:
+        if KLAUDE_AUTH_FILE.exists():
+            KLAUDE_AUTH_FILE.unlink()
+    else:
+        _save_store(store)
+
+
+def list_auth_env() -> dict[str, str]:
+    """List all environment variables in klaude-auth.json 'env' section."""
+    store = _load_store()
+    env_section: Any = store.get("env")
+    if not isinstance(env_section, dict):
+        return {}
+    return {k: str(v) for k, v in env_section.items() if v is not None}

klaude_code/cli/auth_cmd.py

@@ -4,41 +4,81 @@ import datetime
 import webbrowser
 
 import typer
-from prompt_toolkit.styles import Style
 
 from klaude_code.log import log
-from klaude_code.tui.terminal.selector import SelectItem, select_one
-
-_SELECT_STYLE = Style(
-    [
-        ("instruction", "ansibrightblack"),
-        ("pointer", "ansigreen"),
-        ("highlighted", "ansigreen"),
-        ("text", "ansibrightblack"),
-        ("question", "bold"),
-    ]
-)
+from klaude_code.tui.terminal.selector import DEFAULT_PICKER_STYLE, SelectItem, select_one
 
 
 def _select_provider() -> str | None:
     """Display provider selection menu and return selected provider."""
+    from klaude_code.config.builtin_config import SUPPORTED_API_KEYS
+
     items: list[SelectItem[str]] = [
-        SelectItem(title=[("class:text", "Claude Max/Pro Subscription\n")], value="claude", search_text="claude"),
-        SelectItem(title=[("class:text", "ChatGPT Codex Subscription\n")], value="codex", search_text="codex"),
+        SelectItem(
+            title=[("", "Claude Max/Pro Subscription "), ("ansibrightblack", "[OAuth]\n")],
+            value="claude",
+            search_text="claude",
+        ),
+        SelectItem(
+            title=[("", "ChatGPT Codex Subscription "), ("ansibrightblack", "[OAuth]\n")],
+            value="codex",
+            search_text="codex",
+        ),
     ]
+    # Add API key options
+    for key_info in SUPPORTED_API_KEYS:
+        items.append(
+            SelectItem(
+                title=[("", f"{key_info.name} "), ("ansibrightblack", "[API key]\n")],
+                value=key_info.env_var,
+                search_text=key_info.env_var,
+            )
+        )
+
     return select_one(
         message="Select provider to login:",
         items=items,
         pointer="→",
-        style=_SELECT_STYLE,
+        style=DEFAULT_PICKER_STYLE,
         use_search_filter=False,
     )
 
 
+def _configure_api_key(env_var: str) -> None:
+    """Configure a specific API key."""
+    import os
+
+    from klaude_code.auth.env import get_auth_env, set_auth_env
+
+    # Check if already configured
+    current_value = os.environ.get(env_var) or get_auth_env(env_var)
+    if current_value:
+        masked = current_value[:8] + "..." if len(current_value) > 8 else "***"
+        log(f"Current {env_var}: {masked}")
+        if not typer.confirm("Do you want to update it?"):
+            return
+
+    api_key = typer.prompt(f"Enter {env_var}", hide_input=True)
+    if not api_key.strip():
+        log(("Error: API key cannot be empty", "red"))
+        raise typer.Exit(1)
+
+    set_auth_env(env_var, api_key.strip())
+    log((f"{env_var} saved successfully!", "green"))
+
+
+def _build_provider_help() -> str:
+    from klaude_code.config.builtin_config import SUPPORTED_API_KEYS
+
+    # Use first word of name for brevity (e.g., "google" instead of "google gemini")
+    names = ["codex", "claude"] + [k.name.split()[0].lower() for k in SUPPORTED_API_KEYS]
+    return f"Provider name ({', '.join(names)})"
+
+
 def login_command(
-    provider: str | None = typer.Argument(None, help="Provider to login (codex|claude)"),
+    provider: str | None = typer.Argument(None, help=_build_provider_help()),
 ) -> None:
-    """Login to a provider using OAuth."""
+    """Login to a provider or configure API keys."""
     if provider is None:
         provider = _select_provider()
         if provider is None:
@@ -110,8 +150,27 @@ def login_command(
                 log((f"Login failed: {e}", "red"))
                 raise typer.Exit(1) from None
         case _:
-            log((f"Error: Unknown provider '{provider}'. Supported: codex, claude", "red"))
-            raise typer.Exit(1)
+            from klaude_code.config.builtin_config import SUPPORTED_API_KEYS
+
+            # Match by env var (e.g., OPENAI_API_KEY) or name (e.g., openai, google)
+            env_var: str | None = None
+            provider_lower = provider.lower()
+            provider_upper = provider.upper()
+            for key_info in SUPPORTED_API_KEYS:
+                name_lower = key_info.name.lower()
+                # Exact match or starts with (for "google" -> "google gemini")
+                if key_info.env_var == provider_upper or name_lower == provider_lower:
+                    env_var = key_info.env_var
+                    break
+                if name_lower.startswith(provider_lower) or provider_lower in name_lower.split():
+                    env_var = key_info.env_var
+                    break
+
+            if env_var:
+                _configure_api_key(env_var)
+            else:
+                log((f"Error: Unknown provider '{provider}'", "red"))
+                raise typer.Exit(1)
 
 
 def logout_command(
@@ -150,5 +209,14 @@ def logout_command(
 
 def register_auth_commands(app: typer.Typer) -> None:
     """Register auth commands to the given Typer app."""
-    app.command("login")(login_command)
-    app.command("logout")(logout_command)
+    auth_app = typer.Typer(help="Login/logout", invoke_without_command=True)
+
+    @auth_app.callback()
+    def auth_callback(ctx: typer.Context) -> None:
+        """Authentication commands for managing provider logins."""
+        if ctx.invoked_subcommand is None:
+            typer.echo(ctx.get_help())
+
+    auth_app.command("login")(login_command)
+    auth_app.command("logout")(logout_command)
+    app.add_typer(auth_app, name="auth")

klaude_code/cli/config_cmd.py

@@ -11,9 +11,9 @@ from klaude_code.log import log
 
 
 def list_models(
-    show_all: bool = typer.Option(False, "--all", "-a", help="Show all providers including unavailable ones"),
+    show_all: bool = typer.Option(False, "--all", "-a", help="Include unavailable providers"),
 ) -> None:
-    """List all models and providers configuration"""
+    """List available models"""
     from klaude_code.cli.list_model import display_models_and_providers
     from klaude_code.tui.terminal.color import is_light_terminal_background
 
@@ -31,7 +31,7 @@ def list_models(
 
 
 def edit_config() -> None:
-    """Open the configuration file in $EDITOR or default system editor"""
+    """Edit config file"""
     editor = os.environ.get("EDITOR")
 
     # If no EDITOR is set, prioritize TextEdit on macOS
@@ -89,5 +89,5 @@ def edit_config() -> None:
 def register_config_commands(app: typer.Typer) -> None:
     """Register config commands to the given Typer app."""
     app.command("list")(list_models)
-    app.command("config")(edit_config)
-    app.command("conf", hidden=True)(edit_config)
+    app.command("conf")(edit_config)
+    app.command("config", hidden=True)(edit_config)

klaude_code/cli/cost_cmd.py

@@ -5,6 +5,7 @@ from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
 
+import pydantic
 import typer
 from rich.box import Box
 from rich.console import Console
@@ -22,6 +23,7 @@ class ModelUsageStats:
     """Aggregated usage stats for a single model."""
 
     model_name: str
+    provider: str = ""
     input_tokens: int = 0
     output_tokens: int = 0
     cached_tokens: int = 0
@@ -57,8 +59,11 @@ class DailyStats:
             return
 
         model_key = meta.model_name
+        provider = meta.provider or meta.usage.provider or ""
         if model_key not in self.by_model:
-            self.by_model[model_key] = ModelUsageStats(model_name=model_key)
+            self.by_model[model_key] = ModelUsageStats(model_name=model_key, provider=provider)
+        elif not self.by_model[model_key].provider and provider:
+            self.by_model[model_key].provider = provider
 
         self.by_model[model_key].add_usage(meta.usage)
 
@@ -115,6 +120,7 @@ def extract_task_metadata_from_events(events_path: Path) -> list[tuple[str, mode
     """Extract TaskMetadataItem entries from events.jsonl with their dates.
 
     Returns list of (date_str, TaskMetadataItem) tuples.
+    Skips lines that fail pydantic validation.
     """
     results: list[tuple[str, model.TaskMetadataItem]] = []
     try:
@@ -123,7 +129,10 @@ def extract_task_metadata_from_events(events_path: Path) -> list[tuple[str, mode
         return results
 
     for line in content.splitlines():
-        item = decode_jsonl_line(line)
+        try:
+            item = decode_jsonl_line(line)
+        except pydantic.ValidationError:
+            continue
         if isinstance(item, model.TaskMetadataItem):
             date_str = item.created_at.strftime("%Y-%m-%d")
             results.append((date_str, item))
@@ -183,68 +192,123 @@ def render_cost_table(daily_stats: dict[str, DailyStats]) -> Table:
         box=ASCII_HORIZONAL,
     )
 
-    table.add_column("Date", style="cyan", no_wrap=True)
-    table.add_column("Model", no_wrap=True)
-    table.add_column("Input", justify="right", no_wrap=True)
-    table.add_column("Output", justify="right", no_wrap=True)
-    table.add_column("Cache", justify="right", no_wrap=True)
-    table.add_column("Total", justify="right", no_wrap=True)
-    table.add_column("USD", justify="right", no_wrap=True)
-    table.add_column("CNY", justify="right", no_wrap=True)
+    table.add_column("Date", style="cyan")
+    table.add_column("Model", overflow="ellipsis")
+    table.add_column("Input", justify="right")
+    table.add_column("Output", justify="right")
+    table.add_column("Cache", justify="right")
+    table.add_column("Total", justify="right")
+    table.add_column("USD", justify="right")
+    table.add_column("CNY", justify="right")
 
     # Sort dates
     sorted_dates = sorted(daily_stats.keys())
 
-    # Track global totals by model
-    global_by_model: dict[str, ModelUsageStats] = {}
+    # Track global totals by (model, provider)
+    global_by_model: dict[tuple[str, str], ModelUsageStats] = {}
 
     def sort_by_cost(stats: ModelUsageStats) -> tuple[float, float]:
         """Sort key: USD desc, then CNY desc."""
         return (-stats.cost_usd, -stats.cost_cny)
 
-    for date_str in sorted_dates:
-        day = daily_stats[date_str]
-        sorted_models = [s.model_name for s in sorted(day.by_model.values(), key=sort_by_cost)]
+    def render_by_provider(
+        models: dict[str, ModelUsageStats],
+        date_label: str = "",
+        show_subtotal: bool = True,
+    ) -> None:
+        """Render models grouped by provider with tree structure."""
+        # Group models by provider
+        models_by_provider: dict[str, list[ModelUsageStats]] = {}
+        provider_totals: dict[str, ModelUsageStats] = {}
+        for stats in models.values():
+            provider_key = stats.provider or "(unknown)"
+            if provider_key not in models_by_provider:
+                models_by_provider[provider_key] = []
+                provider_totals[provider_key] = ModelUsageStats(model_name=provider_key, provider=provider_key)
+            models_by_provider[provider_key].append(stats)
+            provider_totals[provider_key].input_tokens += stats.input_tokens
+            provider_totals[provider_key].output_tokens += stats.output_tokens
+            provider_totals[provider_key].cached_tokens += stats.cached_tokens
+            provider_totals[provider_key].cost_usd += stats.cost_usd
+            provider_totals[provider_key].cost_cny += stats.cost_cny
+
+        # Sort providers by cost, and models within each provider by cost
+        sorted_providers = sorted(provider_totals.keys(), key=lambda p: sort_by_cost(provider_totals[p]))
+        for provider_key in models_by_provider:
+            models_by_provider[provider_key].sort(key=sort_by_cost)
 
         first_row = True
-        for model_name in sorted_models:
-            stats = day.by_model[model_name]
-            usd_str, cny_str = format_cost_dual(stats.cost_usd, stats.cost_cny)
-
-            # Accumulate to global totals
-            if model_name not in global_by_model:
-                global_by_model[model_name] = ModelUsageStats(model_name=model_name)
-            global_by_model[model_name].input_tokens += stats.input_tokens
-            global_by_model[model_name].output_tokens += stats.output_tokens
-            global_by_model[model_name].cached_tokens += stats.cached_tokens
-            global_by_model[model_name].cost_usd += stats.cost_usd
-            global_by_model[model_name].cost_cny += stats.cost_cny
+        for provider_key in sorted_providers:
+            provider_stats = provider_totals[provider_key]
+            provider_models = models_by_provider[provider_key]
 
+            # Provider row (bold)
+            usd_str, cny_str = format_cost_dual(provider_stats.cost_usd, provider_stats.cost_cny)
             table.add_row(
-                format_date_display(date_str) if first_row else "",
-                f"- {model_name}",
-                format_tokens(stats.input_tokens),
-                format_tokens(stats.output_tokens),
-                format_tokens(stats.cached_tokens),
-                format_tokens(stats.total_tokens),
-                usd_str,
-                cny_str,
+                date_label if first_row else "",
+                f"[bold]{provider_key}[/bold]",
+                f"[bold]{format_tokens(provider_stats.input_tokens)}[/bold]",
+                f"[bold]{format_tokens(provider_stats.output_tokens)}[/bold]",
+                f"[bold]{format_tokens(provider_stats.cached_tokens)}[/bold]",
+                f"[bold]{format_tokens(provider_stats.total_tokens)}[/bold]",
+                f"[bold]{usd_str}[/bold]",
+                f"[bold]{cny_str}[/bold]",
             )
             first_row = False
 
-        # Add subtotal row for this day
-        subtotal = day.get_subtotal()
-        usd_str, cny_str = format_cost_dual(subtotal.cost_usd, subtotal.cost_cny)
-        table.add_row(
-            "",
-            "[cyan]  (subtotal)[/cyan]",
-            f"[cyan]{format_tokens(subtotal.input_tokens)}[/cyan]",
-            f"[cyan]{format_tokens(subtotal.output_tokens)}[/cyan]",
-            f"[cyan]{format_tokens(subtotal.cached_tokens)}[/cyan]",
-            f"[cyan]{format_tokens(subtotal.total_tokens)}[/cyan]",
-            f"[cyan]{usd_str}[/cyan]",
-            f"[cyan]{cny_str}[/cyan]",
-        )
+            # Model rows with tree prefix
+            for i, stats in enumerate(provider_models):
+                is_last = i == len(provider_models) - 1
+                prefix = " └─ " if is_last else " ├─ "
+                usd_str, cny_str = format_cost_dual(stats.cost_usd, stats.cost_cny)
+                table.add_row(
+                    "",
+                    f"[bright_black dim]{prefix}[/bright_black dim]{stats.model_name}",
+                    format_tokens(stats.input_tokens),
+                    format_tokens(stats.output_tokens),
+                    format_tokens(stats.cached_tokens),
+                    format_tokens(stats.total_tokens),
+                    usd_str,
+                    cny_str,
+                )
+
+        # Add subtotal row
+        if show_subtotal:
+            subtotal = ModelUsageStats(model_name="(subtotal)")
+            for stats in models.values():
+                subtotal.input_tokens += stats.input_tokens
+                subtotal.output_tokens += stats.output_tokens
+                subtotal.cached_tokens += stats.cached_tokens
+                subtotal.cost_usd += stats.cost_usd
+                subtotal.cost_cny += stats.cost_cny
+            usd_str, cny_str = format_cost_dual(subtotal.cost_usd, subtotal.cost_cny)
+            table.add_row(
+                "",
+                "[bold](subtotal)[/bold]",
+                f"[bold]{format_tokens(subtotal.input_tokens)}[/bold]",
+                f"[bold]{format_tokens(subtotal.output_tokens)}[/bold]",
+                f"[bold]{format_tokens(subtotal.cached_tokens)}[/bold]",
+                f"[bold]{format_tokens(subtotal.total_tokens)}[/bold]",
+                f"[bold]{usd_str}[/bold]",
+                f"[bold]{cny_str}[/bold]",
+            )
+
+    for date_str in sorted_dates:
+        day = daily_stats[date_str]
+
+        # Accumulate to global totals by (model, provider)
+        for model_name, stats in day.by_model.items():
+            model_key = (model_name, stats.provider or "")
+            if model_key not in global_by_model:
+                global_by_model[model_key] = ModelUsageStats(model_name=model_name, provider=stats.provider)
+            global_by_model[model_key].input_tokens += stats.input_tokens
+            global_by_model[model_key].output_tokens += stats.output_tokens
+            global_by_model[model_key].cached_tokens += stats.cached_tokens
+            global_by_model[model_key].cost_usd += stats.cost_usd
+            global_by_model[model_key].cost_cny += stats.cost_cny
+
+        # Render this day's data grouped by provider
+        render_by_provider(day.by_model, date_label=format_date_display(date_str))
 
         # Add separator between days
         if date_str != sorted_dates[-1]:
@@ -264,27 +328,62 @@ def render_cost_table(daily_stats: dict[str, DailyStats]) -> Table:
     else:
         total_label = "[bold]Total[/bold]"
 
-    # Add per-model totals
-    sorted_global_models = [s.model_name for s in sorted(global_by_model.values(), key=sort_by_cost)]
-    first_total_row = True
-    for model_name in sorted_global_models:
-        # Add empty row before first model to align with Total date range
-        if first_total_row:
-            table.add_row(total_label, "", "", "", "", "", "", "")
-            first_total_row = False
-        stats = global_by_model[model_name]
-        usd_str, cny_str = format_cost_dual(stats.cost_usd, stats.cost_cny)
+    # Group models by provider
+    models_by_provider: dict[str, list[ModelUsageStats]] = {}
+    provider_totals: dict[str, ModelUsageStats] = {}
+    for stats in global_by_model.values():
+        provider_key = stats.provider or "(unknown)"
+        if provider_key not in models_by_provider:
+            models_by_provider[provider_key] = []
+            provider_totals[provider_key] = ModelUsageStats(model_name=provider_key, provider=provider_key)
+        models_by_provider[provider_key].append(stats)
+        provider_totals[provider_key].input_tokens += stats.input_tokens
+        provider_totals[provider_key].output_tokens += stats.output_tokens
+        provider_totals[provider_key].cached_tokens += stats.cached_tokens
+        provider_totals[provider_key].cost_usd += stats.cost_usd
+        provider_totals[provider_key].cost_cny += stats.cost_cny
+
+    # Sort providers by cost, and models within each provider by cost
+    sorted_providers = sorted(provider_totals.keys(), key=lambda p: sort_by_cost(provider_totals[p]))
+    for provider_key in models_by_provider:
+        models_by_provider[provider_key].sort(key=sort_by_cost)
+
+    # Add total label row
+    table.add_row(total_label, "", "", "", "", "", "", "")
+
+    # Render each provider with its models
+    for provider_key in sorted_providers:
+        provider_stats = provider_totals[provider_key]
+        models = models_by_provider[provider_key]
+
+        # Provider row (bold)
+        usd_str, cny_str = format_cost_dual(provider_stats.cost_usd, provider_stats.cost_cny)
         table.add_row(
             "",
-            f"- {model_name}",
-            format_tokens(stats.input_tokens),
-            format_tokens(stats.output_tokens),
-            format_tokens(stats.cached_tokens),
-            format_tokens(stats.total_tokens),
-            usd_str,
-            cny_str,
+            f"[bold]{provider_key}[/bold]",
+            f"[bold]{format_tokens(provider_stats.input_tokens)}[/bold]",
+            f"[bold]{format_tokens(provider_stats.output_tokens)}[/bold]",
+            f"[bold]{format_tokens(provider_stats.cached_tokens)}[/bold]",
+            f"[bold]{format_tokens(provider_stats.total_tokens)}[/bold]",
+            f"[bold]{usd_str}[/bold]",
+            f"[bold]{cny_str}[/bold]",
         )
-        first_total_row = False
+
+        # Model rows with tree prefix
+        for i, stats in enumerate(models):
+            is_last = i == len(models) - 1
+            prefix = " └─ " if is_last else " ├─ "
+            usd_str, cny_str = format_cost_dual(stats.cost_usd, stats.cost_cny)
+            table.add_row(
+                "",
+                f"[bright_black dim]{prefix}[/bright_black dim]{stats.model_name}",
+                format_tokens(stats.input_tokens),
+                format_tokens(stats.output_tokens),
+                format_tokens(stats.cached_tokens),
+                format_tokens(stats.total_tokens),
+                usd_str,
+                cny_str,
+            )
 
     # Add grand total row
     grand_total = ModelUsageStats(model_name="(total)")
@@ -298,7 +397,7 @@ def render_cost_table(daily_stats: dict[str, DailyStats]) -> Table:
     usd_str, cny_str = format_cost_dual(grand_total.cost_usd, grand_total.cost_cny)
     table.add_row(
         "",
-        "[bold]  (total)[/bold]",
+        "[bold](total)[/bold]",
         f"[bold]{format_tokens(grand_total.input_tokens)}[/bold]",
         f"[bold]{format_tokens(grand_total.output_tokens)}[/bold]",
         f"[bold]{format_tokens(grand_total.cached_tokens)}[/bold]",
@@ -313,7 +412,7 @@ def render_cost_table(daily_stats: dict[str, DailyStats]) -> Table:
 def cost_command(
     days: int | None = typer.Option(None, "--days", "-d", help="Limit to last N days"),
 ) -> None:
-    """Display aggregated usage statistics across all sessions."""
+    """Show usage stats"""
     daily_stats = aggregate_all_sessions()
 
     if not daily_stats: