diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a6ef045..1ae65cb8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,8 @@ GitHub Releases page; `0.8.0` is the new starting line. ## Unreleased +- **Quieter `/login`.** Logging in no longer prints a `RuntimeWarning` about an un-awaited `redraw_in_future` coroutine. The prompt redraw throttle now uses a coroutine-free path (`max_render_postpone_time`), eliminating the warning emitted during the login prompt handoff. + ## 0.37.0 (2026-06-07) - **Agent runtime tool visibility hardening.** `PythinkerToolset` now filters the tools advertised to the model by active execution policy, permission profile, root/subagent role, and plan-mode state while preserving execution-time guards as defense in depth. diff --git a/pyproject.toml b/pyproject.toml index 1444cc74..76d9254a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,8 +34,11 @@ dependencies = [ "pyyaml==6.0.3", "rich==15.0.0", "certifi>=2025.10.5", - "click==8.3.0", "pyperclip==1.11.0", + # Pinned: click 8.4.x regresses pyright (`click.Option` typed partially + # unknown, ~95 errors in `make check`). Unpin once that type regression is + # fixed upstream. Mirrored as a Dependabot ignore in .github/dependabot.yml. + "click==8.3.0", "streamingjson==0.0.5", "trafilatura==2.0.0", # lxml is used by trafilatura/htmldate/justext; keep pinned for binary wheels. @@ -73,6 +76,9 @@ dev = [ "pytest>=9.0.3", "pytest-asyncio>=1.3.0", "pytest-cov>=6.0", + # Pinned: ruff 0.15's formatter reflow fails `make check`. Unpin once the + # formatting churn is resolved. Mirrored as a Dependabot ignore in + # .github/dependabot.yml. "ruff>=0.14.10,<0.15", ] diff --git a/src/pythinker_code/models_dev.py b/src/pythinker_code/models_dev.py new file mode 100644 index 00000000..e2fb7af8 --- /dev/null +++ b/src/pythinker_code/models_dev.py @@ -0,0 +1,187 @@ +from __future__ import annotations + +import asyncio +import json +import os +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, cast + +import aiohttp + +from pythinker_code.utils.aiohttp import new_client_session +from pythinker_code.utils.logging import logger + +_MODELS_DEV_URL = "https://models.dev/api.json" +_TTL_SECONDS = 86_400 # 24 hours + +# Canonical providers win over regional/cloud variants when the same +# bare model id appears under multiple providers. +_CANONICAL_PROVIDERS = { + "anthropic", + "openai", + "google", + "deepseek", + "z-ai", + "moonshot", + "minimax", + "meta", + "mistral", + "cohere", + "x-ai", +} + +# Module-level lock: only one network fetch runs at a time. +_refresh_lock = asyncio.Lock() + +# In-process memo: (mtime_ns, catalog_dict) +_catalog_cache: dict[str, Any] = {} + + +@dataclass(frozen=True) +class ModelPrice: + input: float # USD / 1M tokens + output: float + cache_read: float # 0.0 if absent + cache_write: float # 0.0 if absent + + +def _get_cache_path() -> Path: + """Return ~/.pythinker/model-pricing/models-dev.json (or $PYTHINKER_DIR variant).""" + base = Path(os.environ.get("PYTHINKER_DIR") or Path.home() / ".pythinker") + return base / "model-pricing" / "models-dev.json" + + +def _coerce_cost(raw: object) -> float: + """Coerce a cost field to float; returns 0.0 for None, raises on non-numeric.""" + if raw is None: + return 0.0 + if isinstance(raw, (int, float, str)): + return float(raw) + raise TypeError(f"non-numeric cost value: {raw!r}") + + +def _flatten_catalog(raw: dict[str, Any]) -> dict[str, ModelPrice]: + """Flatten provider→models hierarchy into {model_id: ModelPrice}. + + Canonical providers win over regional/compat variants. + Model ids containing '@' (version-tagged) are excluded. + context_over_200k tiered pricing is ignored. + Models with any non-numeric cost field are skipped. + """ + canonical: dict[str, ModelPrice] = {} + fallback: dict[str, ModelPrice] = {} + + for provider_id, provider_data in raw.items(): + if not isinstance(provider_data, dict): + continue + models = cast(dict[str, Any], provider_data).get("models") + if not isinstance(models, dict): + continue + target = canonical if provider_id in _CANONICAL_PROVIDERS else fallback + for model_id, model_data in sorted(cast(dict[str, Any], models).items()): + if "@" in model_id: + continue + if not isinstance(model_data, dict): + continue + cost = cast(dict[str, Any], model_data).get("cost") + if not isinstance(cost, dict): + continue + cost_map = cast(dict[str, Any], cost) + try: + price = ModelPrice( + input=_coerce_cost(cost_map.get("input")), + output=_coerce_cost(cost_map.get("output")), + cache_read=_coerce_cost(cost_map.get("cache_read")), + cache_write=_coerce_cost(cost_map.get("cache_write")), + ) + except (TypeError, ValueError): + continue + if model_id not in target: + target[model_id] = price + + merged = {**fallback, **canonical} + return merged + + +def load_catalog() -> dict[str, ModelPrice]: + """Sync. Return flattened {model_id: ModelPrice} from disk cache. + + Returns {} when no cache file exists or the file is unreadable/corrupt. + Memoises by file mtime_ns — re-parses only after a successful refresh. + """ + cache_path = _get_cache_path() + try: + mtime_ns = cache_path.stat().st_mtime_ns + except OSError: + return {} + + cached = _catalog_cache.get("entry") + if cached is not None and cached[0] == mtime_ns: + return cached[1] # type: ignore[return-value] + + try: + raw = json.loads(cache_path.read_text(encoding="utf-8")) + except Exception: + return {} + + # Valid-but-non-object JSON (e.g. a top-level list) must not escape the + # {} fallback contract — _flatten_catalog assumes a provider mapping. + if not isinstance(raw, dict): + return {} + + result = _flatten_catalog(cast(dict[str, Any], raw)) + _catalog_cache["entry"] = (mtime_ns, result) + return result + + +async def _do_fetch(cache_path: Path) -> bool: + """Fetch models.dev/api.json and write it atomically to *cache_path*.""" + tmp_path = cache_path.with_suffix(".tmp") + try: + async with ( + new_client_session() as session, + session.get( + _MODELS_DEV_URL, + timeout=aiohttp.ClientTimeout(total=10), + raise_for_status=True, + ) as resp, + ): + text = await resp.text() + # Validate it's parseable JSON before writing. + json.loads(text) + cache_path.parent.mkdir(parents=True, exist_ok=True) + tmp_path.write_text(text, encoding="utf-8") + os.replace(tmp_path, cache_path) + return True + except Exception as exc: + tmp_path.unlink(missing_ok=True) + logger.debug("models.dev fetch failed: {error}", error=exc) + return False + + +async def refresh_catalog(*, force: bool = False) -> bool: + """Async. Fetch models.dev/api.json if cache is missing or stale (>24h). + + Returns True when cache is valid (fresh or just refreshed), False on + network/write failure. Never raises. + """ + cache_path = _get_cache_path() + if not force: + try: + age = time.time() - cache_path.stat().st_mtime + if age < _TTL_SECONDS: + return True + except OSError: + pass + + async with _refresh_lock: + if not force: + try: + age = time.time() - cache_path.stat().st_mtime + if age < _TTL_SECONDS: + return True + except OSError: + pass + return await _do_fetch(cache_path) diff --git a/src/pythinker_code/soul/flow_runner.py b/src/pythinker_code/soul/flow_runner.py new file mode 100644 index 00000000..aea780b2 --- /dev/null +++ b/src/pythinker_code/soul/flow_runner.py @@ -0,0 +1,215 @@ +"""Agent-flow execution. + +`FlowRunner` drives a `Flow` graph (and the synthesized "ralph" loop) by feeding +node prompts to a `PythinkerSoul` one turn at a time. It collaborates with the +soul purely through its public turn machinery; the soul owns the agent loop, +`FlowRunner` owns graph traversal. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pythinker_core.message import Message + +from pythinker_code.skill.flow import Flow, FlowEdge, FlowNode, parse_choice +from pythinker_code.soul import MaxStepsReached, wire_send +from pythinker_code.utils.logging import logger +from pythinker_code.wire.types import ContentPart, TurnBegin, TurnEnd + +if TYPE_CHECKING: + from pythinker_code.soul.pythinkersoul import PythinkerSoul, TurnOutcome + +FLOW_COMMAND_PREFIX = "flow:" +DEFAULT_MAX_FLOW_MOVES = 1000 +MAX_INVALID_CHOICE_RETRIES = 3 + + +class FlowRunner: + def __init__( + self, + flow: Flow, + *, + name: str | None = None, + max_moves: int = DEFAULT_MAX_FLOW_MOVES, + ) -> None: + self._flow = flow + self._name = name + self._max_moves = max_moves + + @staticmethod + def ralph_loop( + user_message: Message, + max_ralph_iterations: int, + ) -> FlowRunner: + prompt_content = list(user_message.content) + prompt_text = Message(role="user", content=prompt_content).extract_text(" ").strip() + total_runs = max_ralph_iterations + 1 + if max_ralph_iterations < 0: + total_runs = 1000000000000000 # effectively infinite + + nodes: dict[str, FlowNode] = { + "BEGIN": FlowNode(id="BEGIN", label="BEGIN", kind="begin"), + "END": FlowNode(id="END", label="END", kind="end"), + } + outgoing: dict[str, list[FlowEdge]] = {"BEGIN": [], "END": []} + + nodes["R1"] = FlowNode(id="R1", label=prompt_content, kind="task") + nodes["R2"] = FlowNode( + id="R2", + label=( + f"{prompt_text}. (You are running in an automated loop where the same " + "prompt is fed repeatedly. Only choose STOP when the task is fully complete. " + "Including it will stop further iterations. If you are not 100% sure, " + "choose CONTINUE.)" + ).strip(), + kind="decision", + ) + outgoing["R1"] = [] + outgoing["R2"] = [] + + outgoing["BEGIN"].append(FlowEdge(src="BEGIN", dst="R1", label=None)) + outgoing["R1"].append(FlowEdge(src="R1", dst="R2", label=None)) + outgoing["R2"].append(FlowEdge(src="R2", dst="R2", label="CONTINUE")) + outgoing["R2"].append(FlowEdge(src="R2", dst="END", label="STOP")) + + flow = Flow(nodes=nodes, outgoing=outgoing, begin_id="BEGIN", end_id="END") + max_moves = total_runs + return FlowRunner(flow, max_moves=max_moves) + + async def run(self, soul: PythinkerSoul, args: str) -> None: + if args.strip(): + command = f"/{FLOW_COMMAND_PREFIX}{self._name}" if self._name else "/flow" + logger.warning("Agent flow {command} ignores args: {args}", command=command, args=args) + return + if self._name: + from pythinker_code.telemetry import track + + track("flow_invoked", flow_name=self._name) + + current_id = self._flow.begin_id + moves = 0 + total_steps = 0 + while True: + node = self._flow.nodes[current_id] + edges = self._flow.outgoing.get(current_id, []) + + if node.kind == "end": + logger.info("Agent flow reached END node {node_id}", node_id=current_id) + return + + if node.kind == "begin": + if not edges: + logger.error( + 'Agent flow BEGIN node "{node_id}" has no outgoing edges; stopping.', + node_id=node.id, + ) + return + current_id = edges[0].dst + continue + + if moves >= self._max_moves: + raise MaxStepsReached(total_steps) + next_id, steps_used = await self._execute_flow_node(soul, node, edges) + total_steps += steps_used + if next_id is None: + return + moves += 1 + current_id = next_id + + async def _execute_flow_node( + self, + soul: PythinkerSoul, + node: FlowNode, + edges: list[FlowEdge], + ) -> tuple[str | None, int]: + if not edges: + logger.error( + 'Agent flow node "{node_id}" has no outgoing edges; stopping.', + node_id=node.id, + ) + return None, 0 + + base_prompt = self._build_flow_prompt(node, edges) + prompt = base_prompt + steps_used = 0 + retries = 0 + while True: + result = await self._flow_turn(soul, prompt) + steps_used += result.step_count + if result.stop_reason == "tool_rejected": + logger.error("Agent flow stopped after tool rejection.") + return None, steps_used + + if node.kind != "decision": + return edges[0].dst, steps_used + + choice = ( + parse_choice(result.final_message.extract_text(" ")) + if result.final_message + else None + ) + next_id = self._match_flow_edge(edges, choice) + if next_id is not None: + return next_id, steps_used + + retries += 1 + if retries >= MAX_INVALID_CHOICE_RETRIES: + logger.warning( + "Agent flow: max invalid-choice retries ({n}) reached; stopping.", + n=MAX_INVALID_CHOICE_RETRIES, + ) + return None, steps_used + + options = ", ".join(edge.label or "" for edge in edges) + logger.warning( + "Agent flow invalid choice. Got: {choice}. Available: {options}.", + choice=choice or "", + options=options, + ) + prompt = ( + f"{base_prompt}\n\n" + "Your last response did not include a valid choice. " + "Reply with one of the choices using ...." + ) + + @staticmethod + def _build_flow_prompt(node: FlowNode, edges: list[FlowEdge]) -> str | list[ContentPart]: + if node.kind != "decision": + return node.label + + if not isinstance(node.label, str): + label_text = Message(role="user", content=node.label).extract_text(" ") + else: + label_text = node.label + choices = [edge.label for edge in edges if edge.label] + lines = [ + label_text, + "", + "Available branches:", + *(f"- {choice}" for choice in choices), + "", + "Reply with a choice using ....", + ] + return "\n".join(lines) + + @staticmethod + def _match_flow_edge(edges: list[FlowEdge], choice: str | None) -> str | None: + if not choice: + return None + for edge in edges: + if edge.label == choice: + return edge.dst + return None + + @staticmethod + async def _flow_turn( + soul: PythinkerSoul, + prompt: str | list[ContentPart], + ) -> TurnOutcome: + wire_send(TurnBegin(user_input=prompt)) + try: + res = await soul._turn(Message(role="user", content=prompt)) # type: ignore[reportPrivateUsage] + finally: + wire_send(TurnEnd()) + return res diff --git a/src/pythinker_code/soul/pythinkersoul.py b/src/pythinker_code/soul/pythinkersoul.py index 392d4967..4cc39a46 100644 --- a/src/pythinker_code/soul/pythinkersoul.py +++ b/src/pythinker_code/soul/pythinkersoul.py @@ -41,7 +41,6 @@ ) from pythinker_code.prompt_templates import PromptTemplate, expand_prompt_template from pythinker_code.skill import Skill, read_skill_text_with_local_specialization -from pythinker_code.skill.flow import Flow, FlowEdge, FlowNode, parse_choice from pythinker_code.soul import ( LLMNotSet, LLMNotSupported, @@ -74,6 +73,7 @@ ) from pythinker_code.soul.dynamic_injections.auto_mode import AutoModeInjectionProvider from pythinker_code.soul.dynamic_injections.plan_mode import PlanModeInjectionProvider +from pythinker_code.soul.flow_runner import FLOW_COMMAND_PREFIX, FlowRunner from pythinker_code.soul.message import ( check_message, system, @@ -125,8 +125,6 @@ def type_check(soul: PythinkerSoul): SKILL_COMMAND_PREFIX = "skill:" -FLOW_COMMAND_PREFIX = "flow:" -DEFAULT_MAX_FLOW_MOVES = 1000 def _safe_cwd(fallback: str) -> str: @@ -2028,182 +2026,3 @@ class BackToTheFuture(Exception): def __init__(self, checkpoint_id: int, messages: Sequence[Message]): self.checkpoint_id = checkpoint_id self.messages = messages - - -class FlowRunner: - def __init__( - self, - flow: Flow, - *, - name: str | None = None, - max_moves: int = DEFAULT_MAX_FLOW_MOVES, - ) -> None: - self._flow = flow - self._name = name - self._max_moves = max_moves - - @staticmethod - def ralph_loop( - user_message: Message, - max_ralph_iterations: int, - ) -> FlowRunner: - prompt_content = list(user_message.content) - prompt_text = Message(role="user", content=prompt_content).extract_text(" ").strip() - total_runs = max_ralph_iterations + 1 - if max_ralph_iterations < 0: - total_runs = 1000000000000000 # effectively infinite - - nodes: dict[str, FlowNode] = { - "BEGIN": FlowNode(id="BEGIN", label="BEGIN", kind="begin"), - "END": FlowNode(id="END", label="END", kind="end"), - } - outgoing: dict[str, list[FlowEdge]] = {"BEGIN": [], "END": []} - - nodes["R1"] = FlowNode(id="R1", label=prompt_content, kind="task") - nodes["R2"] = FlowNode( - id="R2", - label=( - f"{prompt_text}. (You are running in an automated loop where the same " - "prompt is fed repeatedly. Only choose STOP when the task is fully complete. " - "Including it will stop further iterations. If you are not 100% sure, " - "choose CONTINUE.)" - ).strip(), - kind="decision", - ) - outgoing["R1"] = [] - outgoing["R2"] = [] - - outgoing["BEGIN"].append(FlowEdge(src="BEGIN", dst="R1", label=None)) - outgoing["R1"].append(FlowEdge(src="R1", dst="R2", label=None)) - outgoing["R2"].append(FlowEdge(src="R2", dst="R2", label="CONTINUE")) - outgoing["R2"].append(FlowEdge(src="R2", dst="END", label="STOP")) - - flow = Flow(nodes=nodes, outgoing=outgoing, begin_id="BEGIN", end_id="END") - max_moves = total_runs - return FlowRunner(flow, max_moves=max_moves) - - async def run(self, soul: PythinkerSoul, args: str) -> None: - if args.strip(): - command = f"/{FLOW_COMMAND_PREFIX}{self._name}" if self._name else "/flow" - logger.warning("Agent flow {command} ignores args: {args}", command=command, args=args) - return - if self._name: - from pythinker_code.telemetry import track - - track("flow_invoked", flow_name=self._name) - - current_id = self._flow.begin_id - moves = 0 - total_steps = 0 - while True: - node = self._flow.nodes[current_id] - edges = self._flow.outgoing.get(current_id, []) - - if node.kind == "end": - logger.info("Agent flow reached END node {node_id}", node_id=current_id) - return - - if node.kind == "begin": - if not edges: - logger.error( - 'Agent flow BEGIN node "{node_id}" has no outgoing edges; stopping.', - node_id=node.id, - ) - return - current_id = edges[0].dst - continue - - if moves >= self._max_moves: - raise MaxStepsReached(total_steps) - next_id, steps_used = await self._execute_flow_node(soul, node, edges) - total_steps += steps_used - if next_id is None: - return - moves += 1 - current_id = next_id - - async def _execute_flow_node( - self, - soul: PythinkerSoul, - node: FlowNode, - edges: list[FlowEdge], - ) -> tuple[str | None, int]: - if not edges: - logger.error( - 'Agent flow node "{node_id}" has no outgoing edges; stopping.', - node_id=node.id, - ) - return None, 0 - - base_prompt = self._build_flow_prompt(node, edges) - prompt = base_prompt - steps_used = 0 - while True: - result = await self._flow_turn(soul, prompt) - steps_used += result.step_count - if result.stop_reason == "tool_rejected": - logger.error("Agent flow stopped after tool rejection.") - return None, steps_used - - if node.kind != "decision": - return edges[0].dst, steps_used - - choice = ( - parse_choice(result.final_message.extract_text(" ")) - if result.final_message - else None - ) - next_id = self._match_flow_edge(edges, choice) - if next_id is not None: - return next_id, steps_used - - options = ", ".join(edge.label or "" for edge in edges) - logger.warning( - "Agent flow invalid choice. Got: {choice}. Available: {options}.", - choice=choice or "", - options=options, - ) - prompt = ( - f"{base_prompt}\n\n" - "Your last response did not include a valid choice. " - "Reply with one of the choices using ...." - ) - - @staticmethod - def _build_flow_prompt(node: FlowNode, edges: list[FlowEdge]) -> str | list[ContentPart]: - if node.kind != "decision": - return node.label - - if not isinstance(node.label, str): - label_text = Message(role="user", content=node.label).extract_text(" ") - else: - label_text = node.label - choices = [edge.label for edge in edges if edge.label] - lines = [ - label_text, - "", - "Available branches:", - *(f"- {choice}" for choice in choices), - "", - "Reply with a choice using ....", - ] - return "\n".join(lines) - - @staticmethod - def _match_flow_edge(edges: list[FlowEdge], choice: str | None) -> str | None: - if not choice: - return None - for edge in edges: - if edge.label == choice: - return edge.dst - return None - - @staticmethod - async def _flow_turn( - soul: PythinkerSoul, - prompt: str | list[ContentPart], - ) -> TurnOutcome: - wire_send(TurnBegin(user_input=prompt)) - res = await soul._turn(Message(role="user", content=prompt)) # type: ignore[reportPrivateUsage] - wire_send(TurnEnd()) - return res diff --git a/src/pythinker_code/tools/file/grep_local.py b/src/pythinker_code/tools/file/grep_local.py index 993f741f..42fba7fe 100644 --- a/src/pythinker_code/tools/file/grep_local.py +++ b/src/pythinker_code/tools/file/grep_local.py @@ -167,23 +167,48 @@ class Params(BaseModel): } +_resolved_rg_path: str | None = None + + def _rg_binary_name() -> str: return "rg.exe" if platform.system() == "Windows" else "rg" -def _find_existing_rg(bin_name: str) -> Path | None: +async def _is_runnable(path: Path) -> bool: + """Return True if the binary at *path* can actually be executed on this platform.""" + try: + proc = await asyncio.create_subprocess_exec( + str(path), + "--version", + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + ) + try: + await asyncio.wait_for(proc.wait(), timeout=5) + except TimeoutError: + proc.kill() + await proc.wait() + return False + return proc.returncode == 0 + except OSError: + return False + + +async def _find_existing_rg(bin_name: str) -> Path | None: + # Explicit override — trust it unconditionally. if env_path := os.getenv("PYTHINKER_RG_PATH"): configured = Path(env_path).expanduser() if configured.is_file(): return configured + # Bundled binaries — verify they can actually run on this platform/arch. share_bin = get_share_dir() / "bin" / bin_name - if share_bin.is_file(): + if share_bin.is_file() and await _is_runnable(share_bin): return share_bin assert pythinker_code.__file__ is not None local_dep = Path(pythinker_code.__file__).parent / "deps" / "bin" / bin_name - if local_dep.is_file(): + if local_dep.is_file() and await _is_runnable(local_dep): return local_dep system_rg = shutil.which("rg") @@ -303,18 +328,28 @@ async def _download_and_install_rg(bin_name: str) -> Path: async def _ensure_rg_path() -> str: + global _resolved_rg_path + if _resolved_rg_path is not None: + p = Path(_resolved_rg_path) + if p.exists() and os.access(str(p), os.X_OK): + return _resolved_rg_path + _resolved_rg_path = None + bin_name = _rg_binary_name() - existing = _find_existing_rg(bin_name) + existing = await _find_existing_rg(bin_name) if existing: - return str(existing) + _resolved_rg_path = str(existing) + return _resolved_rg_path async with _RG_DOWNLOAD_LOCK: - existing = _find_existing_rg(bin_name) + existing = await _find_existing_rg(bin_name) if existing: - return str(existing) + _resolved_rg_path = str(existing) + return _resolved_rg_path downloaded = await _download_and_install_rg(bin_name) - return str(downloaded) + _resolved_rg_path = str(downloaded) + return _resolved_rg_path def _build_rg_args(rg_path: str, params: Params, *, single_threaded: bool = False) -> list[str]: diff --git a/src/pythinker_code/ui/shell/prompt.py b/src/pythinker_code/ui/shell/prompt.py index 3758902f..944c93bc 100644 --- a/src/pythinker_code/ui/shell/prompt.py +++ b/src/pythinker_code/ui/shell/prompt.py @@ -1834,9 +1834,18 @@ def _tip(binding: str, fallback: str, description: str) -> str: # Cap prompt redraws at ~30 fps. Smooth streaming calls ``invalidate()`` on a # fast cadence; without this, prompt_toolkit redraws on *every* invalidate, # which (per its own docs) "could cause a lot of terminal output, which some -# terminals are not able to process" — the classic streaming flicker/lag. With -# it, rapid invalidations coalesce into at most one redraw per interval. -_MIN_REDRAW_INTERVAL_S = 1 / 30 +# terminals are not able to process" — the classic streaming flicker/lag. +# +# We use ``max_render_postpone_time`` rather than ``min_redraw_interval``: the +# latter throttles via an ``async def redraw_in_future`` coroutine that +# ``invalidate()`` schedules onto the loop, and during a prompt-app/loop handoff +# (e.g. ``/login`` swapping prompt sessions) that coroutine can be dropped +# un-awaited, emitting a noisy ``RuntimeWarning``. ``max_render_postpone_time`` +# coalesces rapid invalidations through a coroutine-free path (it batches +# redraws up to this deadline, rendering immediately when the loop is idle), so +# it achieves the same throttling without the leak. See +# tests/ui_and_conv/test_redraw_throttle.py. +_MAX_RENDER_POSTPONE_S = 1 / 30 class CustomPromptSession: @@ -2272,7 +2281,9 @@ def _(event: KeyPressEvent) -> None: # slower terminals (best practice for "invalidate is called a lot"). # prompt_toolkit's renderer is already differential (only emits changed # cells), so this caps frame rate without forcing full repaints. - self._session.app.min_redraw_interval = _MIN_REDRAW_INTERVAL_S + # NB: max_render_postpone_time (not min_redraw_interval) — see the + # constant's definition for why the coroutine-free path matters here. + self._session.app.max_render_postpone_time = _MAX_RENDER_POSTPONE_S self._session.default_buffer.read_only = Condition( lambda: ( (delegate := self._active_prompt_delegate()) is not None diff --git a/src/pythinker_code/ui/shell/stats.py b/src/pythinker_code/ui/shell/stats.py index 3d8c41c7..b9ab80b8 100644 --- a/src/pythinker_code/ui/shell/stats.py +++ b/src/pythinker_code/ui/shell/stats.py @@ -11,6 +11,7 @@ from prompt_toolkit.layout.controls import FormattedTextControl from prompt_toolkit.styles import Style +from pythinker_code.models_dev import refresh_catalog as _refresh_catalog from pythinker_code.ui.shell.console import console from pythinker_code.ui.shell.slash import registry from pythinker_code.ui.shell.stats_collector import ( @@ -38,7 +39,7 @@ } -def _fmt_cost(v: float) -> str: +def fmt_cost(v: float) -> str: if v == 0: return "-" if v < 0.01: @@ -184,7 +185,7 @@ def _pad_left(s: str, w: int) -> str: 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_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"]) @@ -198,7 +199,7 @@ def _pad_left(s: str, w: int) -> str: 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_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"]) @@ -210,7 +211,7 @@ def _pad_left(s: str, w: int) -> str: 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"]) + tot += _pad_left(fmt_cost(cur.total_cost), col_w["cost"]) parts.append(("bold", tot + "\n")) parts.append(("", "\n")) @@ -319,6 +320,7 @@ async def run(self) -> None: @registry.command(name="stats", aliases=["history"]) async def stats(app: Shell, args: str) -> None: """Show usage statistics dashboard (tokens and cost by provider/model).""" + await _refresh_catalog() from pythinker_code.ui.theme import get_tui_tokens as _get_tui_tokens _t = _get_tui_tokens() diff --git a/src/pythinker_code/ui/shell/stats_pricing.py b/src/pythinker_code/ui/shell/stats_pricing.py index 220648bb..8079845e 100644 --- a/src/pythinker_code/ui/shell/stats_pricing.py +++ b/src/pythinker_code/ui/shell/stats_pricing.py @@ -69,28 +69,55 @@ 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. + Resolution order: + 1. Exact match in models.dev catalog (load_catalog) + 2. Longest prefix match in catalog + 3. Exact match in hardcoded _PRICE_TABLE (offline/unknown-model fallback) + 4. Longest prefix match in _PRICE_TABLE + 5. 0.0 """ - 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 + from pythinker_code.models_dev import ModelPrice, load_catalog - 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 + def _apply(p: tuple[float, float, float, float] | ModelPrice) -> float: + if isinstance(p, ModelPrice): + inp, out, cr, cw = p.input, p.output, p.cache_read, p.cache_write + else: + inp, out, cr, cw = p + return ( + usage.input_other * inp + + usage.output * out + + usage.input_cache_read * cr + + usage.input_cache_creation * cw + ) / 1_000_000 + + catalog = load_catalog() + + # 1. Catalog exact + if model in catalog: + return _apply(catalog[model]) + + # 2. Catalog prefix (longest wins) + best_cat: ModelPrice | None = None + best_cat_len = 0 + for key, val in catalog.items(): + if model.startswith(key) and len(key) > best_cat_len: + best_cat = val + best_cat_len = len(key) + if best_cat is not None: + return _apply(best_cat) + + # 3. Hardcoded exact + if model in _PRICE_TABLE: + return _apply(_PRICE_TABLE[model]) + + # 4. Hardcoded prefix (longest 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 not None: + return _apply(best) + + return 0.0 diff --git a/src/pythinker_code/ui/shell/usage.py b/src/pythinker_code/ui/shell/usage.py index 8431943c..f47d7ad5 100644 --- a/src/pythinker_code/ui/shell/usage.py +++ b/src/pythinker_code/ui/shell/usage.py @@ -9,9 +9,12 @@ from pythinker_code.auth.platforms import parse_managed_provider_key from pythinker_code.config import LLMProvider +from pythinker_code.models_dev import refresh_catalog as _refresh_catalog from pythinker_code.soul.pythinkersoul import PythinkerSoul 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 +from pythinker_code.ui.shell.stats_collector import load_all_stats as _load_all_stats_raw from pythinker_code.ui.shell.usage_adapters import ADAPTERS from pythinker_code.ui.shell.usage_adapters.base import UsageAdapter, UsageReport, UsageRow from pythinker_code.ui.shell.usage_adapters.pythinker import to_int as _to_int @@ -30,6 +33,7 @@ from pythinker_code.ui.theme import get_tui_tokens as _get_tui_tokens from pythinker_code.usage_ratelimit_cache import get_cache from pythinker_code.utils.datetime import format_duration +from pythinker_code.utils.logging import logger if TYPE_CHECKING: from pythinker_code.auth.oauth import OAuthManager @@ -199,6 +203,58 @@ def _enrich_with_ratelimit_fallback( return enriched +def _load_cost_stats() -> AllStats | None: + """Return AllStats from wire files, or None if no data or error.""" + try: + stats = _load_all_stats_raw() + if stats.periods["all_time"].total_messages == 0: + return None + return stats + except Exception: + return None + + +def _build_cost_panel(stats: AllStats): + from rich import box as _box + from rich.panel import Panel as _Panel + from rich.table import Table as _Table + + from pythinker_code.ui.shell.stats import fmt_cost + from pythinker_code.ui.theme import tui_rich_style + + t = _Table.grid(padding=(0, 2)) + t.add_column(style=tui_rich_style("info")) + t.add_column(style="bold") + + periods = [ + ("Today", stats.periods["today"].total_cost), + ("This Week", stats.periods["this_week"].total_cost), + ("All time", stats.periods["all_time"].total_cost), + ] + for label, cost in periods: + t.add_row(label, fmt_cost(cost)) + + return _Panel( + t, + title="Session Cost", + border_style=tui_rich_style("border_muted"), + box=_box.ROUNDED, + padding=(0, 2), + expand=False, + ) + + +async def _maybe_print_cost_panel() -> None: + """Print the session cost panel if local usage data exists. Never raises.""" + try: + stats = await asyncio.to_thread(_load_cost_stats) + if stats is None: + return + console.print(_build_cost_panel(stats)) + except Exception as e: + logger.debug("cost panel failed to render: {error}", error=e, exc_info=True) + + @registry.command(aliases=["status", "cost", "/status"]) async def usage(app: Shell, args: str): """Display usage for the current model's provider. @@ -206,6 +262,7 @@ async def usage(app: Shell, args: str): Pass `all` for every provider, or a provider key to filter. """ assert isinstance(app.soul, PythinkerSoul) + await _refresh_catalog() _t = _get_tui_tokens() try: @@ -294,3 +351,5 @@ async def usage(app: Shell, args: str): for report in non_empty_reports: console.print(build_panel(report)) + + await _maybe_print_cost_panel() diff --git a/src/pythinker_code/ui/shell/usage_adapters/__init__.py b/src/pythinker_code/ui/shell/usage_adapters/__init__.py index 35decaa3..5c85d0a4 100644 --- a/src/pythinker_code/ui/shell/usage_adapters/__init__.py +++ b/src/pythinker_code/ui/shell/usage_adapters/__init__.py @@ -1,5 +1,6 @@ from __future__ import annotations +from pythinker_code.ui.shell.usage_adapters.alibaba import AlibabaAdapter from pythinker_code.ui.shell.usage_adapters.anthropic_admin import AnthropicAdminAdapter from pythinker_code.ui.shell.usage_adapters.base import ( UsageAdapter, @@ -27,6 +28,7 @@ # registry has to cover every variant — otherwise `_select_providers` filters # the active provider out and `/usage` falls into the no-adapter branch. ADAPTERS: dict[str, UsageAdapter] = { + AlibabaAdapter.platform_id: AlibabaAdapter(), AnthropicAdminAdapter.platform_id: AnthropicAdminAdapter(), DeepSeekAdapter.platform_id: DeepSeekAdapter(), OpenAIAdminAdapter.platform_id: OpenAIAdminAdapter(), diff --git a/src/pythinker_code/ui/shell/usage_adapters/alibaba.py b/src/pythinker_code/ui/shell/usage_adapters/alibaba.py new file mode 100644 index 00000000..13d2e422 --- /dev/null +++ b/src/pythinker_code/ui/shell/usage_adapters/alibaba.py @@ -0,0 +1,180 @@ +"""Usage adapter for Alibaba DashScope.""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING, Any, cast +from urllib.parse import urlparse + +import aiohttp + +from pythinker_code.auth import ALIBABA_PLATFORM_ID +from pythinker_code.auth.alibaba import ALIBABA_BASE_URL +from pythinker_code.ui.shell.stats_collector import load_all_stats as _load_all_stats +from pythinker_code.ui.shell.usage_adapters.base import UsageReport, UsageRow +from pythinker_code.usage_ratelimit_cache import get_cache +from pythinker_code.utils.aiohttp import new_client_session +from pythinker_code.utils.logging import logger + +if TYPE_CHECKING: + from pythinker_code.auth.oauth import OAuthManager + from pythinker_code.config import LLMProvider + +_QUOTA_PATH = "/api/v1/quotas" +_TIMEOUT = aiohttp.ClientTimeout(total=8, sock_connect=5) + +# DashScope quota API lives on a separate host from the compatible-mode endpoint. +# The "dashscope-us" host is the completion endpoint; quota is on "dashscope-intl". +_DASHSCOPE_HOST_REMAP: dict[str, str] = { + "dashscope-us.aliyuncs.com": "dashscope-intl.aliyuncs.com", +} + + +def _quota_url(base_url: str) -> str: + """Derive the quota API URL from the configured base URL. + + The OpenAI-compatible completion host and the quota API host differ for the + international (US) region — remap known mismatches before constructing the URL. + """ + parsed = urlparse(base_url) + host = _DASHSCOPE_HOST_REMAP.get(parsed.netloc, parsed.netloc) + return f"{parsed.scheme}://{host}{_QUOTA_PATH}" + + +def _parse_quota_response(data: object) -> list[UsageRow]: + """Parse DashScope /api/v1/quotas response into UsageRow list. + + Handles both flat {token_quota, token_used} and {quota_list: [...]} shapes. + Returns [] if the response has an unrecognised shape. + """ + if not isinstance(data, dict): + return [] + data_map = cast(dict[str, Any], data) + rows: list[UsageRow] = [] + payload = data_map.get("data", data_map) + if not isinstance(payload, dict): + return [] + payload_map = cast(dict[str, Any], payload) + + # Flat shape: {token_quota: N, token_used: N} + total = ( + payload_map["token_quota"] + if "token_quota" in payload_map + else payload_map.get("total_quota") + ) + used = ( + payload_map["token_used"] if "token_used" in payload_map else payload_map.get("total_used") + ) + if isinstance(total, (int, float)) and isinstance(used, (int, float)): + rows.append(UsageRow(label="Token quota", used=int(used), limit=int(total), unit="tokens")) + + # List shape: {quota_list: [{quota_name, total_quota, total_used}]} + quota_list = payload_map.get("quota_list") + if isinstance(quota_list, list): + for item in cast(list[Any], quota_list): + if not isinstance(item, dict): + continue + item_map = cast(dict[str, Any], item) + name = item_map.get("quota_name") or item_map.get("quota_type") or "Quota" + t = item_map.get("total_quota") + u = item_map.get("total_used") + if isinstance(t, (int, float)) and isinstance(u, (int, float)): + rows.append( + UsageRow(label=str(name).title(), used=int(u), limit=int(t), unit="tokens") + ) + + return rows + + +class AlibabaAdapter: + platform_id = ALIBABA_PLATFORM_ID + provider_label = "Alibaba DashScope" + requires_admin_key = False + + async def fetch(self, provider: LLMProvider, oauth_mgr: OAuthManager) -> UsageReport: + api_key = provider.api_key.get_secret_value() + base_url = provider.base_url or ALIBABA_BASE_URL + provider_key = f"managed:{ALIBABA_PLATFORM_ID}" + + # --- Local token stats (primary source, always available) --- + local_rows: list[UsageRow] = [] + try: + all_stats = await asyncio.to_thread(_load_all_stats) + for period_name, label in (("today", "Today"), ("all_time", "All time")): + period = all_stats.periods.get(period_name) + if period is None: + continue + prov = period.providers.get(provider_key) + if prov is None or prov.messages == 0: + continue + local_rows.append( + UsageRow( + label=label, + used=prov.tokens, + limit=0, + unit="tokens", + reset_hint=( + f"↑{prov.input_other:,} in ↓{prov.output:,} out ${prov.cost:.4f}" + ), + ) + ) + except Exception as e: + logger.debug("local usage stats unavailable: {error}", error=e, exc_info=True) + + # --- DashScope quota API (best-effort; logs + notes on failure) --- + quota_rows: list[UsageRow] = [] + notes: list[str] = [] + try: + async with ( + new_client_session(timeout=_TIMEOUT) as session, + session.get( + _quota_url(base_url), + headers={"Authorization": f"Bearer {api_key}"}, + ) as resp, + ): + if resp.status == 200: + data = await resp.json(content_type=None) + quota_rows = _parse_quota_response(data) + elif resp.status in (401, 403): + notes.append( + "DashScope quota API: authorization failed — quota data unavailable." + ) + except (aiohttp.ClientError, TimeoutError) as e: + logger.debug("DashScope quota API request failed: {error}", error=e, exc_info=True) + notes.append("DashScope quota API unavailable right now — retry in a moment.") + + # --- Rate-limit headers from last response (if DashScope ever sends them) --- + rl_rows: list[UsageRow] = [] + snap = get_cache().snapshot(provider_key) + if snap is not None: + if snap.requests_limit is not None and snap.requests_remaining is not None: + rl_rows.append( + UsageRow( + label="Requests remaining", + used=snap.requests_remaining, + limit=snap.requests_limit, + unit="requests", + ) + ) + if snap.tokens_limit is not None and snap.tokens_remaining is not None: + rl_rows.append( + UsageRow( + label="Tokens remaining", + used=snap.tokens_remaining, + limit=snap.tokens_limit, + unit="tokens", + ) + ) + + all_rows = local_rows + quota_rows + rl_rows + if not all_rows and not notes: + notes.append("No usage recorded yet. Start a conversation to see token counts here.") + + summary = all_rows[0] if all_rows else None + return UsageReport( + provider_label=self.provider_label, + summary=summary, + limits=all_rows[1:], + notes=notes, + unit_hint="tokens", + ) diff --git a/tasks/decomposition-plan.md b/tasks/decomposition-plan.md new file mode 100644 index 00000000..86981f7d --- /dev/null +++ b/tasks/decomposition-plan.md @@ -0,0 +1,69 @@ +# God-Object Decomposition Plan + +Staged, one-collaborator-per-PR decomposition of the two largest behavioral +classes. **Invariant for every PR: public API frozen, behavior identical, +full test suite green before and after.** Each PR extracts a collaborator that +the host class *delegates to* — no logic rewrites, just relocation behind a +narrow interface. + +Guiding rules: +- Characterize before cutting: confirm the seam is covered by existing tests + *before* moving code. If a seam is thin on tests, add characterization tests + in a separate prior PR. +- One collaborator per PR. Reviewable diffs over big-bang. +- No new abstractions that don't reduce the host class's surface. A 400-line + helper with one caller is the same code in a new file — only worth it if it + is independently testable and shrinks the host's responsibility count. +- Delegation, not duplication: the host keeps thin forwarding methods where the + public API requires them. + +--- + +## Phase A — PythinkerSoul (`src/pythinker_code/soul/pythinkersoul.py`, 2209 lines) + +Seams identified by responsibility cluster (line ranges approximate, current main): + +| PR | Collaborator | Methods to extract (from PythinkerSoul) | Lines | Risk | +|-----|-------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------|------------|------| +| A0 | (characterization only) | Add tests for any thin-covered seam below before moving it (esp. compaction, connection recovery). | n/a | none | +| A1 | `FlowRunner` → own module | Already a *separate class* in this file (lines ~2033-end). Move verbatim to `soul/flow_runner.py`, fix imports. Pure relocation. | ~180 | low | +| A2 | `PlanModeController` | `_bind_plan_mode_tools`, `_ensure_plan_session_id`, `_set_plan_mode`, `get_plan_file_path`, `read_current_plan`, `clear_current_plan`, `toggle_plan_mode*`, `set_plan_mode_from_manual`, `schedule_plan_activation_reminder`, `consume_pending_plan_activation_injection`, `plan_mode` property | ~160 | med | +| A3 | `InjectionManager` | `add_injection_provider`, `rearm_injection`, `_collect_injections`, `_notify_injection_providers_compacted`, `consume_pending_plan_activation_injection` glue | ~70 | med | +| A4 | `SteerQueue` | `steer`, `_consume_pending_steers`, `_inject_steer` + the pending-steer buffer state | ~40 | low | +| A5 | `SlashCommandRegistry` | `available_slash_commands`, `_build_slash_commands`, `_index_slash_commands`, `_find_slash_command`, `_make_prompt_template_runner`, `_make_skill_runner`, `_record_invoked_skill` | ~140 | med | +| A6 | `ConnectionRecovery` | `_is_retryable_error`, `_run_with_connection_recovery`, `_retry_log`, `_emit_step_retry` | ~130 | med | +| A7 | `ContextCompactor` | `_grow_context`, `compact_context`, `_harvest_before_compaction`, `_context_usage` | ~240 | high | + +After A1–A7, the host retains its identity: lifecycle/state (`__init__`, +status, model/agent/runtime/context accessors) and the core loop (`run`, +`_turn`, `_agent_loop`, `_step`) — the irreducible "soul". Target: host class +drops from ~2200 to roughly ~1100 lines. + +**Sequencing rationale:** A1 first (zero-risk warm-up, proves the import-move +mechanics). Then ascending risk. A7 (compaction) last — it is the most +state-entangled (touches context, harvesting, injection-notify) and benefits +from A3 already being extracted. + +**Per-PR verification:** +1. `make check` (ruff + pyright) clean. +2. Full `uv run pytest` green (diff the pass count vs. main — must be ≥). +3. For A7 specifically: add an explicit before/after compaction behavior test + asserting message-count and context-usage parity on a fixed transcript. + +--- + +## Phase B — Shell trio (needs its own seam-discovery pass first) + +`ui/shell/prompt.py` (3591), `ui/shell/__init__.py` (2107), `ui/shell/slash.py` +(1881) are collectively larger than PythinkerSoul and are the bigger target — +but I have **not** yet mapped their seams. Do a discovery PR (read-only, +produces a seam table like Phase A) before committing to extraction PRs. Do not +start Phase B until Phase A lands and stabilizes. + +--- + +## Out of scope (logged, not fixed here) + +- `models_dev.py:65` — `ty` flags `"@" in model_id` (`model_id` typed `object`). + Confirmed `ty` false positive (JSON dict key, always `str`; pyright is fine). + Left as-is; documents why `ty` stays advisory. diff --git a/tests/core/test_wire_file_compat.py b/tests/core/test_wire_file_compat.py new file mode 100644 index 00000000..21caa79a --- /dev/null +++ b/tests/core/test_wire_file_compat.py @@ -0,0 +1,81 @@ +"""Backward-compatibility tests for the wire-file protocol version header. + +Pins the version-detection behaviour of `WireFile`: files written before the +metadata header existed (legacy, headerless) must still be readable and must +report `WIRE_PROTOCOL_LEGACY_VERSION`. A regression here would silently break +old `wire.jsonl` files for every frontend (Shell, Web, Vis, ACP). +""" + +from __future__ import annotations + +import json +import time +from pathlib import Path + +from pythinker_code.wire.file import ( + WireFile, + WireFileMetadata, + WireMessageRecord, + _load_protocol_version, +) +from pythinker_code.wire.protocol import ( + WIRE_PROTOCOL_LEGACY_VERSION, + WIRE_PROTOCOL_VERSION, +) +from pythinker_code.wire.types import TextPart, TurnBegin + + +def _record_line() -> str: + """A single wire record serialized as one jsonl line (no metadata header).""" + record = WireMessageRecord.from_wire_message( + TurnBegin(user_input=[TextPart(text="hello")]), + timestamp=time.time(), + ) + return json.dumps(record.model_dump(mode="json")) + "\n" + + +def test_headerless_file_reports_legacy_version(tmp_path: Path) -> None: + """A pre-header wire file (records only) is detected as the legacy version. + + `WireFile.append_record` always emits a metadata header, so a legacy file + can only be produced by writing the record line directly. + """ + path = tmp_path / "wire.jsonl" + path.write_text(_record_line(), encoding="utf-8") + + assert WireFile(path).version == WIRE_PROTOCOL_LEGACY_VERSION + + +async def test_headerless_records_still_parse(tmp_path: Path) -> None: + """Legacy (headerless) files remain readable end-to-end.""" + path = tmp_path / "wire.jsonl" + path.write_text(_record_line(), encoding="utf-8") + + records = [r async for r in WireFile(path).iter_records()] + assert len(records) == 1 + assert isinstance(records[0].to_wire_message(), TurnBegin) + + +def test_explicit_version_header_round_trips(tmp_path: Path) -> None: + """A file whose header pins a version reports exactly that version.""" + path = tmp_path / "wire.jsonl" + metadata = WireFileMetadata(protocol_version="1.4") + path.write_text( + json.dumps(metadata.model_dump(mode="json")) + "\n" + _record_line(), + encoding="utf-8", + ) + + assert WireFile(path).version == "1.4" + + +def test_new_file_uses_current_version(tmp_path: Path) -> None: + """A not-yet-existing wire file defaults to the current protocol version.""" + path = tmp_path / "does-not-exist.jsonl" + assert WireFile(path).version == WIRE_PROTOCOL_VERSION + + +def test_load_protocol_version_none_for_headerless(tmp_path: Path) -> None: + """The low-level loader returns None when the first line is not metadata.""" + path = tmp_path / "wire.jsonl" + path.write_text(_record_line(), encoding="utf-8") + assert _load_protocol_version(path) is None diff --git a/tests/fixtures/models-dev-subset.json b/tests/fixtures/models-dev-subset.json new file mode 100644 index 00000000..440363f6 --- /dev/null +++ b/tests/fixtures/models-dev-subset.json @@ -0,0 +1,65 @@ +{ + "anthropic": { + "id": "anthropic", + "name": "Anthropic", + "models": { + "claude-sonnet-4-6": { + "id": "claude-sonnet-4-6", + "name": "Claude Sonnet 4.6", + "cost": { + "input": 3.0, + "output": 15.0, + "cache_read": 0.3, + "cache_write": 3.75, + "context_over_200k": { + "input": 6.0, + "output": 22.5 + } + }, + "limit": { "context": 1000000, "output": 64000 } + } + } + }, + "google-vertex-anthropic": { + "id": "google-vertex-anthropic", + "name": "Vertex (Anthropic)", + "models": { + "claude-sonnet-4-6": { + "id": "claude-sonnet-4-6", + "name": "Claude Sonnet 4.6", + "cost": { "input": 3.1, "output": 15.1 }, + "limit": { "context": 1000000 } + }, + "claude-sonnet-4-6@default": { + "id": "claude-sonnet-4-6@default", + "name": "Claude Sonnet 4.6 default", + "cost": { "input": 3.1, "output": 15.1 }, + "limit": { "context": 1000000 } + } + } + }, + "openai": { + "id": "openai", + "name": "OpenAI", + "models": { + "gpt-4o-mini": { + "id": "gpt-4o-mini", + "name": "GPT-4o mini", + "cost": { "input": 0.15, "output": 0.6, "cache_read": 0.08 }, + "limit": { "context": 128000 } + } + } + }, + "unknown-provider": { + "id": "unknown-provider", + "name": "Some Unknown Provider", + "models": { + "some-model": { + "id": "some-model", + "name": "Some Model", + "cost": { "input": 1.0, "output": 2.0 }, + "limit": { "context": 8000 } + } + } + } +} diff --git a/tests/test_models_dev.py b/tests/test_models_dev.py new file mode 100644 index 00000000..cb734e8e --- /dev/null +++ b/tests/test_models_dev.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest + +FIXTURE_PATH = Path(__file__).parent / "fixtures" / "models-dev-subset.json" + + +def _fixture_json() -> str: + return FIXTURE_PATH.read_text() + + +def _fixture_dict() -> dict: + return json.loads(_fixture_json()) + + +# --------------------------------------------------------------------------- +# flatten_catalog +# --------------------------------------------------------------------------- + + +def test_flatten_canonical_provider_wins(): + from pythinker_code.models_dev import _flatten_catalog + + catalog = _fixture_dict() + result = _flatten_catalog(catalog) + # anthropic is canonical; google-vertex-anthropic is not + assert "claude-sonnet-4-6" in result + p = result["claude-sonnet-4-6"] + assert p.input == 3.0 # anthropic price, not vertex (3.1) + + +def test_flatten_skips_versioned_ids(): + from pythinker_code.models_dev import _flatten_catalog + + catalog = _fixture_dict() + result = _flatten_catalog(catalog) + assert "claude-sonnet-4-6@default" not in result + + +def test_flatten_missing_cache_fields_defaults_to_zero(): + from pythinker_code.models_dev import _flatten_catalog + + catalog = _fixture_dict() + result = _flatten_catalog(catalog) + # gpt-4o-mini has no cache_write in fixture + assert result["gpt-4o-mini"].cache_write == 0.0 + # openai gpt-4o-mini has cache_read in fixture + assert result["gpt-4o-mini"].cache_read == 0.08 + + +def test_flatten_unknown_provider_included_as_fallback(): + from pythinker_code.models_dev import _flatten_catalog + + catalog = _fixture_dict() + result = _flatten_catalog(catalog) + # unknown-provider/some-model not overridden by canonical + assert "some-model" in result + + +def test_flatten_context_over_200k_ignored(): + from pythinker_code.models_dev import _flatten_catalog + + catalog = _fixture_dict() + result = _flatten_catalog(catalog) + # Base tier only — input should be 3.0, not 6.0 + assert result["claude-sonnet-4-6"].input == 3.0 + + +def test_flatten_malformed_costs_skipped(): + from pythinker_code.models_dev import _flatten_catalog + + catalog = { + "openai": { + "models": { + "bad-model": {"cost": {"input": "n/a", "output": 1.0}}, + "good-model": {"cost": {"input": 1.0, "output": 2.0}}, + } + } + } + result = _flatten_catalog(catalog) + assert "bad-model" not in result + assert "good-model" in result + assert result["good-model"].input == 1.0 + + +# --------------------------------------------------------------------------- +# load_catalog +# --------------------------------------------------------------------------- + + +def test_load_catalog_empty_when_no_cache_file(tmp_path, monkeypatch): + from pythinker_code import models_dev + + monkeypatch.setattr(models_dev, "_get_cache_path", lambda: tmp_path / "nonexistent.json") + models_dev._catalog_cache.clear() + result = models_dev.load_catalog() + assert result == {} + + +def test_load_catalog_parses_valid_cache(tmp_path, monkeypatch): + from pythinker_code import models_dev + + cache_file = tmp_path / "models-dev.json" + cache_file.write_text(_fixture_json()) + monkeypatch.setattr(models_dev, "_get_cache_path", lambda: cache_file) + models_dev._catalog_cache.clear() + result = models_dev.load_catalog() + assert "claude-sonnet-4-6" in result + assert result["claude-sonnet-4-6"].input == 3.0 + + +def test_load_catalog_returns_empty_on_corrupt_json(tmp_path, monkeypatch): + from pythinker_code import models_dev + + cache_file = tmp_path / "models-dev.json" + cache_file.write_text("{not valid json!!!") + monkeypatch.setattr(models_dev, "_get_cache_path", lambda: cache_file) + models_dev._catalog_cache.clear() + result = models_dev.load_catalog() + assert result == {} + + +def test_load_catalog_returns_empty_on_non_dict_json(tmp_path, monkeypatch): + from pythinker_code import models_dev + + # Valid JSON, but the root is a list — _flatten_catalog would blow up on + # .items(); load_catalog must honour its {} fallback contract instead. + cache_file = tmp_path / "models-dev.json" + cache_file.write_text("[]") + monkeypatch.setattr(models_dev, "_get_cache_path", lambda: cache_file) + models_dev._catalog_cache.clear() + result = models_dev.load_catalog() + assert result == {} + + +def test_load_catalog_memoized_by_mtime(tmp_path, monkeypatch): + from pythinker_code import models_dev + + cache_file = tmp_path / "models-dev.json" + cache_file.write_text(_fixture_json()) + monkeypatch.setattr(models_dev, "_get_cache_path", lambda: cache_file) + models_dev._catalog_cache.clear() + result1 = models_dev.load_catalog() + # Second call with same mtime returns same object + result2 = models_dev.load_catalog() + assert result1 is result2 + + +# --------------------------------------------------------------------------- +# refresh_catalog +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_refresh_catalog_writes_cache(tmp_path, monkeypatch): + from pythinker_code import models_dev + + cache_file = tmp_path / "pricing" / "models-dev.json" + monkeypatch.setattr(models_dev, "_get_cache_path", lambda: cache_file) + models_dev._catalog_cache.clear() + + mock_session = AsyncMock() + mock_resp = AsyncMock() + mock_resp.text = AsyncMock(return_value=_fixture_json()) + mock_resp.__aenter__ = AsyncMock(return_value=mock_resp) + mock_resp.__aexit__ = AsyncMock(return_value=False) + mock_session.get = lambda *a, **kw: mock_resp + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=False) + + with patch("pythinker_code.models_dev.new_client_session", return_value=mock_session): + result = await models_dev.refresh_catalog(force=True) + + assert result is True + assert cache_file.exists() + assert "anthropic" in json.loads(cache_file.read_text()) + + +@pytest.mark.asyncio +async def test_refresh_catalog_noop_within_ttl(tmp_path, monkeypatch): + from pythinker_code import models_dev + + cache_file = tmp_path / "models-dev.json" + cache_file.write_text(_fixture_json()) + # mtime is fresh (just written) + monkeypatch.setattr(models_dev, "_get_cache_path", lambda: cache_file) + models_dev._catalog_cache.clear() + + with patch("pythinker_code.models_dev.new_client_session") as mock_new_session: + result = await models_dev.refresh_catalog(force=False) + + # Cache is fresh — no network I/O should occur + assert result is True + mock_new_session.assert_not_called() + + +@pytest.mark.asyncio +async def test_refresh_catalog_fetches_when_stale(tmp_path, monkeypatch): + import os + + from pythinker_code import models_dev + + cache_file = tmp_path / "models-dev.json" + cache_file.write_text(_fixture_json()) + # Backdate mtime by 25 hours (beyond 24h TTL) + stale_mtime = cache_file.stat().st_mtime - (25 * 3600) + os.utime(cache_file, (stale_mtime, stale_mtime)) + monkeypatch.setattr(models_dev, "_get_cache_path", lambda: cache_file) + models_dev._catalog_cache.clear() + + mock_session = AsyncMock() + mock_resp = AsyncMock() + mock_resp.text = AsyncMock(return_value=_fixture_json()) + mock_resp.__aenter__ = AsyncMock(return_value=mock_resp) + mock_resp.__aexit__ = AsyncMock(return_value=False) + mock_session.get = lambda *a, **kw: mock_resp + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=False) + + with patch( + "pythinker_code.models_dev.new_client_session", return_value=mock_session + ) as mock_new_session: + result = await models_dev.refresh_catalog(force=False) + + # Stale cache — a real network fetch must have been attempted + assert result is True + mock_new_session.assert_called_once() + + +@pytest.mark.asyncio +async def test_refresh_catalog_swallows_network_error(tmp_path, monkeypatch): + import aiohttp + + from pythinker_code import models_dev + + cache_file = tmp_path / "models-dev.json" + monkeypatch.setattr(models_dev, "_get_cache_path", lambda: cache_file) + models_dev._catalog_cache.clear() + + mock_resp = AsyncMock() + mock_resp.__aenter__ = AsyncMock(side_effect=aiohttp.ClientError("connection refused")) + mock_resp.__aexit__ = AsyncMock(return_value=False) + + mock_session = AsyncMock() + mock_session.get = lambda *a, **kw: mock_resp + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=False) + + with patch("pythinker_code.models_dev.new_client_session", return_value=mock_session): + result = await models_dev.refresh_catalog(force=True) + + assert result is False + assert not cache_file.exists() + + +@pytest.mark.asyncio +async def test_refresh_catalog_atomic_write(tmp_path, monkeypatch): + from pythinker_code import models_dev + + cache_file = tmp_path / "models-dev.json" + monkeypatch.setattr(models_dev, "_get_cache_path", lambda: cache_file) + models_dev._catalog_cache.clear() + + mock_session = AsyncMock() + mock_resp = AsyncMock() + mock_resp.text = AsyncMock(return_value=_fixture_json()) + mock_resp.__aenter__ = AsyncMock(return_value=mock_resp) + mock_resp.__aexit__ = AsyncMock(return_value=False) + mock_session.get = lambda *a, **kw: mock_resp + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=False) + + with patch("pythinker_code.models_dev.new_client_session", return_value=mock_session): + await models_dev.refresh_catalog(force=True) + + # No stray .tmp file left behind + tmp_file = cache_file.with_suffix(".tmp") + assert not tmp_file.exists() + assert cache_file.exists() diff --git a/tests/tools/test_grep.py b/tests/tools/test_grep.py index 1f3b2527..c4aa7a0f 100644 --- a/tests/tools/test_grep.py +++ b/tests/tools/test_grep.py @@ -24,18 +24,19 @@ @pytest_asyncio.fixture(scope="module") async def grep_tool() -> Grep: """Create a Grep tool instance when a local ripgrep binary is available.""" - if _find_existing_rg(_rg_binary_name()) is None: + if await _find_existing_rg(_rg_binary_name()) is None: pytest.skip("ripgrep binary is not available in this environment") return Grep() -def test_find_existing_rg_honors_env_path(monkeypatch, tmp_path): +@pytest.mark.asyncio +async def test_find_existing_rg_honors_env_path(monkeypatch, tmp_path): rg_path = tmp_path / _rg_binary_name() rg_path.write_text("fake rg") rg_path.chmod(0o755) monkeypatch.setenv("PYTHINKER_RG_PATH", str(rg_path)) - assert _find_existing_rg(_rg_binary_name()) == rg_path + assert await _find_existing_rg(_rg_binary_name()) == rg_path @pytest.fixture diff --git a/tests/ui/test_usage_cost_panel.py b/tests/ui/test_usage_cost_panel.py new file mode 100644 index 00000000..30b8fb0f --- /dev/null +++ b/tests/ui/test_usage_cost_panel.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from rich.console import Console + +from pythinker_code.ui.shell.stats_collector import AllStats, PeriodStats + + +def _make_all_stats(total_cost: float, messages: int) -> AllStats: + p = PeriodStats() + p.total_cost = total_cost + p.total_messages = messages + return AllStats( + periods={ + "today": p, + "this_week": p, + "last_week": PeriodStats(), + "all_time": p, + }, + insights={}, + ) + + +def _render_panel(panel) -> str: + console = Console(force_terminal=False, width=80) + with console.capture() as cap: + console.print(panel) + return cap.get() + + +def test_build_cost_panel_shows_cost(): + from pythinker_code.ui.shell.usage import _build_cost_panel + + stats = _make_all_stats(total_cost=1.23, messages=5) + panel = _build_cost_panel(stats) + rendered = _render_panel(panel) + assert "$1.23" in rendered + + +def test_build_cost_panel_shows_today_and_week(): + from pythinker_code.ui.shell.usage import _build_cost_panel + + stats = _make_all_stats(total_cost=0.05, messages=3) + panel = _build_cost_panel(stats) + rendered = _render_panel(panel) + assert "Today" in rendered + assert "Week" in rendered + + +def test_build_cost_panel_title(): + from pythinker_code.ui.shell.usage import _build_cost_panel + + stats = _make_all_stats(total_cost=0.01, messages=1) + panel = _build_cost_panel(stats) + assert "Session Cost" in str(panel.title) + + +async def test_usage_prints_cost_panel_when_data_exists(monkeypatch): + from rich.panel import Panel + + from pythinker_code.ui.shell import usage as usage_module + + stats = _make_all_stats(total_cost=2.50, messages=10) + monkeypatch.setattr(usage_module, "_load_cost_stats", lambda: stats) + + printed = [] + monkeypatch.setattr(usage_module.console, "print", lambda *a, **kw: printed.append(a)) + + await usage_module._maybe_print_cost_panel() + + panels = [item[0] for item in printed if item and isinstance(item[0], Panel)] + assert any(p.title == "Session Cost" for p in panels) + + +async def test_usage_omits_cost_panel_when_no_data(monkeypatch): + from pythinker_code.ui.shell import usage as usage_module + + monkeypatch.setattr(usage_module, "_load_cost_stats", lambda: None) + + printed = [] + monkeypatch.setattr(usage_module.console, "print", lambda *a, **kw: printed.append(a)) + + await usage_module._maybe_print_cost_panel() + + assert not printed + + +async def test_usage_omits_cost_panel_on_exception(monkeypatch): + from pythinker_code.ui.shell import usage as usage_module + + def _raise(): + raise RuntimeError("disk error") + + monkeypatch.setattr(usage_module, "_load_cost_stats", _raise) + + printed = [] + monkeypatch.setattr(usage_module.console, "print", lambda *a, **kw: printed.append(a)) + + # Must not raise + await usage_module._maybe_print_cost_panel() + + assert not printed diff --git a/tests/ui/usage_adapters/test_alibaba_adapter.py b/tests/ui/usage_adapters/test_alibaba_adapter.py new file mode 100644 index 00000000..96530f85 --- /dev/null +++ b/tests/ui/usage_adapters/test_alibaba_adapter.py @@ -0,0 +1,379 @@ +"""Tests for AlibabaAdapter.""" + +from __future__ import annotations + +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, MagicMock, patch + +import aiohttp +import pytest +from pydantic import SecretStr + +from pythinker_code.config import LLMProvider +from pythinker_code.ui.shell.stats_collector import AllStats, PeriodStats, ProviderStats +from pythinker_code.ui.shell.usage_adapters.alibaba import ( + AlibabaAdapter, + _parse_quota_response, + _quota_url, +) +from pythinker_code.usage_ratelimit_cache import RateLimitSnapshot + + +@pytest.fixture(autouse=True) +def _patch_local_stats(monkeypatch): + """Return empty AllStats so tests don't read real ~/.pythinker session files.""" + empty = AllStats( + periods={"today": PeriodStats(), "all_time": PeriodStats()}, + insights={}, + ) + monkeypatch.setattr( + "pythinker_code.ui.shell.usage_adapters.alibaba._load_all_stats", + lambda: empty, + ) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_provider( + api_key: str = "test-key", + base_url: str = "https://dashscope-us.aliyuncs.com/compatible-mode/v1", +) -> LLMProvider: + return LLMProvider( + type="openai_legacy", + api_key=SecretStr(api_key), + base_url=base_url, + ) + + +class _StubOAuth: + pass + + +def _make_response(status: int, json_data: object = None) -> MagicMock: + """Build a fake aiohttp response context manager.""" + resp = MagicMock() + resp.status = status + resp.json = AsyncMock(return_value=json_data) + return resp + + +def _make_session(response: MagicMock) -> MagicMock: + """Build a fake aiohttp ClientSession context manager.""" + session = MagicMock() + # session.get(...) is used as an async context manager + get_cm = MagicMock() + get_cm.__aenter__ = AsyncMock(return_value=response) + get_cm.__aexit__ = AsyncMock(return_value=False) + session.get = MagicMock(return_value=get_cm) + return session + + +@asynccontextmanager +async def _fake_new_client_session(session: MagicMock, **_kwargs): + yield session + + +# --------------------------------------------------------------------------- +# Unit tests for pure helpers +# --------------------------------------------------------------------------- + + +def test_quota_url_strips_path() -> None: + # US completion host is remapped to the international quota host + assert _quota_url("https://dashscope-us.aliyuncs.com/compatible-mode/v1") == ( + "https://dashscope-intl.aliyuncs.com/api/v1/quotas" + ) + + +def test_quota_url_china_unchanged() -> None: + # China host has no remap — quota API lives on the same host + assert _quota_url("https://dashscope.aliyuncs.com/compatible-mode/v1") == ( + "https://dashscope.aliyuncs.com/api/v1/quotas" + ) + + +def test_parse_quota_response_flat_shape() -> None: + data = {"token_quota": 1_000_000, "token_used": 123_456} + rows = _parse_quota_response(data) + assert len(rows) == 1 + assert rows[0].label == "Token quota" + assert rows[0].used == 123_456 + assert rows[0].limit == 1_000_000 + assert rows[0].unit == "tokens" + + +def test_parse_quota_response_flat_shape_inside_data_key() -> None: + data = {"data": {"token_quota": 500_000, "token_used": 50_000}} + rows = _parse_quota_response(data) + assert len(rows) == 1 + assert rows[0].used == 50_000 + assert rows[0].limit == 500_000 + + +def test_parse_quota_response_quota_list_shape() -> None: + data = { + "quota_list": [ + {"quota_name": "text", "total_quota": 500_000, "total_used": 50_000}, + ] + } + rows = _parse_quota_response(data) + assert len(rows) == 1 + assert rows[0].label == "Text" + assert rows[0].used == 50_000 + assert rows[0].limit == 500_000 + assert rows[0].unit == "tokens" + + +def test_parse_quota_response_unrecognised_shape_returns_empty() -> None: + assert _parse_quota_response({"foo": "bar"}) == [] + assert _parse_quota_response("not a dict") == [] + assert _parse_quota_response(None) == [] + + +# --------------------------------------------------------------------------- +# Adapter metadata +# --------------------------------------------------------------------------- + + +def test_alibaba_adapter_metadata() -> None: + assert AlibabaAdapter.platform_id == "alibaba" + assert AlibabaAdapter.requires_admin_key is False + assert AlibabaAdapter.provider_label == "Alibaba DashScope" + + +# --------------------------------------------------------------------------- +# Async fetch tests +# --------------------------------------------------------------------------- + + +async def test_quota_api_success() -> None: + """HTTP 200 with flat quota shape produces a correct UsageRow.""" + json_data = {"token_quota": 1_000_000, "token_used": 123_456} + resp = _make_response(200, json_data) + session = _make_session(resp) + + with ( + patch( + "pythinker_code.ui.shell.usage_adapters.alibaba.new_client_session", + side_effect=lambda **kw: _fake_new_client_session(session, **kw), + ), + patch("pythinker_code.ui.shell.usage_adapters.alibaba.get_cache") as mock_cache, + ): + mock_cache.return_value.snapshot.return_value = None + report = await AlibabaAdapter().fetch(_make_provider(), _StubOAuth()) # type: ignore[arg-type] + + assert report.summary is not None + assert report.summary.label == "Token quota" + assert report.summary.used == 123_456 + assert report.summary.limit == 1_000_000 + assert report.summary.unit == "tokens" + assert report.notes == [] + + +async def test_quota_list_shape() -> None: + """HTTP 200 with quota_list shape produces a correct UsageRow.""" + json_data = { + "quota_list": [{"quota_name": "text", "total_quota": 500_000, "total_used": 50_000}] + } + resp = _make_response(200, json_data) + session = _make_session(resp) + + with ( + patch( + "pythinker_code.ui.shell.usage_adapters.alibaba.new_client_session", + side_effect=lambda **kw: _fake_new_client_session(session, **kw), + ), + patch("pythinker_code.ui.shell.usage_adapters.alibaba.get_cache") as mock_cache, + ): + mock_cache.return_value.snapshot.return_value = None + report = await AlibabaAdapter().fetch(_make_provider(), _StubOAuth()) # type: ignore[arg-type] + + assert report.summary is not None + assert report.summary.label == "Text" + assert report.summary.used == 50_000 + assert report.summary.limit == 500_000 + + +async def test_quota_api_404_falls_through() -> None: + """HTTP 404 with no snapshot falls through to the 'no data yet' note.""" + resp = _make_response(404) + session = _make_session(resp) + + with ( + patch( + "pythinker_code.ui.shell.usage_adapters.alibaba.new_client_session", + side_effect=lambda **kw: _fake_new_client_session(session, **kw), + ), + patch("pythinker_code.ui.shell.usage_adapters.alibaba.get_cache") as mock_cache, + ): + mock_cache.return_value.snapshot.return_value = None + report = await AlibabaAdapter().fetch(_make_provider(), _StubOAuth()) # type: ignore[arg-type] + + assert report.summary is None + assert report.limits == [] + assert report.notes # some guidance note is present + + +async def test_quota_api_401_shows_note() -> None: + """HTTP 401 appends an authorization-failure note.""" + resp = _make_response(401) + session = _make_session(resp) + + with ( + patch( + "pythinker_code.ui.shell.usage_adapters.alibaba.new_client_session", + side_effect=lambda **kw: _fake_new_client_session(session, **kw), + ), + patch("pythinker_code.ui.shell.usage_adapters.alibaba.get_cache") as mock_cache, + ): + mock_cache.return_value.snapshot.return_value = None + report = await AlibabaAdapter().fetch(_make_provider(), _StubOAuth()) # type: ignore[arg-type] + + assert any("authorization failed" in n.lower() for n in report.notes) + + +async def test_ratelimit_cache_used() -> None: + """When quota API returns 404, snapshot data fills in rate-limit rows.""" + resp = _make_response(404) + session = _make_session(resp) + + snap = RateLimitSnapshot( + requests_limit=None, + requests_remaining=None, + requests_reset_seconds=None, + tokens_limit=10_000, + tokens_remaining=8_000, + tokens_reset_seconds=None, + captured_at=0.0, + ) + + with ( + patch( + "pythinker_code.ui.shell.usage_adapters.alibaba.new_client_session", + side_effect=lambda **kw: _fake_new_client_session(session, **kw), + ), + patch("pythinker_code.ui.shell.usage_adapters.alibaba.get_cache") as mock_cache, + ): + mock_cache.return_value.snapshot.return_value = snap + report = await AlibabaAdapter().fetch(_make_provider(), _StubOAuth()) # type: ignore[arg-type] + + assert report.summary is not None + assert report.summary.label == "Tokens remaining" + assert report.summary.used == 8_000 + assert report.summary.limit == 10_000 + assert report.summary.unit == "tokens" + + +async def test_ratelimit_requests_row_labeled_remaining() -> None: + """The Requests rate-limit row labels its value as remaining, not consumed.""" + resp = _make_response(404) + session = _make_session(resp) + + snap = RateLimitSnapshot( + requests_limit=1_000, + requests_remaining=750, + requests_reset_seconds=None, + tokens_limit=None, + tokens_remaining=None, + tokens_reset_seconds=None, + captured_at=0.0, + ) + + with ( + patch( + "pythinker_code.ui.shell.usage_adapters.alibaba.new_client_session", + side_effect=lambda **kw: _fake_new_client_session(session, **kw), + ), + patch("pythinker_code.ui.shell.usage_adapters.alibaba.get_cache") as mock_cache, + ): + mock_cache.return_value.snapshot.return_value = snap + report = await AlibabaAdapter().fetch(_make_provider(), _StubOAuth()) # type: ignore[arg-type] + + assert report.summary is not None + assert report.summary.label == "Requests remaining" + assert report.summary.used == 750 + assert report.summary.limit == 1_000 + assert report.summary.unit == "requests" + + +async def test_network_error_falls_through() -> None: + """aiohttp.ClientError during HTTP call is swallowed; returns a UsageReport.""" + + @asynccontextmanager + async def _error_session(**_kw): + session = MagicMock() + get_cm = MagicMock() + get_cm.__aenter__ = AsyncMock(side_effect=aiohttp.ClientConnectionError("network")) + get_cm.__aexit__ = AsyncMock(return_value=False) + session.get = MagicMock(return_value=get_cm) + yield session + + with ( + patch( + "pythinker_code.ui.shell.usage_adapters.alibaba.new_client_session", + side_effect=_error_session, + ), + patch("pythinker_code.ui.shell.usage_adapters.alibaba.get_cache") as mock_cache, + ): + mock_cache.return_value.snapshot.return_value = None + # Must not raise + report = await AlibabaAdapter().fetch(_make_provider(), _StubOAuth()) # type: ignore[arg-type] + + assert report is not None + # A transient fetch failure surfaces an actionable note rather than the + # misleading "No usage recorded yet" fall-through. + assert any("unavailable" in n.lower() for n in report.notes) + + +async def test_local_stats_shown_when_available(monkeypatch) -> None: + """Local ProviderStats are shown as rows even when quota API returns 404.""" + prov = ProviderStats() + prov.messages = 3 + prov.cost = 0.0162 + prov.input_other = 4_200 + prov.output = 1_800 + prov.input_cache_read = 0 + prov.input_cache_creation = 0 + + stats = AllStats( + periods={ + "today": _period_with_provider(prov), + "all_time": _period_with_provider(prov), + }, + insights={}, + ) + monkeypatch.setattr( + "pythinker_code.ui.shell.usage_adapters.alibaba._load_all_stats", + lambda: stats, + ) + + resp = _make_response(404) + session = _make_session(resp) + + with ( + patch( + "pythinker_code.ui.shell.usage_adapters.alibaba.new_client_session", + side_effect=lambda **kw: _fake_new_client_session(session, **kw), + ), + patch("pythinker_code.ui.shell.usage_adapters.alibaba.get_cache") as mock_cache, + ): + mock_cache.return_value.snapshot.return_value = None + report = await AlibabaAdapter().fetch(_make_provider(), _StubOAuth()) # type: ignore[arg-type] + + assert report.summary is not None + assert report.summary.label == "Today" + assert report.summary.used == 6_000 # 4200 + 1800 + assert "4,200" in (report.summary.reset_hint or "") + assert "1,800" in (report.summary.reset_hint or "") + assert report.notes == [] + + +def _period_with_provider(prov: ProviderStats) -> PeriodStats: + period = PeriodStats() + period.total_messages = prov.messages + period.total_cost = prov.cost + period.providers["managed:alibaba"] = prov + return period diff --git a/tests/ui_and_conv/test_redraw_throttle.py b/tests/ui_and_conv/test_redraw_throttle.py new file mode 100644 index 00000000..bd3c49b6 --- /dev/null +++ b/tests/ui_and_conv/test_redraw_throttle.py @@ -0,0 +1,77 @@ +"""Regression test for the `/login` `redraw_in_future` coroutine warning. + +Root cause: setting `Application.min_redraw_interval` activates prompt_toolkit's +*coroutine-based* throttle path in `invalidate()` (`async def redraw_in_future`). +If `invalidate()` fires during an app/loop handoff (as happens when `/login` +swaps prompt sessions), that coroutine is constructed and then dropped +un-awaited -> `RuntimeWarning: coroutine '...redraw_in_future' was never awaited`. + +`max_render_postpone_time` throttles redraws via a coroutine-free path +(`call_soon_threadsafe` only). It never constructs the coroutine, so the +"never awaited" warning is *impossible* on that path — proof by elimination +rather than by racing the garbage collector. + +The tests poke prompt_toolkit internals on purpose: this bug lives exactly at +that boundary, so the regression guard has to exercise it there. +""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Callable, Coroutine +from typing import Any + +from prompt_toolkit.application import Application + + +async def _invalidate_constructs_coroutine(configure: Callable[[Application], None]) -> bool: + """Return True iff `invalidate()` constructs a `redraw_in_future` coroutine. + + The coroutine is the *only* thing that can leak the warning, so whether it is + constructed at all is the deterministic signal for the root cause. + """ + app: Application = Application() + app._is_running = True + app.loop = asyncio.get_running_loop() + app._invalidated = False + app._redraw = lambda render_as_done=True: None # keep the redraw path quiet + # Force the throttle to engage: pretend we just redrew, so the next + # invalidate falls inside the throttle window. + app._last_redraw_time = time.time() + configure(app) + + constructed: list[object] = [] + + def spy(coro: Coroutine[Any, Any, Any]) -> None: + constructed.append(coro) + coro.close() # consume it so the test itself never leaks a coroutine + + app.create_background_task = spy # type: ignore[assignment, method-assign] + + app.invalidate() + await asyncio.sleep(0.02) # let the call_soon_threadsafe callback fire + return bool(constructed) + + +async def test_min_redraw_interval_constructs_redraw_coroutine() -> None: + """Documents the buggy path: min_redraw_interval builds the leak-prone coroutine.""" + + def configure(app: Application) -> None: + app.min_redraw_interval = 10 # large window guarantees the coroutine path + + assert await _invalidate_constructs_coroutine(configure), ( + "expected min_redraw_interval to construct redraw_in_future (the leak source)" + ) + + +async def test_max_render_postpone_time_constructs_no_coroutine() -> None: + """The fix path: max_render_postpone_time never constructs the coroutine.""" + + def configure(app: Application) -> None: + app.max_render_postpone_time = 1 / 30 + + assert not await _invalidate_constructs_coroutine(configure), ( + "max_render_postpone_time must not construct redraw_in_future; " + "no coroutine means the 'never awaited' warning is impossible" + ) diff --git a/tests/ui_and_conv/test_stats_pricing.py b/tests/ui_and_conv/test_stats_pricing.py index 17e4d464..6a91b4ed 100644 --- a/tests/ui_and_conv/test_stats_pricing.py +++ b/tests/ui_and_conv/test_stats_pricing.py @@ -50,3 +50,67 @@ def test_zero_usage_returns_zero(): usage = _usage() cost = get_cost_usd("claude-opus-4-1", usage) assert cost == 0.0 + + +# --------------------------------------------------------------------------- +# catalog-path tests +# --------------------------------------------------------------------------- + + +def test_get_cost_usd_uses_catalog_when_available(monkeypatch): + from pythinker_code import models_dev + from pythinker_code.models_dev import ModelPrice + + fake_catalog = { + "claude-sonnet-4-6": ModelPrice(input=1.0, output=2.0, cache_read=0.1, cache_write=0.2) + } + monkeypatch.setattr(models_dev, "load_catalog", lambda: fake_catalog) + usage = _usage(input_other=1_000_000, output=1_000_000) + cost = get_cost_usd("claude-sonnet-4-6", usage) + assert abs(cost - 3.0) < 0.001 # 1.0 + 2.0 per 1M + + +def test_get_cost_usd_catalog_prefix_match(monkeypatch): + from pythinker_code import models_dev + from pythinker_code.models_dev import ModelPrice + + fake_catalog = { + "claude-sonnet-4-6": ModelPrice(input=9.0, output=9.0, cache_read=0.0, cache_write=0.0) + } + monkeypatch.setattr(models_dev, "load_catalog", lambda: fake_catalog) + usage = _usage(input_other=1_000_000) + # versioned id not in catalog, but prefix matches + cost = get_cost_usd("claude-sonnet-4-6-20251001", usage) + assert abs(cost - 9.0) < 0.001 + + +def test_get_cost_usd_falls_back_to_hardcoded_when_catalog_empty(monkeypatch): + from pythinker_code import models_dev + + monkeypatch.setattr(models_dev, "load_catalog", lambda: {}) + usage = _usage(input_other=1_000_000, output=1_000_000) + # claude-sonnet-4-5 is in _PRICE_TABLE: input=3, output=15 + cost = get_cost_usd("claude-sonnet-4-5", usage) + assert abs(cost - 18.0) < 0.001 + + +def test_get_cost_usd_catalog_beats_hardcoded(monkeypatch): + from pythinker_code import models_dev + from pythinker_code.models_dev import ModelPrice + + # Override a model that IS in _PRICE_TABLE with a different catalog price + fake_catalog = { + "claude-sonnet-4-5": ModelPrice(input=99.0, output=99.0, cache_read=0.0, cache_write=0.0) + } + monkeypatch.setattr(models_dev, "load_catalog", lambda: fake_catalog) + usage = _usage(input_other=1_000_000) + cost = get_cost_usd("claude-sonnet-4-5", usage) + assert abs(cost - 99.0) < 0.001 # catalog wins + + +def test_get_cost_usd_unknown_model_returns_zero(monkeypatch): + from pythinker_code import models_dev + + monkeypatch.setattr(models_dev, "load_catalog", lambda: {}) + usage = _usage(input_other=1_000_000) + assert get_cost_usd("completely-unknown-xyz-model", usage) == 0.0 diff --git a/tests/utils/test_editor.py b/tests/utils/test_editor.py index 7a4a7c13..408a1c58 100644 --- a/tests/utils/test_editor.py +++ b/tests/utils/test_editor.py @@ -97,6 +97,7 @@ def test_returns_none_when_nothing_available(self, monkeypatch: pytest.MonkeyPat def test_empty_configured_is_ignored(self, monkeypatch: pytest.MonkeyPatch): """Empty configured string should be treated as not configured.""" + monkeypatch.delenv("VISUAL", raising=False) monkeypatch.setenv("EDITOR", "nano") assert get_editor_command("") == ["nano"]