import asyncio import base64 from collections.abc import AsyncIterator from pathlib import Path import pytest from pydantic_ai import Agent from pydantic_ai.messages import ( ModelMessage, ModelResponse, SystemPromptPart, TextPart, TextPartDelta, ) from pydantic_ai.models.function import AgentInfo, FunctionModel from pydantic_ai.profiles import ModelProfile from rich.console import Console from rich.json import JSON from rich.text import Text from textual.app import App, ComposeResult from textual.widgets import Markdown from oterm.app.widgets.chat import ( ChatContainer, ChatItem, ToolCallItem, UsageStatus, _format_value, _truncate, ) from oterm.app.widgets.prompt import FlexibleInput from oterm.types import ChatModel, MessageModel from tests._helpers import image_b64, json_server, wait_until class _Host(App): CSS_PATH = "../../src/oterm/app/oterm.tcss" def __init__(self, chat_model: ChatModel, messages: list[MessageModel]): super().__init__() self._chat_model = chat_model self._messages = messages def compose(self) -> ComposeResult: yield ChatContainer(chat_model=self._chat_model, messages=self._messages) def _notifications(app: App) -> list: return list(app._notifications) class TestMount: async def test_renders_info_bar_with_model_name(self, chat_model): chat_model.model = "my-model" app = _Host(chat_model, []) async with app.run_test(): from textual.widgets import Static info = app.query_one("#info", Static) assert "my-model" in str(info.render()) async def test_prompt_is_focused_on_mount(self, chat_model): app = _Host(chat_model, []) async with app.run_test(): assert app.focused is not None assert app.focused.id == "promptArea" class TestLoadMessages: async def test_mounts_chat_item_per_message(self, store, chat_model): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id messages = [ MessageModel(chat_id=chat_id, role="user", text="hi"), MessageModel(chat_id=chat_id, role="assistant", text="hello"), ] app = _Host(chat_model, messages) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.load_messages() await pilot.pause() items = list(container.query(ChatItem)) assert len(items) == 2 assert items[0].author == "user" assert items[1].author == "assistant" async def test_load_is_idempotent(self, store, chat_model): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id messages = [MessageModel(chat_id=chat_id, role="user", text="hi")] app = _Host(chat_model, messages) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.load_messages() await container.load_messages() # second call must not re-mount await pilot.pause() assert len(list(container.query(ChatItem))) == 1 class TestOnSubmit: async def test_empty_input_is_ignored(self, store, chat_model): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) prompt = app.query_one(FlexibleInput) prompt.text = " " await pilot.press("enter") await pilot.pause() assert container.messages == [] async def test_submit_runs_response_task_and_saves_messages( self, store, chat_model ): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str]: yield "hi " yield "there" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = Agent(FunctionModel(stream_function=stream_fn)) prompt = app.query_one(FlexibleInput) prompt.text = "hello" await pilot.press("enter") await wait_until(pilot, lambda: len(container.messages) == 2) assert container.messages[0].role == "user" assert container.messages[0].text == "hello" assert container.messages[1].role == "assistant" assert container.messages[1].text == "hi there" rows = await store.get_messages(chat_id) assert len(rows) == 2 class TestEscapeCancel: async def test_escape_cancels_running_inference(self, store, chat_model): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id started = asyncio.Event() never = asyncio.Event() # stays clear so the stream hangs async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str]: started.set() await never.wait() yield "unreachable" # pragma: no cover app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = Agent(FunctionModel(stream_function=stream_fn)) prompt = app.query_one(FlexibleInput) prompt.text = "slow" await pilot.press("enter") await started.wait() await pilot.press("escape") # Give the cancellation a chance to clean up mounted items. for _ in range(10): await asyncio.sleep(0) await pilot.pause() # Restored prompt text assert prompt.text == "slow" # No messages saved assert container.messages == [] class TestClearChat: async def test_clear_removes_messages_and_ui(self, store, chat_model): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id msgs = [ MessageModel(chat_id=chat_id, role="user", text="q"), MessageModel(chat_id=chat_id, role="assistant", text="a"), ] for m in msgs: m.id = await store.save_message(m) app = _Host(chat_model, msgs) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.load_messages() await pilot.pause() assert len(list(container.query(ChatItem))) == 2 await container.action_clear_chat() await pilot.pause() assert container.messages == [] assert container.pydantic_history == [] assert list(container.query(ChatItem)) == [] rows = await store.get_messages(chat_id) assert rows == [] class TestImages: async def test_image_added_appended_and_token_inserted_at_cursor(self, chat_model): from oterm.app.widgets.image import ImageAdded from oterm.app.widgets.prompt import FlexibleInput, PostableTextArea app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) prompt = container.query_one("#prompt", FlexibleInput) textarea = prompt.query_one("#promptArea", PostableTextArea) textarea.text = "before after" textarea.cursor_location = (0, 7) # between "before " and "after" await pilot.pause() container.post_message(ImageAdded(Path("/tmp/a.png"), "base64data")) await pilot.pause() assert (Path("/tmp/a.png"), "base64data") in container.images assert textarea.text == "before [Image #1] after" async def test_second_image_gets_next_index(self, chat_model): from oterm.app.widgets.image import ImageAdded from oterm.app.widgets.prompt import FlexibleInput, PostableTextArea app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) prompt = container.query_one("#prompt", FlexibleInput) textarea = prompt.query_one("#promptArea", PostableTextArea) await pilot.pause() container.post_message(ImageAdded(Path("/tmp/a.png"), "data1")) await pilot.pause() container.post_message(ImageAdded(Path("/tmp/b.png"), "data2")) await pilot.pause() assert "[Image #1]" in textarea.text assert "[Image #2]" in textarea.text assert textarea.text.index("[Image #1]") < textarea.text.index("[Image #2]") async def test_token_is_highlighted(self, chat_model): from oterm.app.widgets.image import ImageAdded from oterm.app.widgets.prompt import FlexibleInput, PostableTextArea app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) prompt = container.query_one("#prompt", FlexibleInput) textarea = prompt.query_one("#promptArea", PostableTextArea) await pilot.pause() container.post_message(ImageAdded(Path("/tmp/a.png"), "data")) await pilot.pause() highlights = textarea._highlights[0] names = [name for _, _, name in highlights] assert "image-token" in names assert "image-token" in textarea._theme.syntax_styles async def test_backspace_inside_token_deletes_whole_token(self, chat_model): from oterm.app.widgets.image import ImageAdded from oterm.app.widgets.prompt import FlexibleInput, PostableTextArea app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) prompt = container.query_one("#prompt", FlexibleInput) textarea = prompt.query_one("#promptArea", PostableTextArea) await pilot.pause() container.post_message(ImageAdded(Path("/tmp/a.png"), "data")) await pilot.pause() assert textarea.text == "[Image #1] " textarea.cursor_location = (0, 5) await pilot.press("backspace") await pilot.pause() assert "[Image #1]" not in textarea.text assert textarea.text == " " async def test_backspace_just_after_token_deletes_token(self, chat_model): from oterm.app.widgets.image import ImageAdded from oterm.app.widgets.prompt import FlexibleInput, PostableTextArea app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) prompt = container.query_one("#prompt", FlexibleInput) textarea = prompt.query_one("#promptArea", PostableTextArea) await pilot.pause() container.post_message(ImageAdded(Path("/tmp/a.png"), "data")) await pilot.pause() textarea.cursor_location = (0, 10) # right after closing "]" await pilot.press("backspace") await pilot.pause() assert textarea.text == " " async def test_delete_at_token_start_deletes_whole_token(self, chat_model): from oterm.app.widgets.image import ImageAdded from oterm.app.widgets.prompt import FlexibleInput, PostableTextArea app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) prompt = container.query_one("#prompt", FlexibleInput) textarea = prompt.query_one("#promptArea", PostableTextArea) await pilot.pause() container.post_message(ImageAdded(Path("/tmp/a.png"), "data")) await pilot.pause() textarea.cursor_location = (0, 0) await pilot.press("delete") await pilot.pause() assert "[Image #1]" not in textarea.text async def test_backspace_before_token_deletes_one_char(self, chat_model): from oterm.app.widgets.image import ImageAdded from oterm.app.widgets.prompt import FlexibleInput, PostableTextArea app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) prompt = container.query_one("#prompt", FlexibleInput) textarea = prompt.query_one("#promptArea", PostableTextArea) textarea.text = "abc" textarea.cursor_location = (0, 3) await pilot.pause() container.post_message(ImageAdded(Path("/tmp/a.png"), "data")) await pilot.pause() assert textarea.text == "abc[Image #1] " textarea.cursor_location = (0, 3) await pilot.press("backspace") await pilot.pause() assert textarea.text == "ab[Image #1] " async def test_malformed_image_is_skipped_with_notification(self, chat_model): from oterm.app.widgets.chat import build_user_prompt async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str]: yield "ok " yield "response" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = Agent(FunctionModel(stream_function=stream_fn)) good = image_b64() bad = "not-valid-base64!!" user_prompt, skipped = build_user_prompt("look", [good, bad]) assert skipped == 1 chunks = [] async for chunk in container.stream_agent(user_prompt): chunks.append(chunk) text = "".join( p.content_delta for p in chunks if isinstance(p, TextPartDelta) ) assert text == "ok response" await pilot.pause() class TestDecodeImage: def test_jpeg_labeled_jpeg(self): from oterm.app.widgets.chat import _decode_image content = _decode_image(image_b64("JPEG")) assert content is not None assert content.media_type == "image/jpeg" def test_png_labeled_png(self): from oterm.app.widgets.chat import _decode_image content = _decode_image(image_b64("PNG")) assert content is not None assert content.media_type == "image/png" def test_non_image_bytes_return_none(self): from oterm.app.widgets.chat import _decode_image assert _decode_image(base64.b64encode(b"not an image").decode()) is None def test_invalid_b64_returns_none(self): from oterm.app.widgets.chat import _decode_image assert _decode_image("not-valid!!") is None class TestBuildUserPrompt: def test_no_tokens_no_images_returns_text(self): from oterm.app.widgets.chat import build_user_prompt assert build_user_prompt("hello", []) == ("hello", 0) def test_no_tokens_with_images_appends_at_end(self): from pydantic_ai import BinaryContent from oterm.app.widgets.chat import build_user_prompt good = image_b64() prompt, skipped = build_user_prompt("describe", [good]) assert skipped == 0 assert isinstance(prompt, list) assert prompt[0] == "describe" assert isinstance(prompt[1], BinaryContent) def test_empty_text_with_images_sends_only_the_images(self): from pydantic_ai import BinaryContent from oterm.app.widgets.chat import build_user_prompt prompt, skipped = build_user_prompt("", [image_b64()]) assert skipped == 0 assert isinstance(prompt, list) assert len(prompt) == 1 assert isinstance(prompt[0], BinaryContent) def test_token_interleaves_image_at_position(self): from pydantic_ai import BinaryContent from oterm.app.widgets.chat import build_user_prompt good = image_b64() prompt, skipped = build_user_prompt("see [Image #1] please", [good]) assert skipped == 0 assert isinstance(prompt, list) assert prompt[0] == "see " assert isinstance(prompt[1], BinaryContent) assert prompt[2] == " please" def test_unreferenced_image_is_dropped_when_tokens_present(self): from pydantic_ai import BinaryContent from oterm.app.widgets.chat import build_user_prompt good1 = image_b64("PNG") good2 = image_b64("JPEG") prompt, skipped = build_user_prompt("[Image #2] only", [good1, good2]) assert skipped == 0 assert isinstance(prompt, list) # Only image #2 is referenced; image #1 is not appended. binaries = [p for p in prompt if isinstance(p, BinaryContent)] assert len(binaries) == 1 def test_invalid_token_index_preserved_as_literal(self): from oterm.app.widgets.chat import build_user_prompt prompt, skipped = build_user_prompt("[Image #5] hi", []) assert skipped == 0 assert prompt == "[Image #5] hi" def test_no_token_all_invalid_b64_returns_text(self): from oterm.app.widgets.chat import build_user_prompt prompt, skipped = build_user_prompt("describe", ["not-valid!!"]) assert skipped == 1 assert prompt == "describe" def test_token_with_invalid_b64_increments_skipped(self): from oterm.app.widgets.chat import build_user_prompt prompt, skipped = build_user_prompt("see [Image #1]", ["not-valid!!"]) assert skipped == 1 assert prompt == "see [Image #1]" def test_text_ending_at_token_keeps_no_trailing(self): from pydantic_ai import BinaryContent from oterm.app.widgets.chat import build_user_prompt good = image_b64() prompt, skipped = build_user_prompt("see [Image #1]", [good]) assert skipped == 0 assert isinstance(prompt, list) assert prompt[-1].__class__ is BinaryContent # no trailing text part class TestPydanticHistoryRebuild: async def test_user_message_with_token_replays_image_inline( self, store, chat_model ): from pydantic_ai import BinaryContent from pydantic_ai.messages import ModelRequest, UserPromptPart chat_id = await store.save_chat(chat_model) chat_model.id = chat_id good = image_b64() msg = MessageModel( chat_id=chat_id, role="user", text="see [Image #1] please", images=[good], ) msg.id = await store.save_message(msg) app = _Host(chat_model, [msg]) async with app.run_test(): container = app.query_one(ChatContainer) history = container.pydantic_history assert len(history) == 1 req = history[0] assert isinstance(req, ModelRequest) user_part = req.parts[0] assert isinstance(user_part, UserPromptPart) assert isinstance(user_part.content, list) assert any(isinstance(p, BinaryContent) for p in user_part.content) async def test_legacy_user_message_without_tokens_appends_images( self, store, chat_model ): from pydantic_ai import BinaryContent from pydantic_ai.messages import ModelRequest, UserPromptPart chat_id = await store.save_chat(chat_model) chat_model.id = chat_id good = image_b64() msg = MessageModel( chat_id=chat_id, role="user", text="describe", images=[good], ) msg.id = await store.save_message(msg) app = _Host(chat_model, [msg]) async with app.run_test(): container = app.query_one(ChatContainer) req = container.pydantic_history[0] assert isinstance(req, ModelRequest) user_part = req.parts[0] assert isinstance(user_part, UserPromptPart) assert isinstance(user_part.content, list) assert user_part.content[0] == "describe" assert isinstance(user_part.content[1], BinaryContent) async def test_text_only_user_message_replays_as_string(self, store, chat_model): from pydantic_ai.messages import ModelRequest, UserPromptPart chat_id = await store.save_chat(chat_model) chat_model.id = chat_id msg = MessageModel(chat_id=chat_id, role="user", text="hi", images=[]) msg.id = await store.save_message(msg) app = _Host(chat_model, [msg]) async with app.run_test(): container = app.query_one(ChatContainer) req = container.pydantic_history[0] assert isinstance(req, ModelRequest) user_part = req.parts[0] assert isinstance(user_part, UserPromptPart) assert user_part.content == "hi" class TestHistory: async def test_action_history_opens_modal(self, store, chat_model): from oterm.app.prompt_history import PromptHistory chat_id = await store.save_chat(chat_model) chat_model.id = chat_id msgs = [ MessageModel(chat_id=chat_id, role="user", text="first"), MessageModel(chat_id=chat_id, role="assistant", text="answer"), MessageModel(chat_id=chat_id, role="user", text="second"), ] app = _Host(chat_model, msgs) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.action_history() await pilot.pause() assert isinstance(app.screen, PromptHistory) class TestResponseTaskErrors: async def test_cannot_send_when_agent_is_none(self, app_config, chat_model): chat_model.provider = "openai-compat/ghost" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) assert container.agent is None prompt = app.query_one(FlexibleInput) prompt.text = "hello" await pilot.press("enter") await wait_until( pilot, lambda: any("Cannot send" in n.message for n in _notifications(app)), ) assert any("Cannot send" in n.message for n in _notifications(app)) assert container.messages == [] async def test_stream_exception_notifies_and_cleans_up(self, store, chat_model): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str]: raise RuntimeError("boom") yield # pragma: no cover app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = Agent(FunctionModel(stream_function=stream_fn)) prompt = app.query_one(FlexibleInput) prompt.text = "trigger" await pilot.press("enter") await wait_until( pilot, lambda: any( "Unexpected error" in n.message for n in _notifications(app) ), ) assert container.messages == [] assert any("Unexpected error" in n.message for n in _notifications(app)) # UsageStatus is removed when the turn fails. assert list(container.query(UsageStatus)) == [] class TestEditChat: async def test_edit_chat_updates_model_and_agent( self, store, chat_model, monkeypatch ): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id from oterm.types import ChatModel updated = ChatModel( id=chat_id, name="renamed", model=chat_model.model, # store.edit_chat doesn't update model system="be terse", provider="ollama", parameters={"temperature": 0.2}, tools=[], thinking=True, ) app = _Host(chat_model, []) async def fake_push_screen_wait(self, screen): return updated monkeypatch.setattr(type(app), "push_screen_wait", fake_push_screen_wait) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.action_edit_chat() await wait_until(pilot, lambda: container.chat_model.system == "be terse") assert container.chat_model.system == "be terse" assert container.chat_model.thinking is True assert container.chat_model.parameters == {"temperature": 0.2} reloaded = await store.get_chat(chat_id) assert reloaded is not None assert reloaded.system == "be terse" assert reloaded.thinking is True assert reloaded.parameters == {"temperature": 0.2} async def test_edit_chat_cancelled_modal_is_noop( self, store, chat_model, monkeypatch ): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id app = _Host(chat_model, []) async def fake_push_screen_wait(self, screen): return None monkeypatch.setattr(type(app), "push_screen_wait", fake_push_screen_wait) async with app.run_test() as pilot: container = app.query_one(ChatContainer) original_model = container.chat_model.model container.action_edit_chat() for _ in range(5): await pilot.pause() assert container.chat_model.model == original_model class TestToggleThinking: async def test_toggle_enables_thinking_for_session_only( self, store, chat_model, monkeypatch ): from oterm.app.widgets import chat as chat_module from oterm.providers.capabilities import ModelCapabilities chat_id = await store.save_chat(chat_model) chat_model.id = chat_id monkeypatch.setattr( chat_module, "get_capabilities", lambda provider, model: ModelCapabilities(supports_thinking=True), ) app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) assert container.chat_model.thinking is False container.action_toggle_thinking() await pilot.pause() assert container.chat_model.thinking is True assert any("Thinking on" in n.message for n in _notifications(app)) # Session-only: the stored chat is untouched. reloaded = await store.get_chat(chat_id) assert reloaded is not None and reloaded.thinking is False async def test_toggle_disables_thinking_for_session_only( self, store, chat_model, monkeypatch ): from oterm.app.widgets import chat as chat_module from oterm.providers.capabilities import ModelCapabilities chat_model.thinking = True chat_id = await store.save_chat(chat_model) chat_model.id = chat_id monkeypatch.setattr( chat_module, "get_capabilities", lambda provider, model: ModelCapabilities(supports_thinking=True), ) app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.action_toggle_thinking() await pilot.pause() assert container.chat_model.thinking is False assert any("Thinking off" in n.message for n in _notifications(app)) reloaded = await store.get_chat(chat_id) assert reloaded is not None and reloaded.thinking is True async def test_toggle_rebuilds_agent_with_new_thinking( self, store, chat_model, monkeypatch ): from oterm.app.widgets import chat as chat_module from oterm.providers.capabilities import ModelCapabilities chat_id = await store.save_chat(chat_model) chat_model.id = chat_id monkeypatch.setattr( chat_module, "get_capabilities", lambda provider, model: ModelCapabilities(supports_thinking=True), ) thinking_seen: list[bool] = [] def fake_get_agent(*args, thinking, **kwargs): thinking_seen.append(thinking) return None monkeypatch.setattr(chat_module, "get_agent", fake_get_agent) app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) thinking_seen.clear() container.action_toggle_thinking() await pilot.pause() assert thinking_seen == [True] async def test_toggle_on_unsupported_model_notifies_and_noops( self, store, chat_model, monkeypatch ): from oterm.app.widgets import chat as chat_module from oterm.providers.capabilities import ModelCapabilities chat_id = await store.save_chat(chat_model) chat_model.id = chat_id monkeypatch.setattr( chat_module, "get_capabilities", lambda provider, model: ModelCapabilities(supports_thinking=False), ) app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.action_toggle_thinking() await pilot.pause() assert any( "does not support thinking" in n.message for n in _notifications(app) ) assert container.chat_model.thinking is False class TestRenameChat: async def test_rename_chat_updates_store_and_notifies( self, store, chat_model, monkeypatch ): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id from textual.widgets import TabbedContent, TabPane class _Host2(App): def __init__(self): super().__init__() def compose(self): with TabbedContent(): with TabPane("original", id=f"chat-{chat_id}"): yield ChatContainer(chat_model=chat_model, messages=[]) app = _Host2() async def fake_push_screen_wait(self, screen): return "new-name" monkeypatch.setattr(type(app), "push_screen_wait", fake_push_screen_wait) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.action_rename_chat() await wait_until( pilot, lambda: any("renamed" in n.message for n in _notifications(app)), ) reloaded = await store.get_chat(chat_id) assert reloaded and reloaded.name == "new-name" assert any("renamed" in n.message for n in _notifications(app)) async def test_rename_chat_cancelled_is_noop(self, store, chat_model, monkeypatch): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id chat_model.name = "unchanged" await store.edit_chat(chat_model) app = _Host(chat_model, []) async def fake_push_screen_wait(self, screen): return None monkeypatch.setattr(type(app), "push_screen_wait", fake_push_screen_wait) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.action_rename_chat() for _ in range(5): await pilot.pause() reloaded = await store.get_chat(chat_id) assert reloaded is not None assert reloaded.name == "unchanged" class TestHistoryCallback: async def test_selecting_history_entry_fills_prompt(self, store, chat_model): from textual.widgets import OptionList from oterm.app.prompt_history import PromptHistory from oterm.app.widgets.prompt import FlexibleInput chat_id = await store.save_chat(chat_model) chat_model.id = chat_id msgs = [ MessageModel(chat_id=chat_id, role="user", text="previous"), MessageModel(chat_id=chat_id, role="assistant", text="a"), ] app = _Host(chat_model, msgs) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.action_history() await pilot.pause() assert isinstance(app.screen, PromptHistory) option_list = app.screen.query_one(OptionList) option_list.highlighted = 0 await pilot.press("enter") await pilot.pause() prompt = container.query_one(FlexibleInput) assert prompt.text == "previous" async def test_cancelling_history_leaves_prompt_unchanged(self, store, chat_model): from oterm.app.widgets.prompt import FlexibleInput chat_id = await store.save_chat(chat_model) chat_model.id = chat_id msgs = [MessageModel(chat_id=chat_id, role="user", text="prev")] app = _Host(chat_model, msgs) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.action_history() await pilot.pause() await pilot.press("escape") await pilot.pause() prompt = container.query_one(FlexibleInput) assert prompt.text == "" async def test_selecting_history_populates_prompt(self, store, chat_model): from textual.widgets import OptionList from oterm.app.widgets.prompt import FlexibleInput chat_id = await store.save_chat(chat_model) chat_model.id = chat_id msgs = [MessageModel(chat_id=chat_id, role="user", text="line1\nline2")] app = _Host(chat_model, msgs) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.action_history() await pilot.pause() option_list = app.screen.query_one(OptionList) option_list.highlighted = 0 await pilot.press("enter") await pilot.pause() prompt = container.query_one(FlexibleInput) assert prompt.text == "line1\nline2" class TestCopyMessage: async def test_copies_last_message_to_clipboard(self, chat_model): messages = [ MessageModel(chat_id=1, role="user", text="the question"), MessageModel(chat_id=1, role="assistant", text="the answer"), ] app = _Host(chat_model, messages) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.load_messages() copied: list[str] = [] app.copy_to_clipboard = lambda t: copied.append(t) # ty: ignore[invalid-assignment] container.action_copy_message() await pilot.pause() assert copied == ["the answer"] assert any("copied" in n.message.lower() for n in _notifications(app)) async def test_copy_is_a_silent_noop_with_nothing_to_copy(self, chat_model): """No messages, and an assistant item mounted by response_task before its first delta, both copy nothing and say nothing.""" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) copied: list[str] = [] app.copy_to_clipboard = lambda t: copied.append(t) # ty: ignore[invalid-assignment] container.action_copy_message() await pilot.pause() assert copied == [] assert _notifications(app) == [] item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() container.action_copy_message() await pilot.pause() assert copied == [] assert _notifications(app) == [] async def test_clicking_an_empty_message_copies_nothing(self, chat_model): app = _Host(chat_model, []) async with app.run_test(size=(120, 40)) as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() copied: list[str] = [] app.copy_to_clipboard = lambda t: copied.append(t) # ty: ignore[invalid-assignment] await pilot.click(item.query_one(".response", Markdown)) await pilot.pause() assert copied == [] assert _notifications(app) == [] async def test_copies_streaming_text_not_yet_in_messages(self, chat_model): """``messages`` is only appended once streaming finishes, so the copy must read the widget to pick up a response still in flight.""" messages = [MessageModel(chat_id=1, role="assistant", text="stale")] app = _Host(chat_model, messages) async with app.run_test() as pilot: container = app.query_one(ChatContainer) streaming = ChatItem() streaming.author = "assistant" await container.query_one("#messageContainer").mount(streaming) await streaming.append_text("in flight") await pilot.pause() copied: list[str] = [] app.copy_to_clipboard = lambda t: copied.append(t) # ty: ignore[invalid-assignment] container.action_copy_message() await pilot.pause() assert copied == ["in flight"] class TestChatItemClickCopy: async def test_clicking_copies_text_to_clipboard(self, chat_model): app = _Host(chat_model, []) async with app.run_test(size=(120, 40)) as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" item.text = "the answer" await container.query_one("#messageContainer").mount(item) await pilot.pause() copied: list[str] = [] def fake_copy(text): copied.append(text) app.copy_to_clipboard = fake_copy # ty: ignore[invalid-assignment] markdown = item.query_one(".response", Markdown) await pilot.click(markdown) await pilot.pause() assert copied == ["the answer"] class TestChatItemUser: async def test_user_chat_item_renders_with_static(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "user" item.text = "hello" await container.query_one("#messageContainer").mount(item) await pilot.pause() from textual.widgets import Static statics = list(item.query(Static)) assert any(s.has_class("text") for s in statics) # User items don't render a thinking label assert not list(item.query(".thinking-label")) async def test_setting_text_on_user_item_is_noop(self, chat_model): """watch_text short-circuits for user items since they render via Static.""" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "user" await container.query_one("#messageContainer").mount(item) await pilot.pause() item.text = "hi" await pilot.pause() # No Markdown is queried, and no exception is raised. assert not list(item.query(Markdown)) async def test_setting_thinking_on_user_item_is_noop(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "user" await container.query_one("#messageContainer").mount(item) await pilot.pause() item.thinking = "hmm" # should be swallowed without exception await pilot.pause() class TestThinkingCollapse: async def test_thinking_visible_before_response_starts(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() item.thinking = "musing…" await pilot.pause() from textual.widgets import Static label = item.query_one(".thinking-label", Static) body = item.query_one(".thinking-body", Markdown) assert label.display is True assert body.display is True assert "thinking" in str(label.render()) async def test_thinking_collapses_when_response_starts(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() item.thinking = "musing…" await pilot.pause() item.text = "answer" await pilot.pause() from textual.widgets import Static label = item.query_one(".thinking-label", Static) body = item.query_one(".thinking-body", Markdown) assert label.display is True assert body.display is False assert "▸" in str(label.render()) assert "thoughts" in str(label.render()) async def test_thinking_stays_collapsed_on_further_thinking_deltas( self, chat_model ): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() item.thinking = "musing…" item.text = "answer" await pilot.pause() item.thinking = "musing… more" await pilot.pause() body = item.query_one(".thinking-body", Markdown) assert body.display is False async def test_clicking_collapsed_label_re_expands(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() item.thinking = "musing" item.text = "answer" await pilot.pause() from textual.widgets import Static label = item.query_one(".thinking-label", Static) body = item.query_one(".thinking-body", Markdown) assert body.display is False copied: list[str] = [] app.copy_to_clipboard = lambda t: copied.append(t) # ty: ignore[invalid-assignment] await pilot.click(label) await pilot.pause() assert body.display is True assert "▾" in str(label.render()) assert copied == [] await pilot.click(label) await pilot.pause() assert body.display is False assert "▸" in str(label.render()) async def test_clicking_response_still_copies(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" item.thinking = "musing" item.text = "answer" await container.query_one("#messageContainer").mount(item) await pilot.pause() copied: list[str] = [] app.copy_to_clipboard = lambda t: copied.append(t) # ty: ignore[invalid-assignment] response = item.query_one(".response", Markdown) await pilot.click(response) await pilot.pause() assert copied == ["answer"] async def test_clicking_thinking_label_without_text_does_not_toggle( self, chat_model ): from textual.widgets import Static app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() item.thinking = "still thinking" await pilot.pause() label = item.query_one(".thinking-label", Static) assert item.thoughts_collapsed is False copied: list[str] = [] app.copy_to_clipboard = lambda t: copied.append(t) # ty: ignore[invalid-assignment] await pilot.click(label) await pilot.pause() # Click on label returns early: no toggle, no copy. assert item.thoughts_collapsed is False assert copied == [] async def test_clicking_thinking_body_does_not_copy(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" item.thinking = "musing" item.text = "answer" await container.query_one("#messageContainer").mount(item) await pilot.pause() # Surface the body so it can receive a click. item.thoughts_collapsed = False await pilot.pause() copied: list[str] = [] app.copy_to_clipboard = lambda t: copied.append(t) # ty: ignore[invalid-assignment] body = item.query_one(".thinking-body", Markdown) await pilot.click(body) await pilot.pause() assert copied == [] class TestUserItemClick: async def test_clicking_user_item_copies_text(self, chat_model): from textual.widgets import Static app = _Host(chat_model, []) async with app.run_test(size=(120, 40)) as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "user" item.text = "my prompt" await container.query_one("#messageContainer").mount(item) await pilot.pause() copied: list[str] = [] app.copy_to_clipboard = lambda t: copied.append(t) # ty: ignore[invalid-assignment] text_static = next(s for s in item.query(Static) if s.has_class("text")) await pilot.click(text_static) await pilot.pause() assert copied == ["my prompt"] class TestChatItemStreaming: """Streaming path: append_text / append_thinking / finish_stream. These bypass the watch_text / watch_thinking re-render path and write deltas to a Textual ``MarkdownStream`` so long responses render incrementally instead of re-parsing the whole document per token. """ async def test_append_text_updates_text_reactive(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.append_text("hello ") await item.append_text("world") await pilot.pause() # Reactive kept in sync for click-to-copy and chrome refresh. assert item.text == "hello world" async def test_append_text_collapses_thoughts_on_first_delta(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.append_thinking("thinking…") await pilot.pause() assert item.thoughts_collapsed is False await item.append_text("answer") await pilot.pause() assert item.thoughts_collapsed is True from textual.widgets import Static label = item.query_one(".thinking-label", Static) body = item.query_one(".thinking-body", Markdown) assert body.display is False assert "▸" in str(label.render()) async def test_append_thinking_reveals_label_on_first_delta(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() from textual.widgets import Static label = item.query_one(".thinking-label", Static) assert label.display is False await item.append_thinking("musing") await pilot.pause() assert label.display is True assert item.thinking == "musing" async def test_user_item_append_is_noop(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "user" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.append_text("ignored") await item.append_thinking("ignored") await pilot.pause() assert item.text == "" assert item.thinking == "" async def test_finish_stream_idempotent_when_never_streamed(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.finish_stream() await item.finish_stream() # safe to call twice async def test_cancel_streams_safe_after_appends(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.append_text("partial ") item.cancel_streams() # No exception, streams cleared. assert item._response_stream is None assert item._thinking_stream is None async def test_append_thinking_twice_reuses_stream(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.append_thinking("first ") stream = item._thinking_stream assert stream is not None await item.append_thinking("second") await pilot.pause() # Second call must reuse the same MarkdownStream, not recreate one. assert item._thinking_stream is stream assert item.thinking == "first second" async def test_finish_stream_drains_thinking_stream(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.append_thinking("musing") assert item._thinking_stream is not None await item.finish_stream() async def test_finish_stream_re_renders_full_markdown( self, chat_model, monkeypatch ): """Per-delta ``Markdown.append`` leaves ``_last_parsed_line`` at the start of the trailing top-level token (e.g. an open code fence), so the next refresh can drop the streamed blocks. ``finish_stream`` must force a full ``Markdown.update`` with the accumulated text so the widget is in a clean re-parsed state.""" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() deltas = [ "# Summary\n\n", "## Python Example\n\n", "```python\n", 'print("Hello, world!")\n', "```\n", ] for d in deltas: await item.append_text(d) await pilot.pause() response = item.query_one(".response", Markdown) thinking = item.query_one(".thinking-body", Markdown) update_calls: list[tuple[Markdown, str]] = [] original = Markdown.update def _track(self, markdown): update_calls.append((self, markdown)) return original(self, markdown) monkeypatch.setattr(Markdown, "update", _track) await item.finish_stream() await pilot.pause() full = "".join(deltas) assert (response, full) in update_calls assert all(target is not thinking for target, _ in update_calls) async def test_finish_stream_re_renders_thinking(self, chat_model, monkeypatch): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.append_thinking("step 1\n") await item.append_thinking("```py\nx=1\n```\n") await pilot.pause() thinking = item.query_one(".thinking-body", Markdown) update_calls: list[tuple[Markdown, str]] = [] original = Markdown.update def _track(self, markdown): update_calls.append((self, markdown)) return original(self, markdown) monkeypatch.setattr(Markdown, "update", _track) await item.finish_stream() await pilot.pause() assert (thinking, "step 1\n```py\nx=1\n```\n") in update_calls assert item._thinking_stream is None class TestSkippedImageNotify: async def test_response_task_notifies_for_malformed_image(self, store, chat_model): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str]: yield "ok " yield "done" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = Agent(FunctionModel(stream_function=stream_fn)) container.images = [(Path("/tmp/x.png"), "not-valid!!")] prompt = app.query_one(FlexibleInput) prompt.text = "see [Image #1]" await pilot.press("enter") await wait_until(pilot, lambda: len(container.messages) == 2) assert any( "Skipped 1 malformed image" in n.message for n in _notifications(app) ) async def test_regenerate_notifies_for_malformed_image(self, store, chat_model): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id user_msg = MessageModel( chat_id=chat_id, role="user", text="see [Image #1]", images=["not-valid!!"], ) user_msg.id = await store.save_message(user_msg) old_assistant = MessageModel(chat_id=chat_id, role="assistant", text="old") old_assistant.id = await store.save_message(old_assistant) async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str]: yield "new " yield "answer" app = _Host(chat_model, [user_msg, old_assistant]) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.load_messages() container.agent = Agent(FunctionModel(stream_function=stream_fn)) await container.action_regenerate_llm_message() await wait_until(pilot, lambda: container.messages[-1].text == "new answer") assert any( "Skipped 1 malformed image" in n.message for n in _notifications(app) ) class TestThinkingViaResponseTask: async def test_thinking_then_text_streams_through_response_task( self, store, chat_model ): from pydantic_ai.models.function import DeltaThinkingPart chat_id = await store.save_chat(chat_model) chat_model.id = chat_id async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str | dict[int, DeltaThinkingPart]]: yield {0: DeltaThinkingPart(content="weighing… ")} yield {0: DeltaThinkingPart(content="options")} yield "the " yield "answer" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = Agent(FunctionModel(stream_function=stream_fn)) prompt = app.query_one(FlexibleInput) prompt.text = "ask" await pilot.press("enter") await wait_until(pilot, lambda: len(container.messages) == 2) assert container.messages[-1].text == "the answer" items = list(container.query(ChatItem)) assistant = items[-1] assert "weighing… options" in assistant.thinking async def test_tool_call_renders_inside_assistant_item(self, store, chat_model): from pydantic_ai import Tool from pydantic_ai.models.function import DeltaToolCall chat_id = await store.save_chat(chat_model) chat_model.id = chat_id def echo(s: str) -> str: return f"echoed: {s}" async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str | dict[int, DeltaToolCall]]: already_called = any( getattr(m, "parts", None) and any(getattr(p, "part_kind", "") == "tool-return" for p in m.parts) for m in messages ) if already_called: yield "all done" return yield { 0: DeltaToolCall( name="echo", json_args='{"s": "hi"}', tool_call_id="tc-1" ) } app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = Agent( FunctionModel(stream_function=stream_fn), tools=[Tool(echo)] ) prompt = app.query_one(FlexibleInput) prompt.text = "go" await pilot.press("enter") await wait_until(pilot, lambda: len(container.messages) == 2) from textual.widgets import Static assistant = list(container.query(ChatItem))[-1] tool_items = list(assistant.query(ToolCallItem)) assert len(tool_items) == 1 tool_item = tool_items[0] assert tool_item.tool_name == "echo" assert tool_item.tool_call_id == "tc-1" assert tool_item.result == "echoed: hi" header = tool_item.query_one(".tool-call-header", Static) assert "▸ tool call: echo" in str(header.render()) # Click the tool-call to expand; args + result should land in body. await pilot.click(ToolCallItem) await pilot.pause() body_text = _capture(tool_item.query_one(".tool-call-body", Static).content) assert '"s": "hi"' in body_text assert "echoed: hi" in body_text assert "args:" in body_text assert "result:" in body_text assert "▾ tool call: echo" in str(header.render()) async def test_file_part_streamed_through_response_task(self, store, chat_model): """A `FilePart` arrives during streaming and renders as an Image widget.""" from io import BytesIO from PIL import Image as PILImage from pydantic_ai.messages import BinaryImage, FilePart from textual_image.widget import Image as ImageWidget from tests._stream_helpers import make_file_aware_agent buf = BytesIO() PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG") png_bytes = buf.getvalue() chat_id = await store.save_chat(chat_model) chat_model.id = chat_id async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str | FilePart]: yield "before " yield FilePart(content=BinaryImage(data=png_bytes, media_type="image/png")) yield "after" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = make_file_aware_agent(stream_fn) prompt = app.query_one(FlexibleInput) prompt.text = "draw" await pilot.press("enter") await wait_until(pilot, lambda: len(container.messages) == 2) assert container.messages[-1].text == "before after" assistant = list(container.query(ChatItem))[-1] images = list(assistant.query(ImageWidget)) assert len(images) == 1 assert images[0].image is not None # Persisted: the assistant row carries the image as base64. assistant_row = container.messages[-1] assert len(assistant_row.images) == 1 assert base64.b64decode(assistant_row.images[0]) == png_bytes stored = await store.get_messages(chat_id) assistant_rows = [m for m in stored if m.role == "assistant"] assert len(assistant_rows[0].images) == 1 async def test_non_image_file_part_is_ignored_on_stream(self, store, chat_model): """Audio/document FileParts are dropped: no widget, nothing persisted.""" from pydantic_ai.messages import BinaryContent, FilePart from textual_image.widget import Image as ImageWidget from tests._stream_helpers import make_file_aware_agent chat_id = await store.save_chat(chat_model) chat_model.id = chat_id async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str | FilePart]: yield "before " yield FilePart( content=BinaryContent(data=b"ID3\x04audio", media_type="audio/mp3") ) yield "after" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = make_file_aware_agent(stream_fn) prompt = app.query_one(FlexibleInput) prompt.text = "say hi" await pilot.press("enter") await wait_until(pilot, lambda: len(container.messages) == 2) assert container.messages[-1].text == "before after" assistant = list(container.query(ChatItem))[-1] assert list(assistant.query(ImageWidget)) == [] assert container.messages[-1].images == [] async def test_non_image_file_part_is_ignored_on_regenerate( self, store, chat_model ): from pydantic_ai.messages import BinaryContent, FilePart from textual_image.widget import Image as ImageWidget from tests._stream_helpers import make_file_aware_agent chat_id = await store.save_chat(chat_model) chat_model.id = chat_id user_msg = MessageModel(chat_id=chat_id, role="user", text="ask") user_msg.id = await store.save_message(user_msg) old = MessageModel(chat_id=chat_id, role="assistant", text="old") old.id = await store.save_message(old) async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str | FilePart]: yield "regen " yield FilePart( content=BinaryContent(data=b"%PDF-1.4", media_type="application/pdf") ) yield "done" app = _Host(chat_model, [user_msg, old]) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.load_messages() container.agent = make_file_aware_agent(stream_fn) await container.action_regenerate_llm_message() await wait_until(pilot, lambda: container.messages[-1].text == "regen done") assistant = list(container.query(ChatItem))[-1] assert list(assistant.query(ImageWidget)) == [] assert container.messages[-1].images == [] async def test_failing_tool_logs_and_shows_error_in_call_widget( self, store, chat_model ): """A tool raising `ModelRetry` is logged and rendered as `error: …`.""" from pydantic_ai import Tool from pydantic_ai.exceptions import ModelRetry from pydantic_ai.models.function import DeltaToolCall from oterm.log import log_lines chat_id = await store.save_chat(chat_model) chat_model.id = chat_id def boom(s: str) -> str: raise ModelRetry("kaboom") async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str | dict[int, DeltaToolCall]]: already_retried = any( getattr(m, "parts", None) and any(getattr(p, "part_kind", "") == "retry-prompt" for p in m.parts) for m in messages ) if already_retried: yield "gave up" return yield { 0: DeltaToolCall( name="boom", json_args='{"s": "x"}', tool_call_id="tc-b" ) } log_lines.clear() app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = Agent( FunctionModel(stream_function=stream_fn), tools=[Tool(boom)] ) prompt = app.query_one(FlexibleInput) prompt.text = "go" await pilot.press("enter") await wait_until(pilot, lambda: len(container.messages) == 2) assistant = list(container.query(ChatItem))[-1] tool_items = list(assistant.query(ToolCallItem)) assert len(tool_items) == 1 assert tool_items[0].tool_name == "boom" assert isinstance(tool_items[0].result, str) assert tool_items[0].result.startswith("error:") assert "kaboom" in tool_items[0].result messages = [m for _, m in log_lines] assert any("'boom' failed" in m and "kaboom" in m for m in messages) async def test_tool_returning_binary_image_renders_inline(self, store, chat_model): """A tool returning `BinaryImage` mounts an Image widget and persists.""" from io import BytesIO from PIL import Image as PILImage from pydantic_ai import Tool from pydantic_ai.messages import BinaryImage from pydantic_ai.models.function import DeltaToolCall from textual_image.widget import Image as ImageWidget buf = BytesIO() PILImage.new("RGB", (4, 4), "green").save(buf, format="PNG") png_bytes = buf.getvalue() chat_id = await store.save_chat(chat_model) chat_model.id = chat_id def make_image(prompt: str) -> BinaryImage: return BinaryImage(data=png_bytes, media_type="image/png") async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str | dict[int, DeltaToolCall]]: already_called = any( getattr(m, "parts", None) and any(getattr(p, "part_kind", "") == "tool-return" for p in m.parts) for m in messages ) if already_called: yield "done" return yield { 0: DeltaToolCall( name="make_image", json_args='{"prompt": "a square"}', tool_call_id="tc-img", ) } app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = Agent( FunctionModel(stream_function=stream_fn), tools=[Tool(make_image)] ) prompt = app.query_one(FlexibleInput) prompt.text = "draw" await pilot.press("enter") await wait_until(pilot, lambda: len(container.messages) == 2) assistant = list(container.query(ChatItem))[-1] images = list(assistant.query(ImageWidget)) assert len(images) == 1 assert images[0].image is not None tool_items = list(assistant.query(ToolCallItem)) assert len(tool_items) == 1 assert tool_items[0].tool_name == "make_image" assert tool_items[0].result is not None assistant_row = container.messages[-1] assert len(assistant_row.images) == 1 assert base64.b64decode(assistant_row.images[0]) == png_bytes stored = await store.get_messages(chat_id) assistant_rows = [m for m in stored if m.role == "assistant"] assert len(assistant_rows[0].images) == 1 async def test_persisted_assistant_image_renders_on_load(self, store, chat_model): from io import BytesIO from PIL import Image as PILImage from textual_image.widget import Image as ImageWidget buf = BytesIO() PILImage.new("RGB", (4, 4), "blue").save(buf, format="PNG") png_bytes = buf.getvalue() chat_id = await store.save_chat(chat_model) chat_model.id = chat_id user_msg = MessageModel(chat_id=chat_id, role="user", text="draw") user_msg.id = await store.save_message(user_msg) assistant_msg = MessageModel( chat_id=chat_id, role="assistant", text="here you go", images=[base64.b64encode(png_bytes).decode()], ) assistant_msg.id = await store.save_message(assistant_msg) app = _Host(chat_model, [user_msg, assistant_msg]) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.load_messages() await pilot.pause() items = list(container.query(ChatItem)) assert items[-1].author == "assistant" images = list(items[-1].query(ImageWidget)) assert len(images) == 1 assert images[0].image is not None class TestRegenerateCancellation: async def test_cancellation_restores_state(self, store, chat_model): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id user_msg = MessageModel(chat_id=chat_id, role="user", text="ask") user_msg.id = await store.save_message(user_msg) old_assistant = MessageModel(chat_id=chat_id, role="assistant", text="old") old_assistant.id = await store.save_message(old_assistant) async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str]: raise asyncio.CancelledError() yield # pragma: no cover app = _Host(chat_model, [user_msg, old_assistant]) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await container.load_messages() container.agent = Agent(FunctionModel(stream_function=stream_fn)) await container.action_regenerate_llm_message() await wait_until(pilot, lambda: container.messages[-1].text == "old") assert container.messages[-1].text == "old" assert list(container.query(UsageStatus)) == [] class TestToolCallHelpers: def test_truncate_keeps_short_text(self): assert _truncate("short", 10) == "short" def test_truncate_appends_ellipsis(self): assert _truncate("abcdefghij", 5) == "abcde…" @pytest.mark.parametrize( ("value", "expected_type", "expected_substring"), [ ({"a": 1}, JSON, None), ('{"a": 1}', JSON, None), ("hello world", Text, "hello"), (None, Text, "(none)"), (42, Text, "42"), ], ) def test_format_value_dispatches_by_type( self, value, expected_type, expected_substring ): rendered = _format_value(value) assert isinstance(rendered, expected_type) if expected_substring is not None: assert expected_substring in rendered.plain class TestToolCallRendering: async def test_add_tool_call_is_a_no_op_for_user_items(self, chat_model): from pydantic_ai.messages import ToolCallPart app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "user" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.add_tool_call( ToolCallPart(tool_name="x", args="{}", tool_call_id="tc-x") ) assert list(item.query(ToolCallItem)) == [] async def test_add_image_is_a_no_op_for_user_items(self, chat_model): from textual_image.widget import Image as ImageWidget app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "user" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.add_image(b"\x89PNG\r\n") assert list(item.query(ImageWidget)) == [] async def test_clicking_assistant_image_saves_it_to_disk( self, chat_model, tmp_path, monkeypatch ): from io import BytesIO from PIL import Image as PILImage import oterm.config from oterm.app.widgets.chat import AssistantImage monkeypatch.setattr(oterm.config.envConfig, "OTERM_DATA_DIR", tmp_path) buf = BytesIO() PILImage.new("RGB", (4, 4), "lime").save(buf, format="PNG") png_bytes = buf.getvalue() app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.add_image(png_bytes) await pilot.pause() assistant_image = item.query_one(AssistantImage) await item._save_assistant_image(assistant_image) await pilot.pause() saved = list((tmp_path / "downloads").iterdir()) assert len(saved) == 1 assert saved[0].read_bytes() == png_bytes assert saved[0].suffix == ".png" assert any("Image saved" in n.message for n in _notifications(app)) async def test_jpeg_assistant_image_saves_with_jpg_extension( self, chat_model, tmp_path, monkeypatch ): from io import BytesIO from PIL import Image as PILImage import oterm.config from oterm.app.widgets.chat import AssistantImage monkeypatch.setattr(oterm.config.envConfig, "OTERM_DATA_DIR", tmp_path) buf = BytesIO() PILImage.new("RGB", (4, 4), "orange").save(buf, format="JPEG") jpg_bytes = buf.getvalue() app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.add_image(jpg_bytes) await pilot.pause() await item._save_assistant_image(item.query_one(AssistantImage)) saved = list((tmp_path / "downloads").iterdir()) assert len(saved) == 1 assert saved[0].suffix == ".jpg" async def test_chat_item_click_on_image_dispatches_save( self, chat_model, tmp_path, monkeypatch ): from io import BytesIO from PIL import Image as PILImage from textual.events import Click import oterm.config from oterm.app.widgets.chat import AssistantImage monkeypatch.setattr(oterm.config.envConfig, "OTERM_DATA_DIR", tmp_path) buf = BytesIO() PILImage.new("RGB", (4, 4), "navy").save(buf, format="PNG") png_bytes = buf.getvalue() app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.add_image(png_bytes) await pilot.pause() assistant_image = item.query_one(AssistantImage) click = Click( widget=assistant_image, x=0, y=0, delta_x=0, delta_y=0, button=1, shift=False, meta=False, ctrl=False, ) await item.on_click(click) await pilot.pause() saved = list((tmp_path / "downloads").iterdir()) assert len(saved) == 1 assert saved[0].read_bytes() == png_bytes async def test_expand_before_result_shows_only_args(self, chat_model): from pydantic_ai.messages import ToolCallPart from textual.widgets import Static app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) item = ChatItem() item.author = "assistant" await container.query_one("#messageContainer").mount(item) await pilot.pause() await item.add_tool_call( ToolCallPart(tool_name="search", args='{"q": "x"}', tool_call_id="tc-1") ) tool_item = item.query_one(ToolCallItem) tool_item.collapsed = False await pilot.pause() body_text = _capture(tool_item.query_one(".tool-call-body", Static).content) assert '"q": "x"' in body_text assert "args:" in body_text assert "result:" not in body_text def _capture(renderable) -> str: """Render a Rich renderable to plain text without ANSI styling.""" console = Console(width=80, color_system=None) with console.capture() as capture: console.print(renderable) return capture.get() class TestUsageStatus: async def test_zero_tokens_render_only_duration(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) status = UsageStatus() await container.query_one("#messageContainer").mount(status) await pilot.pause() rendered = str(status.render()) assert "↑" not in rendered assert "↓" not in rendered assert rendered.endswith("s") async def test_tick_advances_spinner_frame(self, chat_model): """Direct call so we don't depend on the 0.1s timer firing in time.""" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) status = UsageStatus() await container.query_one("#messageContainer").mount(status) await pilot.pause() before = status._frame status._tick() assert status._frame == (before + 1) % len(UsageStatus.SPINNER_FRAMES) async def test_update_usage_renders_token_arrows(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) status = UsageStatus() await container.query_one("#messageContainer").mount(status) await pilot.pause() status.update_usage(input_tokens=42, output_tokens=7) await pilot.pause() rendered = str(status.render()) assert "↑ 42" in rendered assert "↓ 7" in rendered async def test_update_context_renders_used_window_and_percentage(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) status = UsageStatus() await container.query_one("#messageContainer").mount(status) await pilot.pause() status.update_context(used=1234, window=4096) await pilot.pause() assert "ctx 1.2k / 4.1k (30%)" in str(status.render()) async def test_update_context_renders_million_token_windows(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) status = UsageStatus() await container.query_one("#messageContainer").mount(status) await pilot.pause() status.update_context(used=250_000, window=1_050_000) await pilot.pause() assert "ctx 250.0k / 1.1M (24%)" in str(status.render()) @pytest.mark.parametrize(("used", "variable"), [(600, "warning"), (900, "error")]) async def test_context_colour_follows_how_full_the_window_is( self, chat_model, used, variable ): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) status = UsageStatus() await container.query_one("#messageContainer").mount(status) await pilot.pause() status.update_context(used=used, window=1000) await pilot.pause() segment = next(s for s in status.render_line(0) if "ctx" in s.text) assert segment.style is not None and segment.style.color is not None expected = app.theme_variables[variable].lower() assert segment.style.color.name.lower() == expected async def test_context_below_half_keeps_the_line_colour(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) status = UsageStatus() await container.query_one("#messageContainer").mount(status) await pilot.pause() status.update_context(used=200, window=1000) await pilot.pause() segments = [s for s in status.render_line(0) if s.text.strip()] ctx = next(s for s in segments if "ctx" in s.text) elapsed = next(s for s in segments if s.text.strip().endswith("s")) assert ctx.style is not None and elapsed.style is not None assert ctx.style.color == elapsed.style.color async def test_update_context_without_window_renders_tokens_only(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) status = UsageStatus() await container.query_one("#messageContainer").mount(status) await pilot.pause() status.update_context(used=850, window=None) await pilot.pause() rendered = str(status.render()) assert "ctx 850" in rendered assert "%" not in rendered async def test_finish_drops_spinner_glyph(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) status = UsageStatus() await container.query_one("#messageContainer").mount(status) await pilot.pause() status.update_usage(input_tokens=10, output_tokens=5) status.finish() await pilot.pause() rendered = str(status.render()) assert not any(frame in rendered for frame in UsageStatus.SPINNER_FRAMES) assert "↑ 10" in rendered assert "↓ 5" in rendered async def test_finish_is_idempotent(self, chat_model): app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) status = UsageStatus() await container.query_one("#messageContainer").mount(status) await pilot.pause() status.finish() # Second finish is a no-op and must not raise with the timer gone. status.finish() async def test_status_persists_after_successful_turn(self, store, chat_model): chat_id = await store.save_chat(chat_model) chat_model.id = chat_id async def stream_fn( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str]: yield "first " yield "second" app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = Agent(FunctionModel(stream_function=stream_fn)) prompt = app.query_one(FlexibleInput) prompt.text = "hello" await pilot.press("enter") await wait_until(pilot, lambda: len(container.messages) == 2) statuses = list(container.query(UsageStatus)) assert len(statuses) == 1 # After success, the spinner glyph is gone but the line remains. rendered = str(statuses[0].render()) assert not any(frame in rendered for frame in UsageStatus.SPINNER_FRAMES) async def _stream_hello( messages: list[ModelMessage], info: AgentInfo ) -> AsyncIterator[str]: yield "hello there" async def _run_turn(app: "_Host", agent: Agent) -> "tuple[ChatContainer, str]": """Send one prompt through the real UI and return the finished footer text.""" async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = agent before = len(container.messages) app.query_one(FlexibleInput).text = "hi" await pilot.press("enter") await wait_until(pilot, lambda: len(container.messages) == before + 2) status = list(container.query(UsageStatus))[-1] await wait_until(pilot, lambda: "ctx" in str(status.render())) return container, str(status.render()) def _build_with(monkeypatch, agent: Agent) -> list[dict]: """Have every agent build return ``agent``; returns the arguments of each build.""" import oterm.app.widgets.chat as chat_module builds: list[dict] = [] def get_agent(**kwargs): builds.append(kwargs) return agent monkeypatch.setattr(chat_module, "get_agent", get_agent) return builds def _no_usage(monkeypatch) -> None: """Make FunctionModel report no token usage, as a server that omits it would.""" import pydantic_ai.models.function as function_model from pydantic_ai.usage import RequestUsage monkeypatch.setattr( function_model, "_estimate_usage", lambda *args, **kwargs: RequestUsage() ) monkeypatch.setattr(function_model, "_estimate_string_tokens", lambda text: 0) _LOADED_AT_8K = { "models": [{"name": "test-model", "model": "test-model", "context_length": 8192}] } def _last_response_tokens(container: ChatContainer) -> int: response = container.pydantic_history[-1] assert isinstance(response, ModelResponse) return response.usage.total_tokens class TestContextWindowFooter: async def test_model_profile_window(self, store, chat_model): chat_model.provider = "anthropic" chat_model.id = await store.save_chat(chat_model) agent = Agent( FunctionModel( stream_function=_stream_hello, profile=ModelProfile(context_window=1000), ) ) container, footer = await _run_turn(_Host(chat_model, []), agent) used = _last_response_tokens(container) assert f"ctx {used} / 1.0k ({round(used / 10)}%)" in footer async def test_turn_is_saved_without_waiting_for_the_window_lookup( self, store, chat_model, monkeypatch ): import oterm.config chat_model.id = await store.save_chat(chat_model) ps = { "/api/ps": { "models": [ { "name": "test-model", "model": "test-model", "context_length": 8192, } ] } } with json_server(ps, delay=0.5) as url: monkeypatch.setattr(oterm.config.envConfig, "OLLAMA_URL", url) app = _Host(chat_model, []) agent = Agent(FunctionModel(stream_function=_stream_hello)) _build_with(monkeypatch, agent) async with app.run_test() as pilot: container = app.query_one(ChatContainer) container.agent = agent app.query_one(FlexibleInput).text = "hi" await pilot.press("enter") await wait_until(pilot, lambda: len(container.messages) == 2) status = list(container.query(UsageStatus))[-1] assert not status._streaming assert "ctx" not in str(status.render()) await asyncio.sleep(0.8) await wait_until(pilot, lambda: "/ 8.2k" in str(status.render())) assert "/ 8.2k" in str(status.render()) async def test_ollama_running_context_length(self, store, chat_model, monkeypatch): import oterm.config chat_model.id = await store.save_chat(chat_model) ps = { "/api/ps": { "models": [ { "name": "test-model", "model": "test-model", "context_length": 8192, } ] } } with json_server(ps) as url: monkeypatch.setattr(oterm.config.envConfig, "OLLAMA_URL", url) agent = Agent(FunctionModel(stream_function=_stream_hello)) _build_with(monkeypatch, agent) _, footer = await _run_turn(_Host(chat_model, []), agent) assert "/ 8.2k" in footer async def test_openai_compat_server_window( self, store, chat_model, app_config, monkeypatch ): models = { "/v1/models": { "object": "list", "data": [ { "id": "test-model", "object": "model", "created": 0, "owned_by": "vllm", "max_model_len": 262144, } ], } } with json_server(models) as url: app_config.set("openaiCompatible", {"local": {"base_url": f"{url}/v1"}}) chat_model.provider = "openai-compat/local" chat_model.id = await store.save_chat(chat_model) agent = Agent(FunctionModel(stream_function=_stream_hello)) _build_with(monkeypatch, agent) _, footer = await _run_turn(_Host(chat_model, []), agent) assert "/ 262.1k" in footer async def test_unknown_window_shows_tokens_only(self, store, chat_model): chat_model.provider = "anthropic" chat_model.id = await store.save_chat(chat_model) container, footer = await _run_turn( _Host(chat_model, []), Agent(FunctionModel(stream_function=_stream_hello)) ) assert f"ctx {_last_response_tokens(container)}" in footer assert "%" not in footer def _long_chat(chat_id: int, turns: int) -> list[MessageModel]: messages: list[MessageModel] = [] for i in range(turns): messages.append( MessageModel(chat_id=chat_id, role="user", text=f"question {i} " * 5) ) messages.append( MessageModel(chat_id=chat_id, role="assistant", text=f"answer {i} " * 5) ) return messages def _summarizing_agent(context_window: int) -> Agent: """The chat's own summarize capability, on a scripted model with a known window.""" from oterm.app.widgets.chat import _resolve_tools def reply(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: return ModelResponse(parts=[TextPart(content="the summary")]) _, _, capabilities = _resolve_tools(["summarize"]) return Agent( FunctionModel( reply, stream_function=_stream_hello, profile=ModelProfile(context_window=context_window), ), capabilities=capabilities, ) def _summaries(container: ChatContainer) -> list[str]: return [ part.content for message in container.pydantic_history for part in getattr(message, "parts", []) if isinstance(part, SystemPromptPart) ] class TestSummarize: async def test_long_chat_is_summarized_and_the_user_told(self, store, chat_model): chat_model.provider = "anthropic" chat_model.id = await store.save_chat(chat_model) app = _Host(chat_model, _long_chat(chat_model.id, turns=15)) container, _ = await _run_turn(app, _summarizing_agent(context_window=200)) assert any("the summary" in s for s in _summaries(container)) assert len(container.pydantic_history) < 30 assert any("summarized" in n.message for n in _notifications(app)) async def test_short_chat_is_left_alone(self, store, chat_model): chat_model.provider = "anthropic" chat_model.id = await store.save_chat(chat_model) app = _Host(chat_model, _long_chat(chat_model.id, turns=1)) container, _ = await _run_turn(app, _summarizing_agent(context_window=100_000)) assert _summaries(container) == [] assert not any("summarized" in n.message for n in _notifications(app)) class TestServerWindow: async def test_local_window_is_read_when_the_chat_opens( self, store, chat_model, app_config, monkeypatch ): models = { "/v1/models": { "object": "list", "data": [ { "id": "test-model", "object": "model", "created": 0, "owned_by": "vllm", "max_model_len": 262144, } ], } } with json_server(models) as url: app_config.set("openaiCompatible", {"local": {"base_url": f"{url}/v1"}}) chat_model.provider = "openai-compat/local" chat_model.id = await store.save_chat(chat_model) builds = _build_with( monkeypatch, Agent(FunctionModel(stream_function=_stream_hello)) ) app = _Host(chat_model, []) async with app.run_test() as pilot: await wait_until(pilot, lambda: builds[-1]["context_window"] == 262144) assert [b["context_window"] for b in builds] == [None, 262144] @pytest.mark.parametrize("reports_usage", [True, False]) async def test_window_reported_after_the_first_reply_reaches_the_agent( self, store, chat_model, monkeypatch, reports_usage ): import oterm.config if not reports_usage: _no_usage(monkeypatch) chat_model.id = await store.save_chat(chat_model) ps: dict = {"/api/ps": {"models": []}} with json_server(ps) as url: monkeypatch.setattr(oterm.config.envConfig, "OLLAMA_URL", url) agent = Agent(FunctionModel(stream_function=_stream_hello)) builds = _build_with(monkeypatch, agent) app = _Host(chat_model, []) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await app.workers.wait_for_complete() assert builds[-1]["context_window"] is None ps["/api/ps"] = _LOADED_AT_8K app.query_one(FlexibleInput).text = "hi" await pilot.press("enter") await wait_until(pilot, lambda: builds[-1]["context_window"] == 8192) await app.workers.wait_for_complete() footer = str(list(container.query(UsageStatus))[-1].render()) assert builds[-1]["context_window"] == 8192 assert ("ctx" in footer) is reports_usage async def test_editing_the_chat_forgets_the_server_window( self, store, chat_model, monkeypatch ): import oterm.config chat_id = await store.save_chat(chat_model) chat_model.id = chat_id edited = ChatModel(id=chat_id, model="other-model", provider="ollama") app = _Host(chat_model, []) async def fake_push_screen_wait(self, screen): return edited monkeypatch.setattr(type(app), "push_screen_wait", fake_push_screen_wait) with json_server({"/api/ps": _LOADED_AT_8K}) as url: monkeypatch.setattr(oterm.config.envConfig, "OLLAMA_URL", url) builds = _build_with( monkeypatch, Agent(FunctionModel(stream_function=_stream_hello)) ) async with app.run_test() as pilot: container = app.query_one(ChatContainer) await wait_until(pilot, lambda: builds[-1]["context_window"] == 8192) assert builds[-1]["context_window"] == 8192 container.action_edit_chat() await wait_until(pilot, lambda: builds[-1]["model"] == "other-model") await app.workers.wait_for_complete() assert builds[-1]["model"] == "other-model" assert builds[-1]["context_window"] is None