klaude-code 2.8.1__py3-none-any.whl → 2.9.1__py3-none-any.whl

klaude_code/app/runtime.py

@@ -178,6 +178,7 @@ async def handle_keyboard_interrupt(executor: Executor) -> None:
     log("Bye!")
     session_id = executor.context.current_session_id()
     if session_id and Session.exists(session_id):
-        log(("Resume with:", "dim"), (f"klaude --resume {session_id}", "green"))
+        short_id = Session.shortest_unique_prefix(session_id)
+        log(("Resume with:", "dim"), (f"klaude -r {short_id}", "green"))
     with contextlib.suppress(Exception):
         await executor.submit(op.InterruptOperation(target_session_id=None))

klaude_code/auth/antigravity/oauth.py

@@ -258,42 +258,46 @@ class AntigravityOAuth:
         )
 
     def refresh(self) -> AntigravityAuthState:
-        """Refresh the access token using refresh token."""
-        state = self.token_manager.get_state()
-        if state is None:
-            raise AntigravityNotLoggedInError("Not logged in to Antigravity. Run 'klaude login antigravity' first.")
+        """Refresh the access token using refresh token with file locking.
 
-        data = {
-            "client_id": CLIENT_ID,
-            "client_secret": CLIENT_SECRET,
-            "refresh_token": state.refresh_token,
-            "grant_type": "refresh_token",
-        }
+        Uses file locking to prevent multiple instances from refreshing simultaneously.
+        If another instance has already refreshed, returns the updated state.
+        """
 
-        with httpx.Client() as client:
-            response = client.post(TOKEN_URL, data=data, timeout=30)
+        def do_refresh(current_state: AntigravityAuthState) -> AntigravityAuthState:
+            data = {
+                "client_id": CLIENT_ID,
+                "client_secret": CLIENT_SECRET,
+                "refresh_token": current_state.refresh_token,
+                "grant_type": "refresh_token",
+            }
 
-        if response.status_code != 200:
-            raise AntigravityTokenExpiredError(f"Token refresh failed: {response.text}")
+            with httpx.Client() as client:
+                response = client.post(TOKEN_URL, data=data, timeout=30)
 
-        tokens = response.json()
-        access_token = tokens["access_token"]
-        refresh_token = tokens.get("refresh_token", state.refresh_token)
-        expires_in = tokens.get("expires_in", 3600)
+            if response.status_code != 200:
+                raise AntigravityTokenExpiredError(f"Token refresh failed: {response.text}")
 
-        # Calculate expiry time with 5 minute buffer
-        expires_at = int(time.time()) + expires_in - 300
+            tokens = response.json()
+            access_token = tokens["access_token"]
+            refresh_token = tokens.get("refresh_token", current_state.refresh_token)
+            expires_in = tokens.get("expires_in", 3600)
 
-        new_state = AntigravityAuthState(
-            access_token=access_token,
-            refresh_token=refresh_token,
-            expires_at=expires_at,
-            project_id=state.project_id,
-            email=state.email,
-        )
+            # Calculate expiry time with 5 minute buffer
+            expires_at = int(time.time()) + expires_in - 300
 
-        self.token_manager.save(new_state)
-        return new_state
+            return AntigravityAuthState(
+                access_token=access_token,
+                refresh_token=refresh_token,
+                expires_at=expires_at,
+                project_id=current_state.project_id,
+                email=current_state.email,
+            )
+
+        try:
+            return self.token_manager.refresh_with_lock(do_refresh)
+        except ValueError as e:
+            raise AntigravityNotLoggedInError(str(e)) from e
 
     def ensure_valid_token(self) -> tuple[str, str]:
         """Ensure we have a valid access token, refreshing if needed.
@@ -309,12 +313,3 @@ class AntigravityOAuth:
             state = self.refresh()
 
         return state.access_token, state.project_id
-
-    def get_api_key_json(self) -> str:
-        """Get API key as JSON string for LLM client.
-
-        Returns:
-            JSON string with token and projectId.
-        """
-        access_token, project_id = self.ensure_valid_token()
-        return json.dumps({"token": access_token, "projectId": project_id})

klaude_code/auth/antigravity/token_manager.py

@@ -25,21 +25,3 @@ class AntigravityTokenManager(BaseTokenManager[AntigravityAuthState]):
 
     def _create_state(self, data: dict[str, Any]) -> AntigravityAuthState:
         return AntigravityAuthState.model_validate(data)
-
-    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.antigravity.exceptions import AntigravityNotLoggedInError
-
-            raise AntigravityNotLoggedInError("Not logged in to Antigravity. Run 'klaude login antigravity' first.")
-        return state.access_token
-
-    def get_project_id(self) -> str:
-        """Get project ID, raising if not logged in."""
-        state = self.get_state()
-        if state is None:
-            from klaude_code.auth.antigravity.exceptions import AntigravityNotLoggedInError
-
-            raise AntigravityNotLoggedInError("Not logged in to Antigravity. Run 'klaude login antigravity' first.")
-        return state.project_id

klaude_code/auth/base.py

@@ -3,12 +3,15 @@
 import json
 import time
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from pathlib import Path
 from typing import Any, cast
 
+from filelock import FileLock, Timeout
 from pydantic import BaseModel
 
 KLAUDE_AUTH_FILE = Path.home() / ".klaude" / "klaude-auth.json"
+LOCK_TIMEOUT_SECONDS = 30  # Maximum time to wait for lock acquisition
 
 
 class BaseAuthState(BaseModel):
@@ -99,3 +102,53 @@ class BaseTokenManager[T: BaseAuthState](ABC):
     def clear_cached_state(self) -> None:
         """Clear in-memory cached state to force reload from file on next access."""
         self._state = None
+
+    def _get_lock_file(self) -> Path:
+        """Get the lock file path for this auth file."""
+        return self.auth_file.with_suffix(".lock")
+
+    def refresh_with_lock(self, refresh_fn: Callable[[T], T]) -> T:
+        """Refresh token with file locking to prevent concurrent refresh.
+
+        This prevents multiple instances from simultaneously refreshing the same token.
+        If another instance has already refreshed, returns the updated state.
+
+        Args:
+            refresh_fn: Function that takes current state and returns new state.
+
+        Returns:
+            The new or already-refreshed authentication state.
+
+        Raises:
+            Timeout: If unable to acquire the lock within timeout.
+            ValueError: If not logged in.
+        """
+        lock_file = self._get_lock_file()
+        lock = FileLock(lock_file, timeout=LOCK_TIMEOUT_SECONDS)
+
+        try:
+            with lock:
+                # Re-read file after acquiring lock - another instance may have refreshed
+                self.clear_cached_state()
+                state = self.load()
+
+                if state is None:
+                    raise ValueError(f"Not logged in to {self.storage_key}")
+
+                # Check if token is still expired after re-reading
+                if not state.is_expired():
+                    # Another instance already refreshed, use their result
+                    return state
+
+                # Token still expired, we need to refresh
+                new_state = refresh_fn(state)
+                self.save(new_state)
+                return new_state
+
+        except Timeout:
+            # Lock timeout - try to re-read file in case another instance succeeded
+            self.clear_cached_state()
+            state = self.load()
+            if state and not state.is_expired():
+                return state
+            raise

klaude_code/auth/claude/oauth.py

@@ -125,60 +125,45 @@ class ClaudeOAuth:
             expires_at=int(time.time()) + int(expires_in),
         )
 
-    def _do_refresh_request(self, refresh_token: str) -> httpx.Response:
-        """Send token refresh request to OAuth server."""
-        payload = {
-            "grant_type": "refresh_token",
-            "client_id": CLIENT_ID,
-            "refresh_token": refresh_token,
-        }
-        with httpx.Client() as client:
-            return client.post(
-                TOKEN_URL,
-                json=payload,
-                headers={"Content-Type": "application/json"},
-            )
-
     def refresh(self) -> ClaudeAuthState:
-        """Refresh the access token using refresh token.
+        """Refresh the access token using refresh token with file locking.
 
-        Handles concurrent refresh race conditions by retrying with freshly loaded token
-        if the first attempt fails with invalid_grant error.
+        Uses file locking to prevent multiple instances from refreshing simultaneously.
+        If another instance has already refreshed, returns the updated state.
         """
-        state = self.token_manager.get_state()
-        if state is None:
-            raise ClaudeNotLoggedInError("Not logged in to Claude. Run 'klaude login claude' first.")
 
-        response = self._do_refresh_request(state.refresh_token)
-
-        # Handle race condition: another process may have refreshed the token already
-        if response.status_code != 200 and "invalid_grant" in response.text:
-            # Reload token from file (another process may have updated it)
-            self.token_manager.clear_cached_state()
-            fresh_state = self.token_manager.load()
-            if fresh_state and fresh_state.refresh_token != state.refresh_token:
-                # Token was updated by another process
-                if not fresh_state.is_expired():
-                    # New token is still valid, use it directly
-                    return fresh_state
-                # New token expired, try refreshing with the new refresh_token
-                response = self._do_refresh_request(fresh_state.refresh_token)
-
-        if response.status_code != 200:
-            raise ClaudeAuthError(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)
+        def do_refresh(current_state: ClaudeAuthState) -> ClaudeAuthState:
+            payload = {
+                "grant_type": "refresh_token",
+                "client_id": CLIENT_ID,
+                "refresh_token": current_state.refresh_token,
+            }
+
+            with httpx.Client() as client:
+                response = client.post(
+                    TOKEN_URL,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                )
+
+            if response.status_code != 200:
+                raise ClaudeAuthError(f"Token refresh failed: {response.text}")
+
+            tokens = response.json()
+            access_token = tokens["access_token"]
+            refresh_token = tokens.get("refresh_token", current_state.refresh_token)
+            expires_in = tokens.get("expires_in", 3600)
+
+            return ClaudeAuthState(
+                access_token=access_token,
+                refresh_token=refresh_token,
+                expires_at=int(time.time()) + int(expires_in),
+            )
 
-        new_state = ClaudeAuthState(
-            access_token=access_token,
-            refresh_token=refresh_token,
-            expires_at=int(time.time()) + int(expires_in),
-        )
-        self.token_manager.save(new_state)
-        return new_state
+        try:
+            return self.token_manager.refresh_with_lock(do_refresh)
+        except ValueError as e:
+            raise ClaudeNotLoggedInError(str(e)) from e
 
     def ensure_valid_token(self) -> str:
         """Ensure we have a valid access token, refreshing if needed."""

klaude_code/auth/codex/exceptions.py

@@ -15,7 +15,3 @@ class CodexTokenExpiredError(CodexAuthError):
 
 class CodexOAuthError(CodexAuthError):
     """OAuth flow failed."""
-
-
-class CodexUnsupportedModelError(CodexAuthError):
-    """Model is not supported by codex_oauth protocol."""

klaude_code/auth/codex/oauth.py

@@ -177,43 +177,47 @@ class CodexOAuth:
         )
 
     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
+        """Refresh the access token using refresh token with file locking.
 
-            raise CodexNotLoggedInError("Not logged in to Codex. Run 'klaude login codex' first.")
+        Uses file locking to prevent multiple instances from refreshing simultaneously.
+        If another instance has already refreshed, returns the updated state.
+        """
 
-        data = {
-            "grant_type": "refresh_token",
-            "client_id": CLIENT_ID,
-            "refresh_token": state.refresh_token,
-        }
+        def do_refresh(current_state: CodexAuthState) -> CodexAuthState:
+            data = {
+                "grant_type": "refresh_token",
+                "client_id": CLIENT_ID,
+                "refresh_token": current_state.refresh_token,
+            }
 
-        with httpx.Client() as client:
-            response = client.post(TOKEN_URL, data=data)
+            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
+            if response.status_code != 200:
+                from klaude_code.auth.codex.exceptions import CodexTokenExpiredError
 
-            raise CodexTokenExpiredError(f"Token refresh failed: {response.text}")
+                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)
+            tokens = response.json()
+            access_token = tokens["access_token"]
+            refresh_token = tokens.get("refresh_token", current_state.refresh_token)
+            expires_in = tokens.get("expires_in", 3600)
 
-        account_id = extract_account_id(access_token)
+            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,
-        )
+            return CodexAuthState(
+                access_token=access_token,
+                refresh_token=refresh_token,
+                expires_at=int(time.time()) + expires_in,
+                account_id=account_id,
+            )
+
+        try:
+            return self.token_manager.refresh_with_lock(do_refresh)
+        except ValueError as e:
+            from klaude_code.auth.codex.exceptions import CodexNotLoggedInError
 
-        self.token_manager.save(new_state)
-        return new_state
+            raise CodexNotLoggedInError(str(e)) from e
 
     def ensure_valid_token(self) -> str:
         """Ensure we have a valid access token, refreshing if needed."""

klaude_code/auth/codex/token_manager.py

@@ -24,21 +24,3 @@ class CodexTokenManager(BaseTokenManager[CodexAuthState]):
 
     def _create_state(self, data: dict[str, Any]) -> CodexAuthState:
         return CodexAuthState.model_validate(data)
-
-    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/cost_cmd.py

@@ -34,6 +34,16 @@ class ModelUsageStats:
     def total_tokens(self) -> int:
         return self.input_tokens + self.output_tokens
 
+    @property
+    def non_cached_input_tokens(self) -> int:
+        """Non-cached prompt tokens.
+
+        We store `input_tokens` as the provider-reported prompt token count, which
+        includes cached tokens for providers that support prompt caching.
+        """
+
+        return max(0, self.input_tokens - self.cached_tokens)
+
     def add_usage(self, usage: model.Usage) -> None:
         self.input_tokens += usage.input_tokens
         self.output_tokens += usage.output_tokens
@@ -48,41 +58,99 @@ class ModelUsageStats:
 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.
+@dataclass
+class SubProviderGroup:
+    """Group of models under a sub-provider."""
+
+    name: str
+    models: list[ModelUsageStats]
+    total: ModelUsageStats
+
+
+@dataclass
+class ProviderGroup:
+    """Group of models/sub-providers under a top-level provider."""
+
+    name: str
+    sub_providers: dict[str, SubProviderGroup]  # empty if no sub-providers
+    models: list[ModelUsageStats]  # direct models (when no sub-provider)
+    total: ModelUsageStats
+
 
-    Returns (models_by_provider, provider_totals) where both are sorted by cost desc.
+def _sort_by_cost(stats: ModelUsageStats) -> tuple[float, float]:
+    return (-stats.cost_usd, -stats.cost_cny)
+
+
+def group_models_by_provider(models: dict[ModelKey, ModelUsageStats]) -> dict[str, ProviderGroup]:
+    """Group models by provider with three-level hierarchy.
+
+    Provider strings like "openrouter/Anthropic" are split into:
+    - Top-level: "openrouter"
+    - Sub-provider: "Anthropic"
+
+    Returns dict of ProviderGroup sorted by cost desc.
     """
-    models_by_provider: dict[str, list[ModelUsageStats]] = {}
-    provider_totals: dict[str, ModelUsageStats] = {}
+    provider_groups: dict[str, ProviderGroup] = {}
 
     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
+        provider_raw = stats.provider or "(unknown)"
 
-    def sort_by_cost(stats: ModelUsageStats) -> tuple[float, float]:
-        return (-stats.cost_usd, -stats.cost_cny)
+        # Split provider by first "/"
+        if "/" in provider_raw:
+            parts = provider_raw.split("/", 1)
+            top_provider, sub_provider = parts[0], parts[1]
+        else:
+            top_provider, sub_provider = provider_raw, ""
+
+        # Initialize top-level provider group
+        if top_provider not in provider_groups:
+            provider_groups[top_provider] = ProviderGroup(
+                name=top_provider,
+                sub_providers={},
+                models=[],
+                total=ModelUsageStats(model_name=top_provider),
+            )
+
+        group = provider_groups[top_provider]
+
+        # Accumulate to top-level total
+        group.total.input_tokens += stats.input_tokens
+        group.total.output_tokens += stats.output_tokens
+        group.total.cached_tokens += stats.cached_tokens
+        group.total.cost_usd += stats.cost_usd
+        group.total.cost_cny += stats.cost_cny
+
+        if sub_provider:
+            # Has sub-provider, add to sub-provider group
+            if sub_provider not in group.sub_providers:
+                group.sub_providers[sub_provider] = SubProviderGroup(
+                    name=sub_provider,
+                    models=[],
+                    total=ModelUsageStats(model_name=sub_provider),
+                )
+            sub_group = group.sub_providers[sub_provider]
+            sub_group.models.append(stats)
+            sub_group.total.input_tokens += stats.input_tokens
+            sub_group.total.output_tokens += stats.output_tokens
+            sub_group.total.cached_tokens += stats.cached_tokens
+            sub_group.total.cost_usd += stats.cost_usd
+            sub_group.total.cost_cny += stats.cost_cny
+        else:
+            # No sub-provider, add directly to models
+            group.models.append(stats)
 
-    # 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)
+    # Sort everything by cost
+    for group in provider_groups.values():
+        group.models.sort(key=_sort_by_cost)
+        for sub_group in group.sub_providers.values():
+            sub_group.models.sort(key=_sort_by_cost)
+        # Sort sub-providers by cost
+        group.sub_providers = dict(sorted(group.sub_providers.items(), key=lambda x: _sort_by_cost(x[1].total)))
 
-    # 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}
+    # Sort top-level providers by cost
+    sorted_groups = dict(sorted(provider_groups.items(), key=lambda x: _sort_by_cost(x[1].total)))
 
-    return sorted_models_by_provider, sorted_provider_totals
+    return sorted_groups
 
 
 @dataclass
@@ -223,8 +291,8 @@ def render_cost_table(daily_stats: dict[str, DailyStats]) -> Table:
     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("Output", justify="right")
     table.add_column("Total", justify="right")
     table.add_column("USD", justify="right")
     table.add_column("CNY", justify="right")
@@ -248,9 +316,9 @@ def render_cost_table(daily_stats: dict[str, DailyStats]) -> Table:
         table.add_row(
             date_label,
             model_col,
-            fmt(format_tokens(stats.input_tokens)),
-            fmt(format_tokens(stats.output_tokens)),
+            fmt(format_tokens(stats.non_cached_input_tokens)),
             fmt(format_tokens(stats.cached_tokens)),
+            fmt(format_tokens(stats.output_tokens)),
             fmt(format_tokens(stats.total_tokens)),
             fmt(usd_str),
             fmt(cny_str),
@@ -261,19 +329,40 @@ def render_cost_table(daily_stats: dict[str, DailyStats]) -> Table:
         date_label: str = "",
         show_subtotal: bool = True,
     ) -> None:
-        """Render models grouped by provider with tree structure."""
-        models_by_provider, provider_totals = group_models_by_provider(models)
+        """Render models grouped by provider with three-level tree structure."""
+        provider_groups = group_models_by_provider(models)
 
         first_row = True
-        for provider_key, provider_models in models_by_provider.items():
-            provider_stats = provider_totals[provider_key]
-            add_stats_row(provider_stats, date_label=date_label if first_row else "", bold=True)
+        for group in provider_groups.values():
+            # Top-level provider
+            add_stats_row(group.total, date_label=date_label if first_row else "", bold=True)
             first_row = False
 
-            for i, stats in enumerate(provider_models):
-                is_last = i == len(provider_models) - 1
-                prefix = " └─ " if is_last else " ├─ "
-                add_stats_row(stats, prefix=prefix)
+            if group.sub_providers:
+                # Has sub-providers: render three-level tree
+                sub_list = list(group.sub_providers.values())
+                for sub_idx, sub_group in enumerate(sub_list):
+                    is_last_sub = sub_idx == len(sub_list) - 1
+                    sub_prefix = " └─ " if is_last_sub else " ├─ "
+
+                    # Sub-provider row
+                    add_stats_row(sub_group.total, prefix=sub_prefix, bold=True)
+
+                    # Models under sub-provider
+                    for model_idx, stats in enumerate(sub_group.models):
+                        is_last_model = model_idx == len(sub_group.models) - 1
+                        # Indent based on whether sub-provider is last
+                        if is_last_sub:
+                            model_prefix = "     └─ " if is_last_model else "     ├─ "
+                        else:
+                            model_prefix = " │   └─ " if is_last_model else " │   ├─ "
+                        add_stats_row(stats, prefix=model_prefix)
+            else:
+                # No sub-providers: render two-level tree (direct models)
+                for model_idx, stats in enumerate(group.models):
+                    is_last_model = model_idx == len(group.models) - 1
+                    model_prefix = " └─ " if is_last_model else " ├─ "
+                    add_stats_row(stats, prefix=model_prefix)
 
         if show_subtotal:
             subtotal = ModelUsageStats(model_name="(subtotal)")