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

klaude_code/cli/cost_cmd.py

@@ -45,12 +45,52 @@ class ModelUsageStats:
                 self.cost_usd += usage.total_cost
 
 
+ModelKey = tuple[str, str]  # (model_name, provider)
+
+
+def group_models_by_provider(
+    models: dict[ModelKey, ModelUsageStats],
+) -> tuple[dict[str, list[ModelUsageStats]], dict[str, ModelUsageStats]]:
+    """Group models by provider and compute provider totals.
+
+    Returns (models_by_provider, provider_totals) where both are sorted by cost desc.
+    """
+    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
+
+    def sort_by_cost(stats: ModelUsageStats) -> tuple[float, float]:
+        return (-stats.cost_usd, -stats.cost_cny)
+
+    # Sort providers by cost, and models within each provider
+    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)
+
+    # Rebuild dicts in sorted order
+    sorted_models_by_provider = {p: models_by_provider[p] for p in sorted_providers}
+    sorted_provider_totals = {p: provider_totals[p] for p in sorted_providers}
+
+    return sorted_models_by_provider, sorted_provider_totals
+
+
 @dataclass
 class DailyStats:
     """Aggregated stats for a single day."""
 
     date: str
-    by_model: dict[str, ModelUsageStats] = field(default_factory=lambda: dict[str, ModelUsageStats]())
+    by_model: dict[ModelKey, ModelUsageStats] = field(default_factory=lambda: dict[ModelKey, ModelUsageStats]())
 
     def add_task_metadata(self, meta: model.TaskMetadata, date_str: str) -> None:
         """Add a TaskMetadata to this day's stats."""
@@ -58,26 +98,14 @@ class DailyStats:
         if not meta.usage or not meta.model_name:
             return
 
-        model_key = meta.model_name
         provider = meta.provider or meta.usage.provider or ""
+        model_key: ModelKey = (meta.model_name, provider)
+
         if model_key not in self.by_model:
-            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] = ModelUsageStats(model_name=meta.model_name, provider=provider)
 
         self.by_model[model_key].add_usage(meta.usage)
 
-    def get_subtotal(self) -> ModelUsageStats:
-        """Get subtotal across all models for this day."""
-        subtotal = ModelUsageStats(model_name="(subtotal)")
-        for stats in self.by_model.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
-        return subtotal
-
 
 def iter_all_sessions() -> list[tuple[str, Path]]:
     """Iterate over all sessions across all projects.
@@ -201,78 +229,52 @@ def render_cost_table(daily_stats: dict[str, DailyStats]) -> Table:
     table.add_column("USD", justify="right")
     table.add_column("CNY", justify="right")
 
-    # Sort dates
     sorted_dates = sorted(daily_stats.keys())
+    global_by_model: dict[ModelKey, ModelUsageStats] = {}
+
+    def add_stats_row(stats: ModelUsageStats, date_label: str = "", prefix: str = "", bold: bool = False) -> None:
+        """Add a single stats row to the table."""
+        usd_str, cny_str = format_cost_dual(stats.cost_usd, stats.cost_cny)
+        if prefix:
+            model_col = f"[bright_black dim]{prefix}[/bright_black dim]{stats.model_name}"
+        elif bold:
+            model_col = f"[bold]{stats.model_name}[/bold]"
+        else:
+            model_col = stats.model_name
 
-    # Track global totals by (model, provider)
-    global_by_model: dict[tuple[str, str], ModelUsageStats] = {}
+        def fmt(val: str) -> str:
+            return f"[bold]{val}[/bold]" if bold else val
 
-    def sort_by_cost(stats: ModelUsageStats) -> tuple[float, float]:
-        """Sort key: USD desc, then CNY desc."""
-        return (-stats.cost_usd, -stats.cost_cny)
+        table.add_row(
+            date_label,
+            model_col,
+            fmt(format_tokens(stats.input_tokens)),
+            fmt(format_tokens(stats.output_tokens)),
+            fmt(format_tokens(stats.cached_tokens)),
+            fmt(format_tokens(stats.total_tokens)),
+            fmt(usd_str),
+            fmt(cny_str),
+        )
 
-    def render_by_provider(
-        models: dict[str, ModelUsageStats],
+    def render_grouped(
+        models: dict[ModelKey, 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)
+        models_by_provider, provider_totals = group_models_by_provider(models)
 
         first_row = True
-        for provider_key in sorted_providers:
+        for provider_key, provider_models in models_by_provider.items():
             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(
-                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]",
-            )
+            add_stats_row(provider_stats, date_label=date_label if first_row else "", bold=True)
             first_row = False
 
-            # 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
+                add_stats_row(stats, prefix=prefix)
+
         if show_subtotal:
             subtotal = ModelUsageStats(model_name="(subtotal)")
             for stats in models.values():
@@ -281,43 +283,29 @@ def render_cost_table(daily_stats: dict[str, DailyStats]) -> Table:
                 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]",
-            )
+            add_stats_row(subtotal, bold=True)
 
     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 "")
+        # Accumulate to global totals
+        for model_key, stats in day.by_model.items():
             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] = ModelUsageStats(model_name=stats.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))
+        render_grouped(day.by_model, date_label=format_date_display(date_str))
 
-        # Add separator between days
         if date_str != sorted_dates[-1]:
             table.add_section()
 
-    # Add final section for totals
+    # Total section
     table.add_section()
 
-    # Build date range label for Total
     if sorted_dates:
         first_date = format_date_display(sorted_dates[0])
         last_date = format_date_display(sorted_dates[-1])
@@ -328,64 +316,10 @@ def render_cost_table(daily_stats: dict[str, DailyStats]) -> Table:
     else:
         total_label = "[bold]Total[/bold]"
 
-    # 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_grouped(global_by_model, show_subtotal=False)
 
-    # 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"[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]",
-        )
-
-        # 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
     grand_total = ModelUsageStats(model_name="(total)")
     for stats in global_by_model.values():
         grand_total.input_tokens += stats.input_tokens
@@ -393,18 +327,7 @@ def render_cost_table(daily_stats: dict[str, DailyStats]) -> Table:
         grand_total.cached_tokens += stats.cached_tokens
         grand_total.cost_usd += stats.cost_usd
         grand_total.cost_cny += stats.cost_cny
-
-    usd_str, cny_str = format_cost_dual(grand_total.cost_usd, grand_total.cost_cny)
-    table.add_row(
-        "",
-        "[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]",
-        f"[bold]{format_tokens(grand_total.total_tokens)}[/bold]",
-        f"[bold]{usd_str}[/bold]",
-        f"[bold]{cny_str}[/bold]",
-    )
+    add_stats_row(grand_total, bold=True)
 
     return table
 

klaude_code/cli/list_model.py

@@ -121,6 +121,53 @@ def _get_claude_status_rows() -> list[tuple[Text, Text]]:
     return rows
 
 
+def _get_antigravity_status_rows() -> list[tuple[Text, Text]]:
+    """Get Antigravity OAuth login status as (label, value) tuples for table display."""
+    from klaude_code.auth.antigravity.token_manager import AntigravityTokenManager
+
+    rows: list[tuple[Text, Text]] = []
+    token_manager = AntigravityTokenManager()
+    state = token_manager.get_state()
+
+    if state is None:
+        rows.append(
+            (
+                Text("Status", style=ThemeKey.CONFIG_PARAM_LABEL),
+                Text.assemble(
+                    ("Not logged in", ThemeKey.CONFIG_STATUS_ERROR),
+                    (" (run 'klaude login antigravity' to authenticate)", "dim"),
+                ),
+            )
+        )
+    elif state.is_expired():
+        rows.append(
+            (
+                Text("Status", style=ThemeKey.CONFIG_PARAM_LABEL),
+                Text.assemble(
+                    ("Token expired", ThemeKey.CONFIG_STATUS_ERROR),
+                    (" (will refresh automatically on use; run 'klaude login antigravity' if refresh fails)", "dim"),
+                ),
+            )
+        )
+    else:
+        expires_dt = datetime.datetime.fromtimestamp(state.expires_at, tz=datetime.UTC)
+        email_info = f", email: {state.email}" if state.email else ""
+        rows.append(
+            (
+                Text("Status", style=ThemeKey.CONFIG_PARAM_LABEL),
+                Text.assemble(
+                    ("Logged in", ThemeKey.CONFIG_STATUS_OK),
+                    (
+                        f" (project: {state.project_id}{email_info}, expires: {expires_dt.strftime('%Y-%m-%d %H:%M UTC')})",
+                        "dim",
+                    ),
+                ),
+            )
+        )
+
+    return rows
+
+
 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:
@@ -234,6 +281,9 @@ def _build_provider_info_panel(provider: ProviderConfig, available: bool) -> Quo
     if provider.protocol == LLMClientProtocol.CLAUDE_OAUTH:
         for label, value in _get_claude_status_rows():
             info_table.add_row(label, value)
+    if provider.protocol == LLMClientProtocol.ANTIGRAVITY:
+        for label, value in _get_antigravity_status_rows():
+            info_table.add_row(label, value)
 
     return Quote(
         Group(title, info_table),

klaude_code/cli/main.py

@@ -1,7 +1,10 @@
 import asyncio
 import sys
+from collections.abc import Sequence
+from typing import Any
 
 import typer
+from typer.core import TyperGroup
 
 from klaude_code.cli.auth_cmd import register_auth_commands
 from klaude_code.cli.config_cmd import register_config_commands
@@ -40,7 +43,88 @@ def _build_env_help() -> str:
 
 ENV_HELP = _build_env_help()
 
+
+def _looks_like_flag(token: str) -> bool:
+    return token.startswith("-") and token != "-"
+
+
+def _preprocess_cli_args(args: list[str]) -> list[str]:
+    """Rewrite CLI args to support optional values for selected options.
+
+    Supported rewrites:
+    - --model / -m with no value -> --model-select
+    - --resume / -r with value -> --resume-by-id <value>
+    """
+
+    rewritten: list[str] = []
+    i = 0
+    while i < len(args):
+        token = args[i]
+
+        if token in {"--model", "-m"}:
+            next_token = args[i + 1] if i + 1 < len(args) else None
+            if next_token is None or next_token == "--" or _looks_like_flag(next_token):
+                rewritten.append("--model-select")
+                i += 1
+                continue
+            rewritten.append(token)
+            i += 1
+            continue
+
+        if token.startswith("--model="):
+            value = token.split("=", 1)[1]
+            if value == "":
+                rewritten.append("--model-select")
+            else:
+                rewritten.append(token)
+            i += 1
+            continue
+
+        if token in {"--resume", "-r"}:
+            next_token = args[i + 1] if i + 1 < len(args) else None
+            if next_token is not None and next_token != "--" and not _looks_like_flag(next_token):
+                rewritten.extend(["--resume-by-id", next_token])
+                i += 2
+                continue
+            rewritten.append(token)
+            i += 1
+            continue
+
+        if token.startswith("--resume="):
+            value = token.split("=", 1)[1]
+            rewritten.extend(["--resume-by-id", value])
+            i += 1
+            continue
+
+        rewritten.append(token)
+        i += 1
+
+    return rewritten
+
+
+class _PreprocessingTyperGroup(TyperGroup):
+    def main(
+        self,
+        args: Sequence[str] | None = None,
+        prog_name: str | None = None,
+        complete_var: str | None = None,
+        standalone_mode: bool = True,
+        windows_expand_args: bool = True,
+        **extra: Any,
+    ) -> Any:
+        click_args = _preprocess_cli_args(list(args) if args is not None else sys.argv[1:])
+        return super().main(
+            args=click_args,
+            prog_name=prog_name,
+            complete_var=complete_var,
+            standalone_mode=standalone_mode,
+            windows_expand_args=windows_expand_args,
+            **extra,
+        )
+
+
 app = typer.Typer(
+    cls=_PreprocessingTyperGroup,
     add_completion=False,
     pretty_exceptions_enable=False,
     no_args_is_help=False,
@@ -69,21 +153,27 @@ def main_callback(
         None,
         "--model",
         "-m",
-        help="Select model by name",
+        help="Select model by name; use --model with no value to choose interactively",
         rich_help_panel="LLM",
     ),
     continue_: bool = typer.Option(False, "--continue", "-c", help="Resume latest session"),
-    resume: bool = typer.Option(False, "--resume", "-r", help="Pick a session to resume"),
+    resume: bool = typer.Option(
+        False,
+        "--resume",
+        "-r",
+        help="Resume a session; use --resume <id> to resume directly, or --resume to pick interactively",
+    ),
     resume_by_id: str | None = typer.Option(
         None,
         "--resume-by-id",
         help="Resume session by ID",
+        hidden=True,
     ),
     select_model: bool = typer.Option(
         False,
-        "--select-model",
-        "-s",
-        help="Choose model interactively",
+        "--model-select",
+        help="Choose model interactively (same as --model with no value)",
+        hidden=True,
         rich_help_panel="LLM",
     ),
     debug: bool = typer.Option(
@@ -107,7 +197,7 @@ def main_callback(
     banana: bool = typer.Option(
         False,
         "--banana",
-        help="Image generation mode",
+        help="Image generation mode (alias for --model banana)",
         rich_help_panel="LLM",
     ),
     version: bool = typer.Option(
@@ -130,11 +220,11 @@ def main_callback(
 
         resume_by_id_value = resume_by_id.strip() if resume_by_id is not None else None
         if resume_by_id_value == "":
-            log(("Error: --resume-by-id cannot be empty", "red"))
+            log(("Error: --resume <id> cannot be empty", "red"))
             raise typer.Exit(2)
 
         if resume_by_id_value is not None and (resume or continue_):
-            log(("Error: --resume-by-id cannot be combined with --resume/--continue", "red"))
+            log(("Error: --resume <id> cannot be combined with --continue or interactive --resume", "red"))
             raise typer.Exit(2)
 
         if resume_by_id_value is not None and not Session.exists(resume_by_id_value):
@@ -250,7 +340,7 @@ def main_callback(
         )
 
         if log_path:
-            log(f"Debug log: {log_path}", style="dim")
+            log(f"Debug log: {log_path}", style="red")
 
         asyncio.run(
             run_interactive(