python-codex 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

pycodex/__init__.py

@@ -1,3 +1,7 @@
+from .compat import patch_asyncio
+
+patch_asyncio()
+
 from .agent import AgentLoop
 from .context import ContextConfig, ContextManager
 from .model import (
@@ -60,7 +64,7 @@ from .tools import (
     WriteStdinTool,
 )
 
-def debug(stop: bool = False):
+def debug(stop: 'bool' = False):
 
     import socket
 

pycodex/agent.py

@@ -1,8 +1,7 @@
-from __future__ import annotations
 
 import asyncio
 import json
-from collections.abc import Callable
+from typing import Callable
 
 from .context import ContextManager
 from .model import ModelClient
@@ -19,10 +18,14 @@ from .protocol import (
 )
 from .tools import ToolContext, ToolRegistry
 from .utils import uuid7_string
+import typing
+
+if typing.TYPE_CHECKING:
+    from .utils.session_persist import SessionRolloutRecorder
 
 
 EventHandler = Callable[[AgentEvent], None]
-NOOP_EVENT_HANDLER: EventHandler = lambda _event: None
+NOOP_EVENT_HANDLER: 'EventHandler' = lambda _event: None
 
 
 class TurnInterrupted(RuntimeError):
@@ -40,51 +43,66 @@ class AgentLoop:
 
     def __init__(
         self,
-        model_client: ModelClient,
-        tool_registry: ToolRegistry,
-        context_manager: ContextManager | None = None,
-        parallel_tool_calls: bool = True,
-        event_handler: EventHandler = NOOP_EVENT_HANDLER,
-        initial_history: tuple[ConversationItem, ...] = (),
-    ) -> None:
+        model_client: 'ModelClient',
+        tool_registry: 'ToolRegistry',
+        context_manager: 'typing.Union[ContextManager, None]' = None,
+        parallel_tool_calls: 'bool' = True,
+        event_handler: 'EventHandler' = NOOP_EVENT_HANDLER,
+        initial_history: 'typing.Tuple[ConversationItem, ...]' = (),
+        rollout_recorder: 'typing.Union[SessionRolloutRecorder, None]' = None,
+    ) -> 'None':
         self._model_client = model_client
         self._tool_registry = tool_registry
         self._context_manager = context_manager or ContextManager()
         self._parallel_tool_calls = parallel_tool_calls
         self._event_handler = event_handler
-        self._history: list[ConversationItem] = list(initial_history)
+        self._history: 'typing.List[ConversationItem]' = list(initial_history)
+        self._rollout_recorder = rollout_recorder
         self.interrupt_asap = False
 
     @property
-    def history(self) -> tuple[ConversationItem, ...]:
+    def history(self) -> 'typing.Tuple[ConversationItem, ...]':
         return tuple(self._history)
 
     def set_event_handler(
-        self, event_handler: EventHandler = NOOP_EVENT_HANDLER
-    ) -> None:
+        self, event_handler: 'EventHandler' = NOOP_EVENT_HANDLER
+    ) -> 'None':
         self._event_handler = event_handler
 
+    def replace_history(
+        self,
+        history: 'typing.Iterable[ConversationItem]',
+    ) -> 'None':
+        self._history = list(history)
+
+    def set_rollout_recorder(
+        self,
+        rollout_recorder: 'typing.Union[SessionRolloutRecorder, None]',
+    ) -> 'None':
+        self._rollout_recorder = rollout_recorder
+
     def _raise_if_interrupt_requested(
         self,
-        turn_id: str,
-        iteration: int,
-        output_text: str | None = None,
-    ) -> None:
+        turn_id: 'str',
+        iteration: 'int',
+        output_text: 'typing.Union[str, None]' = None,
+    ) -> 'None':
         if self.interrupt_asap:
             self.interrupt_asap = False
-            payload: dict[str, object] = {"iteration": iteration}
+            payload: 'typing.Dict[str, object]' = {"iteration": iteration}
             if output_text is not None:
                 payload["output_text"] = output_text
             self._emit("turn_interrupted", turn_id, **payload)
             raise TurnInterrupted("turn interrupted")
 
     async def run_turn(
-        self, texts: list[str], turn_id: str | None = None
-    ) -> TurnResult:
+        self, texts: 'typing.List[str]', turn_id: 'typing.Union[str, None]' = None
+    ) -> 'TurnResult':
         turn_id = turn_id or uuid7_string()
         self.interrupt_asap = False
-        for text in texts:
-            self._history.append(UserMessage(text=text))
+        new_user_messages = [UserMessage(text=text) for text in texts]
+        self._history.extend(new_user_messages)
+        self._persist_history_items(new_user_messages)
 
         self._emit(
             "turn_started",
@@ -93,10 +111,8 @@ class AgentLoop:
             user_texts=list(texts),
         )
 
-        last_assistant_message: str | None = None
-        final_response_items: tuple[
-            AssistantMessage | ToolCall | ReasoningItem, ...
-        ] = ()
+        last_assistant_message: 'typing.Union[str, None]' = None
+        final_response_items: 'typing.Tuple[\n    typing.Union[typing.Union[AssistantMessage, ToolCall], ReasoningItem], ...\n]' = ()
 
         iteration = 0
         try:
@@ -132,13 +148,16 @@ class AgentLoop:
                     item_count=len(response.items),
                 )
 
-                tool_calls: list[ToolCall] = []
+                tool_calls: 'typing.List[ToolCall]' = []
+                persisted_response_items: 'typing.List[ConversationItem]' = []
                 for item in response.items:
                     self._history.append(item)
+                    persisted_response_items.append(item)
                     if isinstance(item, AssistantMessage):
                         last_assistant_message = item.text
                     elif isinstance(item, ToolCall):
                         tool_calls.append(item)
+                self._persist_history_items(persisted_response_items)
 
                 if not tool_calls:
                     self._raise_if_interrupt_requested(
@@ -162,7 +181,10 @@ class AgentLoop:
 
                 tool_results = await self._execute_tool_batch(turn_id, tool_calls)
                 self._history.extend(tool_results)
-                self._history.extend(self._build_follow_up_messages(tool_results))
+                self._persist_history_items(tool_results)
+                follow_up_messages = self._build_follow_up_messages(tool_results)
+                self._history.extend(follow_up_messages)
+                self._persist_history_items(follow_up_messages)
                 self._raise_if_interrupt_requested(
                     turn_id,
                     iteration,
@@ -182,11 +204,11 @@ class AgentLoop:
 
     async def _execute_tool_batch(
         self,
-        turn_id: str,
-        tool_calls: list[ToolCall],
-    ) -> list[ToolResult]:
-        results: list[ToolResult] = []
-        parallel_batch: list[ToolCall] = []
+        turn_id: 'str',
+        tool_calls: 'typing.List[ToolCall]',
+    ) -> 'typing.List[ToolResult]':
+        results: 'typing.List[ToolResult]' = []
+        parallel_batch: 'typing.List[ToolCall]' = []
 
         for call in tool_calls:
             can_run_parallel = (
@@ -224,11 +246,16 @@ class AgentLoop:
 
     async def _run_single_tool(
         self,
-        turn_id: str,
-        call: ToolCall,
-        prior_results: tuple[ToolResult, ...] = (),
-    ) -> ToolResult:
-        self._emit("tool_started", turn_id, tool_name=call.name, call_id=call.call_id)
+        turn_id: 'str',
+        call: 'ToolCall',
+        prior_results: 'typing.Tuple[ToolResult, ...]' = (),
+    ) -> 'ToolResult':
+        payload: 'typing.Dict[str, object]' = {
+            "tool_name": call.name,
+            "call_id": call.call_id,
+            "call": call,
+        }
+        self._emit("tool_started", turn_id, **payload)
         result = await self._tool_registry.execute(
             call,
             ToolContext(
@@ -237,32 +264,43 @@ class AgentLoop:
                 collaboration_mode=self._context_manager.collaboration_mode,
             ),
         )
-        payload: dict[str, object] = {
-            "tool_name": call.name,
-            "call_id": call.call_id,
-            "is_error": result.is_error,
-            "call": call,
-            "result": result,
-        }
+        payload["result"] = result
+        payload["is_error"] = result.is_error
         self._emit("tool_completed", turn_id, **payload)
         return result
 
-    def _emit(self, kind: str, turn_id: str, **payload: object) -> None:
+    def _emit(self, kind: 'str', turn_id: 'str', **payload: 'object') -> 'None':
         self._event_handler(
             AgentEvent(kind=kind, turn_id=turn_id, payload=dict(payload))
         )
 
-    def _handle_model_stream_event(self, turn_id: str, event: ModelStreamEvent) -> None:
+    def _persist_history_items(
+        self,
+        items: 'typing.Iterable[ConversationItem]',
+    ) -> 'None':
+        recorder = self._rollout_recorder
+        if recorder is None:
+            return
+        try:
+            recorder.append_history_items(items)
+        except Exception:  # pragma: no cover - persistence should not break turns
+            return
+
+    def _handle_model_stream_event(self, turn_id: 'str', event: 'ModelStreamEvent') -> 'None':
         if event.kind == "assistant_delta":
             self._emit("assistant_delta", turn_id, **event.payload)
         elif event.kind == "tool_call":
             self._emit("tool_called", turn_id, **event.payload)
+        elif event.kind == "token_count":
+            self._emit("token_count", turn_id, **event.payload)
+        elif event.kind == "stream_error":
+            self._emit("stream_error", turn_id, **event.payload)
 
     def _build_follow_up_messages(
         self,
-        tool_results: list[ToolResult],
-    ) -> list[UserMessage]:
-        follow_ups: list[UserMessage] = []
+        tool_results: 'typing.List[ToolResult]',
+    ) -> 'typing.List[UserMessage]':
+        follow_ups: 'typing.List[UserMessage]' = []
         for result in tool_results:
             statuses = None
             if (

pycodex/cli.py

@@ -1,4 +1,3 @@
-from __future__ import annotations
 
 import atexit
 import argparse
@@ -10,23 +9,35 @@ import sys
 import tempfile
 from dataclasses import asdict, replace
 from pathlib import Path
-from typing import Literal, Sequence
+from typing import Sequence
 
 from .agent import AgentLoop
 from .collaboration import DEFAULT_COLLABORATION_MODE, CollaborationMode
+from .compat import Literal
 from .context import ContextManager
 from .model import ResponsesModelClient, ResponsesProviderConfig
 from .portable import bootstrap_called_home, upload_codex_home
 from .protocol import AgentEvent
 from .runtime import AgentRuntime
 from .runtime_services import RuntimeEnvironment, create_runtime_environment
-from .utils import CliSessionView, load_codex_dotenv
+from .utils import CliSessionView, load_codex_dotenv, uuid7_string
+from .utils.compactor import compact_agent_loop
+from .utils.session_persist import (
+    SessionRolloutRecorder,
+    conversation_history_to_turns,
+    list_resumable_sessions,
+    load_resumed_session,
+    resolve_codex_home,
+)
+import typing
 
 EXIT_COMMANDS = {"/exit", "/quit"}
 HISTORY_COMMAND = "/history"
 TITLE_COMMAND = "/title"
 MODEL_COMMAND = "/model"
 QUEUE_COMMAND = "/queue"
+RESUME_COMMAND = "/resume"
+COMPACT_COMMAND = "/compact"
 CliSessionMode = Literal["exec", "tui"]
 LOCAL_RESPONSES_SERVER_API_KEY_ENV = "PYCODEX_LOCAL_RESPONSES_SERVER_KEY"
 CLI_ORIGINATOR = "codex-tui"
@@ -40,7 +51,7 @@ def launch_chat_completion_compat_server(*args, **kwargs):
     return launch_compat_server(*args, **kwargs)
 
 
-def configure_loguru() -> None:
+def configure_loguru() -> 'None':
     try:
         from loguru import logger
     except ImportError:  # pragma: no cover - dependency may be absent in minimal envs
@@ -61,7 +72,7 @@ def configure_loguru() -> None:
         logger.add(sys.stderr, level="DEBUG")
 
 
-def build_parser() -> argparse.ArgumentParser:
+def build_parser() -> 'argparse.ArgumentParser':
     parser = argparse.ArgumentParser(
         prog="pycodex",
         description="Minimal Codex-style local CLI backed by ~/.codex/config.toml.",
@@ -131,11 +142,11 @@ def build_parser() -> argparse.ArgumentParser:
     return parser
 
 
-def should_run_interactive(prompt_parts: Sequence[str], stdin_is_tty: bool) -> bool:
+def should_run_interactive(prompt_parts: 'Sequence[str]', stdin_is_tty: 'bool') -> 'bool':
     return not prompt_parts and stdin_is_tty
 
 
-def resolve_prompt_text(prompt_parts: Sequence[str]) -> str:
+def resolve_prompt_text(prompt_parts: 'Sequence[str]') -> 'str':
     if prompt_parts:
         return " ".join(prompt_parts).strip()
 
@@ -148,8 +159,8 @@ def resolve_prompt_text(prompt_parts: Sequence[str]) -> str:
 
 
 def get_tools(
-    runtime_environment: RuntimeEnvironment | None = None,
-    exec_mode: bool = False,
+    runtime_environment: 'typing.Union[RuntimeEnvironment, None]' = None,
+    exec_mode: 'bool' = False,
 ):
     from .tools import (
         ApplyPatchTool,
@@ -243,7 +254,7 @@ def get_tools(
     return registry
 
 
-def get_subagent_tools(runtime_environment: RuntimeEnvironment | None = None):
+def get_subagent_tools(runtime_environment: 'typing.Union[RuntimeEnvironment, None]' = None):
     from .tools import (
         ApplyPatchTool,
         ExecCommandTool,
@@ -268,13 +279,13 @@ def get_subagent_tools(runtime_environment: RuntimeEnvironment | None = None):
 
 
 def build_runtime(
-    config_path: str,
-    profile: str | None,
-    system_prompt: str | None,
+    config_path: 'str',
+    profile: 'typing.Union[str, None]',
+    system_prompt: 'typing.Union[str, None]',
     client,
-    session_mode: CliSessionMode = "exec",
-    collaboration_mode: CollaborationMode = DEFAULT_COLLABORATION_MODE,
-) -> AgentRuntime:
+    session_mode: 'CliSessionMode' = "exec",
+    collaboration_mode: 'CollaborationMode' = DEFAULT_COLLABORATION_MODE,
+) -> 'AgentRuntime':
     use_tui_context = session_mode == "tui"
     context_manager = ContextManager.from_codex_config(
         config_path,
@@ -283,6 +294,9 @@ def build_runtime(
         collaboration_mode=collaboration_mode,
         include_collaboration_instructions=use_tui_context,
     )
+    session_id = getattr(client, "_session_id", None) or uuid7_string()
+    if hasattr(client, "_session_id"):
+        client._session_id = session_id
     subagent_context_manager = ContextManager.from_codex_config(
         config_path,
         profile,
@@ -292,14 +306,22 @@ def build_runtime(
     runtime_environment = create_runtime_environment()
     runtime_environment.request_user_input_manager.set_handler(None)
     runtime_environment.request_permissions_manager.set_handler(None)
+    rollout_recorder = SessionRolloutRecorder.create(
+        resolve_codex_home(config_path),
+        session_id,
+        context_manager.cwd,
+        getattr(client, "_originator", CLI_ORIGINATOR),
+        getattr(getattr(client, "_config", None), "provider_name", None),
+        context_manager.resolve_base_instructions(),
+    )
 
     def make_subagent_runtime_builder(base_client):
         def build_subagent_runtime(
-            model_override: str | None,
-            reasoning_effort_override: str | None,
+            model_override: 'typing.Union[str, None]',
+            reasoning_effort_override: 'typing.Union[str, None]',
             initial_history=(),
-            session_id: str | None = None,
-        ) -> AgentRuntime:
+            session_id: 'typing.Union[str, None]' = None,
+        ) -> 'AgentRuntime':
             nested_client = base_client.with_overrides(
                 model_override,
                 reasoning_effort_override,
@@ -329,25 +351,28 @@ def build_runtime(
     )
     return AgentRuntime(
         AgentLoop(
-            client, get_tools(runtime_environment, exec_mode=True), context_manager
+            client,
+            get_tools(runtime_environment, exec_mode=True),
+            context_manager,
+            rollout_recorder=rollout_recorder,
         ),
         runtime_environment=runtime_environment,
     )
 
 
-def format_turn_output(result, json_mode: bool) -> str:
+def format_turn_output(result, json_mode: 'bool') -> 'str':
     if json_mode:
         return json.dumps(asdict(result), ensure_ascii=False, indent=2)
     return result.output_text or ""
 
 
 def _build_model_client(
-    config_path: str,
-    profile: str | None,
-    timeout_seconds: float,
-    managed_responses_base_url: str | None = None,
-    vllm_endpoint: str | None = None,
-    use_chat_completion: bool = False,
+    config_path: 'str',
+    profile: 'typing.Union[str, None]',
+    timeout_seconds: 'float',
+    managed_responses_base_url: 'typing.Union[str, None]' = None,
+    vllm_endpoint: 'typing.Union[str, None]' = None,
+    use_chat_completion: 'bool' = False,
 ):
     load_codex_dotenv(config_path)
     provider_config = ResponsesProviderConfig.from_codex_config(
@@ -393,13 +418,13 @@ def _build_model_client(
 
 
 async def prompt_request_user_input(
-    view: CliSessionView,
-    payload: dict[str, object],
-) -> dict[str, object] | None:
+    view: 'CliSessionView',
+    payload: 'typing.Dict[str, object]',
+) -> 'typing.Union[typing.Dict[str, object], None]':
     view.finish_stream()
     view.pause_spinner()
     view.write_line("[request_user_input] waiting for user response")
-    answers: dict[str, dict[str, list[str]]] = {}
+    answers: 'typing.Dict[str, typing.Dict[str, typing.List[str]]]' = {}
     try:
         for question in payload.get("questions", []):
             if not isinstance(question, dict):
@@ -456,9 +481,9 @@ async def prompt_request_user_input(
 
 
 async def prompt_request_permissions(
-    view: CliSessionView,
-    payload: dict[str, object],
-) -> dict[str, object] | None:
+    view: 'CliSessionView',
+    payload: 'typing.Dict[str, object]',
+) -> 'typing.Union[typing.Dict[str, object], None]':
     view.finish_stream()
     view.pause_spinner()
     view.write_line("[request_permissions] user approval required")
@@ -495,14 +520,18 @@ async def prompt_request_permissions(
 
 
 async def run_interactive_session(
-    runtime: AgentRuntime,
-    json_mode: bool,
-) -> int:
+    runtime: 'AgentRuntime',
+    json_mode: 'bool',
+    config_path: 'typing.Union[str, None]' = None,
+) -> 'int':
     worker = asyncio.create_task(runtime.run_forever())
+    context_window_tokens = runtime._agent_loop._context_manager.resolve_model_context_window()
     view = CliSessionView()
+    view.set_context_window_tokens(context_window_tokens)
     model_client = runtime._agent_loop._model_client
+    codex_home = resolve_codex_home(config_path)
     runtime.set_event_handler(view.handle_event)
-    pending_turn_tasks: set[asyncio.Task[None]] = set()
+    pending_turn_tasks: 'typing.Set[asyncio.Task[None]]' = set()
     runtime_environment = runtime.runtime_environment
     if runtime_environment is None:
         runtime_environment = create_runtime_environment()
@@ -514,16 +543,49 @@ async def run_interactive_session(
         lambda payload: prompt_request_permissions(view, payload)
     )
     view.write_line("pycodex interactive mode. Type /exit to quit.")
-    view.write_line("Extra commands: /history, /title, /model")
+    view.write_line("Extra commands: /history, /title, /model, /resume, /compact")
     try:
 
-        def has_pending_turn_tasks() -> bool:
+        def has_pending_turn_tasks() -> 'bool':
             pending_turn_tasks.difference_update(
                 task for task in tuple(pending_turn_tasks) if task.done()
             )
             return bool(pending_turn_tasks)
 
-        async def wait_for_turn_result(future) -> None:
+        async def run_manual_compact() -> 'None':
+            agent_loop = runtime._agent_loop
+            if not agent_loop.history:
+                view.write_line("Nothing to compact.")
+                return
+
+            compact_turn_id = uuid7_string()
+
+            def handle_compact_stream_event(event) -> 'None':
+                if event.kind not in {"token_count", "stream_error"}:
+                    return
+                view.handle_event(
+                    AgentEvent(
+                        kind=event.kind,
+                        turn_id=compact_turn_id,
+                        payload=dict(event.payload),
+                    )
+                )
+
+            view.write_line("Compacting conversation history...")
+            compact_result = await compact_agent_loop(
+                agent_loop,
+                handle_compact_stream_event,
+            )
+            if compact_result is None:
+                view.write_line("Nothing to compact.")
+                return
+            view.load_session_history(
+                getattr(view, "_title", None),
+                conversation_history_to_turns(compact_result.history),
+            )
+            view.write_line(compact_result.display_text())
+
+        async def wait_for_turn_result(future) -> 'None':
             try:
                 result = await future
             except Exception as exc:  # pragma: no cover - defensive surface
@@ -555,6 +617,50 @@ async def run_interactive_session(
             if prompt_text == TITLE_COMMAND:
                 view.show_title()
                 continue
+            if prompt_text == RESUME_COMMAND:
+                sessions = list_resumable_sessions(codex_home)
+                if not sessions:
+                    view.write_line("No resumable sessions found.")
+                    continue
+                view.write_line("Available sessions:")
+                for index, session in enumerate(sessions, start=1):
+                    view.write_line(f"[{index}] {session['preview']}")
+                continue
+            if prompt_text.startswith(f"{RESUME_COMMAND} "):
+                if has_pending_turn_tasks():
+                    view.write_line(
+                        "Cannot resume while work is running or queued."
+                    )
+                    continue
+                resume_target = prompt_text[len(RESUME_COMMAND) :].strip()
+                try:
+                    resumed = load_resumed_session(codex_home, resume_target)
+                    runtime._agent_loop.replace_history(resumed["history"])
+                    if hasattr(model_client, "_session_id"):
+                        model_client._session_id = str(resumed["session_id"])
+                    runtime._agent_loop.set_rollout_recorder(
+                        SessionRolloutRecorder.resume(resumed["rollout_path"])
+                    )
+                    view.load_session_history(
+                        str(resumed["title"]),
+                        tuple(resumed["turns"]),
+                    )
+                    view.write_line(f"Resumed session: {resumed['title']}")
+                    view.show_history()
+                except Exception as exc:  # pragma: no cover - defensive surface
+                    view.show_error(str(exc))
+                continue
+            if prompt_text == COMPACT_COMMAND:
+                if has_pending_turn_tasks():
+                    view.write_line(
+                        "Cannot compact while work is running or queued."
+                    )
+                    continue
+                try:
+                    await run_manual_compact()
+                except Exception as exc:  # pragma: no cover - defensive surface
+                    view.show_error(str(exc))
+                continue
             if prompt_text.startswith(f"{QUEUE_COMMAND} "):
                 queued_text = prompt_text[len(QUEUE_COMMAND) :].strip()
                 if not queued_text:
@@ -618,7 +724,7 @@ async def run_interactive_session(
     return 0
 
 
-async def run_cli(args: argparse.Namespace) -> int:
+async def run_cli(args: 'argparse.Namespace') -> 'int':
     runtime = None
     worker = None
     try:
@@ -628,7 +734,7 @@ async def run_cli(args: argparse.Namespace) -> int:
             raise ValueError("--put does not accept prompt text")
         configure_loguru()
         if args.put is not None:
-            def emit_put_log(message: str) -> None:
+            def emit_put_log(message: 'str') -> 'None':
                 print(message, flush=True)
 
             call_spec = upload_codex_home(args.put, event_handler=emit_put_log)
@@ -662,6 +768,7 @@ async def run_cli(args: argparse.Namespace) -> int:
             return await run_interactive_session(
                 runtime,
                 args.json,
+                args.config,
             )
         else:
             prompt_text = resolve_prompt_text(args.prompt)
@@ -678,7 +785,7 @@ async def run_cli(args: argparse.Namespace) -> int:
             await worker
 
 
-def main(argv: Sequence[str] | None = None) -> int:
+def main(argv: 'typing.Union[Sequence[str], None]' = None) -> 'int':
     raw_args = list(argv) if argv is not None else None
     if raw_args is None:
         raw_args = sys.argv[1:]

pycodex/collaboration.py

@@ -1,13 +1,13 @@
-from __future__ import annotations
 
-from typing import Literal
+from .compat import Literal
+import typing
 
 CollaborationMode = Literal["default", "plan", "execute", "pair_programming"]
 
-DEFAULT_COLLABORATION_MODE: CollaborationMode = "default"
-PLAN_COLLABORATION_MODE: CollaborationMode = "plan"
+DEFAULT_COLLABORATION_MODE: 'CollaborationMode' = "default"
+PLAN_COLLABORATION_MODE: 'CollaborationMode' = "plan"
 
-_MODE_DISPLAY_NAMES: dict[str, str] = {
+_MODE_DISPLAY_NAMES: 'typing.Dict[str, str]' = {
     "default": "Default",
     "plan": "Plan",
     "execute": "Execute",
@@ -15,7 +15,6 @@ _MODE_DISPLAY_NAMES: dict[str, str] = {
 }
 
 
-def collaboration_mode_display_name(mode: str | None) -> str:
+def collaboration_mode_display_name(mode: 'typing.Union[str, None]') -> 'str':
     normalized = (mode or DEFAULT_COLLABORATION_MODE).strip().lower()
     return _MODE_DISPLAY_NAMES.get(normalized, normalized.replace("_", " ").title())
-