diff --git a/README.md b/README.md index e2c69db..ff2e9b1 100644 --- a/README.md +++ b/README.md @@ -117,6 +117,12 @@ Long Claude turns surface an elapsed-time counter and a heartbeat-driven pulse w In Claude mode, workspace files attached as chat context arrive as `@`-mention pointers rather than inlined file contents. Claude's Read tool fetches them on demand, which means images, large files, and notebooks (cell-aware) now work where the older content-injection path silently truncated or skipped them. +#### Usage footer + +Enable **Show usage after each turn** in the Settings dialog to append a small footer under each Claude-mode reply with the turn's duration, token counts (with a cached-input breakdown), and cost — for example `12.3s · 45.2K in (38.1K cached) / 1.2K out · $0.0842`. It is off by default. + +The cost figure comes from the Claude Code CLI, priced from Anthropic's public list rates. It is shown only when NBI is configured with a direct Anthropic API key on the default endpoint, and is omitted otherwise — on a subscription login (where the marginal cost is effectively $0), an enterprise-negotiated contract, or a custom **Base URL** (proxy, gateway, or non-Anthropic endpoint) — since the list-price figure would not match your actual billing. Duration and token counts are always shown. + #### Claude Code launcher tile When the Claude CLI is on `PATH`, the JupyterLab launcher (the panel that opens with new tabs) shows a **Claude Code** tile alongside the standard kernel launchers. Clicking it opens a session picker; search across past transcripts and resume one in a fresh terminal, or start a new session in the file browser's active subdirectory. Session IDs are copyable from the picker for paste into a `claude --resume ` command. diff --git a/notebook_intelligence/claude.py b/notebook_intelligence/claude.py index 262b6a3..7449f03 100644 --- a/notebook_intelligence/claude.py +++ b/notebook_intelligence/claude.py @@ -19,7 +19,7 @@ from notebook_intelligence._version import __version__ as NBI_VERSION import base64 import logging -from claude_agent_sdk import AssistantMessage, PermissionResultAllow, PermissionResultDeny, TextBlock, ToolResultBlock, ToolUseBlock, UserMessage, create_sdk_mcp_server, ClaudeAgentOptions, ClaudeSDKClient, tool +from claude_agent_sdk import AssistantMessage, PermissionResultAllow, PermissionResultDeny, ResultMessage, TextBlock, ToolResultBlock, ToolUseBlock, UserMessage, create_sdk_mcp_server, ClaudeAgentOptions, ClaudeSDKClient, tool from notebook_intelligence.util import ThreadSafeWebSocketConnector, _emit, get_jupyter_root_dir, import_litellm, resolve_claude_cli_path, safe_jupyter_path, terminate_process_tree @@ -190,6 +190,75 @@ def claude_tool_kind(name: str) -> str: return "other" +def _format_token_count(count: int) -> str: + if count >= 1_000_000: + return f"{count / 1_000_000:.1f}M" + if count >= 1_000: + return f"{count / 1_000:.1f}K" + return str(count) + + +def format_result_usage(result: ResultMessage, show_cost: bool = True) -> Optional[str]: + """One-line usage footer for a completed Claude Code turn. + + The SDK ends every turn with a ``ResultMessage`` carrying duration, + token usage, and cost; NBI used to drop it on the floor, leaving + users no signal of what a turn cost short of typing ``/cost``. + Returns a small italic markdown line prefixed with a paragraph + break — the footer is streamed right after the turn's final text + block, and without the separator it would concatenate onto prose + that doesn't end in a blank line (``Done.*12.3s · ...*``), breaking + the emphasis markup. Returns ``None`` when there is nothing + meaningful to show (error results, missing usage) so the caller can + skip the footer entirely. + + ``show_cost`` gates the ``$`` segment. The SDK's ``total_cost_usd`` + is priced from the CLI's built-in public list rates, so it is only + trustworthy for a direct first-party API key; on a subscription + login the marginal cost is really $0, and enterprise/cloud-endpoint + pricing differs. The caller passes ``False`` unless NBI is running + against a direct API key on Anthropic's own endpoint (see + ``_should_show_turn_cost``) so we omit a figure we can't stand + behind, while still showing the always-accurate duration and token + counts. + """ + if getattr(result, "is_error", False): + return None + usage = getattr(result, "usage", None) or {} + parts: list[str] = [] + + duration_ms = getattr(result, "duration_ms", None) + if isinstance(duration_ms, (int, float)) and duration_ms > 0: + parts.append(f"{duration_ms / 1000:.1f}s") + + def _count(key: str) -> int: + value = usage.get(key, 0) + return value if isinstance(value, int) and value > 0 else 0 + + total_in = ( + _count("input_tokens") + + _count("cache_read_input_tokens") + + _count("cache_creation_input_tokens") + ) + output_tokens = _count("output_tokens") + if total_in > 0 or output_tokens > 0: + token_part = f"{_format_token_count(total_in)} in" + cache_read = _count("cache_read_input_tokens") + if cache_read > 0: + token_part += f" ({_format_token_count(cache_read)} cached)" + token_part += f" / {_format_token_count(output_tokens)} out" + parts.append(token_part) + + if show_cost: + cost = getattr(result, "total_cost_usd", None) + if isinstance(cost, (int, float)) and cost > 0: + parts.append(f"${cost:.4f}") + + if not parts: + return None + return f"\n\n*{' · '.join(parts)}*" + + # Cap the diff lines surfaced per tool-call card so a large edit doesn't bloat # the transcript or the wire payload. _MAX_DIFF_LINES = 60 @@ -517,6 +586,63 @@ def _normalize_anthropic_credential(value: Any) -> str | None: return value.strip() or None +def _is_first_party_anthropic_endpoint(base_url: Any) -> bool: + """True when requests go to Anthropic's own API — the only place the + SDK's public-list-price ``total_cost_usd`` is meaningful. + + An unset ``base_url`` means the SDK uses its default + (``api.anthropic.com``). A custom ``base_url`` (proxy, gateway, + Bedrock/Vertex front, OpenAI-compatible endpoint) bills on its own + terms, so cost figures priced off Anthropic's list rates can't be + trusted there. + """ + from urllib.parse import urlparse + + normalized = _normalize_anthropic_credential(base_url) + if normalized is None: + return True + return (urlparse(normalized).hostname or "").lower() == "api.anthropic.com" + + +def _resolve_effective_credential(settings_value: Any, env_var: str) -> str | None: + """Resolve a credential the way the Claude Code subprocess does. + + NBI only overlays ``ANTHROPIC_API_KEY`` / ``ANTHROPIC_BASE_URL`` onto + the CLI's environment when the corresponding ``claude_settings`` field + is non-empty; that env is merged over the server process's own + ``os.environ``, so a blank setting falls through to whatever the + process already had. Mirror that precedence — setting wins, else the + ambient variable — so the cost gate reflects the endpoint/key the CLI + actually uses, not just what the settings panel shows. + """ + resolved = _normalize_anthropic_credential(settings_value) + if resolved is not None: + return resolved + return _normalize_anthropic_credential(os.environ.get(env_var)) + + +def _should_show_turn_cost(claude_settings: dict) -> bool: + """Whether the per-turn footer's ``$`` cost can be trusted. + + Requires both a direct API key (subscription logins report a + meaningless notional cost) and Anthropic's own endpoint (custom + ``base_url`` targets bill differently). Credentials are resolved + against the environment the CLI inherits, not just the settings + panel, so an env-provided key or a custom ``ANTHROPIC_BASE_URL`` is + honored. Either condition failing suppresses the dollar figure while + duration/tokens still render. + """ + api_key = _resolve_effective_credential( + claude_settings.get('api_key'), 'ANTHROPIC_API_KEY' + ) + if api_key is None: + return False + base_url = _resolve_effective_credential( + claude_settings.get('base_url'), 'ANTHROPIC_BASE_URL' + ) + return _is_first_party_anthropic_endpoint(base_url) + + def _create_anthropic_client(api_key: str = None, base_url: str = None) -> "Anthropic": """Create an Anthropic client with normalized credentials and default headers.""" from anthropic import Anthropic @@ -1048,6 +1174,20 @@ async def _client_thread_func(self): status=status, diffs=diffs, )) + elif isinstance(message, ResultMessage): + # End-of-turn accounting from the SDK. + # Opt-in (off by default): the footer + # is persistent per-turn noise, and its + # cost figure is only trustworthy for a + # direct API key against Anthropic's own + # endpoint, so gate the dollar amount on + # both (see _should_show_turn_cost). + claude_settings = self._host.nbi_config.claude_settings + if claude_settings.get('show_turn_usage', False): + show_cost = _should_show_turn_cost(claude_settings) + usage_line = format_result_usage(message, show_cost=show_cost) + if usage_line is not None: + response.stream(MarkdownData(usage_line)) else: pass diff --git a/src/components/settings-panel.tsx b/src/components/settings-panel.tsx index ff5d8a9..4bececa 100644 --- a/src/components/settings-panel.tsx +++ b/src/components/settings-panel.tsx @@ -1195,6 +1195,9 @@ function SettingsPanelComponentClaude(props: any) { const [continueConversation, setContinueConversation] = useState( nbiConfig.claudeSettings.continue_conversation ?? false ); + const [showTurnUsage, setShowTurnUsage] = useState( + nbiConfig.claudeSettings.show_turn_usage ?? false + ); const [claudeModels, setClaudeModels] = useState( nbiConfig.claudeModels ); @@ -1235,7 +1238,8 @@ function SettingsPanelComponentClaude(props: any) { base_url: baseUrl, setting_sources: settingSources, tools: tools, - continue_conversation: continueConversation + continue_conversation: continueConversation, + show_turn_usage: showTurnUsage } }); }; @@ -1250,7 +1254,8 @@ function SettingsPanelComponentClaude(props: any) { baseUrl, settingSources, tools, - continueConversation + continueConversation, + showTurnUsage ]); return ( @@ -1578,6 +1583,27 @@ function SettingsPanelComponentClaude(props: any) { +
+
Usage footer
+
+
+
+
+ { + setShowTurnUsage(!showTurnUsage); + }} + > +
+
+
+
+
+
Claude account
diff --git a/tests/test_claude_turn_usage.py b/tests/test_claude_turn_usage.py new file mode 100644 index 0000000..c706f0f --- /dev/null +++ b/tests/test_claude_turn_usage.py @@ -0,0 +1,189 @@ +# Copyright (c) Mehmet Bektas + +"""Per-turn usage footer built from the SDK's end-of-turn ResultMessage. + +The SDK ends every turn with a ResultMessage carrying duration, token +usage, and cost. NBI used to drop it in the receive loop's silent else +branch, leaving users no signal of what a turn cost short of typing +/cost. +""" + +import pytest +from claude_agent_sdk import ResultMessage + +from notebook_intelligence.claude import ( + format_result_usage, + _should_show_turn_cost, +) + + +@pytest.fixture(autouse=True) +def _clear_anthropic_env(monkeypatch): + """Keep the cost-gate tests hermetic: the resolver now consults the + process ANTHROPIC_* env vars, so clear them unless a test sets them.""" + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) + + +def _result(**overrides) -> ResultMessage: + fields = { + "subtype": "success", + "duration_ms": 12_300, + "duration_api_ms": 11_000, + "is_error": False, + "num_turns": 3, + "session_id": "sess-1", + "total_cost_usd": 0.0842, + "usage": { + "input_tokens": 4_200, + "cache_read_input_tokens": 38_100, + "cache_creation_input_tokens": 2_900, + "output_tokens": 1_200, + }, + } + fields.update(overrides) + return ResultMessage(**fields) + + +class TestFormatResultUsage: + def test_formats_duration_tokens_and_cost(self): + line = format_result_usage(_result()) + assert line == "\n\n*12.3s · 45.2K in (38.1K cached) / 1.2K out · $0.0842*" + + def test_error_results_produce_no_footer(self): + assert format_result_usage(_result(is_error=True)) is None + + def test_missing_usage_and_cost_produce_duration_only(self): + line = format_result_usage(_result(usage=None, total_cost_usd=None)) + assert line == "\n\n*12.3s*" + + def test_nothing_meaningful_returns_none(self): + result = _result(duration_ms=0, usage=None, total_cost_usd=None) + assert format_result_usage(result) is None + + def test_zero_cost_is_omitted(self): + # Subscription (non-API-key) sessions report 0.0 cost; showing + # "$0.0000" would read as a billing claim, so drop the segment. + line = format_result_usage(_result(total_cost_usd=0.0)) + assert "$" not in line + assert "45.2K in" in line + + def test_uncached_turn_omits_cache_annotation(self): + line = format_result_usage( + _result( + usage={ + "input_tokens": 900, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "output_tokens": 150, + } + ) + ) + assert "cached" not in line + assert "900 in / 150 out" in line + + def test_malformed_usage_values_are_ignored(self): + line = format_result_usage( + _result( + usage={ + "input_tokens": "not-a-number", + "output_tokens": None, + "cache_read_input_tokens": -5, + }, + total_cost_usd="free", + ) + ) + # Falls back to just the duration rather than crashing the turn. + assert line == "\n\n*12.3s*" + + def test_million_scale_token_formatting(self): + line = format_result_usage( + _result( + usage={ + "input_tokens": 1_500_000, + "output_tokens": 2_000, + } + ) + ) + assert "1.5M in" in line + + def test_show_cost_false_omits_cost_but_keeps_tokens(self): + # Non-API-key setups (subscription/enterprise/cloud endpoints) + # can't be priced from the CLI's public list rates, so the caller + # suppresses the dollar figure while keeping the accurate + # duration and token counts. + line = format_result_usage(_result(), show_cost=False) + assert "$" not in line + assert line == "\n\n*12.3s · 45.2K in (38.1K cached) / 1.2K out*" + + def test_show_cost_true_is_the_default(self): + assert "$0.0842" in format_result_usage(_result()) + + +class TestShouldShowTurnCost: + def test_direct_key_default_endpoint_shows_cost(self): + assert _should_show_turn_cost({"api_key": "sk-ant-xyz"}) is True + + def test_direct_key_empty_base_url_shows_cost(self): + # Settings-panel saves an unset field as "" rather than None. + assert _should_show_turn_cost({"api_key": "sk-ant-xyz", "base_url": ""}) is True + + def test_direct_key_explicit_anthropic_url_shows_cost(self): + assert _should_show_turn_cost( + {"api_key": "sk-ant-xyz", "base_url": "https://api.anthropic.com"} + ) is True + + def test_custom_base_url_suppresses_cost(self): + # A proxy/gateway/cloud front bills on its own terms, so the SDK's + # public-list-price cost can't be trusted even with an API key. + assert _should_show_turn_cost( + {"api_key": "sk-ant-xyz", "base_url": "https://my-proxy.example.com/v1"} + ) is False + + def test_no_api_key_suppresses_cost(self): + # Subscription login: no key -> notional cost only. + assert _should_show_turn_cost({"api_key": ""}) is False + assert _should_show_turn_cost({}) is False + + +class TestShouldShowTurnCostEnvResolution: + """The CLI subprocess inherits the server's ANTHROPIC_* env vars when the + settings fields are blank, so the cost gate must consult them too.""" + + def test_env_api_key_default_endpoint_shows_cost(self, monkeypatch): + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-env") + monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) + assert _should_show_turn_cost({}) is True + + def test_env_custom_base_url_suppresses_cost(self, monkeypatch): + # Key in settings, but a custom endpoint via env -> proxy billing, + # so the SDK's list-price cost can't be trusted. + monkeypatch.setenv("ANTHROPIC_BASE_URL", "https://gateway.example.com") + assert ( + _should_show_turn_cost({"api_key": "sk-ant-xyz"}) is False + ) + + def test_settings_base_url_overrides_env(self, monkeypatch): + # Settings win over env: a custom setting suppresses cost even when + # the env points at the first-party endpoint. + monkeypatch.setenv("ANTHROPIC_BASE_URL", "https://api.anthropic.com") + assert ( + _should_show_turn_cost( + {"api_key": "sk-ant-xyz", "base_url": "https://proxy.example.com"} + ) + is False + ) + + def test_no_key_anywhere_suppresses_cost(self, monkeypatch): + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + assert _should_show_turn_cost({}) is False + + + def test_footer_starts_with_a_paragraph_break(self): + # The footer is streamed right after the turn's final text block; + # without a separator it concatenates onto prose that doesn't end + # in a blank line ("Done.*12.3s ...*") and breaks the emphasis + # markup. + line = format_result_usage(_result()) + assert line.startswith("\n\n*") + assert ("Done." + line).splitlines()[-1].startswith("*")