codex-autorunner 0.1.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

codex_autorunner/integrations/app_server/supervisor.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 import asyncio
 import logging
+import os
 import time
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, Dict, Optional, Sequence
 
 from ...core.logging_utils import log_event
+from ...core.supervisor_utils import evict_lru_handle_locked, pop_idle_handles_locked
 from ...workspace import canonical_workspace_root, workspace_id_for_path
 from .client import ApprovalHandler, CodexAppServerClient, NotificationHandler
 
@@ -34,8 +36,17 @@ class WorkspaceAppServerSupervisor:
         approval_handler: Optional[ApprovalHandler] = None,
         notification_handler: Optional[NotificationHandler] = None,
         logger: Optional[logging.Logger] = None,
-        auto_restart: bool = True,
+        auto_restart: Optional[bool] = None,
         request_timeout: Optional[float] = None,
+        turn_stall_timeout_seconds: Optional[float] = None,
+        turn_stall_poll_interval_seconds: Optional[float] = None,
+        turn_stall_recovery_min_interval_seconds: Optional[float] = None,
+        max_message_bytes: Optional[int] = None,
+        oversize_preview_bytes: Optional[int] = None,
+        max_oversize_drain_bytes: Optional[int] = None,
+        restart_backoff_initial_seconds: Optional[float] = None,
+        restart_backoff_max_seconds: Optional[float] = None,
+        restart_backoff_jitter_ratio: Optional[float] = None,
         default_approval_decision: str = "cancel",
         max_handles: Optional[int] = None,
         idle_ttl_seconds: Optional[float] = None,
@@ -46,8 +57,27 @@ class WorkspaceAppServerSupervisor:
         self._approval_handler = approval_handler
         self._notification_handler = notification_handler
         self._logger = logger or logging.getLogger(__name__)
-        self._auto_restart = auto_restart
+        disable_restart_env = os.environ.get(
+            "CODEX_DISABLE_APP_SERVER_AUTORESTART_FOR_TESTS"
+        )
+        if disable_restart_env:
+            self._auto_restart = False
+        elif auto_restart is None:
+            self._auto_restart = True
+        else:
+            self._auto_restart = auto_restart
         self._request_timeout = request_timeout
+        self._turn_stall_timeout_seconds = turn_stall_timeout_seconds
+        self._turn_stall_poll_interval_seconds = turn_stall_poll_interval_seconds
+        self._turn_stall_recovery_min_interval_seconds = (
+            turn_stall_recovery_min_interval_seconds
+        )
+        self._max_message_bytes = max_message_bytes
+        self._oversize_preview_bytes = oversize_preview_bytes
+        self._max_oversize_drain_bytes = max_oversize_drain_bytes
+        self._restart_backoff_initial_seconds = restart_backoff_initial_seconds
+        self._restart_backoff_max_seconds = restart_backoff_max_seconds
+        self._restart_backoff_jitter_ratio = restart_backoff_jitter_ratio
         self._default_approval_decision = default_approval_decision
         self._max_handles = max_handles
         self._idle_ttl_seconds = idle_ttl_seconds
@@ -78,7 +108,8 @@ class WorkspaceAppServerSupervisor:
                     last_used_at=handle.last_used_at,
                 )
                 await handle.client.close()
-            except Exception:
+            except Exception as exc:
+                self._logger.debug("Failed to close handle: %s", exc)
                 continue
 
     async def prune_idle(self) -> int:
@@ -100,7 +131,8 @@ class WorkspaceAppServerSupervisor:
                 )
                 await handle.client.close()
                 closed += 1
-            except Exception:
+            except Exception as exc:
+                self._logger.debug("Failed to prune handle: %s", exc)
                 continue
         return closed
 
@@ -129,6 +161,15 @@ class WorkspaceAppServerSupervisor:
                 default_approval_decision=self._default_approval_decision,
                 auto_restart=self._auto_restart,
                 request_timeout=self._request_timeout,
+                turn_stall_timeout_seconds=self._turn_stall_timeout_seconds,
+                turn_stall_poll_interval_seconds=self._turn_stall_poll_interval_seconds,
+                turn_stall_recovery_min_interval_seconds=self._turn_stall_recovery_min_interval_seconds,
+                max_message_bytes=self._max_message_bytes,
+                oversize_preview_bytes=self._oversize_preview_bytes,
+                max_oversize_drain_bytes=self._max_oversize_drain_bytes,
+                restart_backoff_initial_seconds=self._restart_backoff_initial_seconds,
+                restart_backoff_max_seconds=self._restart_backoff_max_seconds,
+                restart_backoff_jitter_ratio=self._restart_backoff_jitter_ratio,
                 notification_handler=self._notification_handler,
                 logger=self._logger,
             )
@@ -157,7 +198,8 @@ class WorkspaceAppServerSupervisor:
                     last_used_at=handle.last_used_at,
                 )
                 await handle.client.close()
-            except Exception:
+            except Exception as exc:
+                self._logger.debug("Failed to close handle: %s", exc)
                 continue
         return handle
 
@@ -173,35 +215,19 @@ class WorkspaceAppServerSupervisor:
             return self._pop_idle_handles_locked()
 
     def _pop_idle_handles_locked(self) -> list[AppServerHandle]:
-        if not self._idle_ttl_seconds or self._idle_ttl_seconds <= 0:
-            return []
-        cutoff = time.monotonic() - self._idle_ttl_seconds
-        stale: list[AppServerHandle] = []
-        for handle in list(self._handles.values()):
-            if handle.last_used_at and handle.last_used_at < cutoff:
-                self._handles.pop(handle.workspace_id, None)
-                stale.append(handle)
-        return stale
+        return pop_idle_handles_locked(
+            self._handles,
+            self._idle_ttl_seconds,
+            self._logger,
+            "app_server",
+            last_used_at_getter=lambda h: h.last_used_at,
+        )
 
     def _evict_lru_handle_locked(self) -> Optional[AppServerHandle]:
-        if not self._max_handles or self._max_handles <= 0:
-            return None
-        if len(self._handles) < self._max_handles:
-            return None
-        lru_handle = min(
-            self._handles.values(),
-            key=lambda handle: handle.last_used_at or 0.0,
-        )
-        log_event(
+        return evict_lru_handle_locked(
+            self._handles,
+            self._max_handles,
             self._logger,
-            logging.INFO,
-            "app_server.handle.evicted",
-            reason="max_handles",
-            workspace_id=lru_handle.workspace_id,
-            workspace_root=str(lru_handle.workspace_root),
-            max_handles=self._max_handles,
-            handle_count=len(self._handles),
-            last_used_at=lru_handle.last_used_at,
+            "app_server",
+            last_used_at_getter=lambda h: h.last_used_at or 0.0,
         )
-        self._handles.pop(lru_handle.workspace_id, None)
-        return lru_handle

codex_autorunner/integrations/telegram/adapter.py

@@ -12,6 +12,16 @@ from ...core.circuit_breaker import CircuitBreaker
 from ...core.exceptions import CodexError, PermanentError, TransientError
 from ...core.logging_utils import log_event
 from ...core.retry import retry_transient
+from .api_schemas import (
+    TelegramAudioSchema,
+    TelegramDocumentSchema,
+    TelegramMessageEntitySchema,
+    TelegramPhotoSizeSchema,
+    TelegramVoiceSchema,
+    parse_callback_query_payload,
+    parse_message_payload,
+    parse_update_payload,
+)
 from .constants import TELEGRAM_CALLBACK_DATA_LIMIT, TELEGRAM_MAX_MESSAGE_LENGTH
 from .retry import _extract_retry_after_seconds
 
@@ -120,6 +130,12 @@ class TelegramMessage:
     voice: Optional[TelegramVoice] = None
     media_group_id: Optional[str] = None
 
+    # Extra metadata used for trigger gating / UX (optional, depends on update payload).
+    chat_type: Optional[str] = None
+    reply_to_message_id: Optional[int] = None
+    reply_to_is_bot: bool = False
+    reply_to_username: Optional[str] = None
+
 
 @dataclass(frozen=True)
 class TelegramCallbackQuery:
@@ -227,12 +243,6 @@ class ReviewCommitCallback:
     sha: str
 
 
-@dataclass(frozen=True)
-class PrFlowStartCallback:
-    slug: str
-    number: int
-
-
 @dataclass(frozen=True)
 class CancelCallback:
     kind: str
@@ -323,111 +333,123 @@ def is_interrupt_alias(text: Optional[str]) -> bool:
 
 
 def parse_update(update: dict[str, Any]) -> Optional[TelegramUpdate]:
-    update_id = update.get("update_id")
-    if not isinstance(update_id, int):
+    try:
+        schema = parse_update_payload(update)
+    except Exception:
         return None
-    message = _parse_message(update_id, update.get("message"), edited=False)
+    message = _parse_message(schema.update_id, schema.message, edited=False)
     if message is None:
-        message = _parse_message(update_id, update.get("edited_message"), edited=True)
-    callback = _parse_callback(update_id, update.get("callback_query"))
+        message = _parse_message(schema.update_id, schema.edited_message, edited=True)
+    callback = _parse_callback(schema.update_id, schema.callback_query)
     if message is None and callback is None:
         return None
-    return TelegramUpdate(update_id=update_id, message=message, callback=callback)
+    return TelegramUpdate(
+        update_id=schema.update_id, message=message, callback=callback
+    )
 
 
 def _parse_message(
     update_id: int, payload: Any, *, edited: bool = False
 ) -> Optional[TelegramMessage]:
-    if not isinstance(payload, dict):
+    schema = parse_message_payload(payload)
+    if schema is None:
         return None
-    message_id = payload.get("message_id")
-    chat = payload.get("chat")
-    if not isinstance(message_id, int) or not isinstance(chat, dict):
-        return None
-    chat_id = chat.get("id")
+
+    chat_id = schema.chat.get("id") if isinstance(schema.chat, dict) else None
     if not isinstance(chat_id, int):
         return None
-    thread_id = payload.get("message_thread_id")
-    if thread_id is not None and not isinstance(thread_id, int):
-        thread_id = None
-    sender = payload.get("from")
-    from_user_id = sender.get("id") if isinstance(sender, dict) else None
+
+    chat_type = schema.chat.get("type") if isinstance(schema.chat, dict) else None
+    if chat_type is not None and not isinstance(chat_type, str):
+        chat_type = None
+
+    reply_to_message_id: Optional[int] = None
+    reply_to_is_bot = False
+    reply_to_username: Optional[str] = None
+    if isinstance(schema.reply_to_message, dict):
+        rmid = schema.reply_to_message.get("message_id")
+        if isinstance(rmid, int):
+            reply_to_message_id = rmid
+        reply_from = schema.reply_to_message.get("from")
+        if isinstance(reply_from, dict):
+            is_bot = reply_from.get("is_bot")
+            if isinstance(is_bot, bool):
+                reply_to_is_bot = is_bot
+            username = reply_from.get("username")
+            if isinstance(username, str):
+                reply_to_username = username
+
+    from_user_id = (
+        schema.from_user.get("id") if isinstance(schema.from_user, dict) else None
+    )
     if from_user_id is not None and not isinstance(from_user_id, int):
         from_user_id = None
-    text = payload.get("text")
-    if text is not None and not isinstance(text, str):
-        text = None
-    caption = payload.get("caption")
-    if caption is not None and not isinstance(caption, str):
-        caption = None
-    entities = _parse_entities(payload.get("entities"))
-    caption_entities = _parse_entities(payload.get("caption_entities"))
-    photos = _parse_photo_sizes(payload.get("photo"))
-    document = _parse_document(payload.get("document"))
-    audio = _parse_audio(payload.get("audio"))
-    voice = _parse_voice(payload.get("voice"))
-    media_group_id = payload.get("media_group_id")
-    if media_group_id is not None and not isinstance(media_group_id, str):
-        media_group_id = None
-    date = payload.get("date")
-    if date is not None and not isinstance(date, int):
-        date = None
-    is_topic_message = bool(payload.get("is_topic_message"))
+
+    entities = _parse_entities(schema.entities)
+    caption_entities = _parse_entities(schema.caption_entities)
+    photos = _parse_photo_sizes(schema.photo)
+    document = _parse_document(schema.document)
+    audio = _parse_audio(schema.audio)
+    voice = _parse_voice(schema.voice)
+
     return TelegramMessage(
         update_id=update_id,
-        message_id=message_id,
+        message_id=schema.message_id,
         chat_id=chat_id,
-        thread_id=thread_id,
+        thread_id=schema.message_thread_id,
         from_user_id=from_user_id,
-        text=text,
-        date=date,
-        is_topic_message=is_topic_message,
+        text=schema.text,
+        date=schema.date,
+        is_topic_message=schema.is_topic_message,
         is_edited=edited,
-        caption=caption,
+        caption=schema.caption,
         entities=entities,
         caption_entities=caption_entities,
         photos=photos,
         document=document,
         audio=audio,
         voice=voice,
-        media_group_id=media_group_id,
+        media_group_id=schema.media_group_id,
+        chat_type=chat_type,
+        reply_to_message_id=reply_to_message_id,
+        reply_to_is_bot=reply_to_is_bot,
+        reply_to_username=reply_to_username,
     )
 
 
 def _parse_callback(update_id: int, payload: Any) -> Optional[TelegramCallbackQuery]:
-    if not isinstance(payload, dict):
-        return None
-    callback_id = payload.get("id")
-    if not isinstance(callback_id, str):
+    schema = parse_callback_query_payload(payload)
+    if schema is None:
         return None
-    sender = payload.get("from")
-    from_user_id = sender.get("id") if isinstance(sender, dict) else None
+
+    from_user_id = (
+        schema.from_user.get("id") if isinstance(schema.from_user, dict) else None
+    )
     if from_user_id is not None and not isinstance(from_user_id, int):
         from_user_id = None
-    data = payload.get("data")
-    if data is not None and not isinstance(data, str):
-        data = None
-    message = payload.get("message")
+
     message_id = None
     chat_id = None
     thread_id = None
-    if isinstance(message, dict):
-        message_id = message.get("message_id")
-        chat = message.get("chat")
+    if isinstance(schema.message, dict):
+        message_id = schema.message.get("message_id")
+        chat = schema.message.get("chat")
         if isinstance(chat, dict):
             chat_id = chat.get("id")
-        thread_id = message.get("message_thread_id")
+        thread_id = schema.message.get("message_thread_id")
+
     if message_id is not None and not isinstance(message_id, int):
         message_id = None
     if chat_id is not None and not isinstance(chat_id, int):
         chat_id = None
     if thread_id is not None and not isinstance(thread_id, int):
         thread_id = None
+
     return TelegramCallbackQuery(
         update_id=update_id,
-        callback_id=callback_id,
+        callback_id=schema.id,
         from_user_id=from_user_id,
-        data=data,
+        data=schema.data,
         message_id=message_id,
         chat_id=chat_id,
         thread_id=thread_id,
@@ -441,26 +463,17 @@ def _parse_photo_sizes(payload: Any) -> tuple[TelegramPhotoSize, ...]:
     for item in payload:
         if not isinstance(item, dict):
             continue
-        file_id = item.get("file_id")
-        if not isinstance(file_id, str) or not file_id:
-            continue
-        file_unique_id = item.get("file_unique_id")
-        if file_unique_id is not None and not isinstance(file_unique_id, str):
-            file_unique_id = None
-        width = item.get("width")
-        height = item.get("height")
-        if not isinstance(width, int) or not isinstance(height, int):
+        try:
+            schema = TelegramPhotoSizeSchema.model_validate(item)
+        except Exception:
             continue
-        file_size = item.get("file_size")
-        if file_size is not None and not isinstance(file_size, int):
-            file_size = None
         sizes.append(
             TelegramPhotoSize(
-                file_id=file_id,
-                file_unique_id=file_unique_id,
-                width=width,
-                height=height,
-                file_size=file_size,
+                file_id=schema.file_id,
+                file_unique_id=schema.file_unique_id,
+                width=schema.width,
+                height=schema.height,
+                file_size=schema.file_size,
             )
         )
     return tuple(sizes)
@@ -469,85 +482,49 @@ def _parse_photo_sizes(payload: Any) -> tuple[TelegramPhotoSize, ...]:
 def _parse_document(payload: Any) -> Optional[TelegramDocument]:
     if not isinstance(payload, dict):
         return None
-    file_id = payload.get("file_id")
-    if not isinstance(file_id, str) or not file_id:
+    try:
+        schema = TelegramDocumentSchema.model_validate(payload)
+    except Exception:
         return None
-    file_unique_id = payload.get("file_unique_id")
-    if file_unique_id is not None and not isinstance(file_unique_id, str):
-        file_unique_id = None
-    file_name = payload.get("file_name")
-    if file_name is not None and not isinstance(file_name, str):
-        file_name = None
-    mime_type = payload.get("mime_type")
-    if mime_type is not None and not isinstance(mime_type, str):
-        mime_type = None
-    file_size = payload.get("file_size")
-    if file_size is not None and not isinstance(file_size, int):
-        file_size = None
     return TelegramDocument(
-        file_id=file_id,
-        file_unique_id=file_unique_id,
-        file_name=file_name,
-        mime_type=mime_type,
-        file_size=file_size,
+        file_id=schema.file_id,
+        file_unique_id=schema.file_unique_id,
+        file_name=schema.file_name,
+        mime_type=schema.mime_type,
+        file_size=schema.file_size,
     )
 
 
 def _parse_audio(payload: Any) -> Optional[TelegramAudio]:
     if not isinstance(payload, dict):
         return None
-    file_id = payload.get("file_id")
-    if not isinstance(file_id, str) or not file_id:
+    try:
+        schema = TelegramAudioSchema.model_validate(payload)
+    except Exception:
         return None
-    file_unique_id = payload.get("file_unique_id")
-    if file_unique_id is not None and not isinstance(file_unique_id, str):
-        file_unique_id = None
-    duration = payload.get("duration")
-    if duration is not None and not isinstance(duration, int):
-        duration = None
-    file_name = payload.get("file_name")
-    if file_name is not None and not isinstance(file_name, str):
-        file_name = None
-    mime_type = payload.get("mime_type")
-    if mime_type is not None and not isinstance(mime_type, str):
-        mime_type = None
-    file_size = payload.get("file_size")
-    if file_size is not None and not isinstance(file_size, int):
-        file_size = None
     return TelegramAudio(
-        file_id=file_id,
-        file_unique_id=file_unique_id,
-        duration=duration,
-        file_name=file_name,
-        mime_type=mime_type,
-        file_size=file_size,
+        file_id=schema.file_id,
+        file_unique_id=schema.file_unique_id,
+        duration=schema.duration,
+        file_name=schema.file_name,
+        mime_type=schema.mime_type,
+        file_size=schema.file_size,
     )
 
 
 def _parse_voice(payload: Any) -> Optional[TelegramVoice]:
     if not isinstance(payload, dict):
         return None
-    file_id = payload.get("file_id")
-    if not isinstance(file_id, str) or not file_id:
+    try:
+        schema = TelegramVoiceSchema.model_validate(payload)
+    except Exception:
         return None
-    file_unique_id = payload.get("file_unique_id")
-    if file_unique_id is not None and not isinstance(file_unique_id, str):
-        file_unique_id = None
-    duration = payload.get("duration")
-    if duration is not None and not isinstance(duration, int):
-        duration = None
-    mime_type = payload.get("mime_type")
-    if mime_type is not None and not isinstance(mime_type, str):
-        mime_type = None
-    file_size = payload.get("file_size")
-    if file_size is not None and not isinstance(file_size, int):
-        file_size = None
     return TelegramVoice(
-        file_id=file_id,
-        file_unique_id=file_unique_id,
-        duration=duration,
-        mime_type=mime_type,
-        file_size=file_size,
+        file_id=schema.file_id,
+        file_unique_id=schema.file_unique_id,
+        duration=schema.duration,
+        mime_type=schema.mime_type,
+        file_size=schema.file_size,
     )
 
 
@@ -558,14 +535,15 @@ def _parse_entities(payload: Any) -> tuple[TelegramMessageEntity, ...]:
     for item in payload:
         if not isinstance(item, dict):
             continue
-        kind = item.get("type")
-        offset = item.get("offset")
-        length = item.get("length")
-        if not isinstance(kind, str):
-            continue
-        if not isinstance(offset, int) or not isinstance(length, int):
+        try:
+            schema = TelegramMessageEntitySchema.model_validate(item)
+        except Exception:
             continue
-        entities.append(TelegramMessageEntity(type=kind, offset=offset, length=length))
+        entities.append(
+            TelegramMessageEntity(
+                type=schema.type, offset=schema.offset, length=schema.length
+            )
+        )
     return tuple(entities)
 
 
@@ -757,12 +735,6 @@ def encode_page_callback(kind: str, page: int) -> str:
     return data
 
 
-def encode_pr_flow_start_callback(slug: str, number: int) -> str:
-    data = f"pr_flow_start:{slug}#{number}"
-    _validate_callback_data(data)
-    return data
-
-
 def encode_compact_callback(action: str) -> str:
     data = f"compact:{action}"
     _validate_callback_data(data)
@@ -786,7 +758,6 @@ def parse_callback_data(
         UpdateCallback,
         UpdateConfirmCallback,
         ReviewCommitCallback,
-        PrFlowStartCallback,
         CancelCallback,
         CompactCallback,
         PageCallback,
@@ -868,16 +839,6 @@ def parse_callback_data(
         if not sha:
             return None
         return ReviewCommitCallback(sha=sha)
-    if data.startswith("pr_flow_start:"):
-        _, _, rest = data.partition(":")
-        if not rest:
-            return None
-        if "#" not in rest:
-            return None
-        slug, _, number_str = rest.partition("#")
-        if not slug or not number_str or not number_str.isdigit():
-            return None
-        return PrFlowStartCallback(slug=slug, number=int(number_str))
     if data.startswith("cancel:"):
         _, _, kind = data.partition(":")
         if not kind:
@@ -1447,6 +1408,7 @@ class TelegramBotClient:
         message_id: int,
         text: str,
         *,
+        message_thread_id: Optional[int] = None,
         reply_markup: Optional[dict[str, Any]] = None,
         parse_mode: Optional[str] = None,
         disable_web_page_preview: bool = True,
@@ -1456,6 +1418,7 @@ class TelegramBotClient:
             logging.INFO,
             "telegram.edit_message",
             chat_id=chat_id,
+            thread_id=message_thread_id,
             message_id=message_id,
             text_len=len(text),
             has_markup=reply_markup is not None,
@@ -1468,6 +1431,8 @@ class TelegramBotClient:
             "text": text,
             "disable_web_page_preview": disable_web_page_preview,
         }
+        if message_thread_id is not None:
+            payload["message_thread_id"] = message_thread_id
         if reply_markup is not None:
             payload["reply_markup"] = reply_markup
         if parse_mode is not None:
@@ -1479,12 +1444,15 @@ class TelegramBotClient:
         self,
         chat_id: Union[int, str],
         message_id: int,
+        *,
+        message_thread_id: Optional[int] = None,
     ) -> bool:
         log_event(
             self._logger,
             logging.INFO,
             "telegram.delete_message",
             chat_id=chat_id,
+            thread_id=message_thread_id,
             message_id=message_id,
         )
         payload: dict[str, Any] = {"chat_id": chat_id, "message_id": message_id}
@@ -1495,6 +1463,9 @@ class TelegramBotClient:
         self,
         callback_query_id: str,
         *,
+        chat_id: Optional[int] = None,
+        thread_id: Optional[int] = None,
+        message_id: Optional[int] = None,
         text: Optional[str] = None,
         show_alert: bool = False,
     ) -> dict[str, Any]:
@@ -1503,6 +1474,9 @@ class TelegramBotClient:
             logging.INFO,
             "telegram.answer_callback",
             callback_query_id=callback_query_id,
+            chat_id=chat_id,
+            thread_id=thread_id,
+            message_id=message_id,
             text_len=len(text) if text else 0,
             show_alert=show_alert,
         )