kimi-cli 0.40__py3-none-any.whl → 0.41__py3-none-any.whl

kimi_cli/soul/globals.py

@@ -0,0 +1,92 @@
+import subprocess
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import NamedTuple
+
+from kimi_cli.config import Config
+from kimi_cli.llm import LLM
+from kimi_cli.metadata import Session
+from kimi_cli.soul.approval import Approval
+from kimi_cli.soul.denwarenji import DenwaRenji
+from kimi_cli.utils.logging import logger
+
+
+class BuiltinSystemPromptArgs(NamedTuple):
+    """Builtin system prompt arguments."""
+
+    KIMI_NOW: str
+    """The current datetime."""
+    KIMI_WORK_DIR: Path
+    """The current working directory."""
+    KIMI_WORK_DIR_LS: str
+    """The directory listing of current working directory."""
+    KIMI_AGENTS_MD: str  # TODO: move to first message from system prompt
+    """The content of AGENTS.md."""
+
+
+def load_agents_md(work_dir: Path) -> str | None:
+    paths = [
+        work_dir / "AGENTS.md",
+        work_dir / "agents.md",
+    ]
+    for path in paths:
+        if path.is_file():
+            logger.info("Loaded agents.md: {path}", path=path)
+            return path.read_text(encoding="utf-8").strip()
+    logger.info("No AGENTS.md found in {work_dir}", work_dir=work_dir)
+    return None
+
+
+def _list_work_dir(work_dir: Path) -> str:
+    if sys.platform == "win32":
+        ls = subprocess.run(
+            ["cmd", "/c", "dir", work_dir],
+            capture_output=True,
+            text=True,
+            encoding="utf-8",
+            errors="replace",
+        )
+    else:
+        ls = subprocess.run(
+            ["ls", "-la", work_dir],
+            capture_output=True,
+            text=True,
+            encoding="utf-8",
+            errors="replace",
+        )
+    return ls.stdout.strip()
+
+
+class AgentGlobals(NamedTuple):
+    """Agent globals."""
+
+    config: Config
+    llm: LLM | None
+    session: Session
+    builtin_args: BuiltinSystemPromptArgs
+    denwa_renji: DenwaRenji
+    approval: Approval
+
+    @classmethod
+    async def create(
+        cls, config: Config, llm: LLM | None, session: Session, yolo: bool
+    ) -> "AgentGlobals":
+        work_dir = Path(session.work_dir.path)
+        # FIXME: do these asynchronously
+        ls_output = _list_work_dir(work_dir)
+        agents_md = load_agents_md(work_dir) or ""
+
+        return cls(
+            config=config,
+            llm=llm,
+            session=session,
+            builtin_args=BuiltinSystemPromptArgs(
+                KIMI_NOW=datetime.now().astimezone().isoformat(),
+                KIMI_WORK_DIR=work_dir,
+                KIMI_WORK_DIR_LS=ls_output,
+                KIMI_AGENTS_MD=agents_md,
+            ),
+            denwa_renji=DenwaRenji(),
+            approval=Approval(yolo=yolo),
+        )

kimi_cli/soul/kimisoul.py

@@ -16,29 +16,28 @@ from kosong.chat_provider import (
 from kosong.tooling import ToolResult
 from tenacity import RetryCallState, retry_if_exception, stop_after_attempt, wait_exponential_jitter
 
-from kimi_cli.agent import Agent, AgentGlobals
 from kimi_cli.config import LoopControl
-from kimi_cli.soul import LLMNotSet, MaxStepsReached, Soul, StatusSnapshot
+from kimi_cli.soul import LLMNotSet, MaxStepsReached, Soul, StatusSnapshot, wire_send
+from kimi_cli.soul.agent import Agent
 from kimi_cli.soul.compaction import SimpleCompaction
 from kimi_cli.soul.context import Context
+from kimi_cli.soul.globals import AgentGlobals
 from kimi_cli.soul.message import system, tool_result_to_messages
-from kimi_cli.soul.wire import (
+from kimi_cli.tools.dmail import NAME as SendDMail_NAME
+from kimi_cli.tools.utils import ToolRejectedError
+from kimi_cli.utils.logging import logger
+from kimi_cli.wire.message import (
     CompactionBegin,
     CompactionEnd,
     StatusUpdate,
     StepBegin,
     StepInterrupted,
-    Wire,
-    current_wire,
 )
-from kimi_cli.tools.dmail import NAME as SendDMail_NAME
-from kimi_cli.tools.utils import ToolRejectedError
-from kimi_cli.utils.logging import logger
 
 RESERVED_TOKENS = 50_000
 
 
-class KimiSoul:
+class KimiSoul(Soul):
     """The soul of Kimi CLI."""
 
     def __init__(
@@ -97,31 +96,27 @@ class KimiSoul:
     async def _checkpoint(self):
         await self._context.checkpoint(self._checkpoint_with_user_message)
 
-    async def run(self, user_input: str, wire: Wire):
+    async def run(self, user_input: str):
         if self._agent_globals.llm is None:
             raise LLMNotSet()
 
         await self._checkpoint()  # this creates the checkpoint 0 on first run
         await self._context.append_message(Message(role="user", content=user_input))
         logger.debug("Appended user message to context")
-        wire_token = current_wire.set(wire)
-        try:
-            await self._agent_loop(wire)
-        finally:
-            current_wire.reset(wire_token)
+        await self._agent_loop()
 
-    async def _agent_loop(self, wire: Wire):
+    async def _agent_loop(self):
         """The main agent loop for one run."""
         assert self._agent_globals.llm is not None
 
         async def _pipe_approval_to_wire():
             while True:
                 request = await self._approval.fetch_request()
-                wire.send(request)
+                wire_send(request)
 
         step_no = 1
         while True:
-            wire.send(StepBegin(step_no))
+            wire_send(StepBegin(step_no))
             approval_task = asyncio.create_task(_pipe_approval_to_wire())
             # FIXME: It's possible that a subagent's approval task steals approval request
             # from the main agent. We must ensure that the Task tool will redirect them
@@ -134,21 +129,21 @@ class KimiSoul:
                     >= self._agent_globals.llm.max_context_size
                 ):
                     logger.info("Context too long, compacting...")
-                    wire.send(CompactionBegin())
+                    wire_send(CompactionBegin())
                     await self.compact_context()
-                    wire.send(CompactionEnd())
+                    wire_send(CompactionEnd())
 
                 logger.debug("Beginning step {step_no}", step_no=step_no)
                 await self._checkpoint()
                 self._denwa_renji.set_n_checkpoints(self._context.n_checkpoints)
-                finished = await self._step(wire)
+                finished = await self._step()
             except BackToTheFuture as e:
                 await self._context.revert_to(e.checkpoint_id)
                 await self._checkpoint()
                 await self._context.append_message(e.messages)
                 continue
             except (ChatProviderError, asyncio.CancelledError):
-                wire.send(StepInterrupted())
+                wire_send(StepInterrupted())
                 # break the agent loop
                 raise
             finally:
@@ -161,7 +156,7 @@ class KimiSoul:
             if step_no > self._loop_control.max_steps_per_run:
                 raise MaxStepsReached(self._loop_control.max_steps_per_run)
 
-    async def _step(self, wire: Wire) -> bool:
+    async def _step(self) -> bool:
         """Run an single step and return whether the run should be stopped."""
         # already checked in `run`
         assert self._agent_globals.llm is not None
@@ -181,8 +176,8 @@ class KimiSoul:
                 self._agent.system_prompt,
                 self._agent.toolset,
                 self._context.history,
-                on_message_part=wire.send,
-                on_tool_result=wire.send,
+                on_message_part=wire_send,
+                on_tool_result=wire_send,
             )
 
         result = await _kosong_step_with_retry()
@@ -190,7 +185,7 @@ class KimiSoul:
         if result.usage is not None:
             # mark the token count for the context before the step
             await self._context.update_token_count(result.usage.input)
-            wire.send(StatusUpdate(status=self.status))
+            wire_send(StatusUpdate(status=self.status))
 
         # wait for all tool results (may be interrupted)
         results = await result.tool_results()

kimi_cli/tools/dmail/__init__.py

@@ -10,7 +10,7 @@ NAME = "SendDMail"
 
 class SendDMail(CallableTool2):
     name: str = NAME
-    description: str = (Path(__file__).parent / "dmail.md").read_text()
+    description: str = (Path(__file__).parent / "dmail.md").read_text(encoding="utf-8")
     params: type[DMail] = DMail
 
     def __init__(self, denwa_renji: DenwaRenji, **kwargs):

kimi_cli/tools/file/glob.md

@@ -14,4 +14,4 @@ Find files and directories using glob patterns. This tool supports standard glob
 
 **Bad example patterns:**
 - `**`, `**/*.py` - Any pattern starting with '**' will be rejected. Because it would recursively search all directories and subdirectories, which is very likely to yield large result that exceeds your context size. Always use more specific patterns like `src/**/*.py` instead.
-- `node_modules/**/*.js` - Although this does not start with '**', it would still highly possible to yield large result because `node_modules` is well-known to contain too many directories and files. Avoid recursivelly searching in such directories, other examples include `venv`, `.venv`, `__pycache__`, `target`. If you really need to search in a dependency, use more specific patterns like `node_modules/react/src/*` instead.
+- `node_modules/**/*.js` - Although this does not start with '**', it would still highly possible to yield large result because `node_modules` is well-known to contain too many directories and files. Avoid recursively searching in such directories, other examples include `venv`, `.venv`, `__pycache__`, `target`. If you really need to search in a dependency, use more specific patterns like `node_modules/react/src/*` instead.

kimi_cli/tools/file/glob.py

@@ -8,7 +8,7 @@ import aiofiles.os
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agent import BuiltinSystemPromptArgs
+from kimi_cli.soul.globals import BuiltinSystemPromptArgs
 from kimi_cli.tools.utils import load_desc
 
 MAX_MATCHES = 1000
@@ -128,7 +128,7 @@ class Glob(CallableTool2[Params]):
             message = (
                 f"Found {len(matches)} matches for pattern `{params.pattern}`."
                 if len(matches) > 0
-                else "No matches found for pattern `{params.pattern}`."
+                else f"No matches found for pattern `{params.pattern}`."
             )
             if len(matches) > MAX_MATCHES:
                 matches = matches[:MAX_MATCHES]

kimi_cli/tools/file/grep.py

@@ -220,7 +220,7 @@ async def _ensure_rg_path() -> str:
 
 class Grep(CallableTool2[Params]):
     name: str = "Grep"
-    description: str = (Path(__file__).parent / "grep.md").read_text()
+    description: str = (Path(__file__).parent / "grep.md").read_text(encoding="utf-8")
     params: type[Params] = Params
 
     @override

kimi_cli/tools/file/patch.py

@@ -6,8 +6,8 @@ import patch_ng
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agent import BuiltinSystemPromptArgs
 from kimi_cli.soul.approval import Approval
+from kimi_cli.soul.globals import BuiltinSystemPromptArgs
 from kimi_cli.tools.file import FileActions
 from kimi_cli.tools.utils import ToolRejectedError
 
@@ -19,7 +19,7 @@ class Params(BaseModel):
 
 class PatchFile(CallableTool2[Params]):
     name: str = "PatchFile"
-    description: str = (Path(__file__).parent / "patch.md").read_text()
+    description: str = (Path(__file__).parent / "patch.md").read_text(encoding="utf-8")
     params: type[Params] = Params
 
     def __init__(self, builtin_args: BuiltinSystemPromptArgs, approval: Approval, **kwargs):

kimi_cli/tools/file/read.py

@@ -5,7 +5,7 @@ import aiofiles
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agent import BuiltinSystemPromptArgs
+from kimi_cli.soul.globals import BuiltinSystemPromptArgs
 from kimi_cli.tools.utils import load_desc, truncate_line
 
 MAX_LINES = 1000

kimi_cli/tools/file/replace.py

@@ -5,8 +5,8 @@ import aiofiles
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agent import BuiltinSystemPromptArgs
 from kimi_cli.soul.approval import Approval
+from kimi_cli.soul.globals import BuiltinSystemPromptArgs
 from kimi_cli.tools.file import FileActions
 from kimi_cli.tools.utils import ToolRejectedError
 
@@ -29,7 +29,7 @@ class Params(BaseModel):
 
 class StrReplaceFile(CallableTool2[Params]):
     name: str = "StrReplaceFile"
-    description: str = (Path(__file__).parent / "replace.md").read_text()
+    description: str = (Path(__file__).parent / "replace.md").read_text(encoding="utf-8")
     params: type[Params] = Params
 
     def __init__(self, builtin_args: BuiltinSystemPromptArgs, approval: Approval, **kwargs):

kimi_cli/tools/file/write.py

@@ -5,8 +5,8 @@ import aiofiles
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agent import BuiltinSystemPromptArgs
 from kimi_cli.soul.approval import Approval
+from kimi_cli.soul.globals import BuiltinSystemPromptArgs
 from kimi_cli.tools.file import FileActions
 from kimi_cli.tools.utils import ToolRejectedError
 
@@ -26,7 +26,7 @@ class Params(BaseModel):
 
 class WriteFile(CallableTool2[Params]):
     name: str = "WriteFile"
-    description: str = (Path(__file__).parent / "write.md").read_text()
+    description: str = (Path(__file__).parent / "write.md").read_text(encoding="utf-8")
     params: type[Params] = Params
 
     def __init__(self, builtin_args: BuiltinSystemPromptArgs, approval: Approval, **kwargs):

kimi_cli/tools/task/__init__.py

@@ -1,17 +1,21 @@
+import asyncio
 from pathlib import Path
 from typing import override
 
 from kosong.tooling import CallableTool2, ToolError, ToolOk, ToolReturnType
 from pydantic import BaseModel, Field
 
-from kimi_cli.agent import Agent, AgentGlobals, AgentSpec, load_agent
-from kimi_cli.soul import MaxStepsReached
+from kimi_cli.agentspec import ResolvedAgentSpec
+from kimi_cli.soul import MaxStepsReached, get_wire_or_none, run_soul
+from kimi_cli.soul.agent import Agent, load_agent
 from kimi_cli.soul.context import Context
+from kimi_cli.soul.globals import AgentGlobals
 from kimi_cli.soul.kimisoul import KimiSoul
-from kimi_cli.soul.wire import ApprovalRequest, Wire, WireMessage, get_wire_or_none
 from kimi_cli.tools.utils import load_desc
 from kimi_cli.utils.message import message_extract_text
 from kimi_cli.utils.path import next_available_rotation
+from kimi_cli.wire import WireUISide
+from kimi_cli.wire.message import ApprovalRequest, WireMessage
 
 # Maximum continuation attempts for task summary
 MAX_CONTINUE_ATTEMPTS = 1
@@ -45,12 +49,11 @@ class Task(CallableTool2[Params]):
     name: str = "Task"
     params: type[Params] = Params
 
-    def __init__(self, agent_spec: AgentSpec, agent_globals: AgentGlobals, **kwargs):
+    def __init__(self, agent_spec: ResolvedAgentSpec, agent_globals: AgentGlobals, **kwargs):
         subagents: dict[str, Agent] = {}
         descs = []
 
         # load all subagents
-        assert agent_spec.subagents is not None, "Task tool expects subagents"
         for name, spec in agent_spec.subagents.items():
             subagents[name] = load_agent(spec.path, agent_globals)
             descs.append(f"- `{name}`: {spec.description}")
@@ -99,6 +102,19 @@ class Task(CallableTool2[Params]):
 
     async def _run_subagent(self, agent: Agent, prompt: str) -> ToolReturnType:
         """Run subagent with optional continuation for task summary."""
+        super_wire = get_wire_or_none()
+        assert super_wire is not None
+
+        def _super_wire_send(msg: WireMessage) -> None:
+            if isinstance(msg, ApprovalRequest):
+                super_wire.soul_side.send(msg)
+            # TODO: visualize subagent behavior by sending other messages in some way
+
+        async def _ui_loop_fn(wire: WireUISide) -> None:
+            while True:
+                msg = await wire.receive()
+                _super_wire_send(msg)
+
         subagent_history_file = await self._get_subagent_history_file()
         context = Context(file_backend=subagent_history_file)
         soul = KimiSoul(
@@ -107,12 +123,9 @@ class Task(CallableTool2[Params]):
             context=context,
             loop_control=self._agent_globals.config.loop_control,
         )
-        wire = get_wire_or_none()
-        assert wire is not None, "Wire is expected to be set"
-        sub_wire = _SubWire(wire)
 
         try:
-            await soul.run(prompt, sub_wire)
+            await run_soul(soul, prompt, _ui_loop_fn, asyncio.Event())
         except MaxStepsReached as e:
             return ToolError(
                 message=(
@@ -135,22 +148,10 @@ class Task(CallableTool2[Params]):
         # Check if response is too brief, if so, run again with continuation prompt
         n_attempts_remaining = MAX_CONTINUE_ATTEMPTS
         if len(final_response) < 200 and n_attempts_remaining > 0:
-            await soul.run(CONTINUE_PROMPT, sub_wire)
+            await run_soul(soul, CONTINUE_PROMPT, _ui_loop_fn, asyncio.Event())
 
             if len(context.history) == 0 or context.history[-1].role != "assistant":
                 return ToolError(message=_error_msg, brief="Failed to run subagent")
             final_response = message_extract_text(context.history[-1])
 
         return ToolOk(output=final_response)
-
-
-class _SubWire(Wire):
-    def __init__(self, super_wire: Wire):
-        super().__init__()
-        self._super_wire = super_wire
-
-    @override
-    def send(self, msg: WireMessage):
-        if isinstance(msg, ApprovalRequest):
-            self._super_wire.send(msg)
-        # TODO: visualize subagent behavior by sending other messages in some way

kimi_cli/tools/task/task.md

@@ -4,7 +4,7 @@ Spawn a subagent to perform a specific task. Subagent will be spawned with a fre
 
 Context isolation is one of the key benefits of using subagents. By delegating tasks to subagents, you can keep your main context clean and focused on the main goal requested by the user.
 
-Here are some scenerios you may want this tool for context isolation:
+Here are some scenarios you may want this tool for context isolation:
 
 - You wrote some code and it did not work as expected. In this case you can spawn a subagent to fix the code, asking the subagent to return how it is fixed. This can potentially benefit because the detailed process of fixing the code may not be relevant to your main goal, and may clutter your context.
 - When you need some latest knowledge of a specific library, framework or technology to proceed with your task, you can spawn a subagent to search on the internet for the needed information and return to you the gathered relevant information, for example code examples, API references, etc. This can avoid ton of irrelevant search results in your own context.

kimi_cli/tools/todo/__init__.py

@@ -16,7 +16,7 @@ class Params(BaseModel):
 
 class SetTodoList(CallableTool2[Params]):
     name: str = "SetTodoList"
-    description: str = (Path(__file__).parent / "set_todo_list.md").read_text()
+    description: str = (Path(__file__).parent / "set_todo_list.md").read_text(encoding="utf-8")
     params: type[Params] = Params
 
     @override

kimi_cli/tools/utils.py

@@ -7,7 +7,7 @@ from kosong.tooling import ToolError, ToolOk
 
 def load_desc(path: Path, substitutions: dict[str, str] | None = None) -> str:
     """Load a tool description from a file, with optional substitutions."""
-    description = path.read_text()
+    description = path.read_text(encoding="utf-8")
     if substitutions:
         description = string.Template(description).substitute(substitutions)
     return description

kimi_cli/tools/web/search.py

@@ -4,8 +4,8 @@ from typing import override
 from kosong.tooling import CallableTool2, ToolReturnType
 from pydantic import BaseModel, Field, ValidationError
 
-import kimi_cli
 from kimi_cli.config import Config
+from kimi_cli.constant import USER_AGENT
 from kimi_cli.soul.toolset import get_current_tool_call_or_none
 from kimi_cli.tools.utils import ToolResultBuilder, load_desc
 from kimi_cli.utils.aiohttp import new_client_session
@@ -66,7 +66,7 @@ class SearchWeb(CallableTool2[Params]):
             session.post(
                 self._base_url,
                 headers={
-                    "User-Agent": kimi_cli.USER_AGENT,
+                    "User-Agent": USER_AGENT,
                     "Authorization": f"Bearer {self._api_key}",
                     "X-Msh-Tool-Call-Id": tool_call.id,
                 },

kimi_cli/ui/__init__.py

@@ -1,69 +0,0 @@
-import asyncio
-import contextlib
-from collections.abc import Callable, Coroutine
-from typing import Any
-
-from kimi_cli.soul import Soul
-from kimi_cli.soul.wire import Wire
-from kimi_cli.utils.logging import logger
-
-type UILoopFn = Callable[[Wire], Coroutine[Any, Any, None]]
-"""A long-running async function to visualize the agent behavior."""
-
-
-class RunCancelled(Exception):
-    """The run was cancelled by the cancel event."""
-
-
-async def run_soul(
-    soul: Soul,
-    user_input: str,
-    ui_loop_fn: UILoopFn,
-    cancel_event: asyncio.Event,
-):
-    """
-    Run the soul with the given user input.
-
-    `cancel_event` is a outside handle that can be used to cancel the run. When the event is set,
-    the run will be gracefully stopped and a `RunCancelled` will be raised.
-
-    Raises:
-        LLMNotSet: When the LLM is not set.
-        ChatProviderError: When the LLM provider returns an error.
-        MaxStepsReached: When the maximum number of steps is reached.
-        RunCancelled: When the run is cancelled by the cancel event.
-    """
-    wire = Wire()
-    logger.debug("Starting UI loop with function: {ui_loop_fn}", ui_loop_fn=ui_loop_fn)
-
-    ui_task = asyncio.create_task(ui_loop_fn(wire))
-    soul_task = asyncio.create_task(soul.run(user_input, wire))
-
-    cancel_event_task = asyncio.create_task(cancel_event.wait())
-    await asyncio.wait(
-        [soul_task, cancel_event_task],
-        return_when=asyncio.FIRST_COMPLETED,
-    )
-
-    try:
-        if cancel_event.is_set():
-            logger.debug("Cancelling the run task")
-            soul_task.cancel()
-            try:
-                await soul_task
-            except asyncio.CancelledError:
-                raise RunCancelled from None
-        else:
-            assert soul_task.done()  # either stop event is set or the run task is done
-            cancel_event_task.cancel()
-            with contextlib.suppress(asyncio.CancelledError):
-                await cancel_event_task
-            soul_task.result()  # this will raise if any exception was raised in the run task
-    finally:
-        logger.debug("Shutting down the visualization loop")
-        # shutting down the event queue should break the visualization loop
-        wire.shutdown()
-        try:
-            await asyncio.wait_for(ui_task, timeout=0.5)
-        except TimeoutError:
-            logger.warning("Visualization loop timed out")

kimi_cli/ui/acp/__init__.py

@@ -12,18 +12,17 @@ from kosong.base.message import (
 from kosong.chat_provider import ChatProviderError
 from kosong.tooling import ToolError, ToolOk, ToolResult
 
-from kimi_cli.soul import LLMNotSet, MaxStepsReached, Soul
-from kimi_cli.soul.wire import (
+from kimi_cli.soul import LLMNotSet, MaxStepsReached, RunCancelled, Soul, run_soul
+from kimi_cli.tools import extract_subtitle
+from kimi_cli.utils.logging import logger
+from kimi_cli.wire import WireUISide
+from kimi_cli.wire.message import (
     ApprovalRequest,
     ApprovalResponse,
     StatusUpdate,
     StepBegin,
     StepInterrupted,
-    Wire,
 )
-from kimi_cli.tools import extract_subtitle
-from kimi_cli.ui import RunCancelled, run_soul
-from kimi_cli.utils.logging import logger
 
 
 class _ToolCallState:
@@ -64,7 +63,7 @@ class _RunState:
         self.cancel_event = asyncio.Event()
 
 
-class ACPAgentImpl:
+class ACPAgent:
     """Implementation of the ACP Agent protocol."""
 
     def __init__(self, soul: Soul, connection: acp.AgentSideConnection):
@@ -172,7 +171,7 @@ class ACPAgentImpl:
             logger.info("Cancelling running prompt")
             self.run_state.cancel_event.set()
 
-    async def _stream_events(self, wire: Wire):
+    async def _stream_events(self, wire: WireUISide):
         try:
             # expect a StepBegin
             assert isinstance(await wire.receive(), StepBegin)
@@ -428,7 +427,7 @@ class ACPServer:
 
         # Create connection - the library handles all JSON-RPC details!
         _ = acp.AgentSideConnection(
-            lambda conn: ACPAgentImpl(self.soul, conn),
+            lambda conn: ACPAgent(self.soul, conn),
             writer,
             reader,
         )