"""Tests for resume-state persistence and token display callbacks.""" from types import SimpleNamespace from typing import Any from langchain_core.messages import AIMessage, HumanMessage from deepagents_code.app import DeepAgentsApp from deepagents_code.resume_state import ( ResumeState, ResumeStateMiddleware, _extract_context_tokens, _extract_model_spec, ) def _runtime(context: dict[str, str | None] | None) -> SimpleNamespace: """Build a stand-in `Runtime` exposing only `.context`.""" return SimpleNamespace(context=context) class TestResumeState: def test_state_has_context_tokens_field(self): """ResumeState declares the `_context_tokens` channel.""" assert "_context_tokens" in ResumeState.__annotations__ def test_state_has_model_spec_field(self): """ResumeState declares the `_model_spec` channel.""" assert "_model_spec" in ResumeState.__annotations__ def test_middleware_exposes_state_schema(self): """ResumeStateMiddleware registers the correct state schema.""" assert ResumeStateMiddleware.state_schema is ResumeState class TestExtractContextTokens: """Tests for `_extract_context_tokens`.""" def test_prefers_input_plus_output(self) -> None: msg = AIMessage( content="hi", usage_metadata={ "input_tokens": 100, "output_tokens": 25, "total_tokens": 200, # deliberately inconsistent }, ) assert _extract_context_tokens(msg) == 125 def test_falls_back_to_total_tokens(self) -> None: msg = AIMessage( content="hi", usage_metadata={ "input_tokens": 0, "output_tokens": 0, "total_tokens": 999, }, ) assert _extract_context_tokens(msg) == 999 def test_returns_none_without_usage_metadata(self) -> None: msg = AIMessage(content="hi") assert _extract_context_tokens(msg) is None def test_returns_none_for_zero_usage(self) -> None: msg = AIMessage( content="hi", usage_metadata={ "input_tokens": 0, "output_tokens": 0, "total_tokens": 0, }, ) assert _extract_context_tokens(msg) is None class TestExtractModelSpec: """Tests for `_extract_model_spec`.""" def test_returns_effective_model_from_context(self) -> None: runtime = _runtime({"effective_model": "anthropic:claude-sonnet-4-5"}) assert _extract_model_spec(runtime) == "anthropic:claude-sonnet-4-5" # ty: ignore def test_returns_none_when_context_missing(self) -> None: assert _extract_model_spec(_runtime(None)) is None # ty: ignore def test_returns_none_when_field_absent(self) -> None: assert _extract_model_spec(_runtime({"model": "x"})) is None # ty: ignore def test_returns_none_for_blank_or_nonstring(self) -> None: assert _extract_model_spec(_runtime({"effective_model": ""})) is None # ty: ignore assert _extract_model_spec(_runtime({"effective_model": None})) is None # ty: ignore def test_returns_none_when_runtime_is_none(self) -> None: assert _extract_model_spec(None) is None # ty: ignore class TestAfterModelHook: """Tests for the `after_model` persistence hook.""" async def test_writes_context_tokens_from_last_ai_message(self) -> None: middleware = ResumeStateMiddleware() state: dict[str, Any] = { "messages": [ HumanMessage(content="hi"), AIMessage( content="response", usage_metadata={ "input_tokens": 1500, "output_tokens": 200, "total_tokens": 1700, }, ), ], } result = middleware.after_model(state, _runtime(None)) # ty: ignore assert result == {"_context_tokens": 1700} async def test_writes_model_spec_from_context(self) -> None: middleware = ResumeStateMiddleware() state: dict[str, Any] = { "messages": [ HumanMessage(content="hi"), AIMessage( content="response", usage_metadata={ "input_tokens": 1500, "output_tokens": 200, "total_tokens": 1700, }, ), ], } runtime = _runtime({"effective_model": "openai:gpt-5.1"}) result = middleware.after_model(state, runtime) # ty: ignore assert result == { "_context_tokens": 1700, "_model_spec": "openai:gpt-5.1", } async def test_writes_model_spec_without_token_usage(self) -> None: """Model spec is recorded even when the AI message reports no usage.""" middleware = ResumeStateMiddleware() state: dict[str, Any] = { "messages": [ HumanMessage(content="hi"), AIMessage(content="no usage info"), ], } runtime = _runtime({"effective_model": "openai:gpt-5.1"}) result = middleware.after_model(state, runtime) # ty: ignore assert result == {"_model_spec": "openai:gpt-5.1"} async def test_returns_none_when_no_ai_message(self) -> None: middleware = ResumeStateMiddleware() state: dict[str, Any] = {"messages": [HumanMessage(content="hi")]} result = middleware.after_model(state, _runtime(None)) # ty: ignore assert result is None async def test_returns_none_when_last_ai_lacks_usage(self) -> None: middleware = ResumeStateMiddleware() state: dict[str, Any] = { "messages": [ HumanMessage(content="hi"), AIMessage(content="no usage info"), ], } result = middleware.after_model(state, _runtime(None)) # ty: ignore assert result is None async def test_handles_empty_messages(self) -> None: middleware = ResumeStateMiddleware() result = middleware.after_model({"messages": []}, _runtime(None)) # ty: ignore assert result is None async def test_skips_intervening_tool_messages(self) -> None: """Picks up the most recent AIMessage even when followed by tool turns.""" from langchain_core.messages import ToolMessage middleware = ResumeStateMiddleware() state: dict[str, Any] = { "messages": [ HumanMessage(content="hi"), AIMessage( content="older", usage_metadata={ "input_tokens": 100, "output_tokens": 10, "total_tokens": 110, }, ), ToolMessage(content="tool out", tool_call_id="t1"), AIMessage( content="newer", usage_metadata={ "input_tokens": 500, "output_tokens": 50, "total_tokens": 550, }, ), ], } result = middleware.after_model(state, _runtime(None)) # ty: ignore assert result == {"_context_tokens": 550} class TestTokenDisplayCallbacks: """Verify the callback-based token tracking that replaced TextualTokenTracker.""" def test_on_tokens_update_sets_cache_and_calls_display(self): """_on_tokens_update should set the local cache and update the status bar.""" display_calls: list[int] = [] class FakeApp: _context_tokens: int = 0 _status_bar = None def _update_tokens(self, count: int) -> None: display_calls.append(count) def _on_tokens_update(self, count: int) -> None: self._context_tokens = count self._update_tokens(count) app = FakeApp() app._on_tokens_update(4200) assert app._context_tokens == 4200 assert display_calls == [4200] def test_show_tokens_restores_cached_value(self): """_show_tokens should re-display the cached value.""" display_calls: list[int] = [] class FakeApp: _context_tokens: int = 1500 def _update_tokens(self, count: int) -> None: display_calls.append(count) def _show_tokens(self) -> None: self._update_tokens(self._context_tokens) app = FakeApp() app._show_tokens() assert display_calls == [1500] def test_show_tokens_preserves_approximate_marker_without_fresh_usage(self): """Turns without usage metadata should not clear a stale-token marker.""" display_calls: list[tuple[int, bool]] = [] def update_tokens(count: int, *, approximate: bool = False) -> None: display_calls.append((count, approximate)) app = SimpleNamespace( _context_tokens=1500, _tokens_approximate=True, _update_tokens=update_tokens, ) DeepAgentsApp._show_tokens(app, approximate=False) # ty: ignore assert app._tokens_approximate is True assert display_calls == [(1500, True)] def test_reset_clears_cache(self): """Resetting (e.g. /clear) should zero the cache and display.""" display_calls: list[int] = [] class FakeApp: _context_tokens: int = 3000 def _update_tokens(self, count: int) -> None: display_calls.append(count) app = FakeApp() app._context_tokens = 0 app._update_tokens(0) assert app._context_tokens == 0 assert display_calls == [0]