Declare caller-owned types for terminal, turn and rendering inputs

src/toad/acp/session_updates.py

@@ -1,118 +1,138 @@
 """Process ACP notifications in order; update declarations own their effects."""
+
 from __future__ import annotations
 
 import asyncio
 from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+from acp.schema import (
+    AgentMessageChunk,
+    AgentPlanUpdate,
+    AgentThoughtChunk,
+    AvailableCommandsUpdate,
+    ConfigOptionUpdate,
+    CurrentModeUpdate,
+    SessionInfoUpdate,
+    SessionNotification,
+    TextContentBlock,
+    ToolCallProgress,
+    ToolCallStart,
+    UsageUpdate,
+    UserMessageChunk,
+)
+from agent_comms.acp_extension import AgentCommsUpdate
+from agent_comms.mro_dispatch import MroDispatch, handles
 
-from agent_comms.declared_family import DeclaredFamily
-
-from toad.core import events
 from toad.acp.client_session import ClientRequestOwner, ClientSessionRequest
 from toad.acp.context_measurement import ContextMeasurement
 from toad.acp.sdk_boundary import ValidateSessionUpdateTask
-from agent_comms.mro_dispatch import MroDispatch, handles
-from acp.schema import (UserMessageChunk, AgentMessageChunk, AgentThoughtChunk,
-    ToolCallStart, ToolCallProgress, AgentPlanUpdate, AvailableCommandsUpdate,
-    CurrentModeUpdate, ConfigOptionUpdate, SessionInfoUpdate, UsageUpdate, TextContentBlock)
+from toad.core import events
 from toad.plan import decode_plan
 from toad.project_link import AcpLogLink
 
+if TYPE_CHECKING:
+    from toad.acp.agent import Agent
+    from toad.acp.comms_updates import CommsUpdateConsumer
+
 
 class SessionUpdateEffect(MroDispatch):
     """Effect registrations consume official SDK declarations, never raw shapes."""
-    def __init__(self, agent, consumer):
+
+    def __init__(self, agent: Agent, consumer: CommsUpdateConsumer):
         self.agent, self.consumer = agent, consumer
 
     @property
     def route(self):
         return self.consumer.route
 
-    def require_supported(self, update):
-        if not any(self.handlers_for(update)):
-            raise ValueError(f'ACP update capability is not supported: {type(update).__name__}')
+    def require_supported(self, notification: SessionNotification):
+        if not any(self.handlers_for(notification.update)):
+            raise ValueError(f"ACP update capability is not supported: {type(notification.update).__name__}")
 
     @handles(UserMessageChunk)
-    def user(self, update):
+    def user(self, update: UserMessageChunk):
         UserContentEffect(self.agent, self.route).dispatch_sync(update.content)
 
     @handles(AgentMessageChunk)
-    def message(self, update):
+    def message(self, update: AgentMessageChunk):
         AgentContentEffect(self.agent, self.route).dispatch_sync(update.content)
 
     @handles(AgentThoughtChunk)
-    def thought(self, update):
+    def thought(self, update: AgentThoughtChunk):
         ThoughtContentEffect(self.agent, self.route).dispatch_sync(update.content)
 
     @handles(ToolCallStart)
-    def tool_start(self, update):
+    def tool_start(self, update: ToolCallStart):
         self.agent.tools.begin(update)
 
     @handles(ToolCallProgress)
-    def tool_progress(self, update):
+    def tool_progress(self, update: ToolCallProgress):
         self.agent.tools.update(update)
 
     @handles(AgentPlanUpdate)
-    def plan(self, update):
+    def plan(self, update: AgentPlanUpdate):
         self.agent.controller.publish_plan(decode_plan(update.entries))
 
     @handles(AvailableCommandsUpdate)
-    def commands(self, update):
+    def commands(self, update: AvailableCommandsUpdate):
         self.agent.controller.publish_commands(update.available_commands)
 
     @handles(CurrentModeUpdate)
-    def mode(self, update):
+    def mode(self, update: CurrentModeUpdate):
         self.agent.configuration.receive_current_mode(update.current_mode_id)
 
     @handles(ConfigOptionUpdate)
-    def config(self, update):
+    def config(self, update: ConfigOptionUpdate):
         self.agent.configuration.receive(update.config_options)
 
     @handles(SessionInfoUpdate)
-    def info(self, update):
-        if 'title' in update.model_fields_set:
+    def info(self, update: SessionInfoUpdate):
+        if "title" in update.model_fields_set:
             self.agent.events.publish(events.SessionInfoUpdate(update.title))
 
     @handles(UsageUpdate)
-    def usage(self, update):
+    def usage(self, update: UsageUpdate):
         self.agent.context_measurement = ContextMeasurement.live(update.used, update.size, update.cost)
         self.agent.update_status_line()
 
 
 class MessageContentEffect(MroDispatch):
-    def __init__(self, agent, route):
+    def __init__(self, agent: Agent, route):
         self.agent, self.route = agent, route
 
     @handles(TextContentBlock)
-    def text(self, content):
+    def text(self, content: TextContentBlock):
         if content.text:
             self.publish(content)
 
     @abstractmethod
-    def publish(self, content): ...
+    def publish(self, content: TextContentBlock): ...
 
 
 class UserContentEffect(MessageContentEffect):
-    def publish(self, content):
+    def publish(self, content: TextContentBlock):
         self.agent.events.publish(events.UserMessage(content.type, content.text))
 
 
 class ThoughtContentEffect(MessageContentEffect):
-    def publish(self, content):
+    def publish(self, content: TextContentBlock):
         self.agent.events.publish(events.Thinking(content.type, content.text))
 
 
 class AgentContentEffect(MessageContentEffect):
-    def publish(self, content):
+    def publish(self, content: TextContentBlock):
         text = content.text
-        if text.startswith('[agent error]'):
-            text += f'\n\n{AcpLogLink.markdown(self.agent.presentation.log_path)}'
+        if text.startswith("[agent error]"):
+            text += f"\n\n{AcpLogLink.markdown(self.agent.presentation.log_path)}"
         from toad.response_delivery import ResponseDelivery
+
         stream = self.agent.presentation.turns.owner.response_stream(ResponseDelivery.from_route(self.route))
         self.agent.events.publish(events.Update(content.type, text, stream))
 
 
 class SessionNotificationOwner(ClientRequestOwner):
-    def __init__(self, agent):
+    def __init__(self, agent: Agent):
         super().__init__(agent)
         self.lock = asyncio.Lock()
 
@@ -122,23 +142,25 @@ async def receive(self, notification, raw_update):
             if authority.retired or not authority.binding.accepts_notification(notification.session_id):
                 return
             validation = await self.agent.controller.task_runner.run(
-                ValidateSessionUpdateTask(notification, raw_update))
+                ValidateSessionUpdateTask(notification, raw_update)
+            )
             if authority.retired:
                 return
             validation.publish(self, notification.session_id, raw_update, notification.field_meta)
 
-    def reject(self, session_id, update, metadata, error):
+    def reject(self, session_id: str, update, metadata, error: str):
         self.agent.log(
             f"[ACP rejected session/update] raw={{'sessionId': {session_id!r}, "
-            f"'update': {update!r}, '_meta': {metadata!r}}}; validation={error}")
+            f"'update': {update!r}, '_meta': {metadata!r}}}; validation={error}"
+        )
         self.agent.events.publish(events.RejectedSessionUpdate())
 
-    def publish(self, session_id, update, facts):
+    def publish(self, session_id: str, update: SessionNotification, facts: tuple[AgentCommsUpdate, ...]):
         metadata = update.update.field_meta
         consumer = self.agent.comms_consumer_class(self.agent, session_id)
         effect = SessionUpdateEffect(self.agent, consumer)
         try:
-            effect.require_supported(update.update)
+            effect.require_supported(update)
         except (TypeError, ValueError) as error:
             self.reject(session_id, update, metadata, str(error))
             return

src/toad/agent.py

@@ -1,5 +1,6 @@
 from abc import abstractmethod
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 from agent_comms.thread_presentation import ThreadPresentation
 from textual.content import Content
@@ -11,6 +12,15 @@
 from toad.core.events import CoreEventStream
 from toad.surface_binding import SurfaceBinding
 
+if TYPE_CHECKING:
+    from agent_comms.session_load import ExistingSessionLoadPlan
+
+    from toad.conversation_turn import TurnOwner
+    from toad.screens.main import MainScreen
+    from toad.session_observation import SessionObservation
+    from toad.session_presentation import OperationalSessionSources
+    from toad.widgets.conversation import Conversation, ConversationSessionBinding
+
 
 class AgentBase(AgentConnection):
     """Base class for an 'agent'."""
@@ -35,7 +45,7 @@ def attach_view(self, view) -> None:
     def matches_agent(self, agent) -> bool:
         return self is agent
 
-    def cancel_allowed(self, turn):
+    def cancel_allowed(self, turn: TurnOwner):
         return turn.can_cancel
 
     def matches_session(self, session_id):
@@ -44,14 +54,16 @@ def matches_session(self, session_id):
     def borrowed_processes(self):
         return (self.process,)
 
-    async def reconnect_session(self, source, plan) -> None:
+    async def reconnect_session(
+        self, source: MainScreen, plan: ExistingSessionLoadPlan
+    ) -> None:
         await self.session.reconnect(plan)
 
-    async def change_project(self, view, path) -> None:
+    async def change_project(self, view: Conversation, path: Path) -> None:
         project = await self.update_project(str(path))
         view.feedback.flash(Content(f"Project changed to {project}"), style="success")
 
-    def request_compaction(self, view, instructions) -> None:
+    def request_compaction(self, view: Conversation, instructions) -> None:
         view.compact_context(instructions)
 
     def begin_submission(self, owner, submission):
@@ -69,49 +81,51 @@ def accepts_failed_input(self, message):
     def send_now(self, owner):
         return owner.send_bound_now(self)
 
-    def invalidate_observation(self, observation) -> None:
+    def invalidate_observation(self, observation: SessionObservation) -> None:
         observation.schedule(self)
 
     @property
     def automatic_titles(self) -> bool:
         return self.observed_thread() is None
 
-    async def project_changed(self, path) -> None:
+    async def project_changed(self, path: Path) -> None:
         from toad.db import DB
 
         self.project_root_path = path
         if (session_pk := self.session.pk) is not None:
             await DB().session_update_project(session_pk, path)
 
-    def refresh_context(self, view) -> None:
+    def refresh_context(self, view: Conversation) -> None:
         view.status = self.context_measurement.status()
 
     async def compact_context(self, instructions) -> None:
         await self.controller.compact_context(instructions)
 
-    async def cancel_turn(self, view) -> None:
+    async def cancel_turn(self, view: Conversation) -> None:
         await view.cancel_bound_turn(self)
 
     def accepts_message(self, message) -> bool:
         return message.publisher is self and message.event.session_id == self.session_id
 
-    def refresh_projection(self, view) -> None:
+    def refresh_projection(self, view: ConversationSessionBinding) -> None:
         view.status = self.context_measurement.status()
         view.turns.bound()
         view.submissions.publish_pending()
         view.delivery_observation.invalidate()
 
-    def start_session(self, view) -> None:
+    def start_session(self, view: ConversationSessionBinding) -> None:
         if self.ready:
             view.session_ready()
 
-    def start_agent(self, sources, view) -> None:
+    def start_agent(
+        self, sources: OperationalSessionSources, view: Conversation
+    ) -> None:
         """A connected agent is not created a second time."""
 
-    def update_view(self, view) -> None:
+    def update_view(self, view: ConversationSessionBinding) -> None:
         view.apply_bound_agent(self)
 
-    def connected(self, view) -> None:
+    def connected(self, view: ConversationSessionBinding) -> None:
         view.agent_connected(self)
 
     @property
@@ -125,7 +139,7 @@ def current_turn(self):
     def attach_surface(self, binding: SurfaceBinding) -> None:
         binding.close()
 
-    async def retire_surface(self, surface) -> None:
+    async def retire_surface(self, surface: Conversation) -> None:
         await self.stop()
 
     def observed_thread(self) -> tuple[str, str] | None: