diff --git a/CHANGELOG.md b/CHANGELOG.md index bda2ea4f..e49f99cf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,10 @@ GitHub Releases page; `0.8.0` is the new starting line. ## Unreleased +- **`/stats` usage dashboard.** New slash command opens an interactive prompt_toolkit TUI showing token and cost breakdown by provider/model across Today / This Week / Last Week / All Time. Powered by a static pricing table (`get_cost_usd`) and a session collector that walks `~/.pythinker/sessions/` wire files. `StatusUpdate` wire events now carry optional `model_name` / `provider_key` fields for per-step attribution. +- **Z AI provider auth.** Login/logout via API key, model discovery, and OAuth selector wired into the TUI and `refresh_managed_models`. +- **Moonshot provider auth.** Login/logout via API key, model discovery (Kimi K2.x catalog), OAuth selector wired into the TUI and `refresh_managed_models`. + ## 0.33.0 (2026-06-03) ### What changed in this release diff --git a/src/pythinker_code/auth/__init__.py b/src/pythinker_code/auth/__init__.py index 6856885a..5035f8b9 100644 --- a/src/pythinker_code/auth/__init__.py +++ b/src/pythinker_code/auth/__init__.py @@ -5,21 +5,25 @@ OPENAI_CHATGPT_PLATFORM_ID = "openai-chatgpt" OPENCODE_GO_PLATFORM_ID = "opencode-go" MINIMAX_PLATFORM_ID = "minimax" +MOONSHOT_PLATFORM_ID = "moonshot" DEEPSEEK_PLATFORM_ID = "deepseek" ANTHROPIC_PLATFORM_ID = "anthropic" OPENROUTER_PLATFORM_ID = "openrouter" LM_STUDIO_PLATFORM_ID = "lm-studio" OLLAMA_PLATFORM_ID = "ollama" +ZAI_PLATFORM_ID = "z-ai" __all__ = [ "ANTHROPIC_PLATFORM_ID", "DEEPSEEK_PLATFORM_ID", "LM_STUDIO_PLATFORM_ID", "MINIMAX_PLATFORM_ID", + "MOONSHOT_PLATFORM_ID", "OLLAMA_PLATFORM_ID", "OPENAI_API_PLATFORM_ID", "OPENAI_CHATGPT_PLATFORM_ID", "OPENCODE_GO_PLATFORM_ID", "OPENROUTER_PLATFORM_ID", "PYTHINKER_CODE_PLATFORM_ID", + "ZAI_PLATFORM_ID", ] diff --git a/src/pythinker_code/auth/moonshot.py b/src/pythinker_code/auth/moonshot.py new file mode 100644 index 00000000..4652776a --- /dev/null +++ b/src/pythinker_code/auth/moonshot.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +import os +from collections.abc import AsyncIterator +from dataclasses import dataclass +from typing import Any, cast + +import aiohttp +from pydantic import SecretStr + +from pythinker_code.auth import MOONSHOT_PLATFORM_ID +from pythinker_code.auth.oauth import OAuthEvent +from pythinker_code.auth.platforms import managed_model_key, managed_provider_key +from pythinker_code.config import Config, LLMModel, LLMProvider, save_config +from pythinker_code.thinking import apply_login_thinking_defaults +from pythinker_code.utils.aiohttp import new_client_session + +MOONSHOT_BASE_URL = "https://api.moonshot.ai/v1" +MOONSHOT_PROVIDER_KEY = managed_provider_key(MOONSHOT_PLATFORM_ID) +MOONSHOT_DEFAULT_MODEL_ALIAS = managed_model_key(MOONSHOT_PLATFORM_ID, "kimi-k2.6") + + +@dataclass(frozen=True, slots=True) +class MoonshotModel: + model_id: str + alias_suffix: str + display_name: str + provider_key: str = MOONSHOT_PROVIDER_KEY + max_context_size: int = 262_144 + + @property + def alias(self) -> str: + return f"{MOONSHOT_PLATFORM_ID}/{self.alias_suffix}" + + +MOONSHOT_MODELS: tuple[MoonshotModel, ...] = ( + MoonshotModel("kimi-k2.6", "kimi-k2.6", "Kimi K2.6"), + MoonshotModel("kimi-k2.5", "kimi-k2.5", "Kimi K2.5"), + MoonshotModel("kimi-k2-thinking", "kimi-k2-thinking", "Kimi K2 Thinking"), +) + + +def get_moonshot_api_key_from_env() -> str | None: + raw = os.environ.get("MOONSHOT_API_KEY", "") + if not raw: + return None + return raw.strip() + + +def _apply_moonshot_config( + config: Config, + api_key: SecretStr, + models: tuple[MoonshotModel, ...] = MOONSHOT_MODELS, +) -> None: + config.providers[MOONSHOT_PROVIDER_KEY] = LLMProvider( + type="openai_legacy", + base_url=MOONSHOT_BASE_URL, + api_key=api_key, + ) + + provider_keys = {MOONSHOT_PROVIDER_KEY} + for key, model in list(config.models.items()): + if model.provider in provider_keys: + del config.models[key] + + for model in models: + config.models[model.alias] = LLMModel( + provider=model.provider_key, + model=model.model_id, + max_context_size=model.max_context_size, + display_name=model.display_name, + ) + + fallback = next( + (m.alias for m in models), + next(iter(config.models), ""), + ) + if MOONSHOT_DEFAULT_MODEL_ALIAS in config.models: + config.default_model = MOONSHOT_DEFAULT_MODEL_ALIAS + else: + config.default_model = fallback + apply_login_thinking_defaults(config, thinking=False, effort="off") + + +def _model_by_id() -> dict[str, MoonshotModel]: + return {model.model_id: model for model in MOONSHOT_MODELS} + + +def _parse_discovered_models(data: object) -> tuple[MoonshotModel, ...] | None: + """Return parsed models, or None if the payload is structurally invalid.""" + if not isinstance(data, dict): + return None + d = cast(dict[str, Any], data) + items = d.get("data") + if not isinstance(items, list): + return None + catalog = _model_by_id() + results: list[MoonshotModel] = [] + for raw_item in items: # pyright: ignore[reportUnknownVariableType] + if not isinstance(raw_item, dict): + continue + item = cast(dict[str, Any], raw_item) + model_id = item.get("id") + if not isinstance(model_id, str) or model_id not in catalog: + continue + base = catalog[model_id] + ctx = item.get("context_length") + max_ctx = base.max_context_size + if isinstance(ctx, int) and ctx > 0: + max_ctx = ctx + display_name = base.display_name + api_name = item.get("display_name") + if isinstance(api_name, str) and api_name: + display_name = api_name + results.append( + MoonshotModel( + model_id=base.model_id, + alias_suffix=base.alias_suffix, + display_name=display_name, + provider_key=base.provider_key, + max_context_size=max_ctx, + ) + ) + return tuple(results) + + +async def _discover_moonshot_models( + api_key: str, +) -> tuple[MoonshotModel, ...] | None: + async with ( + new_client_session() as session, + session.get( + f"{MOONSHOT_BASE_URL}/models", + headers={"Authorization": f"Bearer {api_key}"}, + ) as resp, + ): + resp.raise_for_status() + payload: object = await resp.json() + return _parse_discovered_models(payload) + + +async def login_moonshot_api_key( + config: Config, api_key: str | None = None +) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Login requires the default config file; restart without --config/--config-file.", + ) + return + + resolved_key = (api_key or get_moonshot_api_key_from_env() or "").strip() + if not resolved_key: + yield OAuthEvent("error", "Moonshot API key is required.") + return + + models = MOONSHOT_MODELS + try: + discovered = await _discover_moonshot_models(resolved_key) + if discovered is not None and discovered: + models = discovered + except aiohttp.ClientResponseError as exc: + if exc.status in {401, 403}: + yield OAuthEvent("error", "Invalid Moonshot API key; the key was not saved.") + return + yield OAuthEvent( + "info", + "Moonshot model listing is unavailable; using the built-in model list.", + ) + except (aiohttp.ClientError, TimeoutError, ValueError): + yield OAuthEvent( + "info", + "Moonshot model listing is unavailable; using the built-in model list.", + ) + + _apply_moonshot_config(config, SecretStr(resolved_key), models=models) + save_config(config) + yield OAuthEvent("success", f"Moonshot configured with model {config.default_model}.") + + +async def logout_moonshot(config: Config) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Logout requires the default config file; restart without --config/--config-file.", + ) + return + + provider_keys = {MOONSHOT_PROVIDER_KEY} + config.providers.pop(MOONSHOT_PROVIDER_KEY, None) + for key, model in list(config.models.items()): + if model.provider in provider_keys: + del config.models[key] + + if config.default_model not in config.models: + config.default_model = next(iter(config.models), "") + save_config(config) + yield OAuthEvent("success", "Logged out of Moonshot successfully.") diff --git a/src/pythinker_code/auth/platforms.py b/src/pythinker_code/auth/platforms.py index 77fec38e..e40bddcf 100644 --- a/src/pythinker_code/auth/platforms.py +++ b/src/pythinker_code/auth/platforms.py @@ -243,6 +243,12 @@ async def refresh_managed_models(config: Config) -> bool: apply_opencode_go_models, refresh_opencode_go_models, ) + from pythinker_code.auth.z_ai import ( + ZAI_PROVIDER_KEY, + ZaiModel, + apply_z_ai_models, + refresh_z_ai_models, + ) managed_providers = { key: provider for key, provider in config.providers.items() if is_managed_provider_key(key) @@ -258,9 +264,9 @@ async def refresh_managed_models(config: Config) -> bool: # generic `managed:` path can't express OpenCode Go's # two-provider split, and MiniMax's provider key intentionally includes # the wire-shape suffix (`managed:minimax-anthropic`). - if ( - provider_key in OPENCODE_GO_PROVIDER_KEYS - or provider_key == MINIMAX_ANTHROPIC_PROVIDER_KEY + if provider_key in OPENCODE_GO_PROVIDER_KEYS or provider_key in ( + MINIMAX_ANTHROPIC_PROVIDER_KEY, + ZAI_PROVIDER_KEY, ): continue platform_id = parse_managed_provider_key(provider_key) @@ -417,6 +423,14 @@ async def refresh_managed_models(config: Config) -> bool: if minimax_models is not None and apply_minimax_models(config, minimax_models): changed = True + z_ai_models: tuple[ZaiModel, ...] | None = None + try: + z_ai_models = await refresh_z_ai_models(config) + except (aiohttp.ClientError, TimeoutError, ValueError) as exc: + logger.warning("Failed to refresh Z AI models: {error}", error=exc) + if z_ai_models is not None and apply_z_ai_models(config, z_ai_models): + changed = True + if changed: config_for_save = load_config() save_changed = False @@ -427,6 +441,8 @@ async def refresh_managed_models(config: Config) -> bool: save_changed = True if minimax_models is not None and apply_minimax_models(config_for_save, minimax_models): save_changed = True + if z_ai_models is not None and apply_z_ai_models(config_for_save, z_ai_models): + save_changed = True if save_changed: save_config(config_for_save) return changed diff --git a/src/pythinker_code/auth/z_ai.py b/src/pythinker_code/auth/z_ai.py new file mode 100644 index 00000000..15be1464 --- /dev/null +++ b/src/pythinker_code/auth/z_ai.py @@ -0,0 +1,315 @@ +from __future__ import annotations + +import os +from collections.abc import AsyncIterator, Mapping +from dataclasses import dataclass +from typing import Any, cast + +import aiohttp +from pydantic import SecretStr + +from pythinker_code.auth import ZAI_PLATFORM_ID +from pythinker_code.auth.oauth import OAuthEvent +from pythinker_code.auth.platforms import managed_model_key, managed_provider_key +from pythinker_code.config import Config, LLMModel, LLMProvider, save_config +from pythinker_code.thinking import apply_login_thinking_defaults +from pythinker_code.utils.aiohttp import new_client_session + +ZAI_BASE_URL = "https://api.z.ai/api/anthropic" +ZAI_MODELS_URL = "https://api.z.ai/api/anthropic/v1/models" +ZAI_PROVIDER_KEY = managed_provider_key(ZAI_PLATFORM_ID) +ZAI_DEFAULT_MODEL_ALIAS = managed_model_key(ZAI_PLATFORM_ID, "glm-5.1") +ZAI_MODEL_DISCOVERY_TIMEOUT = aiohttp.ClientTimeout(total=15, sock_connect=8, sock_read=10) + + +@dataclass(frozen=True, slots=True) +class ZaiModel: + model_id: str + alias_suffix: str + display_name: str + provider_key: str = ZAI_PROVIDER_KEY + max_context_size: int = 131_072 + + @property + def alias(self) -> str: + return f"{ZAI_PLATFORM_ID}/{self.alias_suffix}" + + +ZAI_MODELS: tuple[ZaiModel, ...] = ( + ZaiModel("glm-5.1", "glm-5.1", "GLM-5.1", max_context_size=204_800), + ZaiModel("glm-5", "glm-5", "GLM-5"), + ZaiModel("glm-5-turbo", "glm-5-turbo", "GLM-5-Turbo"), + ZaiModel("glm-4.7", "glm-4.7", "GLM-4.7"), + ZaiModel("glm-4.5-air", "glm-4.5-air", "GLM-4.5-Air", max_context_size=98_304), +) + + +def get_z_ai_api_key_from_env() -> str | None: + value = os.getenv("ZAI_API_KEY") + if value and value.strip(): + return value.strip() + return None + + +def _is_supported_z_ai_model(model_id: str) -> bool: + return model_id.lower().startswith("glm-") + + +def _model_by_id() -> dict[str, ZaiModel]: + return {model.model_id: model for model in ZAI_MODELS} + + +def _derive_alias_suffix(model_id: str) -> str: + return model_id.lower().strip() + + +def _derive_display_name(model_id: str) -> str: + parts = model_id.split("-") + return "-".join(p.upper() if p.lower() == "glm" else p.capitalize() for p in parts) + + +def _to_positive_int(value: Any) -> int | None: + if isinstance(value, bool): + return None + try: + parsed = int(value) + except (TypeError, ValueError): + return None + return parsed if parsed > 0 else None + + +def _context_size_from_item(item: Mapping[str, Any], fallback: int) -> int: + for key in ("context_length", "max_context_length", "context_window"): + parsed = _to_positive_int(item.get(key)) + if parsed is not None: + return parsed + return fallback + + +def _display_name_from_item(item: Mapping[str, Any], fallback: str) -> str: + for key in ("display_name", "name"): + value = item.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + return fallback + + +def _parse_discovered_models(data: object) -> tuple[ZaiModel, ...] | None: + """Return parsed models, or None if the payload is structurally invalid.""" + if not isinstance(data, dict): + return None + raw_items = cast(dict[str, Any], data).get("data") + if not isinstance(raw_items, list): + return None + + known = _model_by_id() + seen: set[str] = set() + result: list[ZaiModel] = [] + for raw_item in cast(list[Any], raw_items): + if not isinstance(raw_item, Mapping): + continue + item = cast(Mapping[str, Any], raw_item) + model_id = item.get("id") + if not isinstance(model_id, str) or not model_id.strip(): + continue + model_id = model_id.strip() + if model_id in seen or not _is_supported_z_ai_model(model_id): + continue + seen.add(model_id) + + current = known.get(model_id) + alias_suffix = current.alias_suffix if current else _derive_alias_suffix(model_id) + display_name = _display_name_from_item( + item, + current.display_name if current else _derive_display_name(model_id), + ) + max_context_size = _context_size_from_item( + item, + current.max_context_size if current else 131_072, + ) + result.append( + ZaiModel( + model_id=model_id, + alias_suffix=alias_suffix, + display_name=display_name, + provider_key=current.provider_key if current else ZAI_PROVIDER_KEY, + max_context_size=max_context_size, + ) + ) + return tuple(result) + + +async def _discover_z_ai_models(api_key: str) -> tuple[ZaiModel, ...] | None: + async with ( + new_client_session(timeout=ZAI_MODEL_DISCOVERY_TIMEOUT) as session, + session.get( + ZAI_MODELS_URL, + headers={"x-api-key": api_key}, + raise_for_status=True, + ) as response, + ): + payload = await response.json(content_type=None) + return _parse_discovered_models(payload) + + +def _apply_z_ai_config( + config: Config, + api_key: SecretStr, + models: tuple[ZaiModel, ...] = ZAI_MODELS, +) -> None: + config.providers[ZAI_PROVIDER_KEY] = LLMProvider( + type="anthropic", + base_url=ZAI_BASE_URL, + api_key=api_key, + ) + + provider_keys = {ZAI_PROVIDER_KEY} + for key, model in list(config.models.items()): + if model.provider in provider_keys: + del config.models[key] + + for model in models: + config.models[model.alias] = LLMModel( + provider=model.provider_key, + model=model.model_id, + max_context_size=model.max_context_size, + display_name=model.display_name, + ) + + fallback = next( + (m.alias for m in models), + next(iter(config.models), ""), + ) + if ZAI_DEFAULT_MODEL_ALIAS in config.models: + config.default_model = ZAI_DEFAULT_MODEL_ALIAS + else: + config.default_model = fallback + apply_login_thinking_defaults(config, thinking=False, effort="off") + + +async def login_z_ai_api_key( + config: Config, api_key: str | None = None +) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Login requires the default config file; restart without --config/--config-file.", + ) + return + + resolved_key = (api_key or get_z_ai_api_key_from_env() or "").strip() + if not resolved_key: + yield OAuthEvent("error", "Z AI API key is required.") + return + + models = ZAI_MODELS + try: + discovered = await _discover_z_ai_models(resolved_key) + if discovered is not None and discovered: + models = discovered + except aiohttp.ClientResponseError as exc: + if exc.status in {401, 403}: + yield OAuthEvent("error", "Invalid Z AI API key; the key was not saved.") + return + yield OAuthEvent( + "info", + "Z AI model listing is unavailable; using the built-in model list.", + ) + except (aiohttp.ClientError, TimeoutError, ValueError): + yield OAuthEvent( + "info", + "Z AI model listing is unavailable; using the built-in model list.", + ) + + _apply_z_ai_config(config, SecretStr(resolved_key), models=models) + save_config(config) + yield OAuthEvent("success", f"Z AI configured with model {config.default_model}.") + + +async def logout_z_ai(config: Config) -> AsyncIterator[OAuthEvent]: + if not config.is_from_default_location: + yield OAuthEvent( + "error", + "Logout requires the default config file; restart without --config/--config-file.", + ) + return + + provider_keys = {ZAI_PROVIDER_KEY} + config.providers.pop(ZAI_PROVIDER_KEY, None) + for key, model in list(config.models.items()): + if model.provider in provider_keys: + del config.models[key] + + if config.default_model not in config.models: + config.default_model = next(iter(config.models), "") + save_config(config) + yield OAuthEvent("success", "Logged out of Z AI successfully.") + + +def apply_z_ai_models(config: Config, models: tuple[ZaiModel, ...]) -> bool: + """Upsert the live Z AI catalog and prune models no longer returned. + + Preserves user preferences unless the selected Z AI model disappeared. + """ + changed = False + aliases: list[str] = [] + for model in models: + alias = model.alias + aliases.append(alias) + existing = config.models.get(alias) + if existing is None: + config.models[alias] = LLMModel( + provider=model.provider_key, + model=model.model_id, + max_context_size=model.max_context_size, + display_name=model.display_name, + ) + changed = True + continue + if existing.provider != model.provider_key: + existing.provider = model.provider_key + changed = True + if existing.model != model.model_id: + existing.model = model.model_id + changed = True + if existing.max_context_size != model.max_context_size: + existing.max_context_size = model.max_context_size + changed = True + if existing.display_name != model.display_name: + existing.display_name = model.display_name + changed = True + + alias_set = set(aliases) + removed_default = False + for alias, model_cfg in list(config.models.items()): + if model_cfg.provider != ZAI_PROVIDER_KEY: + continue + if alias in alias_set: + continue + del config.models[alias] + if config.default_model == alias: + removed_default = True + changed = True + + if removed_default: + config.default_model = aliases[0] if aliases else next(iter(config.models), "") + changed = True + elif config.default_model and config.default_model not in config.models: + config.default_model = next(iter(config.models), "") + changed = True + return changed + + +def _z_ai_api_key(config: Config) -> str | None: + provider = config.providers.get(ZAI_PROVIDER_KEY) + if provider is None: + return None + value = provider.api_key.get_secret_value().strip() + return value or None + + +async def refresh_z_ai_models(config: Config) -> tuple[ZaiModel, ...] | None: + api_key = _z_ai_api_key(config) + if api_key is None: + return None + return await _discover_z_ai_models(api_key) diff --git a/src/pythinker_code/llm.py b/src/pythinker_code/llm.py index d738149b..737989df 100644 --- a/src/pythinker_code/llm.py +++ b/src/pythinker_code/llm.py @@ -436,6 +436,10 @@ def derive_model_capabilities(model: LLMModel) -> set[ModelCapability]: # switch in create_llm(). if _is_kimi_k2_model(model.model): capabilities.add("thinking") + # kimi-k2-thinking is Moonshot's thinking-only variant; unlike the + # hybrid K2.5/K2.6 it cannot be switched off. + if "thinking" in model_name: + capabilities.add("always_thinking") # Models with "thinking" in their name are always-thinking models elif "thinking" in model_name or "reason" in model_name: capabilities.update(("thinking", "always_thinking")) diff --git a/src/pythinker_code/soul/pythinkersoul.py b/src/pythinker_code/soul/pythinkersoul.py index 5f18d6ec..8640ade0 100644 --- a/src/pythinker_code/soul/pythinkersoul.py +++ b/src/pythinker_code/soul/pythinkersoul.py @@ -1509,8 +1509,16 @@ async def _pythinker_core_step_with_retry() -> StepResult: input_tokens=usage.input if usage else "?", output_tokens=usage.output if usage else "?", ) + _step_model_name: str | None = chat_provider.model_name + _step_provider_key: str | None = None + if self._runtime.llm.model_config is not None: + _step_provider_key = self._runtime.llm.model_config.provider status_update = StatusUpdate( - token_usage=usage, message_id=result.id, plan_mode=self._plan_mode + token_usage=usage, + message_id=result.id, + model_name=_step_model_name, + provider_key=_step_provider_key, + plan_mode=self._plan_mode, ) if usage is not None: # mark the token count for the context before the step diff --git a/src/pythinker_code/thinking.py b/src/pythinker_code/thinking.py index 5b8f2677..0d30f963 100644 --- a/src/pythinker_code/thinking.py +++ b/src/pythinker_code/thinking.py @@ -63,9 +63,17 @@ def apply_login_thinking_defaults( cross-session user preference. ``create_llm`` clamps effort to the chosen model's capabilities at use-time, so a previously-set value is always safe to keep; only an unset (``None``) effort is initialized here. + + Additionally, a user who explicitly enabled thinking (``default_thinking=True``) + on the legacy boolean path is not silently downgraded when a provider defaults + to ``thinking=False`` — their preference is preserved. """ if config.default_thinking_effort is not None: return + # Preserve an explicit user preference for thinking when the provider default + # would downgrade it. The user can still override per-session. + if config.default_thinking and not thinking: + return config.default_thinking = thinking config.default_thinking_effort = effort diff --git a/src/pythinker_code/ui/shell/oauth.py b/src/pythinker_code/ui/shell/oauth.py index 8cdd9185..adee6b9b 100644 --- a/src/pythinker_code/ui/shell/oauth.py +++ b/src/pythinker_code/ui/shell/oauth.py @@ -12,11 +12,13 @@ DEEPSEEK_PLATFORM_ID, LM_STUDIO_PLATFORM_ID, MINIMAX_PLATFORM_ID, + MOONSHOT_PLATFORM_ID, OLLAMA_PLATFORM_ID, OPENAI_API_PLATFORM_ID, OPENAI_CHATGPT_PLATFORM_ID, OPENCODE_GO_PLATFORM_ID, OPENROUTER_PLATFORM_ID, + ZAI_PLATFORM_ID, ) from pythinker_code.auth.anthropic_direct import ( ANTHROPIC_PROVIDER_KEY, @@ -38,6 +40,11 @@ login_minimax_api_key, logout_minimax, ) +from pythinker_code.auth.moonshot import ( + MOONSHOT_PROVIDER_KEY, + login_moonshot_api_key, + logout_moonshot, +) from pythinker_code.auth.oauth import OAuthEvent from pythinker_code.auth.ollama import ( OLLAMA_PROVIDER_KEY, @@ -62,6 +69,11 @@ logout_openrouter, ) from pythinker_code.auth.platforms import managed_provider_key +from pythinker_code.auth.z_ai import ( + ZAI_PROVIDER_KEY, + login_z_ai_api_key, + logout_z_ai, +) from pythinker_code.cli import Reload from pythinker_code.ui.shell.console import console from pythinker_code.ui.shell.selectors.oauth import ( @@ -126,6 +138,8 @@ async def _prompt_api_key(label: str) -> str | None: OAuthProviderEntry(id="opencode-go", name="OpenCode Go", auth_type="api_key"), OAuthProviderEntry(id="minimax", name="MiniMax", auth_type="api_key"), OAuthProviderEntry(id="deepseek", name="DeepSeek", auth_type="api_key"), + OAuthProviderEntry(id="z-ai", name="Z AI", auth_type="api_key"), + OAuthProviderEntry(id="moonshot", name="Moonshot", auth_type="api_key"), OAuthProviderEntry(id="anthropic", name="Anthropic", auth_type="api_key"), OAuthProviderEntry(id="openrouter", name="OpenRouter", auth_type="api_key"), OAuthProviderEntry(id="lm-studio", name="LM Studio", auth_type="api_key"), @@ -148,6 +162,8 @@ async def _prompt_api_key(label: str) -> str | None: "opencode-go": (OPENCODE_GO_OPENAI_PROVIDER_KEY, OPENCODE_GO_ANTHROPIC_PROVIDER_KEY), "minimax": (MINIMAX_ANTHROPIC_PROVIDER_KEY,), "deepseek": (DEEPSEEK_PROVIDER_KEY,), + "z-ai": (ZAI_PROVIDER_KEY,), + "moonshot": (MOONSHOT_PROVIDER_KEY,), "anthropic": (ANTHROPIC_PROVIDER_KEY,), "openrouter": (OPENROUTER_PROVIDER_KEY,), "lm-studio": (LM_STUDIO_PROVIDER_KEY,), @@ -161,6 +177,8 @@ async def _prompt_api_key(label: str) -> str | None: OAuthProviderEntry(id="opencode-go", name="OpenCode Go", auth_type="api_key"), OAuthProviderEntry(id="minimax", name="MiniMax", auth_type="api_key"), OAuthProviderEntry(id="deepseek", name="DeepSeek", auth_type="api_key"), + OAuthProviderEntry(id="z-ai", name="Z AI", auth_type="api_key"), + OAuthProviderEntry(id="moonshot", name="Moonshot", auth_type="api_key"), OAuthProviderEntry(id="anthropic", name="Anthropic", auth_type="api_key"), OAuthProviderEntry(id="openrouter", name="OpenRouter", auth_type="api_key"), OAuthProviderEntry(id="lm-studio", name="LM Studio", auth_type="api_key"), @@ -238,6 +256,20 @@ async def login(app: Shell, args: str) -> None: return ok = await _render_oauth_events(login_deepseek_api_key(soul.runtime.config, api_key)) provider = DEEPSEEK_PLATFORM_ID + elif mode == "z-ai": + api_key = await _prompt_api_key("Z AI") + if not api_key: + console.print(f"[{_t.error}]No Z AI API key entered.[/]") + return + ok = await _render_oauth_events(login_z_ai_api_key(soul.runtime.config, api_key)) + provider = ZAI_PLATFORM_ID + elif mode == "moonshot": + api_key = await _prompt_api_key("Moonshot") + if not api_key: + console.print(f"[{_t.error}]No Moonshot API key entered.[/]") + return + ok = await _render_oauth_events(login_moonshot_api_key(soul.runtime.config, api_key)) + provider = MOONSHOT_PLATFORM_ID elif mode == "anthropic": api_key = await _prompt_api_key("Anthropic") if not api_key: @@ -261,8 +293,8 @@ async def login(app: Shell, args: str) -> None: else: console.print( f"[{_t.error}]Usage: /login " - "[browser|headless|api-key|opencode-go|minimax|deepseek|anthropic|openrouter|" - "lm-studio|ollama][/]" + "[browser|headless|api-key|opencode-go|minimax|deepseek|z-ai|moonshot|anthropic|" + "openrouter|lm-studio|ollama][/]" ) return if not ok: @@ -316,6 +348,10 @@ async def logout(app: Shell, args: str) -> None: ok = await _render_oauth_events(logout_anthropic(config)) elif mode == "deepseek": ok = await _render_oauth_events(logout_deepseek(config)) + elif mode == "z-ai": + ok = await _render_oauth_events(logout_z_ai(config)) + elif mode == "moonshot": + ok = await _render_oauth_events(logout_moonshot(config)) elif mode == "minimax": ok = await _render_oauth_events(logout_minimax(config)) elif mode in ("opencode-go", "opencode", "go"): @@ -333,8 +369,8 @@ async def logout(app: Shell, args: str) -> None: else: console.print( f"[{_t.error}]Usage: /logout " - "[openai|opencode-go|minimax|deepseek|anthropic|openrouter|lm-studio|ollama|" - "github-feedback][/]" + "[openai|opencode-go|minimax|deepseek|z-ai|moonshot|anthropic|openrouter|lm-studio|" + "ollama|github-feedback][/]" ) return if not ok: diff --git a/src/pythinker_code/ui/shell/slash.py b/src/pythinker_code/ui/shell/slash.py index 0028b19a..7e2a1ed1 100644 --- a/src/pythinker_code/ui/shell/slash.py +++ b/src/pythinker_code/ui/shell/slash.py @@ -1875,6 +1875,7 @@ async def fork(app: Shell, args: str): export_import, # noqa: F401 # type: ignore[reportUnusedImport] oauth, # noqa: F401 # type: ignore[reportUnusedImport] setup, # noqa: F401 # type: ignore[reportUnusedImport] + stats, # noqa: F401 # type: ignore[reportUnusedImport] update, # noqa: F401 # type: ignore[reportUnusedImport] usage, # noqa: F401 # type: ignore[reportUnusedImport] ) diff --git a/src/pythinker_code/ui/shell/stats.py b/src/pythinker_code/ui/shell/stats.py new file mode 100644 index 00000000..3d8c41c7 --- /dev/null +++ b/src/pythinker_code/ui/shell/stats.py @@ -0,0 +1,340 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from typing import TYPE_CHECKING + +from prompt_toolkit.application import Application +from prompt_toolkit.formatted_text import StyleAndTextTuples +from prompt_toolkit.key_binding import KeyBindings, KeyPressEvent +from prompt_toolkit.layout import HSplit, Layout, Window +from prompt_toolkit.layout.controls import FormattedTextControl +from prompt_toolkit.styles import Style + +from pythinker_code.ui.shell.console import console +from pythinker_code.ui.shell.slash import registry +from pythinker_code.ui.shell.stats_collector import ( + AllStats, + PeriodStats, + ProviderStats, + load_all_stats, +) + +if TYPE_CHECKING: + from pythinker_code.ui.shell import Shell + +type _LineFn = Callable[[str, str], None] + +# --------------------------------------------------------------------------- +# Formatting helpers +# --------------------------------------------------------------------------- + +_TAB_KEYS = ["today", "this_week", "last_week", "all_time"] +_TAB_LABELS = { + "today": "Today", + "this_week": "This Week", + "last_week": "Last Week", + "all_time": "All Time", +} + + +def _fmt_cost(v: float) -> str: + if v == 0: + return "-" + if v < 0.01: + return f"${v:.4f}" + if v < 1: + return f"${v:.2f}" + if v < 10: + return f"${v:.2f}" + return f"${v:.1f}" + + +def _fmt_tokens(n: int) -> str: + if n == 0: + return "-" + if n < 1_000: + return str(n) + if n < 10_000: + return f"{n / 1000:.1f}k" + if n < 1_000_000: + return f"{n // 1000}k" + return f"{n / 1_000_000:.1f}M" + + +def _fmt_num(n: int) -> str: + if n == 0: + return "-" + return f"{n:,}" + + +# --------------------------------------------------------------------------- +# Dashboard model +# --------------------------------------------------------------------------- + + +class StatsApp: + """Interactive usage statistics dashboard.""" + + def __init__(self, data: AllStats) -> None: + self._data = data + self._periods = data.periods + self._tab_idx = 0 + self._view = "table" + self._selected_idx = 0 + self._expanded: set[str] = set() + self._app = self._build_app() + + @property + def _tab(self) -> str: + return _TAB_KEYS[self._tab_idx] + + @property + def _current(self) -> PeriodStats: + return self._periods[self._tab] + + def _providers_sorted(self) -> list[tuple[str, ProviderStats]]: + return sorted( + self._current.providers.items(), + key=lambda kv: kv[1].cost, + reverse=True, + ) + + def _render(self) -> StyleAndTextTuples: + parts: StyleAndTextTuples = [] + + def line(text: str, style: str = "") -> None: + parts.append((style, text)) + parts.append(("", "\n")) + + def txt(text: str, style: str = "") -> None: + parts.append((style, text)) + + # Title + title = "Usage Insights" if self._view == "insights" else "Usage Statistics" + line(title, "bold ansicyan") + line("") + + # Tabs + tab_parts: list[str] = [] + for i, key in enumerate(_TAB_KEYS): + label = _TAB_LABELS[key] + if i == self._tab_idx: + tab_parts.append(f"[{label}]") + else: + tab_parts.append(f" {label} ") + line(" ".join(tab_parts), "ansiblue") + line("") + + cur = self._current + + if self._view == "insights": + self._render_insights(parts, line, txt, cur) + else: + self._render_table(parts, line, txt, cur) + + # Help line + if self._view == "insights": + line("[Tab/←→] period [v] table view [q] close", "ansigray") + else: + line( + "[Tab/←→] period [↑↓] select [Enter] expand [v] insights [q] close", + "ansigray", + ) + + return parts + + def _render_table( + self, + parts: StyleAndTextTuples, + line: _LineFn, + txt: _LineFn, + cur: PeriodStats, + ) -> None: + col_w = {"sessions": 9, "msgs": 9, "cost": 9, "tokens": 9, "in": 8, "out": 8} + name_w = 26 + + def _pad_right(s: str, w: int) -> str: + return s[:w].ljust(w) + + def _pad_left(s: str, w: int) -> str: + return s[:w].rjust(w) + + # Header + hdr = _pad_right("Provider / Model", name_w) + hdr += _pad_left("Sessions", col_w["sessions"]) + hdr += _pad_left("Msgs", col_w["msgs"]) + hdr += _pad_left("Cost", col_w["cost"]) + hdr += _pad_left("Tokens", col_w["tokens"]) + hdr += _pad_left("↑In", col_w["in"]) + hdr += _pad_left("↓Out", col_w["out"]) + parts.append(("ansigray", hdr + "\n")) + parts.append(("ansigray", "─" * (name_w + sum(col_w.values())) + "\n")) + + providers = self._providers_sorted() + if not providers: + parts.append(("ansigray", " No usage data for this period\n")) + else: + for i, (pname, pstats) in enumerate(providers): + is_sel = i == self._selected_idx + is_exp = pname in self._expanded + arrow = "▾" if is_exp else "▸" + style = "bold ansicyan" if is_sel else "" + + row = f"{arrow} {_pad_right(pname, name_w - 2)}" + row += _pad_left(_fmt_num(len(pstats.sessions)), col_w["sessions"]) + row += _pad_left(_fmt_num(pstats.messages), col_w["msgs"]) + row += _pad_left(_fmt_cost(pstats.cost), col_w["cost"]) + row += _pad_left(_fmt_tokens(pstats.tokens), col_w["tokens"]) + in_tokens = pstats.input_other + pstats.input_cache_creation + row += _pad_left(_fmt_tokens(in_tokens), col_w["in"]) + row += _pad_left(_fmt_tokens(pstats.output), col_w["out"]) + parts.append((style, row + "\n")) + + if is_exp: + for mname, mstats in sorted( + pstats.models.items(), key=lambda kv: kv[1].cost, reverse=True + ): + mrow = " " + _pad_right(mname, name_w - 4) + mrow += _pad_left(_fmt_num(len(mstats.sessions)), col_w["sessions"]) + mrow += _pad_left(_fmt_num(mstats.messages), col_w["msgs"]) + mrow += _pad_left(_fmt_cost(mstats.cost), col_w["cost"]) + mrow += _pad_left(_fmt_tokens(mstats.tokens), col_w["tokens"]) + m_in = mstats.input_other + mstats.input_cache_creation + mrow += _pad_left(_fmt_tokens(m_in), col_w["in"]) + mrow += _pad_left(_fmt_tokens(mstats.output), col_w["out"]) + parts.append(("ansigray", mrow + "\n")) + + # Totals + parts.append(("ansigray", "─" * (name_w + sum(col_w.values())) + "\n")) + tot = _pad_right("Total", name_w) + tot += _pad_left(_fmt_num(cur.total_sessions), col_w["sessions"]) + tot += _pad_left(_fmt_num(cur.total_messages), col_w["msgs"]) + tot += _pad_left(_fmt_cost(cur.total_cost), col_w["cost"]) + parts.append(("bold", tot + "\n")) + parts.append(("", "\n")) + + def _render_insights( + self, parts: StyleAndTextTuples, line: _LineFn, txt: _LineFn, cur: PeriodStats + ) -> None: + insights = self._data.insights.get(self._tab) + if cur.total_messages == 0: + parts.append(("ansigray", " No usage recorded for this period.\n\n")) + return + if cur.total_cost == 0 or insights is None or not insights.insights: + parts.append( + ("ansigray", " No cost data available (models not yet priced or no sessions).\n\n") + ) + return + label = _TAB_LABELS[self._tab] + parts.append(("ansigray", f" {label} · weighted by cost (USD)\n\n")) + for insight in insights.insights: + pct_fmt = ".0f" if insight.percent >= 10 else ".1f" + pct_str = f"{insight.percent:{pct_fmt}}%" + parts.append(("bold ansicyan", f" {pct_str} ")) + parts.append(("", insight.headline + "\n")) + parts.append(("ansigray", f" {insight.advice}\n\n")) + + def _build_app(self) -> Application[None]: + kb = KeyBindings() + + @kb.add("q") + @kb.add("escape") + def _quit(event: KeyPressEvent) -> None: + event.app.exit() + + @kb.add("tab") + @kb.add("right") + def _next_tab(event: KeyPressEvent) -> None: + self._tab_idx = (self._tab_idx + 1) % len(_TAB_KEYS) + self._selected_idx = 0 + event.app.invalidate() + + @kb.add("s-tab") + @kb.add("left") + def _prev_tab(event: KeyPressEvent) -> None: + self._tab_idx = (self._tab_idx - 1) % len(_TAB_KEYS) + self._selected_idx = 0 + event.app.invalidate() + + @kb.add("up") + def _up(event: KeyPressEvent) -> None: + if self._view == "table" and self._selected_idx > 0: + self._selected_idx -= 1 + event.app.invalidate() + + @kb.add("down") + def _down(event: KeyPressEvent) -> None: + if self._view == "table": + providers = self._providers_sorted() + if self._selected_idx < len(providers) - 1: + self._selected_idx += 1 + event.app.invalidate() + + @kb.add("enter") + @kb.add("space") + def _toggle_expand(event: KeyPressEvent) -> None: + if self._view == "table": + providers = self._providers_sorted() + if providers and self._selected_idx < len(providers): + pname = providers[self._selected_idx][0] + if pname in self._expanded: + self._expanded.discard(pname) + else: + self._expanded.add(pname) + event.app.invalidate() + + @kb.add("v") + def _toggle_view(event: KeyPressEvent) -> None: + self._view = "insights" if self._view == "table" else "table" + event.app.invalidate() + + # Mark handlers as used + _ = (_quit, _next_tab, _prev_tab, _up, _down, _toggle_expand, _toggle_view) + + ctrl = FormattedTextControl(self._render, focusable=False) + layout = Layout(HSplit([Window(content=ctrl)])) + + return Application( + layout=layout, + key_bindings=kb, + full_screen=False, + style=Style.from_dict( + { + "": "bg:#1e1e1e fg:#d4d4d4", + } + ), + mouse_support=False, + ) + + async def run(self) -> None: + await self._app.run_async() + + +# --------------------------------------------------------------------------- +# Slash command +# --------------------------------------------------------------------------- + + +@registry.command(name="stats", aliases=["history"]) +async def stats(app: Shell, args: str) -> None: + """Show usage statistics dashboard (tokens and cost by provider/model).""" + from pythinker_code.ui.theme import get_tui_tokens as _get_tui_tokens + + _t = _get_tui_tokens() + + try: + data = await asyncio.to_thread(load_all_stats) + except Exception as e: + from pythinker_code.utils.logging import logger as _logger + + _logger.exception("Failed to load stats: {error}", error=e) + console.print(f"[{_t.error}]Failed to load stats: {e}[/]") + return + + if data.periods["all_time"].total_messages == 0: + console.print(f"[{_t.warning}]No usage data found in ~/.pythinker/sessions/[/]") + return + + dashboard = StatsApp(data) + await dashboard.run() diff --git a/src/pythinker_code/ui/shell/stats_collector.py b/src/pythinker_code/ui/shell/stats_collector.py new file mode 100644 index 00000000..7986fd49 --- /dev/null +++ b/src/pythinker_code/ui/shell/stats_collector.py @@ -0,0 +1,402 @@ +from __future__ import annotations + +import json +import os +from collections.abc import Generator +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any, cast + +from pythinker_core.chat_provider import TokenUsage + +from pythinker_code.ui.shell.stats_pricing import get_cost_usd +from pythinker_code.utils.logging import logger + + +def _set_str() -> set[str]: + return set() + + +def _dict_model_stats() -> dict[str, ModelStats]: + return {} + + +def _dict_provider_stats() -> dict[str, ProviderStats]: + return {} + + +def _list_insight() -> list[Insight]: + return [] + + +@dataclass(slots=True) +class StepRecord: + session_id: str + timestamp: float + model_name: str + provider_key: str + input_other: int + output: int + input_cache_read: int + input_cache_creation: int + + @property + def total_tokens(self) -> int: + return self.input_other + self.output + self.input_cache_read + self.input_cache_creation + + @property + def cost_usd(self) -> float: + usage = TokenUsage( + input_other=self.input_other, + output=self.output, + input_cache_read=self.input_cache_read, + input_cache_creation=self.input_cache_creation, + ) + return get_cost_usd(self.model_name, usage) + + +@dataclass(slots=True) +class ModelStats: + messages: int = 0 + cost: float = 0.0 + input_other: int = 0 + output: int = 0 + input_cache_read: int = 0 + input_cache_creation: int = 0 + sessions: set[str] = field(default_factory=_set_str) + + def add(self, step: StepRecord) -> None: + self.messages += 1 + self.cost += step.cost_usd + self.input_other += step.input_other + self.output += step.output + self.input_cache_read += step.input_cache_read + self.input_cache_creation += step.input_cache_creation + self.sessions.add(step.session_id) + + @property + def tokens(self) -> int: + return self.input_other + self.output + self.input_cache_read + self.input_cache_creation + + +@dataclass(slots=True) +class ProviderStats: + messages: int = 0 + cost: float = 0.0 + input_other: int = 0 + output: int = 0 + input_cache_read: int = 0 + input_cache_creation: int = 0 + sessions: set[str] = field(default_factory=_set_str) + models: dict[str, ModelStats] = field(default_factory=_dict_model_stats) + + def add(self, step: StepRecord) -> None: + self.messages += 1 + self.cost += step.cost_usd + self.input_other += step.input_other + self.output += step.output + self.input_cache_read += step.input_cache_read + self.input_cache_creation += step.input_cache_creation + self.sessions.add(step.session_id) + m = self.models.setdefault(step.model_name, ModelStats()) + m.add(step) + + @property + def tokens(self) -> int: + return self.input_other + self.output + self.input_cache_read + self.input_cache_creation + + +@dataclass(slots=True) +class PeriodStats: + total_messages: int = 0 + total_cost: float = 0.0 + total_sessions: int = 0 + providers: dict[str, ProviderStats] = field(default_factory=_dict_provider_stats) + _sessions: set[str] = field(default_factory=_set_str) + + def add(self, step: StepRecord) -> None: + self.total_messages += 1 + self.total_cost += step.cost_usd + self._sessions.add(step.session_id) + self.total_sessions = len(self._sessions) + p = self.providers.setdefault(step.provider_key, ProviderStats()) + p.add(step) + + +@dataclass(slots=True) +class Insight: + percent: float + headline: str + advice: str + + +@dataclass(slots=True) +class PeriodInsights: + insights: list[Insight] = field(default_factory=_list_insight) + + +@dataclass(slots=True) +class UsagePeriod: + stats: PeriodStats = field(default_factory=PeriodStats) + insights: PeriodInsights = field(default_factory=PeriodInsights) + + +@dataclass(slots=True) +class AllStats: + periods: dict[str, PeriodStats] + insights: dict[str, PeriodInsights] + + +def get_sessions_root() -> Path: + """Return the path to ~/.pythinker/sessions/.""" + agent_dir = os.environ.get("PYTHINKER_DIR") or os.path.join( + os.path.expanduser("~"), ".pythinker" + ) + return Path(agent_dir) / "sessions" + + +def collect_session_files(sessions_root: Path) -> list[Path]: + """Recursively collect all wire.jsonl files under sessions_root.""" + result: list[Path] = [] + if not sessions_root.is_dir(): + return result + for wd_dir in sessions_root.iterdir(): + if not wd_dir.is_dir(): + continue + for sess_dir in wd_dir.iterdir(): + if not sess_dir.is_dir(): + continue + _collect_from_session_dir(sess_dir, result) + return sorted(result) + + +def _collect_from_session_dir(sess_dir: Path, result: list[Path]) -> None: + wire = sess_dir / "wire.jsonl" + if wire.is_file(): + result.append(wire) + subagents = sess_dir / "subagents" + if subagents.is_dir(): + for agent_dir in subagents.iterdir(): + if agent_dir.is_dir(): + sub_wire = agent_dir / "wire.jsonl" + if sub_wire.is_file(): + result.append(sub_wire) + + +def parse_wire_file( + wire_path: Path, + session_id: str, + seen_hashes: set[str], +) -> Generator[StepRecord, None, None]: + """Parse one wire.jsonl and yield StepRecords for each StatusUpdate.""" + try: + with wire_path.open(encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + raw: object = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(raw, dict): + continue + obj = cast(dict[str, Any], raw) + raw_msg: object = obj.get("message") + if not isinstance(raw_msg, dict): + continue + msg = cast(dict[str, Any], raw_msg) + if msg.get("type") != "StatusUpdate": + continue + raw_payload: object = msg.get("payload") + if not isinstance(raw_payload, dict): + continue + payload = cast(dict[str, Any], raw_payload) + raw_tu: object = payload.get("token_usage") + if not isinstance(raw_tu, dict): + continue + tu = cast(dict[str, Any], raw_tu) + + input_other = int(tu.get("input_other", 0)) + output = int(tu.get("output", 0)) + cache_read = int(tu.get("input_cache_read", 0)) + cache_write = int(tu.get("input_cache_creation", 0)) + ts = float(obj.get("timestamp", 0)) + + # Include individual token fields so distinct StatusUpdate events + # with the same timestamp and equal totals don't collide. + h = f"{session_id}:{ts}:{input_other}:{output}:{cache_read}:{cache_write}" + if h in seen_hashes: + continue + seen_hashes.add(h) + + model_name: str = str(payload.get("model_name") or "unknown") + provider_key: str = str(payload.get("provider_key") or "unknown") + + yield StepRecord( + session_id=session_id, + timestamp=ts, + model_name=model_name, + provider_key=provider_key, + input_other=input_other, + output=output, + input_cache_read=cache_read, + input_cache_creation=cache_write, + ) + except OSError as exc: + logger.warning("Failed to read wire file {path}: {error}", path=wire_path, error=exc) + + +def _period_boundaries() -> tuple[float, float, float]: + """Return (today_start, week_start, last_week_start) as UTC timestamps.""" + now = datetime.now(UTC) + today_start = now.replace(hour=0, minute=0, second=0, microsecond=0) + days_since_monday = now.weekday() # Monday=0 + week_start = today_start - timedelta(days=days_since_monday) + last_week_start = week_start - timedelta(days=7) + return ( + today_start.timestamp(), + week_start.timestamp(), + last_week_start.timestamp(), + ) + + +def compute_period_stats(steps: list[StepRecord]) -> dict[str, PeriodStats]: + """Bin steps into Today / This Week / Last Week / All Time.""" + today_ts, week_ts, last_week_ts = _period_boundaries() + + periods: dict[str, PeriodStats] = { + "today": PeriodStats(), + "this_week": PeriodStats(), + "last_week": PeriodStats(), + "all_time": PeriodStats(), + } + + for step in steps: + ts = step.timestamp + periods["all_time"].add(step) + if ts >= today_ts: + periods["today"].add(step) + if ts >= week_ts: + periods["this_week"].add(step) + elif ts >= last_week_ts: + periods["last_week"].add(step) + + return periods + + +_PARALLEL_WINDOW_S = 120.0 +_PARALLEL_SESSION_THRESHOLD = 4 +_LARGE_CONTEXT_THRESHOLD = 150_000 +_LARGE_UNCACHED_THRESHOLD = 100_000 +_LONG_SESSION_S = 8 * 3600 +_TOP_SESSION_COUNT = 5 +_MIN_PERCENT = 1.0 + + +def compute_insights(steps: list[StepRecord]) -> PeriodInsights: + """Compute cost-weighted usage insights for a period.""" + if not steps: + return PeriodInsights() + + total_cost = sum(s.cost_usd for s in steps) + if total_cost <= 0: + return PeriodInsights() + + candidates: list[Insight] = [] + + # Large context + large_ctx_cost = sum( + s.cost_usd + for s in steps + if (s.input_other + s.input_cache_read + s.input_cache_creation) > _LARGE_CONTEXT_THRESHOLD + ) + candidates.append( + Insight( + percent=(large_ctx_cost / total_cost) * 100, + headline=f"of your cost was at >{_LARGE_CONTEXT_THRESHOLD // 1000}k context", + advice=( + "Longer sessions are more expensive even when cached. " + "/compact mid-task, /clear when switching to new tasks." + ), + ) + ) + + # Large uncached prompt + uncached_cost = sum( + s.cost_usd + for s in steps + if (s.input_other + s.input_cache_creation) > _LARGE_UNCACHED_THRESHOLD + ) + candidates.append( + Insight( + percent=(uncached_cost / total_cost) * 100, + headline=( + f"of your cost came from" + f" >{_LARGE_UNCACHED_THRESHOLD // 1000}k-token uncached prompts" + ), + advice=( + "Uncached input is expensive. " + "/compact before stepping away keeps the cold-start small." + ), + ) + ) + + # Top-N session concentration + session_costs: dict[str, float] = {} + for s in steps: + session_costs[s.session_id] = session_costs.get(s.session_id, 0.0) + s.cost_usd + if len(session_costs) > _TOP_SESSION_COUNT: + top_cost = sum(sorted(session_costs.values(), reverse=True)[:_TOP_SESSION_COUNT]) + candidates.append( + Insight( + percent=(top_cost / total_cost) * 100, + headline=f"of your cost came from your top {_TOP_SESSION_COUNT} sessions", + advice="A small number of sessions drives most of your spend.", + ) + ) + + insights = [i for i in candidates if i.percent >= _MIN_PERCENT] + insights.sort(key=lambda i: i.percent, reverse=True) + return PeriodInsights(insights=insights) + + +def load_all_stats() -> AllStats: + """Load and aggregate all pythinker session usage with insights.""" + root = get_sessions_root() + wire_files = collect_session_files(root) + seen_hashes: set[str] = set() + all_steps: list[StepRecord] = [] + + for wire_path in wire_files: + # sessions///wire.jsonl → parents[1]/ + # sessions///subagents//wire.jsonl + # parents[1] = "subagents", so walk up to parents[3]/parents[2] + if wire_path.parent.parent.name == "subagents": + session_id = f"{wire_path.parents[3].name}/{wire_path.parents[2].name}" + else: + session_id = f"{wire_path.parents[1].name}/{wire_path.parents[0].name}" + for step in parse_wire_file(wire_path, session_id, seen_hashes): + all_steps.append(step) + + periods = compute_period_stats(all_steps) + + today_ts, week_ts, last_week_ts = _period_boundaries() + period_steps: dict[str, list[StepRecord]] = { + "today": [], + "this_week": [], + "last_week": [], + "all_time": list(all_steps), + } + for step in all_steps: + ts = step.timestamp + if ts >= today_ts: + period_steps["today"].append(step) + if ts >= week_ts: + period_steps["this_week"].append(step) + elif ts >= last_week_ts: + period_steps["last_week"].append(step) + + insights = {k: compute_insights(v) for k, v in period_steps.items()} + return AllStats(periods=periods, insights=insights) diff --git a/src/pythinker_code/ui/shell/stats_pricing.py b/src/pythinker_code/ui/shell/stats_pricing.py new file mode 100644 index 00000000..220648bb --- /dev/null +++ b/src/pythinker_code/ui/shell/stats_pricing.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +from pythinker_core.chat_provider import TokenUsage + +# Prices in USD per million tokens. +# Source: Pi's models.generated.ts (blackbox/pi-main/packages/ai/src/models.generated.ts) +# Format: {model_id: (input, output, cache_read, cache_write)} +_PRICE_TABLE: dict[str, tuple[float, float, float, float]] = { + # Anthropic — direct API + "claude-3-haiku-20240307": (0.25, 1.25, 0.03, 0.3), + "claude-3-sonnet-20240229": (3.0, 15.0, 0.3, 0.3), + "claude-3-opus-20240229": (15.0, 75.0, 1.5, 18.75), + "claude-3-5-haiku-20241022": (0.8, 4.0, 0.08, 1.0), + "claude-3-5-haiku-latest": (0.8, 4.0, 0.08, 1.0), + "claude-3-5-sonnet-20240620": (3.0, 15.0, 0.3, 3.75), + "claude-3-5-sonnet-20241022": (3.0, 15.0, 0.3, 3.75), + "claude-3-7-sonnet-20250219": (3.0, 15.0, 0.3, 3.75), + "claude-sonnet-4-20250514": (3.0, 15.0, 0.3, 3.75), + "claude-sonnet-4-0": (3.0, 15.0, 0.3, 3.75), + "claude-sonnet-4": (3.0, 15.0, 0.3, 3.75), + "claude-sonnet-4-5": (3.0, 15.0, 0.3, 3.75), + "claude-sonnet-4-5-20250929": (3.0, 15.0, 0.3, 3.75), + "claude-sonnet-4-6": (3.0, 15.0, 0.3, 3.75), + "claude-opus-4-20250514": (15.0, 75.0, 1.5, 18.75), + "claude-opus-4-0": (15.0, 75.0, 1.5, 18.75), + "claude-opus-4-1": (15.0, 75.0, 1.5, 18.75), + "claude-opus-4-1-20250805": (15.0, 75.0, 1.5, 18.75), + "claude-opus-4-5": (5.0, 25.0, 0.5, 6.25), + "claude-opus-4-5-20251101": (5.0, 25.0, 0.5, 6.25), + "claude-haiku-4-5": (1.0, 5.0, 0.1, 1.25), + "claude-haiku-4-5-20251001": (1.0, 5.0, 0.1, 1.25), + # OpenAI GPT-5 family + "gpt-5": (2.5, 15.0, 0.25, 0.0), + "gpt-5.5": (2.5, 15.0, 0.25, 0.0), + "gpt-5-chat-latest": (2.5, 15.0, 0.25, 0.0), + "gpt-5-mini": (0.75, 4.5, 0.075, 0.0), + "gpt-5.4-mini": (0.75, 4.5, 0.075, 0.0), + "gpt-5-nano": (0.75, 4.5, 0.075, 0.0), + "gpt-5-pro": (5.0, 30.0, 0.5, 0.0), + "gpt-5.5-pro": (5.0, 30.0, 0.5, 0.0), + "gpt-4o": (2.5, 10.0, 0.25, 0.0), + "gpt-4o-mini": (0.15, 0.6, 0.075, 0.0), + # DeepSeek (via opencode-go or direct) + "deepseek-v4-flash": (0.14, 0.28, 0.0028, 0.0), + "deepseek-v4-pro": (0.435, 0.87, 0.003625, 0.0), + "deepseek-chat": (0.27, 1.1, 0.0, 0.0), + "deepseek-reasoner": (0.55, 2.19, 0.55, 0.0), + # GLM (Z.AI / OpenCode-Go) + "glm-5": (1.0, 3.2, 0.2, 0.0), + "glm-5.1": (1.4, 4.4, 0.26, 0.0), + "glm-5-turbo": (0.5, 1.5, 0.1, 0.0), + "glm-4.7": (0.5, 1.5, 0.1, 0.0), + "glm-4.5-air": (0.3, 1.0, 0.06, 0.0), + # Kimi (opencode-go) + "kimi-k2.5": (0.6, 3.0, 0.08, 0.0), + "kimi-k2.6": (0.95, 4.0, 0.16, 0.0), + # MiniMax (opencode-go / anthropic shape) + "minimax-m2.5": (0.3, 1.2, 0.06, 0.0), + "minimax-m2.7": (0.3, 1.2, 0.06, 0.0), + # Gemini (Google) + "gemini-2.0-flash": (0.1, 0.4, 0.025, 0.0), + "gemini-2.0-flash-lite": (0.075, 0.3, 0.0, 0.0), + "gemini-2.5-flash": (0.3, 2.5, 0.03, 0.0), + "gemini-2.5-flash-lite": (0.1, 0.4, 0.01, 0.0), + "gemini-2.5-pro": (1.25, 10.0, 0.125, 0.0), +} + + +def get_cost_usd(model: str, usage: TokenUsage) -> float: + """Return estimated USD cost for one LLM step. + + Looks up the exact model ID first; falls back to prefix matching + for versioned aliases (e.g. ``claude-sonnet-4-5-20250929`` → + ``claude-sonnet-4-5``). Returns 0.0 for unknown models. + """ + pricing = _PRICE_TABLE.get(model) + if pricing is None: + # Prefix fallback: longest matching key wins. + best: tuple[float, float, float, float] | None = None + best_len = 0 + for key, val in _PRICE_TABLE.items(): + if model.startswith(key) and len(key) > best_len: + best = val + best_len = len(key) + if best is None: + return 0.0 + pricing = best + + inp, out, cr, cw = pricing + total = ( + usage.input_other * inp + + usage.output * out + + usage.input_cache_read * cr + + usage.input_cache_creation * cw + ) / 1_000_000 + return total diff --git a/src/pythinker_code/wire/types.py b/src/pythinker_code/wire/types.py index c9e12428..ad62fd8a 100644 --- a/src/pythinker_code/wire/types.py +++ b/src/pythinker_code/wire/types.py @@ -230,6 +230,10 @@ class StatusUpdate(BaseModel): """The token usage statistics of the current step.""" message_id: str | None = None """The message ID of the current step.""" + model_name: str | None = None + """The model ID that produced this step (e.g. 'claude-sonnet-4-5').""" + provider_key: str | None = None + """The provider config key for this step (e.g. 'managed:openai-chatgpt').""" plan_mode: bool | None = None """Whether plan mode (read-only) is active. None means no change.""" mcp_status: MCPStatusSnapshot | None = None diff --git a/tests/auth/test_z_ai_auth.py b/tests/auth/test_z_ai_auth.py new file mode 100644 index 00000000..b029029d --- /dev/null +++ b/tests/auth/test_z_ai_auth.py @@ -0,0 +1,432 @@ +from __future__ import annotations + +import aiohttp +import pytest +from multidict import CIMultiDict, CIMultiDictProxy +from pydantic import SecretStr +from yarl import URL + +from pythinker_code.config import Config + + +def _request_info(url: str) -> aiohttp.RequestInfo: + return aiohttp.RequestInfo( + url=URL(url), + method="GET", + headers=CIMultiDictProxy(CIMultiDict()), + real_url=URL(url), + ) + + +def test_z_ai_model_catalog_contains_five_models(): + from pythinker_code.auth.z_ai import ZAI_MODELS + + aliases = {model.alias for model in ZAI_MODELS} + assert aliases == { + "z-ai/glm-5.1", + "z-ai/glm-5", + "z-ai/glm-5-turbo", + "z-ai/glm-4.7", + "z-ai/glm-4.5-air", + } + + api_ids = {m.alias: m.model_id for m in ZAI_MODELS} + assert api_ids == { + "z-ai/glm-5.1": "glm-5.1", + "z-ai/glm-5": "glm-5", + "z-ai/glm-5-turbo": "glm-5-turbo", + "z-ai/glm-4.7": "glm-4.7", + "z-ai/glm-4.5-air": "glm-4.5-air", + } + + assert all(m.provider_key == "managed:z-ai" for m in ZAI_MODELS) + + +def test_z_ai_glm51_has_200k_context(): + from pythinker_code.auth.z_ai import ZAI_MODELS + + glm51 = next(m for m in ZAI_MODELS if m.model_id == "glm-5.1") + assert glm51.max_context_size == 204_800 + + +def test_z_ai_glm45air_has_96k_context(): + from pythinker_code.auth.z_ai import ZAI_MODELS + + air = next(m for m in ZAI_MODELS if m.model_id == "glm-4.5-air") + assert air.max_context_size == 98_304 + + +def test_z_ai_env_key_uses_zai_api_key(monkeypatch): + from pythinker_code.auth.z_ai import get_z_ai_api_key_from_env + + monkeypatch.delenv("ZAI_API_KEY", raising=False) + assert get_z_ai_api_key_from_env() is None + + monkeypatch.setenv("ZAI_API_KEY", " zai-key ") + assert get_z_ai_api_key_from_env() == "zai-key" + + monkeypatch.setenv("ZAI_API_KEY", "") + assert get_z_ai_api_key_from_env() is None + + +@pytest.mark.parametrize( + "payload, expected_aliases", + [ + # Structurally invalid → None (no prune should happen) + (None, None), + ({}, None), + ({"data": "not a list"}, None), + # Valid structure, no matching models → empty tuple + ({"data": [{"context_length": 1000}]}, set()), + ({"data": [{"id": "unknown-model-xyz"}]}, set()), + # Valid structure with known models + ({"data": [{"id": "glm-5.1"}]}, {"z-ai/glm-5.1"}), + ({"data": [{"id": "glm-5-turbo"}]}, {"z-ai/glm-5-turbo"}), + ( + {"data": [{"id": "glm-4.7"}, {"id": "glm-4.5-air"}]}, + {"z-ai/glm-4.7", "z-ai/glm-4.5-air"}, + ), + ], +) +def test_parse_discovered_z_ai_models_handles_payloads(payload, expected_aliases): + from pythinker_code.auth.z_ai import _parse_discovered_models + + result = _parse_discovered_models(payload) + if expected_aliases is None: + assert result is None + else: + assert result is not None + assert {m.alias for m in result} == expected_aliases + + +def test_parse_discovered_z_ai_models_uses_context_length_when_positive(): + from pythinker_code.auth.z_ai import _parse_discovered_models + + payload = { + "data": [ + {"id": "glm-5.1", "context_length": 400_000}, + {"id": "glm-4.5-air", "context_length": -5}, + {"id": "glm-4.7", "context_length": "bogus"}, + ] + } + result = _parse_discovered_models(payload) + assert result is not None + by_id = {m.model_id: m for m in result} + assert by_id["glm-5.1"].max_context_size == 400_000 + assert by_id["glm-4.5-air"].max_context_size == 98_304 # fallback to hardcoded + assert by_id["glm-4.7"].max_context_size == 131_072 # fallback to hardcoded + + +def test_parse_discovered_z_ai_models_accepts_unknown_glm_future_models(): + from pythinker_code.auth.z_ai import _parse_discovered_models + + payload = {"data": [{"id": "glm-6.0", "context_length": 512_000}]} + result = _parse_discovered_models(payload) + assert result is not None + assert len(result) == 1 + assert result[0].model_id == "glm-6.0" + assert result[0].alias_suffix == "glm-6.0" + assert result[0].max_context_size == 512_000 + + +def test_parse_discovered_z_ai_models_deduplicates(): + from pythinker_code.auth.z_ai import _parse_discovered_models + + payload = {"data": [{"id": "glm-5.1"}, {"id": "glm-5.1"}]} + result = _parse_discovered_models(payload) + assert result is not None + assert len(result) == 1 + + +def test_apply_z_ai_config_writes_provider_and_default(): + from pythinker_code.auth.z_ai import ( + ZAI_BASE_URL, + ZAI_DEFAULT_MODEL_ALIAS, + ZAI_PROVIDER_KEY, + _apply_z_ai_config, + ) + + config = Config(is_from_default_location=True) + _apply_z_ai_config(config, SecretStr("zai-test")) + + assert set(config.providers) == {ZAI_PROVIDER_KEY} + provider = config.providers[ZAI_PROVIDER_KEY] + assert provider.type == "anthropic" + assert provider.base_url == ZAI_BASE_URL + assert provider.api_key.get_secret_value() == "zai-test" + assert config.models["z-ai/glm-5.1"].provider == ZAI_PROVIDER_KEY + assert config.models["z-ai/glm-5.1"].model == "glm-5.1" + assert config.models["z-ai/glm-5.1"].max_context_size == 204_800 + assert config.default_model == ZAI_DEFAULT_MODEL_ALIAS + + +def test_apply_z_ai_config_replaces_existing_z_ai_models(): + from pythinker_code.auth.z_ai import ( + ZAI_PROVIDER_KEY, + ZaiModel, + _apply_z_ai_config, + ) + + config = Config(is_from_default_location=True) + _apply_z_ai_config(config, SecretStr("zai-test")) + + new_models = (ZaiModel("glm-5.1", "glm-5.1", "GLM-5.1 New", max_context_size=300_000),) + _apply_z_ai_config(config, SecretStr("zai-test-2"), models=new_models) + + z_ai_aliases = [a for a, m in config.models.items() if m.provider == ZAI_PROVIDER_KEY] + assert z_ai_aliases == ["z-ai/glm-5.1"] + assert config.models["z-ai/glm-5.1"].max_context_size == 300_000 + + +@pytest.mark.asyncio +async def test_login_z_ai_saves_static_models_when_discovery_fails(monkeypatch, tmp_path): + from pythinker_code.auth.z_ai import login_z_ai_api_key + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + + async def fake_discover(api_key: str): + assert api_key == "zai-test" + raise aiohttp.ClientConnectionError("models unavailable") + + monkeypatch.setattr("pythinker_code.auth.z_ai._discover_z_ai_models", fake_discover) + + events = [event async for event in login_z_ai_api_key(config, "zai-test")] + + assert [e.type for e in events] == ["info", "success"] + assert config.default_model == "z-ai/glm-5.1" + assert "z-ai/glm-5-turbo" in config.models + assert (tmp_path / "config.toml").exists() + + +@pytest.mark.asyncio +async def test_login_z_ai_falls_back_on_non_auth_response_error(monkeypatch, tmp_path): + from pythinker_code.auth.z_ai import login_z_ai_api_key + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + + async def fake_discover(api_key: str): + raise aiohttp.ClientResponseError( + _request_info("https://api.z.ai/api/anthropic/v1/models"), + (), + status=503, + message="Service Unavailable", + ) + + monkeypatch.setattr("pythinker_code.auth.z_ai._discover_z_ai_models", fake_discover) + + events = [event async for event in login_z_ai_api_key(config, "zai-test")] + + assert [e.type for e in events] == ["info", "success"] + assert config.default_model == "z-ai/glm-5.1" + + +@pytest.mark.asyncio +async def test_login_z_ai_rejects_401(monkeypatch, tmp_path): + from pythinker_code.auth.z_ai import login_z_ai_api_key + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + + async def fake_discover(api_key: str): + raise aiohttp.ClientResponseError( + _request_info("https://api.z.ai/api/anthropic/v1/models"), + (), + status=401, + message="Unauthorized", + ) + + monkeypatch.setattr("pythinker_code.auth.z_ai._discover_z_ai_models", fake_discover) + + events = [event async for event in login_z_ai_api_key(config, "bad-key")] + + assert events[-1].type == "error" + assert "Invalid Z AI API key" in events[-1].message + assert config.providers == {} + assert config.models == {} + + +@pytest.mark.asyncio +async def test_login_z_ai_uses_discovered_context_length(monkeypatch, tmp_path): + from pythinker_code.auth.z_ai import ZaiModel, login_z_ai_api_key + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + + async def fake_discover(api_key: str): + return (ZaiModel("glm-5.1", "glm-5.1", "GLM-5.1", max_context_size=512_000),) + + monkeypatch.setattr("pythinker_code.auth.z_ai._discover_z_ai_models", fake_discover) + + events = [event async for event in login_z_ai_api_key(config, "zai-test")] + + assert events[-1].type == "success" + assert config.models["z-ai/glm-5.1"].max_context_size == 512_000 + + +@pytest.mark.asyncio +async def test_login_z_ai_requires_key(tmp_path): + from pythinker_code.auth.z_ai import login_z_ai_api_key + + config = Config(is_from_default_location=True) + + events = [event async for event in login_z_ai_api_key(config, "")] + + assert events[-1].type == "error" + assert events[-1].message == "Z AI API key is required." + + +@pytest.mark.asyncio +async def test_logout_z_ai_removes_only_z_ai(monkeypatch, tmp_path): + from pythinker_code.auth.z_ai import ( + ZAI_PROVIDER_KEY, + _apply_z_ai_config, + logout_z_ai, + ) + from pythinker_code.config import LLMModel, LLMProvider + + monkeypatch.setenv("PYTHINKER_SHARE_DIR", str(tmp_path)) + config = Config(is_from_default_location=True) + config.providers["managed:openai"] = LLMProvider( + type="openai_responses", + base_url="https://api.openai.com/v1", + api_key=SecretStr("sk-test"), + ) + config.models["openai/gpt-5.2"] = LLMModel( + provider="managed:openai", + model="gpt-5.2", + max_context_size=400_000, + ) + _apply_z_ai_config(config, SecretStr("zai-test")) + + events = [event async for event in logout_z_ai(config)] + + assert events[-1].type == "success" + assert ZAI_PROVIDER_KEY not in config.providers + assert "z-ai/glm-5.1" not in config.models + assert "managed:openai" in config.providers + assert "openai/gpt-5.2" in config.models + assert config.default_model == "openai/gpt-5.2" + + +@pytest.mark.asyncio +async def test_logout_z_ai_rejects_non_default_config_location(): + from pythinker_code.auth.z_ai import logout_z_ai + + config = Config(is_from_default_location=False) + + events = [event async for event in logout_z_ai(config)] + + assert events[-1].type == "error" + assert "default config file" in events[-1].message + assert config.providers == {} + + +def test_apply_z_ai_models_prunes_stale_models_and_preserves_user_default(): + from pythinker_code.auth.z_ai import ( + ZAI_PROVIDER_KEY, + ZaiModel, + _apply_z_ai_config, + apply_z_ai_models, + ) + + config = Config(is_from_default_location=True) + _apply_z_ai_config(config, SecretStr("zai-test")) + config.default_model = "z-ai/glm-5.1" + + discovered = ( + ZaiModel("glm-5.1", "glm-5.1", "GLM-5.1", max_context_size=400_000), + ZaiModel("glm-6.0", "glm-6.0", "GLM-6.0", max_context_size=512_000), + ) + + changed = apply_z_ai_models(config, discovered) + + assert changed is True + z_ai_aliases = {a for a, m in config.models.items() if m.provider == ZAI_PROVIDER_KEY} + assert z_ai_aliases == {"z-ai/glm-5.1", "z-ai/glm-6.0"} + assert config.models["z-ai/glm-5.1"].max_context_size == 400_000 + assert config.default_model == "z-ai/glm-5.1" + + +def test_apply_z_ai_models_reassigns_default_when_it_disappears(): + from pythinker_code.auth.z_ai import ( + ZaiModel, + _apply_z_ai_config, + apply_z_ai_models, + ) + + config = Config(is_from_default_location=True) + _apply_z_ai_config( + config, + SecretStr("zai-test"), + models=(ZaiModel("glm-5.1", "glm-5.1", "GLM-5.1"),), + ) + config.default_model = "z-ai/glm-5.1" + + changed = apply_z_ai_models( + config, + (ZaiModel("glm-5-turbo", "glm-5-turbo", "GLM-5-Turbo"),), + ) + + assert changed is True + assert "z-ai/glm-5.1" not in config.models + assert config.default_model == "z-ai/glm-5-turbo" + + +def test_apply_z_ai_models_returns_false_for_noop(): + from pythinker_code.auth.z_ai import ( + ZAI_MODELS, + _apply_z_ai_config, + apply_z_ai_models, + ) + + config = Config(is_from_default_location=True) + _apply_z_ai_config(config, SecretStr("zai-test")) + + assert apply_z_ai_models(config, ZAI_MODELS) is False + + +def test_apply_z_ai_config_defaults_thinking_off(): + """Z.ai is wired to its Anthropic-compatible endpoint, which (verified + empirically 2026-06-03 against glm-5.1) defaults thinking OFF and honors + `thinking: {"type": "disabled"}`. The login default of effort="off" + therefore matches the endpoint's native behavior. If the provider type or + base_url ever moves to the OpenAI-compatible endpoint (which defaults + thinking ON for GLM-5.x), these defaults must be re-evaluated. + """ + from pythinker_code.auth.z_ai import _apply_z_ai_config + + config = Config(is_from_default_location=True) + assert config.default_thinking_effort is None + _apply_z_ai_config(config, SecretStr("zai-test")) + + assert config.providers["managed:z-ai"].type == "anthropic" + assert config.default_thinking is False # unchanged from Config default + assert config.default_thinking_effort == "off" + + +def test_apply_z_ai_config_preserves_legacy_thinking_true(): + """Users with legacy default_thinking=True must not be silently downgraded.""" + from pythinker_code.auth.z_ai import _apply_z_ai_config + + config = Config(is_from_default_location=True) + config.default_thinking = True + # effort unset — the legacy path + assert config.default_thinking_effort is None + _apply_z_ai_config(config, SecretStr("zai-test")) + + assert config.default_thinking is True + assert config.default_thinking_effort is None + + +def test_apply_z_ai_config_preserves_existing_effort_choice(): + from pythinker_code.auth.z_ai import _apply_z_ai_config + + config = Config(is_from_default_location=True) + config.default_thinking = True + config.default_thinking_effort = "medium" + _apply_z_ai_config(config, SecretStr("zai-test")) + + assert config.default_thinking is True + assert config.default_thinking_effort == "medium" diff --git a/tests/core/test_create_llm.py b/tests/core/test_create_llm.py index 0cb2eb13..7e1d7b02 100644 --- a/tests/core/test_create_llm.py +++ b/tests/core/test_create_llm.py @@ -752,3 +752,38 @@ def test_clone_llm_with_model_alias_preserves_kimi_thinking_disabled(): assert cloned.chat_provider._generation_kwargs.get("extra_body") == { # pyright: ignore[reportPrivateUsage] "thinking": {"type": "disabled"} } + + +def test_derive_model_capabilities_marks_kimi_k2_thinking_as_always_on(): + model = LLMModel( + provider="openai-compatible", + model="kimi-k2-thinking", + max_context_size=262_144, + capabilities=None, + ) + assert derive_model_capabilities(model) == {"thinking", "always_thinking"} + + +def test_create_llm_kimi_k2_thinking_ignores_thinking_off(): + provider = LLMProvider( + type="openai_legacy", + base_url="https://api.example.com/v1", + api_key=SecretStr("test-key"), + ) + model = LLMModel( + provider="kimi-provider", + model="kimi-k2-thinking", + max_context_size=262_144, + capabilities=None, + ) + llm = create_llm(provider, model, thinking=False) + assert llm is not None + assert isinstance(llm.chat_provider, OpenAILegacy) + assert llm.capabilities == {"thinking", "always_thinking"} + # Always-thinking: "off" is overridden to the default effort, and the + # Kimi body switch must say enabled, never disabled. + assert llm.thinking is True + assert llm.thinking_effort == "high" + assert llm.chat_provider._generation_kwargs.get("extra_body") == { # pyright: ignore[reportPrivateUsage] + "thinking": {"type": "enabled"} + } diff --git a/tests/core/test_wire_message.py b/tests/core/test_wire_message.py index 5076f7bc..5583e914 100644 --- a/tests/core/test_wire_message.py +++ b/tests/core/test_wire_message.py @@ -173,6 +173,8 @@ async def test_wire_message_serde(): "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": None, + "provider_key": None, "plan_mode": None, "mcp_status": { "loading": True, diff --git a/tests/ui_and_conv/test_stats_collector.py b/tests/ui_and_conv/test_stats_collector.py new file mode 100644 index 00000000..3ca93f1a --- /dev/null +++ b/tests/ui_and_conv/test_stats_collector.py @@ -0,0 +1,171 @@ +from __future__ import annotations + +import json +from datetime import UTC, datetime +from pathlib import Path + +from pythinker_code.ui.shell.stats_collector import ( + StepRecord, + collect_session_files, + compute_period_stats, + get_sessions_root, + parse_wire_file, +) + + +def _make_wire(tmp_path: Path, records: list[dict]) -> Path: + p = tmp_path / "wire.jsonl" + lines = [json.dumps({"type": "metadata", "protocol_version": "1.9"})] + for r in records: + lines.append(json.dumps(r)) + p.write_text("\n".join(lines)) + return p + + +def _status_update( + ts: float, + input_other: int, + output: int, + cache_read: int = 0, + cache_write: int = 0, + model_name: str | None = None, + provider_key: str | None = None, +) -> dict: + return { + "timestamp": ts, + "message": { + "type": "StatusUpdate", + "payload": { + "token_usage": { + "input_other": input_other, + "output": output, + "input_cache_read": cache_read, + "input_cache_creation": cache_write, + }, + "model_name": model_name, + "provider_key": provider_key, + }, + }, + } + + +def test_parse_wire_extracts_steps(tmp_path): + now = datetime.now(UTC).timestamp() + wire = _make_wire( + tmp_path, + [ + _status_update(now, 1000, 200), + _status_update(now + 1, 2000, 400), + ], + ) + session_id = "test-session" + seen = set() + steps = list(parse_wire_file(wire, session_id, seen)) + assert len(steps) == 2 + assert steps[0].input_other == 1000 + assert steps[1].output == 400 + + +def test_parse_wire_deduplicates(tmp_path): + now = datetime.now(UTC).timestamp() + # Same timestamp and total_tokens → duplicate + record = _status_update(now, 1000, 200) + wire = _make_wire(tmp_path, [record, record]) + seen = set() + steps = list(parse_wire_file(wire, "s", seen)) + assert len(steps) == 1 + + +def test_parse_wire_unknown_model_defaults(tmp_path): + now = datetime.now(UTC).timestamp() + wire = _make_wire(tmp_path, [_status_update(now, 500, 100)]) + seen = set() + steps = list(parse_wire_file(wire, "s", seen)) + assert steps[0].model_name == "unknown" + assert steps[0].provider_key == "unknown" + + +def test_compute_period_stats_today(tmp_path): + now = datetime.now(UTC).timestamp() + steps = [ + StepRecord( + session_id="s1", + timestamp=now, + model_name="claude-sonnet-4-5", + provider_key="anthropic", + input_other=1000, + output=200, + input_cache_read=0, + input_cache_creation=0, + ), + ] + stats = compute_period_stats(steps) + assert stats["all_time"].total_messages == 1 + assert stats["today"].total_messages == 1 + assert "anthropic" in stats["all_time"].providers + + +def test_compute_period_stats_excludes_old(tmp_path): + old_ts = datetime(2020, 1, 1, tzinfo=UTC).timestamp() + now = datetime.now(UTC).timestamp() + steps = [ + StepRecord( + session_id="s1", + timestamp=old_ts, + model_name="m", + provider_key="p", + input_other=100, + output=50, + input_cache_read=0, + input_cache_creation=0, + ), + StepRecord( + session_id="s2", + timestamp=now, + model_name="m", + provider_key="p", + input_other=200, + output=100, + input_cache_read=0, + input_cache_creation=0, + ), + ] + stats = compute_period_stats(steps) + assert stats["all_time"].total_messages == 2 + assert stats["today"].total_messages == 1 + + +def test_get_sessions_root_path(monkeypatch, tmp_path): + monkeypatch.setenv("PYTHINKER_DIR", str(tmp_path)) + root = get_sessions_root() + assert root.name == "sessions" + assert root.parent == tmp_path + + +def test_collect_session_files_finds_wires(tmp_path): + # Structure: tmp/sessions/wdhash/sessid/wire.jsonl + sess_dir = tmp_path / "sessions" / "abc123" / "sess1" + sess_dir.mkdir(parents=True) + w = sess_dir / "wire.jsonl" + w.write_text('{"type":"metadata","protocol_version":"1.9"}\n') + # Subagent wire + sub_dir = sess_dir / "subagents" / "agent1" + sub_dir.mkdir(parents=True) + sw = sub_dir / "wire.jsonl" + sw.write_text('{"type":"metadata","protocol_version":"1.9"}\n') + files = collect_session_files(tmp_path / "sessions") + assert w in files + assert sw in files + + +def test_load_all_stats_returns_all_stats(tmp_path, monkeypatch): + from pythinker_code.ui.shell.stats_collector import AllStats, load_all_stats + + monkeypatch.setattr( + "pythinker_code.ui.shell.stats_collector.get_sessions_root", + lambda: tmp_path / "sessions", + ) + result = load_all_stats() + assert isinstance(result, AllStats) + assert "all_time" in result.periods + assert "all_time" in result.insights diff --git a/tests/ui_and_conv/test_stats_pricing.py b/tests/ui_and_conv/test_stats_pricing.py new file mode 100644 index 00000000..17e4d464 --- /dev/null +++ b/tests/ui_and_conv/test_stats_pricing.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +from pythinker_core.chat_provider import TokenUsage + +from pythinker_code.ui.shell.stats_pricing import get_cost_usd + + +def _usage(input_other=0, output=0, cache_read=0, cache_write=0) -> TokenUsage: + return TokenUsage( + input_other=input_other, + output=output, + input_cache_read=cache_read, + input_cache_creation=cache_write, + ) + + +def test_known_model_cost(): + # claude-sonnet-4-5: input=$3/M, output=$15/M + usage = _usage(input_other=1_000_000, output=1_000_000) + cost = get_cost_usd("claude-sonnet-4-5", usage) + assert abs(cost - 18.0) < 0.001 + + +def test_cache_read_cost(): + usage = _usage(cache_read=1_000_000) + cost = get_cost_usd("claude-sonnet-4-5", usage) + assert abs(cost - 0.3) < 0.001 + + +def test_cache_write_cost(): + usage = _usage(cache_write=1_000_000) + cost = get_cost_usd("claude-sonnet-4-5", usage) + assert abs(cost - 3.75) < 0.001 + + +def test_unknown_model_returns_zero(): + usage = _usage(input_other=1_000_000, output=1_000_000) + cost = get_cost_usd("totally-unknown-model-xyz", usage) + assert cost == 0.0 + + +def test_prefix_match_fallback(): + # "claude-sonnet-4-5-20251999" is NOT in _PRICE_TABLE, so must hit prefix fallback + usage = _usage(input_other=1_000_000, output=1_000_000) + cost = get_cost_usd("claude-sonnet-4-5-20251999", usage) + assert cost > 0.0 + + +def test_zero_usage_returns_zero(): + usage = _usage() + cost = get_cost_usd("claude-opus-4-1", usage) + assert cost == 0.0 diff --git a/tests_e2e/test_wire_approvals_tools.py b/tests_e2e/test_wire_approvals_tools.py index 3a87a936..9ecea309 100644 --- a/tests_e2e/test_wire_approvals_tools.py +++ b/tests_e2e/test_wire_approvals_tools.py @@ -118,6 +118,8 @@ def test_shell_approval_approve(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -183,6 +185,8 @@ def test_shell_approval_approve(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -259,6 +263,8 @@ def test_shell_approval_reject(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -390,6 +396,8 @@ def test_approve_for_session(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -459,6 +467,8 @@ def test_approve_for_session(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -494,6 +504,8 @@ def test_approve_for_session(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -537,6 +549,8 @@ def test_approve_for_session(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -610,6 +624,8 @@ def test_yolo_skips_approval(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -653,6 +669,8 @@ def test_yolo_skips_approval(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -923,6 +941,8 @@ def test_display_block_todo(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -963,6 +983,8 @@ def test_display_block_todo(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -1046,6 +1068,8 @@ def test_tool_call_part_streaming(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -1086,6 +1110,8 @@ def test_tool_call_part_streaming(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -1159,6 +1185,8 @@ def test_default_agent_missing_tool(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -1192,6 +1220,8 @@ def test_default_agent_missing_tool(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -1279,6 +1309,8 @@ def test_custom_agent_exclude_tool(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -1312,6 +1344,8 @@ def test_custom_agent_exclude_tool(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, diff --git a/tests_e2e/test_wire_config.py b/tests_e2e/test_wire_config.py index b2737b97..6abee0c6 100644 --- a/tests_e2e/test_wire_config.py +++ b/tests_e2e/test_wire_config.py @@ -79,6 +79,8 @@ def test_config_string(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -169,6 +171,8 @@ def test_model_override(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "provider-b", "plan_mode": False, "mcp_status": None, }, diff --git a/tests_e2e/test_wire_prompt.py b/tests_e2e/test_wire_prompt.py index cb8c8144..d422ab97 100644 --- a/tests_e2e/test_wire_prompt.py +++ b/tests_e2e/test_wire_prompt.py @@ -89,6 +89,8 @@ def test_basic_prompt_events(tmp_path) -> None: "input_cache_creation": 0, }, "message_id": "scripted-1", + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -293,6 +295,8 @@ def test_max_steps_reached(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -377,6 +381,8 @@ def test_status_update_fields(tmp_path) -> None: "input_cache_creation": 0, }, "message_id": "scripted-1", + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -475,6 +481,8 @@ def test_concurrent_prompt_error(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -540,6 +548,8 @@ def test_concurrent_prompt_error(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, diff --git a/tests_e2e/test_wire_protocol.py b/tests_e2e/test_wire_protocol.py index 53a149a3..992eb123 100644 --- a/tests_e2e/test_wire_protocol.py +++ b/tests_e2e/test_wire_protocol.py @@ -464,6 +464,8 @@ def handle_request(msg: dict[str, Any]) -> dict[str, Any]: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -511,6 +513,8 @@ def handle_request(msg: dict[str, Any]) -> dict[str, Any]: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -563,6 +567,8 @@ def test_prompt_without_initialize(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, diff --git a/tests_e2e/test_wire_sessions.py b/tests_e2e/test_wire_sessions.py index bd3a1e3d..c2e57b5f 100644 --- a/tests_e2e/test_wire_sessions.py +++ b/tests_e2e/test_wire_sessions.py @@ -220,6 +220,8 @@ def test_clear_context_rotates(tmp_path) -> None: "max_context_tokens": 100000, "token_usage": None, "message_id": None, + "model_name": None, + "provider_key": None, "plan_mode": None, "mcp_status": None, }, @@ -305,6 +307,8 @@ def test_manual_compact(tmp_path) -> None: "max_context_tokens": 100000, "token_usage": None, "message_id": None, + "model_name": None, + "provider_key": None, "plan_mode": None, "mcp_status": None, }, @@ -461,6 +465,8 @@ def test_replay_streams_wire_history(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -504,6 +510,8 @@ def test_replay_streams_wire_history(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, diff --git a/tests_e2e/test_wire_skills_mcp.py b/tests_e2e/test_wire_skills_mcp.py index 0a8e12df..0ea92be7 100644 --- a/tests_e2e/test_wire_skills_mcp.py +++ b/tests_e2e/test_wire_skills_mcp.py @@ -116,6 +116,8 @@ def test_skill_prompt_injects_skill_text(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -207,6 +209,8 @@ def test_flow_skill(tmp_path) -> None: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -305,6 +309,8 @@ def ping(text: str) -> str: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": None, + "provider_key": None, "plan_mode": None, "mcp_status": { "loading": True, @@ -325,6 +331,8 @@ def ping(text: str) -> str: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": None, + "provider_key": None, "plan_mode": None, "mcp_status": { "loading": False, @@ -361,6 +369,8 @@ def ping(text: str) -> str: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, }, @@ -421,6 +431,8 @@ def ping(text: str) -> str: "max_context_tokens": None, "token_usage": None, "message_id": None, + "model_name": "scripted_echo", + "provider_key": "scripted_provider", "plan_mode": False, "mcp_status": None, },