refactor: tighten runtime lifecycle and remove maintenance debt

This commit is contained in:
Kazuhiro Sera
2026-07-01 13:47:14 +09:00
parent 655c805c92
commit 5c87f14fbd
34 changed files with 642 additions and 620 deletions
+8 -1
View File
@@ -119,6 +119,11 @@ select = [
"I", # isort
"B", # flake8-bugbear
"C4", # flake8-comprehensions
"DTZ005", # datetime.now() without a timezone
"G004", # logging statement uses f-string
"RUF006", # unowned asyncio tasks
"RUF012", # mutable class attributes without ClassVar
"RUF100", # unused noqa directives
"UP", # pyupgrade
]
isort = { combine-as-imports = true, known-first-party = ["agents"] }
@@ -127,7 +132,9 @@ isort = { combine-as-imports = true, known-first-party = ["agents"] }
convention = "google"
[tool.ruff.lint.per-file-ignores]
"examples/**/*.py" = ["E501"]
"examples/**/*.py" = ["DTZ005", "E501", "G004", "RUF006", "RUF012", "RUF100"]
"examples/**/*.ipynb" = ["RUF100"]
"tests/**/*.py" = ["RUF006", "RUF012", "RUF100"]
[tool.mypy]
strict = true
@@ -532,7 +532,7 @@ def _validate_default_run_context_thread_id_suffix(value: str) -> str:
def _parse_tool_input(parameters_model: type[BaseModel], input_json: str) -> BaseModel:
try:
json_data = json.loads(input_json) if input_json else {}
except Exception as exc: # noqa: BLE001
except Exception as exc:
if _debug.DONT_LOG_TOOL_DATA:
logger.debug("Invalid JSON input for codex tool")
else:
@@ -933,7 +933,7 @@ def _store_thread_id_in_run_context(
try:
setattr(context, key, thread_id)
except Exception as exc: # noqa: BLE001
except Exception as exc:
raise UserError(
f'Unable to store Codex thread_id in run context field "{key}". '
"Use a mutable dict context or set a writable attribute."
@@ -965,7 +965,7 @@ def _set_pydantic_context_value(context: BaseModel, key: str, value: str) -> boo
if key in model_fields:
try:
setattr(context, key, value)
except Exception: # noqa: BLE001
except Exception:
return False
return True
@@ -974,7 +974,7 @@ def _set_pydantic_context_value(context: BaseModel, key: str, value: str) -> boo
return True
except ValueError:
pass
except Exception: # noqa: BLE001
except Exception:
return False
state = getattr(context, "__dict__", None)
@@ -150,7 +150,7 @@ class Thread:
) from exc
try:
parsed = _parse_event(item)
except Exception as exc: # noqa: BLE001
except Exception as exc:
raise RuntimeError(f"Failed to parse event: {item}") from exc
if isinstance(parsed, ThreadStartedEvent):
# Capture the thread id so callers can resume later.
@@ -268,7 +268,7 @@ class AdvancedSQLiteSession(SQLiteSession):
# Only update turn-level usage - session usage is aggregated on demand
await self._update_turn_usage_internal(current_turn, result.context_wrapper.usage)
except Exception as e:
self._logger.error(f"Failed to store usage for session {self.session_id}: {e}")
self._logger.error("Failed to store usage for session %s: %s", self.session_id, e)
def _get_next_turn_number(self, branch_id: str) -> int:
"""Get the next turn number for a specific branch.
@@ -491,7 +491,7 @@ class AdvancedSQLiteSession(SQLiteSession):
deleted_count = cursor.rowcount
if deleted_count:
self._logger.info(f"Cleaned up {deleted_count} orphaned messages")
self._logger.info("Cleaned up %s orphaned messages", deleted_count)
return deleted_count
def _classify_message_type(self, item: TResponseInputItem) -> str:
@@ -639,7 +639,11 @@ class AdvancedSQLiteSession(SQLiteSession):
self._current_branch_id = branch_name
self._logger.debug(
f"Created branch '{branch_name}' from turn {turn_number} ('{turn_content}') in '{old_branch}'" # noqa: E501
"Created branch '%s' from turn %s ('%s') in '%s'",
branch_name,
turn_number,
turn_content,
old_branch,
)
return branch_name
@@ -697,7 +701,7 @@ class AdvancedSQLiteSession(SQLiteSession):
old_branch = self._current_branch_id
self._current_branch_id = branch_id
self._logger.info(f"Switched from branch '{old_branch}' to '{branch_id}'")
self._logger.info("Switched from branch '%s' to '%s'", old_branch, branch_id)
async def delete_branch(self, branch_id: str, force: bool = False) -> None:
"""Delete a branch and all its associated data.
@@ -778,8 +782,11 @@ class AdvancedSQLiteSession(SQLiteSession):
)
self._logger.info(
f"Deleted branch '{branch_id}': {structure_deleted} message entries, "
f"{usage_deleted} usage entries, {orphaned_messages_deleted} orphaned messages"
"Deleted branch '%s': %s message entries, %s usage entries, %s orphaned messages",
branch_id,
structure_deleted,
usage_deleted,
orphaned_messages_deleted,
)
async def list_branches(self) -> list[dict[str, Any]]:
@@ -1305,7 +1312,7 @@ class AdvancedSQLiteSession(SQLiteSession):
try:
input_details_json = json.dumps(usage_data.input_tokens_details.__dict__)
except (TypeError, ValueError) as e:
self._logger.warning(f"Failed to serialize input tokens details: {e}")
self._logger.warning("Failed to serialize input tokens details: %s", e)
input_details_json = None
if (
@@ -1315,7 +1322,7 @@ class AdvancedSQLiteSession(SQLiteSession):
try:
output_details_json = json.dumps(usage_data.output_tokens_details.__dict__)
except (TypeError, ValueError) as e:
self._logger.warning(f"Failed to serialize output tokens details: {e}")
self._logger.warning("Failed to serialize output tokens details: %s", e)
output_details_json = None
with closing(conn.cursor()) as cursor:
@@ -35,7 +35,7 @@ import json
import threading
import weakref
from datetime import datetime, timezone
from typing import Any
from typing import Any, ClassVar
from ._optional_imports import raise_optional_dependency_error
@@ -97,8 +97,8 @@ class MongoDBSession(SessionABC):
# one across loops raises RuntimeError. create_index is idempotent, so
# we only need the threading lock to guard the boolean done flag — no
# async coordination is required.
_init_state: dict[int, dict[tuple[str, str, str], bool]] = {}
_init_guard: threading.Lock = threading.Lock()
_init_state: ClassVar[dict[int, dict[tuple[str, str, str], bool]]] = {}
_init_guard: ClassVar[threading.Lock] = threading.Lock()
session_settings: SessionSettings | None = None
+82
View File
@@ -0,0 +1,82 @@
from __future__ import annotations
from ...sandbox.entries.mounts.patterns import RcloneMountPattern
from ...sandbox.errors import MountConfigError
from ...sandbox.session.base_sandbox_session import BaseSandboxSession
_APT = "DEBIAN_FRONTEND=noninteractive DEBCONF_NOWARNINGS=yes apt-get -o Dpkg::Use-Pty=0"
_RCLONE_CHECK = "command -v rclone >/dev/null 2>&1 || test -x /usr/local/bin/rclone"
_INSTALL_RCLONE_COMMANDS = (
f"{_APT} update -qq",
f"{_APT} install -y -qq curl unzip ca-certificates",
"curl -fsSL https://rclone.org/install.sh | bash",
)
async def ensure_rclone(session: BaseSandboxSession) -> None:
rclone = await session.exec("sh", "-lc", _RCLONE_CHECK, shell=False)
if rclone.ok():
return
apt = await session.exec("sh", "-lc", "command -v apt-get >/dev/null 2>&1", shell=False)
if not apt.ok():
raise MountConfigError(
message="rclone is not installed and apt-get is unavailable; preinstall rclone",
context={"package": "rclone"},
)
for command in _INSTALL_RCLONE_COMMANDS:
install = await session.exec(
"sh",
"-lc",
command,
shell=False,
timeout=300,
user="root",
)
if not install.ok():
raise MountConfigError(
message="failed to install rclone",
context={"package": "rclone", "exit_code": install.exit_code},
)
rclone = await session.exec("sh", "-lc", _RCLONE_CHECK, shell=False)
if not rclone.ok():
raise MountConfigError(
message="rclone was installed but is still not available on PATH",
context={"package": "rclone"},
)
async def _default_user_ids(session: BaseSandboxSession) -> tuple[str, str] | None:
result = await session.exec("sh", "-lc", "id -u; id -g", shell=False, timeout=30)
if not result.ok():
return None
lines = result.stdout.decode("utf-8", errors="replace").splitlines()
if len(lines) < 2 or not lines[0].isdigit() or not lines[1].isdigit():
return None
return lines[0], lines[1]
def _append_option(args: list[str], option: str, *values: str) -> None:
if option not in args:
args.extend([option, *values])
async def rclone_pattern_for_session(
session: BaseSandboxSession,
pattern: RcloneMountPattern,
) -> RcloneMountPattern:
if pattern.mode != "fuse":
return pattern
extra_args = list(pattern.extra_args)
_append_option(extra_args, "--allow-other")
user_ids = await _default_user_ids(session)
if user_ids is not None:
uid, gid = user_ids
_append_option(extra_args, "--uid", uid)
_append_option(extra_args, "--gid", gid)
return pattern.model_copy(update={"extra_args": extra_args})
@@ -44,6 +44,7 @@ from ....sandbox.session import SandboxSession, SandboxSessionState
from ....sandbox.session.base_sandbox_session import BaseSandboxSession
from ....sandbox.session.dependencies import Dependencies
from ....sandbox.session.manager import Instrumentation
from ....sandbox.session.pty_output import collect_pty_output
from ....sandbox.session.pty_types import (
PTY_PROCESSES_MAX,
PTY_PROCESSES_WARNING,
@@ -52,7 +53,6 @@ from ....sandbox.session.pty_types import (
clamp_pty_yield_time_ms,
process_id_to_prune_from_meta,
resolve_pty_write_yield_time_ms,
truncate_text_by_tokens,
)
from ....sandbox.session.runtime_helpers import RESOLVE_WORKSPACE_PATH_HELPER, RuntimeHelperScript
from ....sandbox.session.sandbox_client import BaseSandboxClient
@@ -959,34 +959,14 @@ class BlaxelSandboxSession(BaseSandboxSession):
yield_time_ms: int,
max_output_tokens: int | None,
) -> tuple[bytes, int | None]:
deadline = time.monotonic() + (yield_time_ms / 1000)
output = bytearray()
while True:
async with entry.output_lock:
while entry.output_chunks:
output.extend(entry.output_chunks.popleft())
if time.monotonic() >= deadline:
break
if entry.done:
async with entry.output_lock:
while entry.output_chunks:
output.extend(entry.output_chunks.popleft())
break
remaining_s = deadline - time.monotonic()
if remaining_s <= 0:
break
try:
await asyncio.wait_for(entry.output_notify.wait(), timeout=remaining_s)
except asyncio.TimeoutError:
break
entry.output_notify.clear()
text = output.decode("utf-8", errors="replace")
truncated, original_token_count = truncate_text_by_tokens(text, max_output_tokens)
return truncated.encode("utf-8", errors="replace"), original_token_count
return await collect_pty_output(
output_chunks=entry.output_chunks,
output_lock=entry.output_lock,
output_notify=entry.output_notify,
is_done=lambda: entry.done,
yield_time_ms=yield_time_ms,
max_output_tokens=max_output_tokens,
)
async def _finalize_pty_update(
self,
@@ -43,6 +43,7 @@ from ....sandbox.session import SandboxSession, SandboxSessionState
from ....sandbox.session.base_sandbox_session import BaseSandboxSession
from ....sandbox.session.dependencies import Dependencies
from ....sandbox.session.manager import Instrumentation
from ....sandbox.session.pty_output import collect_pty_output
from ....sandbox.session.pty_types import (
PTY_PROCESSES_MAX,
PTY_PROCESSES_WARNING,
@@ -51,7 +52,6 @@ from ....sandbox.session.pty_types import (
clamp_pty_yield_time_ms,
process_id_to_prune_from_meta,
resolve_pty_write_yield_time_ms,
truncate_text_by_tokens,
)
from ....sandbox.session.runtime_helpers import RESOLVE_WORKSPACE_PATH_HELPER, RuntimeHelperScript
from ....sandbox.session.sandbox_client import BaseSandboxClient, BaseSandboxClientOptions
@@ -388,6 +388,7 @@ class _DaytonaPtySessionEntry:
last_used: float = field(default_factory=time.monotonic)
done: bool = False
exit_code: int | None = None
worker_task: asyncio.Task[None] | None = None
class DaytonaSandboxSession(BaseSandboxSession):
@@ -672,7 +673,7 @@ class DaytonaSandboxSession(BaseSandboxSession):
timeout=exec_timeout,
)
entry.pty_handle = pty_handle
asyncio.create_task(self._run_pty_waiter(entry))
entry.worker_task = asyncio.create_task(self._run_pty_waiter(entry))
await asyncio.wait_for(pty_handle.wait_for_connection(), timeout=exec_timeout)
await asyncio.wait_for(
pty_handle.send_input(cmd_str + "\n"),
@@ -699,7 +700,7 @@ class DaytonaSandboxSession(BaseSandboxSession):
timeout=exec_timeout,
)
entry.cmd_id = resp.cmd_id
asyncio.create_task(
entry.worker_task = asyncio.create_task(
self._run_session_reader(
entry,
daytona_session_id,
@@ -885,36 +886,14 @@ class DaytonaSandboxSession(BaseSandboxSession):
yield_time_ms: int,
max_output_tokens: int | None,
) -> tuple[bytes, int | None]:
deadline = time.monotonic() + (yield_time_ms / 1000)
output = bytearray()
while True:
async with entry.output_lock:
while entry.output_chunks:
output.extend(entry.output_chunks.popleft())
if time.monotonic() >= deadline:
break
if entry.done:
async with entry.output_lock:
while entry.output_chunks:
output.extend(entry.output_chunks.popleft())
break
remaining_s = deadline - time.monotonic()
if remaining_s <= 0:
break
try:
await asyncio.wait_for(entry.output_notify.wait(), timeout=remaining_s)
except asyncio.TimeoutError:
break
entry.output_notify.clear()
text = output.decode("utf-8", errors="replace")
truncated, original_token_count = truncate_text_by_tokens(text, max_output_tokens)
return truncated.encode("utf-8", errors="replace"), original_token_count
return await collect_pty_output(
output_chunks=entry.output_chunks,
output_lock=entry.output_lock,
output_notify=entry.output_notify,
is_done=lambda: entry.done,
yield_time_ms=yield_time_ms,
max_output_tokens=max_output_tokens,
)
def _prune_pty_sessions_if_needed(self) -> _DaytonaPtySessionEntry | None:
if len(self._pty_sessions) < PTY_PROCESSES_MAX:
@@ -937,6 +916,13 @@ class DaytonaSandboxSession(BaseSandboxSession):
except Exception:
pass
worker_task = entry.worker_task
entry.worker_task = None
if worker_task is not None and worker_task is not asyncio.current_task():
if not worker_task.done():
worker_task.cancel()
await asyncio.gather(worker_task, return_exceptions=True)
async def read(self, path: Path | str, *, user: str | User | None = None) -> io.IOBase:
error_path = posix_path_as_path(coerce_posix_path(path))
if user is not None:
+4 -69
View File
@@ -10,14 +10,11 @@ from ....sandbox.entries.mounts.patterns import RcloneMountPattern
from ....sandbox.errors import MountConfigError
from ....sandbox.materialization import MaterializedFile
from ....sandbox.session.base_sandbox_session import BaseSandboxSession
_APT = "DEBIAN_FRONTEND=noninteractive DEBCONF_NOWARNINGS=yes apt-get -o Dpkg::Use-Pty=0"
_RCLONE_CHECK = "command -v rclone >/dev/null 2>&1 || test -x /usr/local/bin/rclone"
_INSTALL_RCLONE_COMMANDS = (
f"{_APT} update -qq",
f"{_APT} install -y -qq curl unzip ca-certificates",
"curl -fsSL https://rclone.org/install.sh | bash",
from .._rclone import (
ensure_rclone as _ensure_rclone,
rclone_pattern_for_session as _rclone_pattern_for_session,
)
_FUSE_ALLOW_OTHER = (
"chmod a+rw /dev/fuse && "
"touch /etc/fuse.conf && "
@@ -55,68 +52,6 @@ async def _ensure_fuse_support(session: BaseSandboxSession) -> None:
)
async def _ensure_rclone(session: BaseSandboxSession) -> None:
rclone = await session.exec("sh", "-lc", _RCLONE_CHECK, shell=False)
if rclone.ok():
return
apt = await session.exec("sh", "-lc", "command -v apt-get >/dev/null 2>&1", shell=False)
if not apt.ok():
raise MountConfigError(
message="rclone is not installed and apt-get is unavailable; preinstall rclone",
context={"package": "rclone"},
)
for command in _INSTALL_RCLONE_COMMANDS:
install = await session.exec("sh", "-lc", command, shell=False, timeout=300, user="root")
if not install.ok():
raise MountConfigError(
message="failed to install rclone",
context={"package": "rclone", "exit_code": install.exit_code},
)
rclone = await session.exec("sh", "-lc", _RCLONE_CHECK, shell=False)
if not rclone.ok():
raise MountConfigError(
message="rclone was installed but is still not available on PATH",
context={"package": "rclone"},
)
async def _default_user_ids(session: BaseSandboxSession) -> tuple[str, str] | None:
result = await session.exec("sh", "-lc", "id -u; id -g", shell=False, timeout=30)
if not result.ok():
return None
lines = result.stdout.decode("utf-8", errors="replace").splitlines()
if len(lines) < 2 or not lines[0].isdigit() or not lines[1].isdigit():
return None
return lines[0], lines[1]
def _append_option(args: list[str], option: str, *values: str) -> None:
if option not in args:
args.extend([option, *values])
async def _rclone_pattern_for_session(
session: BaseSandboxSession,
pattern: RcloneMountPattern,
) -> RcloneMountPattern:
if pattern.mode != "fuse":
return pattern
extra_args = list(pattern.extra_args)
_append_option(extra_args, "--allow-other")
user_ids = await _default_user_ids(session)
if user_ids is not None:
uid, gid = user_ids
_append_option(extra_args, "--uid", uid)
_append_option(extra_args, "--gid", gid)
return pattern.model_copy(update={"extra_args": extra_args})
def _assert_e2b_session(session: BaseSandboxSession) -> None:
if type(session).__name__ != "E2BSandboxSession":
raise MountConfigError(
@@ -380,15 +380,6 @@ async def _sandbox_write_file(
)
async def _sandbox_remove_file(
sandbox: object,
path: str,
*,
request_timeout: float | None = None,
) -> object:
return await _as_sandbox_api(sandbox).files.remove(path, request_timeout=request_timeout)
async def _sandbox_make_dir(
sandbox: object,
path: str,
@@ -10,14 +10,12 @@ from ....sandbox.entries.mounts.patterns import RcloneMountPattern
from ....sandbox.errors import MountConfigError
from ....sandbox.materialization import MaterializedFile
from ....sandbox.session.base_sandbox_session import BaseSandboxSession
from .._rclone import (
ensure_rclone as _ensure_rclone,
rclone_pattern_for_session as _rclone_pattern_for_session,
)
_APT = "DEBIAN_FRONTEND=noninteractive DEBCONF_NOWARNINGS=yes apt-get -o Dpkg::Use-Pty=0"
_RCLONE_CHECK = "command -v rclone >/dev/null 2>&1 || test -x /usr/local/bin/rclone"
_INSTALL_RCLONE_COMMANDS = (
f"{_APT} update -qq",
f"{_APT} install -y -qq curl unzip ca-certificates",
"curl -fsSL https://rclone.org/install.sh | bash",
)
_INSTALL_FUSE_COMMANDS = (
f"{_APT} update -qq",
f"{_APT} install -y -qq fuse3",
@@ -100,68 +98,6 @@ async def _ensure_fuse_support(session: BaseSandboxSession) -> None:
)
async def _ensure_rclone(session: BaseSandboxSession) -> None:
rclone = await session.exec("sh", "-lc", _RCLONE_CHECK, shell=False)
if rclone.ok():
return
apt = await session.exec("sh", "-lc", "command -v apt-get >/dev/null 2>&1", shell=False)
if not apt.ok():
raise MountConfigError(
message="rclone is not installed and apt-get is unavailable; preinstall rclone",
context={"package": "rclone"},
)
for command in _INSTALL_RCLONE_COMMANDS:
install = await session.exec("sh", "-lc", command, shell=False, timeout=300, user="root")
if not install.ok():
raise MountConfigError(
message="failed to install rclone",
context={"package": "rclone", "exit_code": install.exit_code},
)
rclone = await session.exec("sh", "-lc", _RCLONE_CHECK, shell=False)
if not rclone.ok():
raise MountConfigError(
message="rclone was installed but is still not available on PATH",
context={"package": "rclone"},
)
async def _default_user_ids(session: BaseSandboxSession) -> tuple[str, str] | None:
result = await session.exec("sh", "-lc", "id -u; id -g", shell=False, timeout=30)
if not result.ok():
return None
lines = result.stdout.decode("utf-8", errors="replace").splitlines()
if len(lines) < 2 or not lines[0].isdigit() or not lines[1].isdigit():
return None
return lines[0], lines[1]
def _append_option(args: list[str], option: str, *values: str) -> None:
if option not in args:
args.extend([option, *values])
async def _rclone_pattern_for_session(
session: BaseSandboxSession,
pattern: RcloneMountPattern,
) -> RcloneMountPattern:
if pattern.mode != "fuse":
return pattern
extra_args = list(pattern.extra_args)
_append_option(extra_args, "--allow-other")
user_ids = await _default_user_ids(session)
if user_ids is not None:
uid, gid = user_ids
_append_option(extra_args, "--uid", uid)
_append_option(extra_args, "--gid", gid)
return pattern.model_copy(update={"extra_args": extra_args})
def _assert_runloop_session(session: BaseSandboxSession) -> None:
if type(session).__name__ != "RunloopSandboxSession":
raise MountConfigError(
+3 -2
View File
@@ -148,8 +148,9 @@ class ToolOutputTrimmer:
if trimmed_count > 0:
logger.debug(
f"ToolOutputTrimmer: trimmed {trimmed_count} tool output(s), "
f"saved ~{chars_saved} chars"
"ToolOutputTrimmer: trimmed %s tool output(s), saved ~%s chars",
trimmed_count,
chars_saved,
)
return _ModelInputData(input=new_items, instructions=model_data.instructions)
@@ -196,14 +196,18 @@ class OpenAIResponsesCompactionSession(SessionABC, OpenAIResponsesCompactionAwar
if not should_compact:
logger.debug(
f"skip: decision hook declined compaction for {self._response_id} "
f"(mode={resolved_mode})"
"skip: decision hook declined compaction for %s (mode=%s)",
self._response_id,
resolved_mode,
)
return
self._deferred_response_id = None
logger.debug(
f"compact: start for {self._response_id} using {self.model} (mode={resolved_mode})"
"compact: start for %s using %s (mode=%s)",
self._response_id,
self.model,
resolved_mode,
)
compact_kwargs: dict[str, Any] = {"model": self.model}
@@ -228,9 +232,11 @@ class OpenAIResponsesCompactionSession(SessionABC, OpenAIResponsesCompactionAwar
self._session_items = output_items
logger.debug(
f"compact: done for {self._response_id} "
f"(mode={resolved_mode}, output={len(output_items)}, "
f"candidates={len(self._compaction_candidate_items)})"
"compact: done for %s (mode=%s, output=%s, candidates=%s)",
self._response_id,
resolved_mode,
len(output_items),
len(self._compaction_candidate_items),
)
async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
@@ -367,7 +373,9 @@ class OpenAIResponsesCompactionSession(SessionABC, OpenAIResponsesCompactionAwar
self._session_items = history
logger.debug(
f"candidates: initialized (history={len(history)}, candidates={len(candidates)})"
"candidates: initialized (history=%s, candidates=%s)",
len(history),
len(candidates),
)
return (candidates[:], history[:])
+10 -117
View File
@@ -1,121 +1,16 @@
from __future__ import annotations
import time
from collections.abc import Iterator, Mapping
from email.utils import parsedate_to_datetime
from typing import Any
import httpx
from openai import APIConnectionError, APIStatusError, APITimeoutError
from openai import APIConnectionError, APITimeoutError
from ..retry import ModelRetryAdvice, ModelRetryAdviceRequest, ModelRetryNormalizedError
def _iter_error_chain(error: Exception) -> Iterator[Exception]:
current: Exception | None = error
seen: set[int] = set()
while current is not None and id(current) not in seen:
seen.add(id(current))
yield current
next_error = current.__cause__ or current.__context__
current = next_error if isinstance(next_error, Exception) else None
def _header_lookup(headers: Any, key: str) -> str | None:
normalized_key = key.lower()
if isinstance(headers, httpx.Headers):
value = headers.get(key)
return value if isinstance(value, str) else None
if isinstance(headers, Mapping):
for header_name, header_value in headers.items():
if str(header_name).lower() == normalized_key and isinstance(header_value, str):
return header_value
return None
def _get_header_value(error: Exception, key: str) -> str | None:
for candidate in _iter_error_chain(error):
response = getattr(candidate, "response", None)
if isinstance(response, httpx.Response):
header_value = _header_lookup(response.headers, key)
if header_value is not None:
return header_value
for attr_name in ("headers", "response_headers"):
header_value = _header_lookup(getattr(candidate, attr_name, None), key)
if header_value is not None:
return header_value
return None
def _parse_retry_after_ms(value: str | None) -> float | None:
if value is None:
return None
try:
parsed = float(value) / 1000.0
except ValueError:
return None
return parsed if parsed >= 0 else None
def _parse_retry_after(value: str | None) -> float | None:
if value is None:
return None
try:
parsed = float(value)
except ValueError:
parsed = None
if parsed is not None:
return parsed if parsed >= 0 else None
try:
retry_datetime = parsedate_to_datetime(value)
except (TypeError, ValueError, IndexError):
return None
return max(retry_datetime.timestamp() - time.time(), 0.0)
def _get_status_code(error: Exception) -> int | None:
for candidate in _iter_error_chain(error):
if isinstance(candidate, APIStatusError):
return candidate.status_code
status_code = getattr(candidate, "status_code", None)
if isinstance(status_code, int):
return status_code
status = getattr(candidate, "status", None)
if isinstance(status, int):
return status
return None
def _get_request_id(error: Exception) -> str | None:
for candidate in _iter_error_chain(error):
request_id = getattr(candidate, "request_id", None)
if isinstance(request_id, str):
return request_id
return None
def _get_error_code(error: Exception) -> str | None:
for candidate in _iter_error_chain(error):
error_code = getattr(candidate, "code", None)
if isinstance(error_code, str):
return error_code
body = getattr(candidate, "body", None)
if isinstance(body, Mapping):
nested_error = body.get("error")
if isinstance(nested_error, Mapping):
nested_code = nested_error.get("code")
if isinstance(nested_code, str):
return nested_code
body_code = body.get("code")
if isinstance(body_code, str):
return body_code
return None
from ._retry_runtime import (
get_error_code as _get_error_code,
get_error_header as _get_header_value,
get_request_id as _get_request_id,
get_retry_after,
get_status_code as _get_status_code,
iter_error_chain as _iter_error_chain,
)
def _is_stateful_request(request: ModelRetryAdviceRequest) -> bool:
@@ -163,9 +58,7 @@ def get_openai_retry_advice(request: ModelRetryAdviceRequest) -> ModelRetryAdvic
reason=str(error),
)
retry_after = _parse_retry_after_ms(_get_header_value(error, "retry-after-ms"))
if retry_after is None:
retry_after = _parse_retry_after(_get_header_value(error, "retry-after"))
retry_after = get_retry_after(error)
normalized = _build_normalized_error(error, retry_after=retry_after)
stateful_request = _is_stateful_request(request)
+117 -1
View File
@@ -1,8 +1,14 @@
from __future__ import annotations
from collections.abc import Iterator
import time
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from contextvars import ContextVar
from email.utils import parsedate_to_datetime
from typing import Any
import httpx
from openai import APIStatusError
_DISABLE_PROVIDER_MANAGED_RETRIES: ContextVar[bool] = ContextVar(
"disable_provider_managed_retries",
@@ -38,3 +44,113 @@ def websocket_pre_event_retries_disabled(disabled: bool) -> Iterator[None]:
def should_disable_websocket_pre_event_retries() -> bool:
return _DISABLE_WEBSOCKET_PRE_EVENT_RETRIES.get()
def iter_error_chain(error: Exception) -> Iterator[Exception]:
current: Exception | None = error
seen: set[int] = set()
while current is not None and id(current) not in seen:
seen.add(id(current))
yield current
next_error = current.__cause__ or current.__context__
current = next_error if isinstance(next_error, Exception) else None
def header_lookup(headers: Any, key: str) -> str | None:
normalized_key = key.lower()
if isinstance(headers, httpx.Headers):
value = headers.get(key)
return value if isinstance(value, str) else None
if isinstance(headers, Mapping):
for header_name, header_value in headers.items():
if str(header_name).lower() == normalized_key and isinstance(header_value, str):
return header_value
return None
def get_error_header(error: Exception, key: str) -> str | None:
for candidate in iter_error_chain(error):
response = getattr(candidate, "response", None)
if isinstance(response, httpx.Response):
header_value = header_lookup(response.headers, key)
if header_value is not None:
return header_value
for attr_name in ("headers", "response_headers"):
header_value = header_lookup(getattr(candidate, attr_name, None), key)
if header_value is not None:
return header_value
return None
def parse_retry_after_ms(value: str | None) -> float | None:
if value is None:
return None
try:
parsed = float(value) / 1000.0
except ValueError:
return None
return parsed if parsed >= 0 else None
def parse_retry_after_value(value: str | None) -> float | None:
if value is None:
return None
try:
parsed = float(value)
except ValueError:
parsed = None
if parsed is not None:
return parsed if parsed >= 0 else None
try:
retry_datetime = parsedate_to_datetime(value)
except (TypeError, ValueError, IndexError):
return None
return max(retry_datetime.timestamp() - time.time(), 0.0)
def get_retry_after(error: Exception) -> float | None:
retry_after = parse_retry_after_ms(get_error_header(error, "retry-after-ms"))
if retry_after is not None:
return retry_after
return parse_retry_after_value(get_error_header(error, "retry-after"))
def get_status_code(error: Exception) -> int | None:
for candidate in iter_error_chain(error):
if isinstance(candidate, APIStatusError):
return candidate.status_code
for attr_name in ("status_code", "status"):
value = getattr(candidate, attr_name, None)
if isinstance(value, int):
return value
return None
def get_request_id(error: Exception) -> str | None:
for candidate in iter_error_chain(error):
request_id = getattr(candidate, "request_id", None)
if isinstance(request_id, str):
return request_id
return None
def get_error_code(error: Exception) -> str | None:
for candidate in iter_error_chain(error):
error_code = getattr(candidate, "code", None)
if isinstance(error_code, str):
return error_code
body = getattr(candidate, "body", None)
if isinstance(body, Mapping):
nested_error = body.get("error")
if isinstance(nested_error, Mapping):
nested_code = nested_error.get("code")
if isinstance(nested_code, str):
return nested_code
body_code = body.get("code")
if isinstance(body_code, str):
return body_code
return None
+2 -9
View File
@@ -9,7 +9,7 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequenc
from contextvars import ContextVar
from dataclasses import asdict, dataclass, is_dataclass
from enum import Enum
from typing import TYPE_CHECKING, Any, Literal, TypedDict, TypeGuard, cast, get_args, overload
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast, overload
import httpx
from openai import AsyncOpenAI, NotGiven, Omit, omit
@@ -88,9 +88,6 @@ _HEADERS = {"User-Agent": _USER_AGENT}
_HEADERS_OVERRIDE: ContextVar[dict[str, str] | None] = ContextVar(
"openai_responses_headers_override", default=None
)
_RESPONSE_INCLUDABLE_VALUES = frozenset(
value for value in get_args(ResponseIncludable) if isinstance(value, str)
)
class _NamespaceToolParam(TypedDict):
@@ -132,10 +129,6 @@ def _require_responses_tool_param(value: object) -> ResponsesToolParam:
return cast(ResponsesToolParam, value)
def _is_response_includable(value: object) -> TypeGuard[ResponseIncludable]:
return isinstance(value, str) and value in _RESPONSE_INCLUDABLE_VALUES
def _coerce_response_includables(values: Sequence[str]) -> list[ResponseIncludable]:
includables: list[ResponseIncludable] = []
for value in values:
@@ -221,7 +214,7 @@ class OpenAIResponsesWebSocketOptions(TypedDict):
class _ResponseStreamWithRequestId:
"""Wrap an SDK event stream and retain the originating request ID."""
_TERMINAL_EVENT_TYPES = {
_TERMINAL_EVENT_TYPES: ClassVar[set[str]] = {
"response.completed",
"response.failed",
"response.incomplete",
+3 -3
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import time
from dataclasses import dataclass
from datetime import datetime
from ._util import calculate_audio_length_ms
from .config import RealtimeAudioFormat
@@ -9,7 +9,7 @@ from .config import RealtimeAudioFormat
@dataclass
class ModelAudioState:
initial_received_time: datetime
initial_received_time: float
audio_length_ms: float
@@ -35,7 +35,7 @@ class ModelAudioTracker:
self._last_audio_item = new_key
if new_key not in self._states:
self._states[new_key] = ModelAudioState(datetime.now(), ms)
self._states[new_key] = ModelAudioState(time.monotonic(), ms)
else:
self._states[new_key].audio_length_ms += ms
+2 -4
View File
@@ -6,9 +6,9 @@ import inspect
import json
import math
import os
import time
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import Annotated, Any, Literal, TypeAlias, cast
import pydantic
@@ -890,9 +890,7 @@ class OpenAIRealtimeWebSocketModel(RealtimeModel):
item_id, item_content_index = last_audio_item_id
audio_state = self._audio_state_tracker.get_state(item_id, item_content_index)
if audio_state:
elapsed_ms = (
datetime.now() - audio_state.initial_received_time
).total_seconds() * 1000
elapsed_ms = (time.monotonic() - audio_state.initial_received_time) * 1000
return {
"current_item_id": item_id,
"current_item_content_index": item_content_index,
+34 -37
View File
@@ -499,6 +499,10 @@ class RealtimeSession(RealtimeModelListener):
"""Put an event into the queue."""
await self._event_queue.put(event)
def _put_event_nowait(self, event: RealtimeSessionEvent) -> None:
"""Put an event into the unbounded queue from a synchronous callback."""
self._event_queue.put_nowait(event)
async def _function_needs_approval(
self, function_tool: FunctionTool, tool_call: RealtimeModelToolCallEvent
) -> bool:
@@ -1282,20 +1286,23 @@ class RealtimeSession(RealtimeModelListener):
exception = task.exception()
if exception:
# Create an exception event instead of raising
asyncio.create_task(
self._put_event(
RealtimeError(
info=self._event_info,
error={"message": f"Guardrail task failed: {str(exception)}"},
)
self._put_event_nowait(
RealtimeError(
info=self._event_info,
error={"message": f"Guardrail task failed: {str(exception)}"},
)
)
def _cleanup_guardrail_tasks(self) -> None:
for task in self._guardrail_tasks:
if not task.done():
task.cancel()
self._guardrail_tasks.clear()
@staticmethod
async def _cancel_and_wait_for_tasks(tasks: set[asyncio.Task[Any]]) -> None:
observed: set[asyncio.Task[Any]] = set()
while new_tasks := tuple(task for task in tasks if task not in observed):
observed.update(new_tasks)
for task in new_tasks:
if not task.done():
task.cancel()
await asyncio.gather(*new_tasks, return_exceptions=True)
tasks.difference_update(observed)
def _enqueue_tool_call_task(
self,
@@ -1335,17 +1342,14 @@ class RealtimeSession(RealtimeModelListener):
exception.call_id,
exc_info=exception,
)
asyncio.create_task(
self._put_event(
RealtimeError(
info=self._event_info,
error={
"message": (
"Tool output send failed; cached output will be retried: "
f"{exception}"
)
},
)
self._put_event_nowait(
RealtimeError(
info=self._event_info,
error={
"message": (
f"Tool output send failed; cached output will be retried: {exception}"
)
},
)
)
return
@@ -1355,21 +1359,13 @@ class RealtimeSession(RealtimeModelListener):
if self._stored_exception is None:
self._stored_exception = exception
asyncio.create_task(
self._put_event(
RealtimeError(
info=self._event_info,
error={"message": f"Tool call task failed: {exception}"},
)
self._put_event_nowait(
RealtimeError(
info=self._event_info,
error={"message": f"Tool call task failed: {exception}"},
)
)
def _cleanup_tool_call_tasks(self) -> None:
for task in self._tool_call_tasks:
if not task.done():
task.cancel()
self._tool_call_tasks.clear()
def _wake_event_iterators(self) -> None:
for _ in range(self._event_iterator_waiters):
self._event_queue.put_nowait(_REALTIME_SESSION_CLOSED_SENTINEL)
@@ -1380,9 +1376,10 @@ class RealtimeSession(RealtimeModelListener):
self._wake_event_iterators()
return
# Cancel and cleanup guardrail tasks
self._cleanup_guardrail_tasks()
self._cleanup_tool_call_tasks()
# Cancel each observed task once, await its finalizer, then rescan for tasks
# created while the observed batch was unwinding.
await self._cancel_and_wait_for_tasks(self._guardrail_tasks)
await self._cancel_and_wait_for_tasks(self._tool_call_tasks)
# Remove ourselves as a listener
self._model.remove_listener(self)
+2 -2
View File
@@ -145,8 +145,8 @@ def _mark_retry_capabilities(
retries_safe_transport_errors: bool,
retries_all_transient_errors: bool,
) -> RetryPolicy:
setattr(policy, _RETRIES_SAFE_TRANSPORT_ERRORS_ATTR, retries_safe_transport_errors) # noqa: B010
setattr(policy, _RETRIES_ALL_TRANSIENT_ERRORS_ATTR, retries_all_transient_errors) # noqa: B010
setattr(policy, _RETRIES_SAFE_TRANSPORT_ERRORS_ATTR, retries_safe_transport_errors)
setattr(policy, _RETRIES_ALL_TRANSIENT_ERRORS_ATTR, retries_all_transient_errors)
return policy
+8 -113
View File
@@ -2,18 +2,21 @@ from __future__ import annotations
import asyncio
import random
import time
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping
from email.utils import parsedate_to_datetime
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping
from inspect import isawaitable
from typing import Any
import httpx
from openai import APIConnectionError, APIStatusError, APITimeoutError, BadRequestError
from openai import APIConnectionError, APITimeoutError, BadRequestError
from ..items import ModelResponse, TResponseStreamEvent
from ..logger import logger
from ..models._retry_runtime import (
get_error_code as _get_error_code,
get_request_id as _get_request_id,
get_retry_after as _get_retry_after,
get_status_code as _get_status_code,
iter_error_chain as _iter_error_chain,
provider_managed_retries_disabled,
websocket_pre_event_retries_disabled,
)
@@ -44,120 +47,12 @@ COMPATIBILITY_CONVERSATION_LOCKED_RETRIES = 3
_RETRY_SAFE_STREAM_EVENT_TYPES = frozenset({"response.created", "response.in_progress"})
def _iter_error_chain(error: Exception) -> Iterator[Exception]:
current: Exception | None = error
seen: set[int] = set()
while current is not None and id(current) not in seen:
seen.add(id(current))
yield current
next_error = current.__cause__ or current.__context__
current = next_error if isinstance(next_error, Exception) else None
def _is_conversation_locked_error(error: Exception) -> bool:
return (
isinstance(error, BadRequestError) and getattr(error, "code", "") == "conversation_locked"
)
def _get_header_value(headers: Any, key: str) -> str | None:
normalized_key = key.lower()
if isinstance(headers, httpx.Headers):
value = headers.get(key)
return value if isinstance(value, str) else None
if isinstance(headers, Mapping):
for header_name, header_value in headers.items():
if str(header_name).lower() == normalized_key and isinstance(header_value, str):
return header_value
return None
def _extract_headers(error: Exception) -> httpx.Headers | Mapping[str, str] | None:
for candidate in _iter_error_chain(error):
response = getattr(candidate, "response", None)
if isinstance(response, httpx.Response):
return response.headers
for attr_name in ("headers", "response_headers"):
headers = getattr(candidate, attr_name, None)
if isinstance(headers, httpx.Headers | Mapping):
return headers
return None
def _parse_retry_after(headers: httpx.Headers | Mapping[str, str] | None) -> float | None:
if headers is None:
return None
retry_after_ms = _get_header_value(headers, "retry-after-ms")
if retry_after_ms is not None:
try:
parsed_ms = float(retry_after_ms) / 1000.0
except ValueError:
parsed_ms = None
if parsed_ms is not None and parsed_ms >= 0:
return parsed_ms
retry_after = _get_header_value(headers, "retry-after")
if retry_after is None:
return None
try:
parsed_seconds = float(retry_after)
except ValueError:
parsed_seconds = None
if parsed_seconds is not None:
return parsed_seconds if parsed_seconds >= 0 else None
try:
retry_datetime = parsedate_to_datetime(retry_after)
except (TypeError, ValueError, IndexError):
return None
return max(retry_datetime.timestamp() - time.time(), 0.0)
def _get_status_code(error: Exception) -> int | None:
for candidate in _iter_error_chain(error):
if isinstance(candidate, APIStatusError):
return candidate.status_code
for attr_name in ("status_code", "status"):
value = getattr(candidate, attr_name, None)
if isinstance(value, int):
return value
return None
def _get_error_code(error: Exception) -> str | None:
for candidate in _iter_error_chain(error):
error_code = getattr(candidate, "code", None)
if isinstance(error_code, str):
return error_code
body = getattr(candidate, "body", None)
if isinstance(body, Mapping):
nested_error = body.get("error")
if isinstance(nested_error, Mapping):
nested_code = nested_error.get("code")
if isinstance(nested_code, str):
return nested_code
body_code = body.get("code")
if isinstance(body_code, str):
return body_code
return None
def _get_request_id(error: Exception) -> str | None:
for candidate in _iter_error_chain(error):
request_id = getattr(candidate, "request_id", None)
if isinstance(request_id, str):
return request_id
return None
def _is_abort_like_error(error: Exception) -> bool:
if isinstance(error, asyncio.CancelledError):
return True
@@ -211,7 +106,7 @@ def _normalize_retry_error(
error_code=_get_error_code(error),
message=str(error),
request_id=_get_request_id(error),
retry_after=_parse_retry_after(_extract_headers(error)),
retry_after=_get_retry_after(error),
is_abort=_is_abort_like_error(error),
is_network_error=_is_network_like_error(error),
is_timeout=any(
+24 -32
View File
@@ -49,6 +49,7 @@ from ..session import SandboxSession, SandboxSessionState
from ..session.base_sandbox_session import BaseSandboxSession
from ..session.dependencies import Dependencies
from ..session.manager import Instrumentation
from ..session.pty_output import collect_pty_output
from ..session.pty_types import (
PTY_PROCESSES_MAX,
PTY_PROCESSES_WARNING,
@@ -57,7 +58,6 @@ from ..session.pty_types import (
clamp_pty_yield_time_ms,
process_id_to_prune_from_meta,
resolve_pty_write_yield_time_ms,
truncate_text_by_tokens,
)
from ..session.runtime_helpers import RESOLVE_WORKSPACE_PATH_HELPER, RuntimeHelperScript
from ..session.sandbox_client import BaseSandboxClient, BaseSandboxClientOptions
@@ -164,6 +164,7 @@ class DockerSandboxSession(BaseSandboxSession):
_pty_lock: asyncio.Lock
_pty_processes: dict[int, _DockerPtyProcessEntry]
_reserved_pty_process_ids: set[int]
_cleanup_tasks: set[asyncio.Task[None]]
state: DockerSandboxSessionState
_ARCHIVE_STAGING_DIR: Path = posix_path_as_path(
@@ -185,6 +186,7 @@ class DockerSandboxSession(BaseSandboxSession):
self._pty_lock = asyncio.Lock()
self._pty_processes = {}
self._reserved_pty_process_ids = set()
self._cleanup_tasks = set()
@classmethod
def from_state(
@@ -764,6 +766,12 @@ class DockerSandboxSession(BaseSandboxSession):
# If the container is already gone/stopped, ignore.
pass
async def _after_stop(self) -> None:
await self._wait_for_cleanup_tasks()
async def _after_shutdown(self) -> None:
await self._wait_for_cleanup_tasks()
@staticmethod
def _start_exec_socket(*, api: Any, exec_id: str, tty: bool = False) -> _DockerExecSocket:
if not all(
@@ -1068,36 +1076,14 @@ class DockerSandboxSession(BaseSandboxSession):
yield_time_ms: int,
max_output_tokens: int | None,
) -> tuple[bytes, int | None]:
deadline = time.monotonic() + (yield_time_ms / 1000)
output = bytearray()
while True:
async with entry.output_lock:
while entry.output_chunks:
output.extend(entry.output_chunks.popleft())
if time.monotonic() >= deadline:
break
if entry.output_closed.is_set():
async with entry.output_lock:
while entry.output_chunks:
output.extend(entry.output_chunks.popleft())
break
remaining_s = deadline - time.monotonic()
if remaining_s <= 0:
break
try:
await asyncio.wait_for(entry.output_notify.wait(), timeout=remaining_s)
except asyncio.TimeoutError:
break
entry.output_notify.clear()
text = output.decode("utf-8", errors="replace")
truncated_text, original_token_count = truncate_text_by_tokens(text, max_output_tokens)
return truncated_text.encode("utf-8", errors="replace"), original_token_count
return await collect_pty_output(
output_chunks=entry.output_chunks,
output_lock=entry.output_lock,
output_notify=entry.output_notify,
is_done=entry.output_closed.is_set,
yield_time_ms=yield_time_ms,
max_output_tokens=max_output_tokens,
)
async def _finalize_pty_update(
self,
@@ -1282,7 +1268,13 @@ class DockerSandboxSession(BaseSandboxSession):
def _schedule_rm_best_effort(self, path: Path) -> None:
loop = asyncio.get_running_loop()
loop.create_task(self._rm_best_effort(path))
task = loop.create_task(self._rm_best_effort(path))
self._cleanup_tasks.add(task)
task.add_done_callback(self._cleanup_tasks.discard)
async def _wait_for_cleanup_tasks(self) -> None:
while cleanup_tasks := tuple(self._cleanup_tasks):
await asyncio.gather(*cleanup_tasks, return_exceptions=True)
def _workspace_archive_stream(
self,
+29 -34
View File
@@ -45,6 +45,7 @@ from ..session import SandboxSession, SandboxSessionState
from ..session.base_sandbox_session import BaseSandboxSession
from ..session.dependencies import Dependencies
from ..session.manager import Instrumentation
from ..session.pty_output import collect_pty_output
from ..session.pty_types import (
PTY_PROCESSES_MAX,
PTY_PROCESSES_WARNING,
@@ -53,7 +54,6 @@ from ..session.pty_types import (
clamp_pty_yield_time_ms,
process_id_to_prune_from_meta,
resolve_pty_write_yield_time_ms,
truncate_text_by_tokens,
)
from ..session.sandbox_client import BaseSandboxClient, BaseSandboxClientOptions
from ..session.workspace_payloads import coerce_write_payload
@@ -70,6 +70,7 @@ _DEFAULT_WORKSPACE_PREFIX = "sandbox-local-"
_DEFAULT_MANIFEST_ROOT = cast(str, Manifest.model_fields["root"].default)
_PTY_READ_CHUNK_BYTES = 16_384
_PTY_CHILD_SIGNAL_DEFAULTS = (signal.SIGINT, signal.SIGQUIT)
_PTY_FD_CLOSE_GRACE_SECONDS = 0.1
logger = logging.getLogger(__name__)
@@ -130,6 +131,7 @@ class UnixLocalSandboxSession(BaseSandboxSession):
_pty_lock: asyncio.Lock
_pty_processes: dict[int, _UnixPtyProcessEntry]
_reserved_pty_process_ids: set[int]
_fd_close_tasks: set[asyncio.Task[None]]
def __init__(self, *, state: UnixLocalSandboxSessionState) -> None:
self.state = state
@@ -137,6 +139,7 @@ class UnixLocalSandboxSession(BaseSandboxSession):
self._pty_lock = asyncio.Lock()
self._pty_processes = {}
self._reserved_pty_process_ids = set()
self._fd_close_tasks = set()
@classmethod
def from_state(cls, state: UnixLocalSandboxSessionState) -> "UnixLocalSandboxSession":
@@ -192,10 +195,14 @@ class UnixLocalSandboxSession(BaseSandboxSession):
)
async def _after_shutdown(self) -> None:
await self._wait_for_fd_close_tasks()
# Best-effort: mark session not running. We intentionally do not delete the workspace
# directory here; cleanup is handled by the Client.delete().
self._running = False
async def _after_stop(self) -> None:
await self._wait_for_fd_close_tasks()
async def _resolve_exposed_port(self, port: int) -> ExposedPortEndpoint:
return ExposedPortEndpoint(host="127.0.0.1", port=port, tls=False)
@@ -476,36 +483,14 @@ class UnixLocalSandboxSession(BaseSandboxSession):
yield_time_ms: int,
max_output_tokens: int | None,
) -> tuple[bytes, int | None]:
deadline = time.monotonic() + (yield_time_ms / 1000)
output = bytearray()
while True:
async with entry.output_lock:
while entry.output_chunks:
output.extend(entry.output_chunks.popleft())
if time.monotonic() >= deadline:
break
if entry.output_closed.is_set():
async with entry.output_lock:
while entry.output_chunks:
output.extend(entry.output_chunks.popleft())
break
remaining_s = deadline - time.monotonic()
if remaining_s <= 0:
break
try:
await asyncio.wait_for(entry.output_notify.wait(), timeout=remaining_s)
except asyncio.TimeoutError:
break
entry.output_notify.clear()
text = output.decode("utf-8", errors="replace")
truncated_text, original_token_count = truncate_text_by_tokens(text, max_output_tokens)
return truncated_text.encode("utf-8", errors="replace"), original_token_count
return await collect_pty_output(
output_chunks=entry.output_chunks,
output_lock=entry.output_lock,
output_notify=entry.output_notify,
is_done=entry.output_closed.is_set,
yield_time_ms=yield_time_ms,
max_output_tokens=max_output_tokens,
)
async def _finalize_pty_update(
self,
@@ -564,9 +549,9 @@ class UnixLocalSandboxSession(BaseSandboxSession):
if entry.tty:
if primary_fd is not None:
# On macOS we have observed os.close() on the PTY master fd block while a
# background reader thread is still inside os.read(). Close it off-thread so
# session teardown remains best-effort and non-blocking.
asyncio.create_task(asyncio.to_thread(_close_fd_quietly, primary_fd))
# background reader thread is still inside os.read(). Keep the close task owned
# by the session without making PTY termination wait indefinitely for it.
self._schedule_fd_close(primary_fd)
entry.output_closed.set()
entry.output_notify.set()
return
@@ -577,6 +562,16 @@ class UnixLocalSandboxSession(BaseSandboxSession):
if entry.wait_task is not None:
await asyncio.gather(entry.wait_task, return_exceptions=True)
def _schedule_fd_close(self, fd: int) -> None:
task = asyncio.create_task(asyncio.to_thread(_close_fd_quietly, fd))
self._fd_close_tasks.add(task)
task.add_done_callback(self._fd_close_tasks.discard)
async def _wait_for_fd_close_tasks(self) -> None:
tasks = tuple(self._fd_close_tasks)
if tasks:
await asyncio.wait(tasks, timeout=_PTY_FD_CLOSE_GRACE_SECONDS)
def _confined_exec_command(
self,
*,
+50
View File
@@ -0,0 +1,50 @@
from __future__ import annotations
import asyncio
import time
from collections import deque
from collections.abc import Callable
from .pty_types import truncate_text_by_tokens
async def collect_pty_output(
*,
output_chunks: deque[bytes],
output_lock: asyncio.Lock,
output_notify: asyncio.Event,
is_done: Callable[[], bool],
yield_time_ms: int,
max_output_tokens: int | None,
) -> tuple[bytes, int | None]:
"""Collect and truncate PTY output until the deadline or provider completion."""
deadline = time.monotonic() + (yield_time_ms / 1000)
output = bytearray()
while True:
async with output_lock:
while output_chunks:
output.extend(output_chunks.popleft())
if time.monotonic() >= deadline:
break
if is_done():
async with output_lock:
while output_chunks:
output.extend(output_chunks.popleft())
break
remaining_s = deadline - time.monotonic()
if remaining_s <= 0:
break
try:
await asyncio.wait_for(output_notify.wait(), timeout=remaining_s)
except asyncio.TimeoutError:
break
output_notify.clear()
text = output.decode("utf-8", errors="replace")
truncated, original_token_count = truncate_text_by_tokens(text, max_output_tokens)
return truncated.encode("utf-8", errors="replace"), original_token_count
+37
View File
@@ -1526,6 +1526,43 @@ class TestDaytonaSandbox:
assert entry.done is False
assert entry.exit_code is None
@pytest.mark.asyncio
async def test_terminate_pty_entry_awaits_worker_finalizer(
self,
monkeypatch: pytest.MonkeyPatch,
) -> None:
daytona_module = _load_daytona_module(monkeypatch)
sandbox = _FakeDaytonaSandbox()
state = daytona_module.DaytonaSandboxSessionState(
manifest=Manifest(root=daytona_module.DEFAULT_DAYTONA_WORKSPACE_ROOT),
snapshot=NoopSnapshot(id="snapshot"),
sandbox_id=sandbox.id,
)
session = daytona_module.DaytonaSandboxSession.from_state(state, sandbox=sandbox)
entry = daytona_module._DaytonaPtySessionEntry( # noqa: SLF001
daytona_session_id="session-123",
pty_handle=object(),
tty=False,
cmd_id="cmd-123",
)
finalizer_finished = asyncio.Event()
async def worker() -> None:
try:
await asyncio.Event().wait()
finally:
await asyncio.sleep(0)
finalizer_finished.set()
entry.worker_task = asyncio.create_task(worker())
await asyncio.sleep(0)
await session._terminate_pty_entry(entry) # noqa: SLF001
assert finalizer_finished.is_set()
assert entry.worker_task is None
assert sandbox.process.delete_session_calls == ["session-123"]
# ---------------------------------------------------------------------------
# DaytonaCloudBucketMountStrategy tests
+4 -2
View File
@@ -16,12 +16,14 @@ import pytest
from pydantic import Field, PrivateAttr
import agents.extensions.sandbox.e2b.sandbox as e2b_module
from agents.extensions.sandbox._rclone import (
ensure_rclone as _ensure_rclone,
rclone_pattern_for_session as _rclone_pattern_for_session,
)
from agents.extensions.sandbox.e2b.mounts import (
E2BCloudBucketMountStrategy,
_assert_e2b_session,
_ensure_fuse_support,
_ensure_rclone,
_rclone_pattern_for_session,
)
from agents.extensions.sandbox.e2b.sandbox import (
E2BSandboxClient,
@@ -135,7 +135,7 @@ def test_runloop_session_guard_accepts_correct_type() -> None:
@pytest.mark.asyncio
async def test_runloop_ensure_rclone_installs_with_root_apt() -> None:
from agents.extensions.sandbox.runloop.mounts import _ensure_rclone
from agents.extensions.sandbox._rclone import ensure_rclone
session = _FakeRunloopMountSession(
[
@@ -147,7 +147,7 @@ async def test_runloop_ensure_rclone_installs_with_root_apt() -> None:
]
)
await _ensure_rclone(session)
await ensure_rclone(session)
assert session.exec_calls[:2] == [
"sh -lc command -v rclone >/dev/null 2>&1 || test -x /usr/local/bin/rclone",
@@ -215,10 +215,10 @@ async def test_runloop_ensure_fuse_installs_missing_fusermount() -> None:
@pytest.mark.asyncio
async def test_runloop_rclone_pattern_adds_fuse_access_args() -> None:
from agents.extensions.sandbox.runloop.mounts import _rclone_pattern_for_session
from agents.extensions.sandbox._rclone import rclone_pattern_for_session
session = _FakeRunloopMountSession([_exec_ok(stdout=b"1000\n1000\n")])
pattern = await _rclone_pattern_for_session(session, RcloneMountPattern(mode="fuse"))
pattern = await rclone_pattern_for_session(session, RcloneMountPattern(mode="fuse"))
assert pattern.extra_args == ["--allow-other", "--uid", "1000", "--gid", "1000"]
+28 -8
View File
@@ -12,16 +12,17 @@ from email.utils import format_datetime
import httpx
from agents.models._openai_retry import (
_get_error_code,
_get_header_value,
_get_status_code,
_header_lookup,
_parse_retry_after,
_parse_retry_after_ms,
get_openai_retry_advice,
from agents.models._openai_retry import get_openai_retry_advice
from agents.models._retry_runtime import (
get_error_code as _get_error_code,
get_error_header as _get_header_value,
get_status_code as _get_status_code,
header_lookup as _header_lookup,
parse_retry_after_ms as _parse_retry_after_ms,
parse_retry_after_value as _parse_retry_after,
)
from agents.retry import ModelRetryAdviceRequest
from agents.run_internal.model_retry import _normalize_retry_error
class _HeaderError(Exception):
@@ -99,6 +100,25 @@ def test_get_error_code_from_body_mapping() -> None:
assert _get_error_code(Exception("none")) is None
def test_provider_and_runner_retry_normalization_share_metadata() -> None:
class _RetryableError(Exception):
status_code = 429
request_id = "req_test"
body = {"error": {"code": "rate_limit_exceeded"}}
headers = {"retry-after-ms": "1500"}
error = _RetryableError("slow down")
advice = get_openai_retry_advice(_make_request(error))
runner_normalized = _normalize_retry_error(error, None)
assert advice is not None
assert advice.normalized is not None
assert advice.normalized.status_code == runner_normalized.status_code
assert advice.normalized.error_code == runner_normalized.error_code
assert advice.normalized.request_id == runner_normalized.request_id
assert advice.normalized.retry_after == runner_normalized.retry_after
def test_advice_unsafe_to_replay() -> None:
error = Exception("cannot replay")
error.unsafe_to_replay = True # type: ignore[attr-defined]
+3 -3
View File
@@ -1,6 +1,6 @@
import asyncio
import json
from datetime import datetime, timedelta
import time
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import AsyncMock, Mock, patch
@@ -815,7 +815,7 @@ class TestEventHandlingRobustness(TestOpenAIRealtimeWebSocketModel):
model._audio_state_tracker.on_audio_delta("i1", 0, b"a" * 48_000)
state = model._audio_state_tracker.get_state("i1", 0)
assert state is not None
state.initial_received_time = datetime.now() - timedelta(seconds=5)
state.initial_received_time = time.monotonic() - 5
monkeypatch.setattr(
model,
@@ -846,7 +846,7 @@ class TestEventHandlingRobustness(TestOpenAIRealtimeWebSocketModel):
model._audio_state_tracker.on_audio_delta("i1", 0, b"a" * 48_000)
state = model._audio_state_tracker.get_state("i1", 0)
assert state is not None
state.initial_received_time = datetime.now() - timedelta(seconds=5)
state.initial_received_time = time.monotonic() - 5
model._ongoing_response = True
monkeypatch.setattr(
+11 -1
View File
@@ -1,4 +1,4 @@
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, patch
import pytest
@@ -135,6 +135,16 @@ class TestPlaybackTracker:
expected_length = (8 / (24_000 * 2)) * 1000
assert state.audio_length_ms == pytest.approx(expected_length, rel=0, abs=1e-6)
def test_audio_state_uses_monotonic_timestamp(self):
tracker = ModelAudioTracker()
with patch("agents.realtime._default_tracker.time.monotonic", return_value=42.5):
tracker.on_audio_delta("item_1", 0, b"test")
state = tracker.get_state("item_1", 0)
assert state is not None
assert state.initial_received_time == 42.5
def test_state_cleanup_on_interruption(self):
"""Test both trackers properly reset state on interruption."""
+60 -3
View File
@@ -322,14 +322,71 @@ async def test_on_guardrail_task_done_emits_error_event():
session._on_guardrail_task_done(task)
# Allow event task to enqueue
await asyncio.sleep(0.01)
# Should have a RealtimeError queued
err = await session._event_queue.get()
assert isinstance(err, RealtimeError)
@pytest.mark.asyncio
async def test_close_awaits_task_finalizers_without_recancelling():
session = RealtimeSession(_DummyModel(), RealtimeAgent(name="agent"), None)
finalizer_started = asyncio.Event()
release_finalizer = asyncio.Event()
finalizer_finished = asyncio.Event()
async def background_task() -> None:
try:
await asyncio.Event().wait()
finally:
finalizer_started.set()
await release_finalizer.wait()
finalizer_finished.set()
task = asyncio.create_task(background_task())
session._guardrail_tasks.add(task)
await asyncio.sleep(0)
close_task = asyncio.create_task(session.close())
await finalizer_started.wait()
await asyncio.sleep(0)
assert not close_task.done()
release_finalizer.set()
await close_task
assert finalizer_finished.is_set()
assert session._guardrail_tasks == set()
@pytest.mark.asyncio
async def test_close_rescans_for_tasks_created_during_unwind():
session = RealtimeSession(_DummyModel(), RealtimeAgent(name="agent"), None)
spawned_task_finished = asyncio.Event()
async def spawned_task() -> None:
try:
await asyncio.Event().wait()
finally:
spawned_task_finished.set()
async def original_task() -> None:
try:
await asyncio.Event().wait()
finally:
task = asyncio.create_task(spawned_task())
session._tool_call_tasks.add(task)
await asyncio.sleep(0)
task = asyncio.create_task(original_task())
session._tool_call_tasks.add(task)
await asyncio.sleep(0)
await session.close()
assert spawned_task_finished.is_set()
assert session._tool_call_tasks == set()
@pytest.mark.asyncio
async def test_get_handoffs_async_is_enabled(monkeypatch):
# Agent includes both a direct Handoff and a RealtimeAgent (auto-converted)
+7 -12
View File
@@ -249,16 +249,11 @@ class TestSessionExceptions:
session = RealtimeSession(fake_model, fake_agent, None)
# Add some fake guardrail tasks
fake_task1 = Mock()
fake_task1.done.return_value = False
fake_task1.cancel = Mock()
fake_task2 = Mock()
fake_task2.done.return_value = True
fake_task2.cancel = Mock()
session._guardrail_tasks = {fake_task1, fake_task2}
# Add one pending and one completed guardrail task.
pending_task = asyncio.create_task(asyncio.Event().wait())
completed_task = asyncio.create_task(asyncio.sleep(0))
await completed_task
session._guardrail_tasks = {pending_task, completed_task}
fake_model.set_next_events([exception_event])
@@ -268,8 +263,8 @@ class TestSessionExceptions:
pass
# Verify guardrail tasks were properly cleaned up
fake_task1.cancel.assert_called_once()
fake_task2.cancel.assert_not_called() # Already done
assert pending_task.cancelled()
assert not completed_task.cancelled()
assert len(session._guardrail_tasks) == 0
@pytest.mark.asyncio
+2 -1
View File
@@ -677,9 +677,10 @@ async def test_docker_persist_workspace_defers_stage_cleanup_until_archive_close
assert session.stage_cleanup_calls == []
_ = archive.read()
await asyncio.sleep(0)
await session._wait_for_cleanup_tasks()
assert session.stage_cleanup_calls == [session.last_staging_parent]
assert session._cleanup_tasks == set()
def test_docker_start_exec_socket_closes_underlying_http_response() -> None:
+38
View File
@@ -1,7 +1,10 @@
from __future__ import annotations
import asyncio
import signal
from pathlib import Path
from types import SimpleNamespace
from typing import cast
import pytest
@@ -11,6 +14,7 @@ from agents.sandbox.sandboxes.unix_local import (
UnixLocalSandboxClient,
UnixLocalSandboxSession,
UnixLocalSandboxSessionState,
_UnixPtyProcessEntry,
)
from agents.sandbox.snapshot import NoopSnapshot
from agents.sandbox.types import ExecResult, User
@@ -37,6 +41,40 @@ class _RecordingUnixLocalSession(UnixLocalSandboxSession):
class TestUnixLocalPty:
@pytest.mark.asyncio
async def test_tty_fd_close_is_owned_without_blocking_termination(
self,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
session = _RecordingUnixLocalSession(tmp_path)
close_started = asyncio.Event()
release_close = asyncio.Event()
async def blocked_to_thread(*args: object, **kwargs: object) -> None:
_ = (args, kwargs)
close_started.set()
await release_close.wait()
monkeypatch.setattr(asyncio, "to_thread", blocked_to_thread)
process = cast(
asyncio.subprocess.Process,
SimpleNamespace(returncode=0, pid=None),
)
entry = _UnixPtyProcessEntry(process=process, tty=True, primary_fd=123)
await asyncio.wait_for(session._terminate_pty_entry(entry), timeout=0.5)
await close_started.wait()
assert len(session._fd_close_tasks) == 1
await asyncio.wait_for(session._wait_for_fd_close_tasks(), timeout=0.5)
assert len(session._fd_close_tasks) == 1
release_close.set()
await asyncio.gather(*session._fd_close_tasks)
await asyncio.sleep(0)
assert session._fd_close_tasks == set()
@pytest.mark.asyncio
async def test_pty_exec_write_poll_and_unknown_session_errors(self, tmp_path: Path) -> None:
client = UnixLocalSandboxClient()