From ca8cdaa4aea4e0efd597b3dab8b465c8fe2050ed Mon Sep 17 00:00:00 2001 From: Stefan Gersmann Date: Wed, 16 Sep 2026 17:40:31 +0200 Subject: [PATCH] fix(sdk): make cleanup reliable and preserve typed tool arguments --- .github/workflows/release-published.yml | 24 +- codex/_file_utils.py | 18 - codex/app_server/_async_threads.py | 18 +- codex/app_server/_session.py | 37 +- codex/app_server/_sync_services.py | 350 ++----------- codex/app_server/transports.py | 2 - codex/dynamic_tools.py | 3 +- codex/output_schema_file.py | 35 -- codex/thread.py | 12 +- scripts/generate_protocol_types.py | 13 +- scripts/postprocess_schema_titles.py | 378 -------------- tests/test_api_features.py | 111 +++++ tests/test_app_server_async_client.py | 107 +++- tests/test_app_server_client.py | 582 ++++++++++++---------- tests/test_app_server_session.py | 274 +++++++++- tests/test_app_server_transports.py | 12 +- tests/test_config_overrides.py | 33 -- tests/test_dynamic_tools.py | 37 +- tests/test_protocol_generation_scripts.py | 151 +++--- tests/test_release_workflows.py | 95 ++++ 20 files changed, 1103 insertions(+), 1189 deletions(-) delete mode 100644 codex/_file_utils.py delete mode 100644 codex/output_schema_file.py delete mode 100644 scripts/postprocess_schema_titles.py diff --git a/.github/workflows/release-published.yml b/.github/workflows/release-published.yml index 00a5fa7..0dcccc7 100644 --- a/.github/workflows/release-published.yml +++ b/.github/workflows/release-published.yml @@ -118,18 +118,25 @@ jobs: git push origin HEAD:"${DEFAULT_BRANCH}" fi + - name: Resolve build SHA + id: build_sha + shell: bash + run: | + set -euo pipefail + SHA=$(git rev-parse HEAD) + echo "sha=$SHA" >> "$GITHUB_OUTPUT" + echo "Using build SHA: $SHA" + - name: Create new semver tag and repoint release if: ${{ steps.resolve.outputs.mode == 'bump' }} shell: bash env: PLACEHOLDER_TAG: ${{ steps.ctx.outputs.tag }} VERSION: ${{ steps.resolve.outputs.version }} + BUILD_SHA: ${{ steps.build_sha.outputs.sha }} run: | - DEFAULT_BRANCH="${{ github.event.repository.default_branch }}" - git fetch origin "$DEFAULT_BRANCH" --depth=1 - NEW_SHA=$(git rev-parse "origin/${DEFAULT_BRANCH}") REAL_TAG="v${VERSION}" - git tag -fa "$REAL_TAG" "$NEW_SHA" -m "Release $REAL_TAG" + git tag -fa "$REAL_TAG" "$BUILD_SHA" -m "Release $REAL_TAG" git push origin "refs/tags/${REAL_TAG}" --force # Delete placeholder tag if it exists if git rev-parse -q --verify "refs/tags/${PLACEHOLDER_TAG}" >/dev/null; then @@ -156,15 +163,6 @@ jobs: }); core.info(`Updated release ${rel.id} to tag ${realTag}`); - - name: Resolve build SHA - id: build_sha - shell: bash - run: | - set -euo pipefail - SHA=$(git rev-parse HEAD) - echo "sha=$SHA" >> "$GITHUB_OUTPUT" - echo "Using build SHA: $SHA" - build-wheels: name: Build native wheels needs: prepare diff --git a/codex/_file_utils.py b/codex/_file_utils.py deleted file mode 100644 index 3d307ab..0000000 --- a/codex/_file_utils.py +++ /dev/null @@ -1,18 +0,0 @@ -from __future__ import annotations - -import os -import tempfile -from pathlib import Path - - -def atomic_write_text(path: Path, text: str, *, encoding: str = "utf-8") -> None: - path.parent.mkdir(parents=True, exist_ok=True) - fd, temp_path_str = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) - temp_path = Path(temp_path_str) - try: - with os.fdopen(fd, "w", encoding=encoding) as handle: - handle.write(text) - temp_path.replace(path) - except Exception: - temp_path.unlink(missing_ok=True) - raise diff --git a/codex/app_server/_async_threads.py b/codex/app_server/_async_threads.py index 6375485..61383c3 100644 --- a/codex/app_server/_async_threads.py +++ b/codex/app_server/_async_threads.py @@ -140,7 +140,6 @@ def __init__( self._item_index: dict[str, int] = {} self._text_deltas: list[str] = [] self._retryable_error_notifications: list[protocol.ErrorNotificationModel] = [] - self._done = False self._closed = False @classmethod @@ -191,14 +190,16 @@ def __aiter__(self) -> AsyncTurnStream: return self async def __anext__(self) -> Notification: - if self._done: - await self.close() + if self._closed: raise StopAsyncIteration - notification = await self._subscription.next() + try: + notification = await self._subscription.next() + except Exception: + await self.close() + raise self._apply(notification) if isinstance(notification, protocol.ErrorNotificationModel): if not notification.params.willRetry: - self._done = True await self.close() error = notification.params.error message = error.message @@ -207,15 +208,14 @@ async def __anext__(self) -> Notification: raise AppServerTurnError(message, error=error) self._retryable_error_notifications.append(notification) if isinstance(notification, protocol.TurnCompletedNotificationModel): - self._done = True + await self.close() return notification async def wait(self) -> AsyncTurnStream: """Consume the stream to completion and return `self`.""" try: - if not self._done: - async for _ in self: - pass + async for _ in self: + pass self._require_terminal_turn() finally: await self.close() diff --git a/codex/app_server/_session.py b/codex/app_server/_session.py index 028e2f8..23e1696 100644 --- a/codex/app_server/_session.py +++ b/codex/app_server/_session.py @@ -89,7 +89,7 @@ def __init__( self._transport = transport self._initialize_options = initialize_options or AppServerInitializeOptions() self._started = False - self._closed = False + self._close_task: asyncio.Task[Exception | None] | None = None self._next_request_id = 0 self._pending: dict[int | str, asyncio.Future[object]] = {} self._request_handlers: dict[str, _RegisteredHandler] = {} @@ -101,7 +101,7 @@ def __init__( self._initialize_result: InitializeResult | None = None async def start(self) -> InitializeResult: - if self._closed: + if self._close_task is not None: raise AppServerClosedError("app-server client is closed") if self._started: if self._initialize_result is None: @@ -126,9 +126,15 @@ async def start(self) -> InitializeResult: return result async def close(self) -> None: - if self._closed: + if self._close_task is None: + self._close_task = asyncio.create_task(self._close()) + elif self._close_task.done(): return - self._closed = True + close_error = await asyncio.shield(self._close_task) + if close_error is not None: + raise close_error + + async def _close(self) -> Exception | None: close_error: Exception | None = None if self._reader_error_reported else self._reader_error if self._reader_task is not None: if not self._reader_task.done(): @@ -150,8 +156,7 @@ async def close(self) -> None: for sink in list(self._notification_sinks): await sink.queue.put(None) self._notification_sinks.clear() - if close_error is not None: - raise close_error + return close_error async def _close_for_start_failure(self) -> Exception | None: try: @@ -182,9 +187,6 @@ async def request( await self._ensure_started_or_starting() request_id_value = self._next_request_id self._next_request_id += 1 - loop = asyncio.get_running_loop() - future: asyncio.Future[object] = loop.create_future() - self._pending[request_id_value] = future message: JsonObject = {"id": request_id_value, "method": method} if params is not None: serialized = serialize_value(params) @@ -193,8 +195,17 @@ async def request( f"Request params must serialize to an object, got {type(serialized).__name__}" ) message["params"] = cast(JsonObject, serialized) - await self._transport.send(message) - return await self._await_future(future) + future: asyncio.Future[object] = asyncio.get_running_loop().create_future() + self._pending[request_id_value] = future + try: + await self._transport.send(message) + return await self._await_future(future) + finally: + self._pending.pop(request_id_value, None) + if not future.done(): + future.cancel() + elif not future.cancelled(): + future.exception() async def request_typed( self, @@ -227,7 +238,7 @@ def subscribe_notifications( return _AsyncNotificationSubscription(sink, sink.queue, lambda: self._remove_sink(sink)) async def _ensure_started_or_starting(self) -> None: - if self._closed: + if self._close_task is not None: raise AppServerClosedError("app-server client is closed") if self._reader_task is None: raise AppServerClosedError("app-server client is not started") @@ -290,6 +301,8 @@ def _reader_failure(self) -> Exception: return self._reader_error if self._reader_task is None: return AppServerClosedError("app-server client is not started") + if self._reader_task.cancelled(): + return AppServerClosedError("app-server reader was cancelled") task_exception = self._reader_task.exception() if task_exception is not None: if isinstance(task_exception, Exception): diff --git a/codex/app_server/_sync_services.py b/codex/app_server/_sync_services.py index 293b83d..4c3cbfb 100644 --- a/codex/app_server/_sync_services.py +++ b/codex/app_server/_sync_services.py @@ -1,8 +1,23 @@ from __future__ import annotations from collections.abc import Callable, Coroutine, Mapping, Sequence -from typing import Any, Protocol - +from typing import Any + +from codex.app_server._async_services import ( + AsyncAccountClient, + AsyncAppsClient, + AsyncCommandClient, + AsyncConfigClient, + AsyncEnvironmentClient, + AsyncExternalAgentConfigClient, + AsyncFeedbackClient, + AsyncFsClient, + AsyncMcpServersClient, + AsyncModelsClient, + AsyncSkillsClient, + AsyncThreadSectionsClient, + AsyncWindowsSandboxClient, +) from codex.app_server._sync_support import _SyncRunner from codex.app_server.models import ( AccountCancelLoginResult, @@ -33,315 +48,10 @@ from codex.protocol import types as protocol -class _AsyncModelsClientLike(Protocol): - async def list( - self, - *, - cursor: str | None = None, - include_hidden: bool | None = None, - limit: int | None = None, - ) -> list[ModelInfo]: ... - - async def list_page( - self, - *, - cursor: str | None = None, - include_hidden: bool | None = None, - limit: int | None = None, - ) -> ModelListResult: ... - - -class _AsyncAppsClientLike(Protocol): - async def list( - self, - *, - cursor: str | None = None, - force_refetch: bool | None = None, - limit: int | None = None, - thread_id: str | None = None, - ) -> list[protocol.AppInfo]: ... - - async def list_page( - self, - *, - cursor: str | None = None, - force_refetch: bool | None = None, - limit: int | None = None, - thread_id: str | None = None, - ) -> AppListResult: ... - - -class _AsyncThreadSectionsClientLike(Protocol): - async def list( - self, - *, - cursor: str | None = None, - limit: int | None = None, - ) -> list[protocol.ThreadSection]: ... - - async def list_page( - self, - *, - cursor: str | None = None, - limit: int | None = None, - ) -> protocol.ThreadSectionListResponse: ... - - async def create( - self, - *, - name: str, - appearance: protocol.ThreadSectionAppearance | None = None, - ) -> protocol.ThreadSection: ... - - async def rename(self, *, section_id: str, name: str) -> protocol.ThreadSection: ... - - async def delete(self, *, section_id: str) -> EmptyResult: ... - - -class _AsyncSkillsClientLike(Protocol): - def input(self, *, name: str, path: str) -> protocol.SkillUserInput: ... - - async def list( - self, - *, - cwds: Sequence[str] | None = None, - force_reload: bool | None = None, - ) -> list[SkillsListEntry]: ... - - async def list_page( - self, - *, - cwds: Sequence[str] | None = None, - force_reload: bool | None = None, - ) -> SkillsListResult: ... - - async def reload(self, *, cwds: Sequence[str] | None = None) -> Sequence[SkillsListEntry]: ... - - async def write_config(self, *, path: str, enabled: bool) -> SkillsConfigWriteResult: ... - - async def write_skill( - self, - *, - name: str, - directory: str, - instructions: str | bytes, - reload_cwds: Sequence[str] | None = None, - ) -> protocol.SkillUserInput: ... - - -class _AsyncFsClientLike(Protocol): - async def create_directory( - self, - *, - path: str, - recursive: bool | None = True, - ) -> protocol.FsCreateDirectoryResponse: ... - - async def write_file( - self, - *, - path: str, - data: str | bytes, - encoding: str = "utf-8", - ) -> protocol.FsWriteFileResponse: ... - - -class _AsyncEnvironmentClientLike(Protocol): - async def info(self, *, environment_id: str) -> protocol.EnvironmentInfoResponse: ... - - -class _AsyncAccountClientLike(Protocol): - async def read(self, *, refresh_token: bool | None = None) -> AccountReadResult: ... - - async def login_api_key(self, *, api_key: str) -> ApiKeyLoginResult: ... - - async def login_chatgpt(self) -> ChatGptLoginResult: ... - - async def login_chatgpt_tokens( - self, - *, - access_token: str, - chatgpt_account_id: str, - chatgpt_plan_type: protocol.PlanType | None = None, - ) -> ChatGptAuthTokensLoginResult: ... - - async def cancel_login(self, *, login_id: str) -> AccountCancelLoginResult: ... - - async def logout(self) -> EmptyResult: ... - - async def read_rate_limits(self) -> AccountRateLimitsResult: ... - - -class _AsyncConfigClientLike(Protocol): - async def read( - self, - *, - cwd: str | None = None, - include_layers: bool | None = None, - ) -> ConfigReadResult: ... - - async def reload_mcp_servers(self) -> EmptyResult: ... - - async def write_value( - self, - *, - key_path: str, - value: Any, - merge_strategy: protocol.MergeStrategy, - expected_version: str | None = None, - file_path: str | None = None, - ) -> ConfigWriteResult: ... - - async def batch_write( - self, - *, - edits: Sequence[protocol.ConfigEdit], - expected_version: str | None = None, - file_path: str | None = None, - ) -> ConfigWriteResult: ... - - async def read_requirements(self) -> ConfigRequirementsReadResult: ... - - -class _AsyncMcpServersClientLike(Protocol): - async def set_enabled_tools( - self, - *, - name: str, - tools: Sequence[str], - plugin_id: str | None = None, - reload: bool = True, - file_path: str | None = None, - ) -> ConfigWriteResult: ... - - async def set_disabled_tools( - self, - *, - name: str, - tools: Sequence[str], - plugin_id: str | None = None, - reload: bool = True, - file_path: str | None = None, - ) -> ConfigWriteResult: ... - - async def oauth_login( - self, - *, - client_registration: protocol.McpServerOauthClientRegistration | None = None, - name: str, - scopes: Sequence[str] | None = None, - thread_id: str | None = None, - timeout_seconds: int | None = None, - ) -> McpServerOauthLoginResult: ... - - async def list( - self, - *, - cursor: str | None = None, - limit: int | None = None, - ) -> list[McpServerStatus]: ... - - async def list_page( - self, - *, - cursor: str | None = None, - limit: int | None = None, - ) -> McpServerStatusListResult: ... - - -class _AsyncFeedbackClientLike(Protocol): - async def upload( - self, - *, - classification: str, - include_logs: bool, - extra_log_files: Sequence[str] | None = None, - reason: str | None = None, - thread_id: str | None = None, - ) -> FeedbackUploadResult: ... - - -class _AsyncCommandClientLike(Protocol): - async def execute( - self, - *, - command: Sequence[str], - cwd: str | None = None, - disable_output_cap: bool | None = None, - disable_timeout: bool | None = None, - env: Mapping[str, object | None] | None = None, - output_bytes_cap: int | None = None, - permission_profile: str | None = None, - process_id: str | None = None, - sandbox_policy: protocol.SandboxPolicy | None = None, - size: protocol.CommandExecTerminalSize | None = None, - stream_stdin: bool | None = None, - stream_stdout_stderr: bool | None = None, - timeout_ms: int | None = None, - tty: bool | None = None, - ) -> CommandExecResult: ... - - async def write_stdin( - self, - *, - process_id: str, - close_stdin: bool | None = None, - delta_base64: str | None = None, - ) -> EmptyResult: ... - - async def resize_terminal( - self, - *, - process_id: str, - size: protocol.CommandExecTerminalSize, - ) -> EmptyResult: ... - - async def terminate_process(self, *, process_id: str) -> EmptyResult: ... - - -class _AsyncExternalAgentConfigClientLike(Protocol): - async def detect( - self, - *, - cwds: Sequence[str] | None = None, - include_home: bool | None = None, - max_session_age_days: int | None = None, - max_sessions: int | None = None, - migration_source: str | None = None, - ) -> ExternalAgentConfigDetectResult: ... - - async def import_items( - self, - *, - migration_items: Sequence[protocol.ExternalAgentConfigMigrationItem], - migration_source: str | None = None, - provider_id: str | None = None, - source: str | None = None, - ) -> ExternalAgentConfigImportResult: ... - - async def record_history( - self, - *, - item_type_results: Sequence[ - protocol.ExternalAgentConfigImportHistoryRecordTypeResultParams - ], - provider_id: str, - ) -> ExternalAgentConfigImportResult: ... - - -class _AsyncWindowsSandboxClientLike(Protocol): - async def setup_start( - self, - *, - mode: protocol.WindowsSandboxSetupMode, - cwd: str | None = None, - ) -> WindowsSandboxSetupStartResult: ... - - class _ModelsClient(_SyncRunner): def __init__( self, - async_client: _AsyncModelsClientLike, + async_client: AsyncModelsClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -381,7 +91,7 @@ def list_page( class _AppsClient(_SyncRunner): def __init__( self, - async_client: _AsyncAppsClientLike, + async_client: AsyncAppsClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -425,7 +135,7 @@ def list_page( class _ThreadSectionsClient(_SyncRunner): def __init__( self, - async_client: _AsyncThreadSectionsClientLike, + async_client: AsyncThreadSectionsClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -465,7 +175,7 @@ def delete(self, *, section_id: str) -> EmptyResult: class _SkillsClient(_SyncRunner): def __init__( self, - async_client: _AsyncSkillsClientLike, + async_client: AsyncSkillsClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -527,7 +237,7 @@ def write_skill( class _FsClient(_SyncRunner): def __init__( self, - async_client: _AsyncFsClientLike, + async_client: AsyncFsClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -554,7 +264,7 @@ def write_file( class _EnvironmentClient(_SyncRunner): def __init__( self, - async_client: _AsyncEnvironmentClientLike, + async_client: AsyncEnvironmentClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -567,7 +277,7 @@ def info(self, *, environment_id: str) -> protocol.EnvironmentInfoResponse: class _AccountClient(_SyncRunner): def __init__( self, - async_client: _AsyncAccountClientLike, + async_client: AsyncAccountClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -610,7 +320,7 @@ def read_rate_limits(self) -> AccountRateLimitsResult: class _ConfigClient(_SyncRunner): def __init__( self, - async_client: _AsyncConfigClientLike, + async_client: AsyncConfigClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -668,7 +378,7 @@ def read_requirements(self) -> ConfigRequirementsReadResult: class _McpServersClient(_SyncRunner): def __init__( self, - async_client: _AsyncMcpServersClientLike, + async_client: AsyncMcpServersClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -755,7 +465,7 @@ def list_page( class _FeedbackClient(_SyncRunner): def __init__( self, - async_client: _AsyncFeedbackClientLike, + async_client: AsyncFeedbackClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -784,7 +494,7 @@ def upload( class _CommandClient(_SyncRunner): def __init__( self, - async_client: _AsyncCommandClientLike, + async_client: AsyncCommandClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -875,7 +585,7 @@ def terminate_process(self, *, process_id: str) -> EmptyResult: class _ExternalAgentConfigClient(_SyncRunner): def __init__( self, - async_client: _AsyncExternalAgentConfigClientLike, + async_client: AsyncExternalAgentConfigClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) @@ -936,7 +646,7 @@ def record_history( class _WindowsSandboxClient(_SyncRunner): def __init__( self, - async_client: _AsyncWindowsSandboxClientLike, + async_client: AsyncWindowsSandboxClient, runner: Callable[[Coroutine[Any, Any, Any]], Any], ) -> None: super().__init__(runner) diff --git a/codex/app_server/transports.py b/codex/app_server/transports.py index 9c78adb..93f5925 100644 --- a/codex/app_server/transports.py +++ b/codex/app_server/transports.py @@ -60,7 +60,6 @@ def __init__(self, options: AppServerProcessOptions | None = None) -> None: self._options = options or AppServerProcessOptions() self._process: asyncio.subprocess.Process | None = None self._stderr_task: asyncio.Task[None] | None = None - self._stderr_lines: list[str] = [] async def start(self) -> None: if self._process is not None: @@ -96,7 +95,6 @@ async def _drain_stderr(self, stderr: asyncio.StreamReader) -> None: line = await _readline_with_limit_error(stderr, stream_name="stderr") if line == b"": break - self._stderr_lines.append(line.decode("utf-8", errors="replace").rstrip()) async def send(self, message: JsonObject) -> None: if self._process is None or self._process.stdin is None: diff --git a/codex/dynamic_tools.py b/codex/dynamic_tools.py index 4381e10..9661bd0 100644 --- a/codex/dynamic_tools.py +++ b/codex/dynamic_tools.py @@ -142,8 +142,7 @@ async def dispatch( ) validated = tool.input_model.model_validate(request.params.arguments) - arguments = validated.model_dump(mode="python") - result = tool.callable(**arguments) + result = tool.callable(**dict(validated)) if inspect.isawaitable(result): result = await cast(Awaitable[object], result) return _normalize_tool_result(result) diff --git a/codex/output_schema_file.py b/codex/output_schema_file.py deleted file mode 100644 index a6fcce2..0000000 --- a/codex/output_schema_file.py +++ /dev/null @@ -1,35 +0,0 @@ -from __future__ import annotations - -import json -import shutil -import tempfile -from dataclasses import dataclass -from pathlib import Path - -from codex._file_utils import atomic_write_text -from codex.output_schema import OutputSchemaInput, normalize_output_schema - - -@dataclass(slots=True, frozen=True) -class OutputSchemaFile: - schema_path: str | None - schema_dir: str | None - - def cleanup(self) -> None: - if self.schema_dir is not None: - shutil.rmtree(self.schema_dir, ignore_errors=True) - - -def create_output_schema_file(schema: OutputSchemaInput | None) -> OutputSchemaFile: - normalized_schema = normalize_output_schema(schema) - if normalized_schema is None: - return OutputSchemaFile(schema_path=None, schema_dir=None) - - schema_dir = Path(tempfile.mkdtemp(prefix="codex-output-schema-")) - schema_path = schema_dir / "schema.json" - try: - atomic_write_text(schema_path, json.dumps(normalized_schema)) - except Exception: - shutil.rmtree(schema_dir, ignore_errors=True) - raise - return OutputSchemaFile(schema_path=str(schema_path), schema_dir=str(schema_dir)) diff --git a/codex/thread.py b/codex/thread.py index 5c20d8b..f3c15d1 100644 --- a/codex/thread.py +++ b/codex/thread.py @@ -58,6 +58,7 @@ def __init__( ) -> None: self._stream = stream self._thread_id = thread_id + self._owner: Thread | None = None self._closed = False self._interrupt_requested = False self._watcher = _SignalWatcher(self, signal) @@ -66,7 +67,11 @@ def __iter__(self) -> CodexTurnStream: return self def __next__(self) -> BaseModel: - notification: BaseModel = next(self._stream) + try: + notification: BaseModel = next(self._stream) + except BaseException: + self._watcher.stop() + raise if self.final_turn is not None: self._watcher.stop() return notification @@ -174,7 +179,10 @@ def run( raise ThreadRunError("Turn aborted: interrupted") thread = self._ensure_thread() stream = thread.run(input, _to_app_server_turn_options(effective_turn_options)) - return CodexTurnStream(stream, thread_id=thread.id, signal=signal) + turn_stream = CodexTurnStream(stream, thread_id=thread.id, signal=signal) + # The bound client factory keeps temporary Codex owners alive. + turn_stream._owner = self + return turn_stream def run_text( self, diff --git a/scripts/generate_protocol_types.py b/scripts/generate_protocol_types.py index d6a98cf..3443a30 100644 --- a/scripts/generate_protocol_types.py +++ b/scripts/generate_protocol_types.py @@ -209,17 +209,20 @@ def main() -> int: output_path = Path(args.output) output_path.parent.mkdir(parents=True, exist_ok=True) - with tempfile.TemporaryDirectory(prefix="codex-app-server-schema-") as temp_dir: + with tempfile.TemporaryDirectory( + prefix="codex-app-server-schema-", dir=output_path.parent + ) as temp_dir: schema_dir = Path(temp_dir) / "schemas" + candidate_path = Path(temp_dir) / "types.py" schema_path = export_protocol_schema( codex_bin=args.codex_bin, schema_dir=schema_dir, experimental=args.experimental, ) - generate_protocol_models(schema_path=schema_path, output_path=output_path) - append_extra_protocol_models(schema_dir=schema_dir, output_path=output_path) - - postprocess_protocol_models(output_path) + generate_protocol_models(schema_path=schema_path, output_path=candidate_path) + append_extra_protocol_models(schema_dir=schema_dir, output_path=candidate_path) + postprocess_protocol_models(candidate_path) + candidate_path.replace(output_path) return 0 diff --git a/scripts/postprocess_schema_titles.py b/scripts/postprocess_schema_titles.py deleted file mode 100644 index 2cf088b..0000000 --- a/scripts/postprocess_schema_titles.py +++ /dev/null @@ -1,378 +0,0 @@ -#!/usr/bin/env python3 -from __future__ import annotations - -import argparse -import json -import os -import tempfile -from collections.abc import Callable -from pathlib import Path - -type SchemaNode = dict[str, object] - -TARGETS = [ - ("EventMsg", "type"), - ("ClientRequest", "method"), - ("ServerRequest", "method"), - ("ServerNotification", "method"), - ("InputItem", "type"), -] - - -def atomic_write_text(path: Path, text: str, *, encoding: str = "utf-8") -> None: - path.parent.mkdir(parents=True, exist_ok=True) - fd, temp_path_str = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) - temp_path = Path(temp_path_str) - try: - with os.fdopen(fd, "w", encoding=encoding) as handle: - handle.write(text) - temp_path.replace(path) - except Exception: - temp_path.unlink(missing_ok=True) - raise - - -def camelize(s: str) -> str: - parts = [p for p in s.replace("-", "_").replace(" ", "_").split("_") if p] - return "".join(p[:1].upper() + p[1:] for p in parts) - - -def _definitions_node(schema: SchemaNode) -> tuple[dict[str, object], str] | None: - for key in ("definitions", "$defs"): - defs = schema.get(key) - if isinstance(defs, dict): - return defs, key - return None - - -def _tag_variants(node: SchemaNode) -> list[SchemaNode] | None: - one_of = node.get("oneOf") or node.get("anyOf") - if not isinstance(one_of, list): - return None - return [variant for variant in one_of if isinstance(variant, dict)] - - -def _tag_value(properties: object, tag_key: str) -> str | None: - if not isinstance(properties, dict): - return None - tag = properties.get(tag_key) - if not isinstance(tag, dict): - return None - enum = tag.get("enum") - if isinstance(enum, list) and enum and isinstance(enum[0], str): - return enum[0] - const = tag.get("const") - if isinstance(const, str): - return const - return None - - -def _walk_schema(node: object, visit: Callable[[SchemaNode], None]) -> None: - if isinstance(node, dict): - visit(node) - for key in ("items", "additionalProperties", "not"): - child = node.get(key) - if isinstance(child, dict): - _walk_schema(child, visit) - for key in ("anyOf", "oneOf", "allOf"): - children = node.get(key) - if isinstance(children, list): - for child in children: - _walk_schema(child, visit) - for key in ("definitions", "$defs", "patternProperties"): - children = node.get(key) - if isinstance(children, dict): - for child in children.values(): - _walk_schema(child, visit) - return - if isinstance(node, list): - for child in node: - _walk_schema(child, visit) - - -def _nullable(prop_schema: object) -> bool: - if not isinstance(prop_schema, dict): - return False - field_type = prop_schema.get("type") - if field_type == "null": - return True - if isinstance(field_type, list) and "null" in field_type: - return True - for key in ("anyOf", "oneOf"): - variants = prop_schema.get(key) - if not isinstance(variants, list): - continue - if any(isinstance(variant, dict) and variant.get("type") == "null" for variant in variants): - return True - return False - - -def _normalize_numeric_type(value: object) -> tuple[object, bool]: - if value == "number": - return "integer", True - if not isinstance(value, list) or "number" not in value: - return value, False - normalized = ["integer" if item == "number" else item for item in value] - return _dedupe_preserve_order(normalized), True - - -def _replace_number_with_integer(node: SchemaNode) -> bool: - changed = False - normalized_type, type_changed = _normalize_numeric_type(node.get("type")) - if type_changed: - node["type"] = normalized_type - changed = True - for key in ("anyOf", "oneOf"): - variants = node.get(key) - if not isinstance(variants, list): - continue - for variant in variants: - if not isinstance(variant, dict): - continue - normalized_variant_type, variant_changed = _normalize_numeric_type(variant.get("type")) - if variant_changed: - variant["type"] = normalized_variant_type - changed = True - return changed - - -def _duration_union(description: object) -> list[SchemaNode]: - description_field = {"description": description} if isinstance(description, str) else {} - return [ - {"type": "string", **description_field}, - { - "type": "object", - "properties": { - "secs": {"type": "integer"}, - "nanos": {"type": "integer"}, - }, - "required": ["secs", "nanos"], - "additionalProperties": False, - **description_field, - }, - ] - - -def add_titles(schema: dict) -> tuple[bool, int]: - definitions = _definitions_node(schema) - if definitions is None: - return (False, 0) - defs, base_key = definitions - changed = False - added = 0 - for name, tag_key in TARGETS: - node = defs.get(name) - if not isinstance(node, dict): - continue - variants = _tag_variants(node) - if variants is None: - continue - one_of = node.get("oneOf") or node.get("anyOf") - if not isinstance(one_of, list): - continue - for index, variant in enumerate(one_of): - if not isinstance(variant, dict): - continue - tag_value = _tag_value(variant.get("properties"), tag_key) - if tag_value is None: - continue - title = f"{name}_{camelize(tag_value)}" - variant["title"] = title - if title not in defs: - defs[title] = variant - changed = True - added += 1 - one_of[index] = {"$ref": f"#/{base_key}/{title}"} - return changed, added - - -def relax_required_for_nullables(schema: dict) -> tuple[bool, int]: - """Recursively remove nullable properties from 'required' arrays. - - Applies to the whole schema tree, not just top-level $defs/definitions, to - capture inline object schemas generated within oneOf/anyOf branches. - """ - changed = False - count = 0 - - def visit(node: SchemaNode) -> None: - nonlocal changed, count - properties = node.get("properties") - required = node.get("required") - if not isinstance(properties, dict) or not isinstance(required, list): - return - new_required = [name for name in required if not _nullable(properties.get(name))] - if len(new_required) == len(required): - return - node["required"] = new_required - changed = True - count += len(required) - len(new_required) - - _walk_schema(schema, visit) - return changed, count - - -def enforce_request_id_integer(schema: dict) -> bool: - # Force RequestId to be string|integer (not number) so Python maps to str|int - defs = schema.get("definitions") or schema.get("$defs") - if not isinstance(defs, dict): - return False - node = defs.get("RequestId") - if not isinstance(node, dict): - return False - current = node.get("type") - desired = ["string", "integer"] - if current != desired: - node["type"] = desired - # remove other conflicting keys if any - for k in ("anyOf", "oneOf"): - if k in node: - node.pop(k) - return True - return False - - -def enforce_exec_exit_code_integer(schema: dict) -> bool: - # Force ExecCommandEndEvent.exit_code to integer - defs = schema.get("definitions") or schema.get("$defs") - if not isinstance(defs, dict): - return False - node = defs.get("ExecCommandEndEvent") - if not isinstance(node, dict): - return False - props = node.get("properties") - if not isinstance(props, dict): - return False - exit_node = props.get("exit_code") - if not isinstance(exit_node, dict): - return False - if exit_node.get("type") != "integer": - exit_node["type"] = "integer" - return True - return False - - -def _dedupe_preserve_order(items: list[str]) -> list[str]: - """Return a new list with duplicates removed, preserving order.""" - seen: set[str] = set() - out: list[str] = [] - for x in items: - if x not in seen: - seen.add(x) - out.append(x) - return out - - -INTEGER_FIELDS = { - # Exec - "exit_code", - # Token usage counters - "input_tokens", - "cached_input_tokens", - "output_tokens", - "reasoning_output_tokens", - "total_tokens", - # Context window capacity - "model_context_window", - # History identifiers and counters - "log_id", - "history_log_id", - "history_entry_count", - "offset", -} - - -def enforce_integer_fields(schema: dict) -> int: - """Walk the JSON Schema and coerce selected numeric fields to integer. - - Applies to both hoisted $defs and inline subschemas to avoid mismatches - between event structs and EventMsg wrappers. - """ - changed = 0 - - def visit(node: SchemaNode) -> None: - nonlocal changed - properties = node.get("properties") - if not isinstance(properties, dict): - return - for name, sub_schema in properties.items(): - if ( - name in INTEGER_FIELDS - and isinstance(sub_schema, dict) - and _replace_number_with_integer(sub_schema) - ): - changed += 1 - - _walk_schema(schema, visit) - return changed - - -def enforce_duration_union(schema: dict) -> int: - """Allow duration fields to be either string or {secs,nanos} object. - - Some upstream emitters serialize Rust `Duration` as an object - `{secs, nanos}` while the TypeScript schema uses `string`. - To tolerate both without breaking older clients, convert any - property named `duration` that is currently `type: string` into - a `oneOf: [string, {secs:int, nanos:int}]`. - Applies recursively across the schema tree. - """ - changed = 0 - - def visit(node: SchemaNode) -> None: - nonlocal changed - properties = node.get("properties") - if not isinstance(properties, dict): - return - duration = properties.get("duration") - if not isinstance(duration, dict) or duration.get("type") != "string": - return - description = duration.get("description") - duration.clear() - duration["oneOf"] = _duration_union(description) - changed += 1 - - _walk_schema(schema, visit) - return changed - - -def main() -> int: - parser = argparse.ArgumentParser( - description="Post-process generated JSON Schema: add titles/hoist, integer coercions, and optional tweaks", - ) - parser.add_argument( - "schema", - nargs="?", - default=Path(".generated/schema/protocol.schema.json"), - type=Path, - help="Path to protocol.schema.json", - ) - parser.add_argument( - "--relax-nullable-required", - action="store_true", - help="If set, remove nullable properties from 'required' to make them optional in Python.", - ) - args = parser.parse_args() - - path: Path = args.schema - data = json.loads(path.read_text()) - t_changed, t_added = add_titles(data) - r_changed = False - r_count = 0 - if args.relax_nullable_required: - r_changed, r_count = relax_required_for_nullables(data) - id_fixed = enforce_request_id_integer(data) - exit_fixed = enforce_exec_exit_code_integer(data) - coerced = enforce_integer_fields(data) - durations = enforce_duration_union(data) - if t_changed or r_changed or id_fixed or exit_fixed or coerced or durations: - atomic_write_text(path, json.dumps(data, indent=2, ensure_ascii=False) + "\n") - print( - f"Schema postprocess: titles+hoist added={t_added}, relaxed_required={r_count if args.relax_nullable_required else 0}, " - f"requestId_fixed={'yes' if id_fixed else 'no'}, exit_code_fixed={'yes' if exit_fixed else 'no'}, integers_coerced={coerced}, durations_patched={durations} in {path.name}" - ) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/test_api_features.py b/tests/test_api_features.py index 3619063..3cdf932 100644 --- a/tests/test_api_features.py +++ b/tests/test_api_features.py @@ -1,6 +1,8 @@ from __future__ import annotations +import gc import threading +import weakref from collections.abc import Sequence from typing import Any @@ -10,6 +12,7 @@ from codex import ( Codex, CodexOptions, + CodexTurnStream, ThreadResumeOptions, ThreadStartOptions, TurnOptions, @@ -519,6 +522,55 @@ def lookup_ticket(id: str) -> str: ] +@pytest.mark.parametrize("use_temporary_thread", [False, True]) +def test_stream_keeps_temporary_owner_alive_until_released( + monkeypatch: pytest.MonkeyPatch, use_temporary_thread: bool +) -> None: + class CloseSensitiveStream(_FakeAppTurnStream): + def __next__(self) -> BaseModel: + if fake_client.closed: + raise RuntimeError("client closed during streaming") + return super().__next__() + + app_stream = CloseSensitiveStream( + [_item_completed_notification("answer"), _turn_completed_notification()] + ) + fake_client = _FakeAppServerClient(_FakeAppThread("thr-1", [app_stream])) + _patch_connect_stdio(monkeypatch, fake_client=fake_client, capture={}) + + if use_temporary_thread: + stream = Codex().start_thread().run("hello") + else: + stream = Codex().run("hello") + stream_reference = weakref.ref(stream) + + gc.collect() + stream.wait() + assert stream.final_text == "answer" + assert not fake_client.closed + stream.close() + assert app_stream.closed + + del stream + gc.collect() + assert stream_reference() is None + assert fake_client.closed + + +def test_explicit_codex_close_closes_client_with_reachable_stream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_client = _FakeAppServerClient(_FakeAppThread("thr-1", [_FakeAppTurnStream([])])) + _patch_connect_stdio(monkeypatch, fake_client=fake_client, capture={}) + client = Codex() + stream = client.run("hello") + + client.close() + + assert fake_client.closed + assert stream.thread_id == "thr-1" + + def test_run_preserves_usage_when_turn_completed_omits_usage( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -795,6 +847,65 @@ def test_run_turn_signal_interrupts_in_flight_turn(monkeypatch: pytest.MonkeyPat _ = fake_stream +@pytest.mark.parametrize( + "error", + [ + StopIteration(), + AppServerTurnError("turn failed"), + RuntimeError("reader failed"), + KeyboardInterrupt(), + ], +) +def test_stream_iteration_failure_stops_signal_watcher(error: BaseException) -> None: + class FailingStream(_FakeAppTurnStream): + def __next__(self) -> BaseModel: + raise error + + stream = CodexTurnStream(FailingStream([]), thread_id="thr-1", signal=threading.Event()) + watcher = stream._watcher._thread + assert watcher is not None + try: + assert watcher.is_alive() + with pytest.raises(type(error)) as exc_info: + next(stream) + assert exc_info.value is error + watcher.join(timeout=1) + assert not watcher.is_alive() + finally: + stream.close() + + +def test_stream_signal_watcher_survives_retryable_events_until_terminal_delivery() -> None: + retryable_error = protocol.ErrorNotificationModel.model_validate( + { + "method": "error", + "params": { + "threadId": "thr-1", + "turnId": "turn-1", + "willRetry": True, + "error": {"message": "temporary outage"}, + }, + } + ) + notifications = [_turn_started_notification(), retryable_error, _turn_completed_notification()] + app_stream = _FakeAppTurnStream(notifications) + stream = CodexTurnStream(app_stream, thread_id="thr-1", signal=threading.Event()) + watcher = stream._watcher._thread + assert watcher is not None + try: + for notification in notifications[:-1]: + assert next(stream) is notification + assert watcher.is_alive() + assert next(stream) is notifications[-1] + watcher.join(timeout=1) + assert not watcher.is_alive() + stream.close() + stream.close() + assert app_stream.closed + finally: + stream.close() + + def test_run_raises_thread_run_error_for_failed_turn(monkeypatch: pytest.MonkeyPatch) -> None: fake_thread = _FakeAppThread( "thr-1", diff --git a/tests/test_app_server_async_client.py b/tests/test_app_server_async_client.py index 668f78a..5f48ebf 100644 --- a/tests/test_app_server_async_client.py +++ b/tests/test_app_server_async_client.py @@ -5,6 +5,8 @@ import pytest from codex.app_server._async_client import AsyncEventsClient, AsyncTurnStream +from codex.app_server._session import _AsyncNotificationSubscription, _NotificationSink +from codex.app_server._sync_threads import TurnStream from codex.app_server.errors import AppServerProtocolError, AppServerTurnError from codex.app_server.models import ReviewResult from codex.protocol import types as protocol @@ -372,23 +374,114 @@ def test_async_turn_stream_apply_replaces_existing_item_state() -> None: _ = stream.final_text -def test_async_turn_stream_wait_returns_immediately_when_done() -> None: +def test_async_turn_stream_terminal_delivery_unsubscribes_before_returning() -> None: async def scenario() -> None: - subscription = _FakeSubscription() + sink = _NotificationSink() + active_sinks = [sink] + subscription = _AsyncNotificationSubscription( + sink, sink.queue, lambda: active_sinks.remove(sink) + ) + terminal = protocol.TurnCompletedNotificationModel.model_validate( + { + "method": "turn/completed", + "params": {"threadId": "thr-1", "turn": _turn_payload()}, + } + ) + sink.queue.put_nowait(terminal) stream = AsyncTurnStream( _FakeThread(), # type: ignore[arg-type] - subscription, # type: ignore[arg-type] + subscription, protocol.Turn.model_validate(_turn_payload(status="inProgress")), ) - stream._done = True - stream.final_turn = protocol.Turn.model_validate(_turn_payload(status="completed")) - assert await stream.wait() is stream - assert subscription.closed is True + assert await anext(stream) is terminal + assert active_sinks == [] + assert stream.final_turn is terminal.params.turn + for _ in range(2): + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(anext(stream), timeout=1) + assert await asyncio.wait_for(stream.wait(), timeout=1) is stream + await stream.close() + + asyncio.run(scenario()) + + +@pytest.mark.parametrize("ending", ["close", "eof", "reader_error"]) +def test_async_turn_stream_remains_exhausted_without_terminal_result(ending: str) -> None: + async def scenario() -> None: + sink = _NotificationSink() + active_sinks = [sink] + subscription = _AsyncNotificationSubscription( + sink, sink.queue, lambda: active_sinks.remove(sink) + ) + stream = AsyncTurnStream( + _FakeThread(), # type: ignore[arg-type] + subscription, + protocol.Turn.model_validate(_turn_payload(status="inProgress")), + ) + if ending == "close": + await stream.close() + elif ending == "eof": + sink.queue.put_nowait(None) + else: + error = AppServerProtocolError("reader failed") + sink.queue.put_nowait(error) + with pytest.raises(AppServerProtocolError) as exc_info: + await anext(stream) + assert exc_info.value is error + + for _ in range(2): + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(anext(stream), timeout=1) + assert active_sinks == [] + with pytest.raises(ValueError, match="No terminal turn is available yet"): + await asyncio.wait_for(stream.wait(), timeout=1) + assert stream.final_turn is None + await stream.close() asyncio.run(scenario()) +@pytest.mark.parametrize("complete_turn", [False, True]) +def test_sync_turn_stream_preserves_exhaustion_and_completion(complete_turn: bool) -> None: + sink = _NotificationSink() + active_sinks = [sink] + subscription = _AsyncNotificationSubscription( + sink, sink.queue, lambda: active_sinks.remove(sink) + ) + async_stream = AsyncTurnStream( + _FakeThread(), # type: ignore[arg-type] + subscription, + protocol.Turn.model_validate(_turn_payload(status="inProgress")), + ) + with asyncio.Runner() as runner: + stream = TurnStream( + async_stream, lambda coro: runner.run(asyncio.wait_for(coro, timeout=1)) + ) + if complete_turn: + terminal = protocol.TurnCompletedNotificationModel.model_validate( + { + "method": "turn/completed", + "params": {"threadId": "thr-1", "turn": _turn_payload()}, + } + ) + sink.queue.put_nowait(terminal) + assert next(stream) is terminal + else: + stream.close() + assert active_sinks == [] + + for _ in range(2): + with pytest.raises(StopIteration): + next(stream) + if complete_turn: + assert stream.wait() is stream + else: + with pytest.raises(ValueError, match="No terminal turn is available yet"): + stream.wait() + stream.close() + + def test_async_turn_stream_wait_closes_subscription_after_terminal_notification() -> None: class _CompletedSubscription(_FakeSubscription): def __init__(self) -> None: diff --git a/tests/test_app_server_client.py b/tests/test_app_server_client.py index ae14a73..71f8fff 100644 --- a/tests/test_app_server_client.py +++ b/tests/test_app_server_client.py @@ -3,6 +3,9 @@ import asyncio import concurrent.futures import inspect +import subprocess +import sys +import textwrap import threading import time from collections.abc import Callable @@ -545,36 +548,71 @@ async def scenario() -> None: ), ) - await client.start() - thread = await client.start_thread() + async with client: + thread = await client.start_thread() - assert transport.started is True - initialize_request = transport.wait_for_method("initialize") - assert initialize_request["params"]["clientInfo"] == { - "name": "pytest-client", - "title": "Pytest Client", - "version": "1.2.3", - } - assert initialize_request["params"]["capabilities"] == { - "experimentalApi": True, - "extensions": {"openai/form": {}}, - "optOutNotificationMethods": ["item/agentMessage/delta"], - } - assert transport.wait_for_method("initialized") == {"method": "initialized", "params": {}} - assert transport.wait_for_method("thread/start") == { - "id": 1, - "method": "thread/start", - "params": {}, - } - assert thread.id == "thr-1" - assert isinstance(thread.snapshot, protocol.Thread) - assert thread.snapshot.cwd.root == "/repo" - - await client.close() + assert transport.started is True + initialize_request = transport.wait_for_method("initialize") + assert initialize_request["params"]["clientInfo"] == { + "name": "pytest-client", + "title": "Pytest Client", + "version": "1.2.3", + } + assert initialize_request["params"]["capabilities"] == { + "experimentalApi": True, + "extensions": {"openai/form": {}}, + "optOutNotificationMethods": ["item/agentMessage/delta"], + } + assert transport.wait_for_method("initialized") == { + "method": "initialized", + "params": {}, + } + assert transport.wait_for_method("thread/start") == { + "id": 1, + "method": "thread/start", + "params": {}, + } + assert thread.id == "thr-1" + assert isinstance(thread.snapshot, protocol.Thread) + assert thread.snapshot.cwd.root == "/repo" asyncio.run(scenario()) +def test_async_client_context_closes_transport_after_assertion_failure() -> None: + # Bound the whole process: an unclosed Queue.get worker can stall asyncio.run teardown. + subprocess.run( + [ + sys.executable, + "-c", + textwrap.dedent("""\ + import asyncio + import sys + sys.path.insert(0, "tests") + from test_app_server_client import AsyncAppServerClient, ScriptedTransport + + transport = ScriptedTransport() + + async def scenario(): + async with AsyncAppServerClient(transport): + raise AssertionError("scenario failed") + + try: + asyncio.run(scenario()) + except AssertionError as exc: + assert str(exc) == "scenario failed" + assert transport.closed + else: + raise AssertionError("scenario failure was hidden") + """), + ], + check=True, + timeout=10, + capture_output=True, + text=True, + ) + + def test_app_server_thread_start_options_serialize_with_camel_case_aliases() -> None: params = AppServerThreadStartOptions( allow_provider_model_fallback=True, @@ -2246,252 +2284,254 @@ def windows_sandbox_setup_start(message: JsonObject) -> JsonObject: ) transport.responses["windowsSandbox/setupStart"] = windows_sandbox_setup_start client = AsyncAppServerClient(transport) - await client.start() - - models = await client.models.list(limit=20, include_hidden=False) - model_page = await client.models.list_page(limit=20, include_hidden=False) - apps = await client.apps.list( - cursor="cursor-1", - force_refetch=True, - limit=10, - thread_id="thr-1", - ) - app_page = await client.apps.list_page( - cursor="cursor-1", - force_refetch=True, - limit=10, - thread_id="thr-1", - ) - skills = await client.skills.list( - cwds=["/repo"], - force_reload=True, - ) - skills_result = await client.skills.list_page( - cwds=["/repo"], - force_reload=True, - ) - skill_config = await client.skills.write_config( - path="/repo/.codex/skills/skill-creator/SKILL.md", - enabled=False, - ) - account = await client.account.read(refresh_token=True) - api_key_login = await client.account.login_api_key(api_key="sk-test") - chatgpt_login = await client.account.login_chatgpt() - chatgpt_tokens_login = await client.account.login_chatgpt_tokens( - access_token="access-token", - chatgpt_account_id="acct-1", - chatgpt_plan_type=protocol.PlanType("enterprise"), - ) - canceled_login = await client.account.cancel_login(login_id="login-1") - logout_result = await client.account.logout() - rate_limits = await client.account.read_rate_limits() - config = await client.config.read(cwd="/repo", include_layers=True) - write_result = await client.config.write_value( - key_path="model", - value="gpt-5.4", - merge_strategy="replace", - expected_version="v1", - ) - batch_result = await client.config.batch_write( - edits=[ - protocol.ConfigEdit( - keyPath="apps.demo.enabled", - mergeStrategy="upsert", - value=True, - ) - ], - file_path="/home/user/.codex/config.toml", - ) - requirements = await client.config.read_requirements() - reload_result = await client.config.reload_mcp_servers() - oauth_result = await client.mcp_servers.oauth_login( - client_registration=protocol.McpServerOauthClientRegistration("dcr"), - name="github", - scopes=["repo"], - thread_id="thr-1", - timeout_seconds=30, - ) - mcp_status = await client.mcp_servers.list(cursor="cursor-2", limit=5) - mcp_status_page = await client.mcp_servers.list_page(cursor="cursor-2", limit=5) - mcp_status_alias = await client.mcp_servers.list_status(cursor="cursor-2", limit=5) - mcp_status_page_alias = await client.mcp_servers.list_status_page( - cursor="cursor-2", - limit=5, - ) - feedback = await client.feedback.upload( - classification="bug", - include_logs=True, - extra_log_files=["/tmp/app.log"], - reason="Needs follow-up", - thread_id="thr-1", - ) - command = await client.command.execute( - command=["git", "status"], - cwd="/repo", - sandbox_policy=protocol.WorkspaceWriteSandboxPolicy( - type="workspaceWrite", - networkAccess=True, - ), - timeout_ms=5000, - ) - created_dir = await client.fs.create_directory( - path="/repo/.codex/skills/generated", - recursive=True, - ) - wrote_file = await client.fs.write_file( - path="/repo/.codex/skills/generated/SKILL.md", - data="# Generated Skill\n", - ) - environment = await client.environment.info(environment_id="environment-1") - generated_skill = await client.skills.write_skill( - name="generated", - directory="/repo/.codex/skills/generated", - instructions="# Generated Skill\n", - reload_cwds=["/repo"], - ) - detected = await client.external_agent_config.detect( - cwds=["/repo"], - include_home=True, - max_session_age_days=30, - max_sessions=100, - migration_source="claude", - ) - import_result = await client.external_agent_config.import_items( - migration_items=[ - protocol.ExternalAgentConfigMigrationItem( - itemType="AGENTS_MD", - description="Import CLAUDE.md", - cwd="/repo", - ) - ], - migration_source="claude", - provider_id="provider-1", - source="pytest", - ) - recorded_history = await client.external_agent_config.record_history( - item_type_results=[ - protocol.ExternalAgentConfigImportHistoryRecordTypeResultParams( - itemType="AGENTS_MD", - successes=[], - failures=[], - ) - ], - provider_id="provider-1", - ) - windows_setup = await client.windows_sandbox.setup_start(mode="elevated", cwd="C:/repo") - - assert models[0].display_name == "GPT-5.4" - assert models[0].additional_speed_tiers == ["flex"] - assert models[0].model_specialty == "coding" - assert models[0].multi_agent_version == protocol.MultiAgentVersion("v2") - assert models[0].upgrade_info is not None - assert models[0].upgrade_info.retirement_at == 1800000000 - assert model_page.data[0].display_name == "GPT-5.4" - assert model_page.data[0].additional_speed_tiers == ["flex"] - assert apps[0].id == "demo-app" - assert app_page.data[0].id == "demo-app" - assert skills[0].cwd == "/repo" - assert skills[0].errors[0].message == "missing dependency" - assert skills[0].errors[0].path == "/repo/.codex/skills/broken/SKILL.md" - assert skills[0].skills[0].dependencies is not None - assert skills[0].skills[0].dependencies.tools[0].value == "git" - assert skills[0].skills[0].interface is not None - assert skills[0].skills[0].interface.display_name == "Skill Creator" - assert skills[0].skills[0].plugin_id == "plugin-1" - assert skills[0].skills[0].short_description == "Create or update skills" - assert skills_result.data[0].cwd == "/repo" - assert skill_config.effective_enabled is False - assert account.account is not None - assert account.account.type == "chatgpt" - assert api_key_login.type == "apiKey" - assert chatgpt_login.login_id == "login-1" - assert chatgpt_tokens_login.type == "chatgptAuthTokens" - assert canceled_login.status == "canceled" - assert logout_result == EmptyResult() - assert rate_limits.rate_limits.limitId == "codex" - assert rate_limits.account_id == "acct-1" - assert rate_limits.rate_limit_reset_credits is not None - assert rate_limits.rate_limit_reset_credits.availableCount == 2 - assert rate_limits.rate_limit_upsell == {"banner_type": "usage", "dismissed_at": None} - assert config.config.model == "gpt-5.4" - assert write_result.version == "v2" - assert batch_result.version == "v3" - assert requirements.requirements is not None - assert ( - requirements.requirements.additional_developer_instructions == "Follow managed policy." - ) - assert requirements.requirements.allow_browser_and_computer_use is True - assert requirements.requirements.auto_review is not None - assert requirements.requirements.auto_review.ignoreRules == ["safe-command"] - assert requirements.requirements.chatgpt_base_url == "https://chatgpt.example.com" - assert requirements.requirements.cli_auth_credentials_store is not None - assert requirements.requirements.cli_auth_credentials_store.root == "file" - assert requirements.requirements.in_app_browser is not None - assert requirements.requirements.in_app_browser.allowExternalBrowserSettingsImport is False - assert requirements.requirements.allowed_sandbox_modes is not None - assert [mode.root for mode in requirements.requirements.allowed_sandbox_modes] == [ - "read-only", - "workspace-write", - ] - assert requirements.requirements.allowed_approvals_reviewers is not None - assert [ - reviewer.root for reviewer in requirements.requirements.allowed_approvals_reviewers - ] == ["user"] - assert requirements.requirements.allowed_web_search_modes is not None - assert [mode.root for mode in requirements.requirements.allowed_web_search_modes] == [ - "disabled", - "live", - ] - assert requirements.requirements.enforce_residency is not None - assert requirements.requirements.enforce_residency.root == "us" - assert requirements.requirements.feature_requirements == {"personality": {"required": True}} - assert requirements.requirements.network is not None - assert requirements.requirements.network.enabled is True - assert requirements.requirements.network.allowedDomains == ["api.openai.com"] - assert requirements.requirements.network.deniedDomains == ["example.invalid"] - assert requirements.requirements.network.managedAllowedDomainsOnly is True - assert reload_result == EmptyResult() - assert oauth_result.authorization_url == "https://example.com/oauth" - assert mcp_status[0].name == "github" - assert isinstance(mcp_status[0].auth_status, protocol.McpAuthStatus) - assert mcp_status[0].auth_status.root == "oAuth" - assert mcp_status[0].plugin_id == "plugin-1" - assert mcp_status[0].runtime_status == protocol.McpServerConnectionStatus("connected") - assert isinstance(mcp_status[0].tools["repo_status"], protocol.Tool) - assert mcp_status[0].tools["repo_status"].field_meta == {"origin": "pytest"} - assert mcp_status[0].tools["repo_status"].inputSchema == { - "type": "object", - "properties": {}, - } - assert mcp_status[0].tools["repo_status"].outputSchema == {"type": "object"} - assert isinstance(mcp_status[0].resources[0], protocol.Resource) - assert mcp_status[0].resources[0].field_meta == {"origin": "pytest"} - assert mcp_status[0].resources[0].mimeType == "text/markdown" - assert mcp_status[0].resources[0].uri == "file:///repo/README.md" - assert isinstance(mcp_status[0].resource_templates[0], protocol.ResourceTemplate) - assert mcp_status[0].resource_templates[0].uriTemplate == "file:///repo/{path}" - assert mcp_status_page.data[0].name == "github" - assert mcp_status_alias[0].name == "github" - assert mcp_status_page_alias.data[0].name == "github" - assert feedback.thread_id == "thr-feedback" - assert command.exit_code == 0 - assert isinstance(created_dir, protocol.FsCreateDirectoryResponse) - assert isinstance(wrote_file, protocol.FsWriteFileResponse) - assert environment.cwd == protocol.PathUri("file:///repo") - assert environment.shell == protocol.EnvironmentShellInfo(name="zsh", path="/bin/zsh") - assert generated_skill == protocol.SkillUserInput( - type=protocol.SkillUserInputType("skill"), - name="generated", - path="/repo/.codex/skills/generated/SKILL.md", - ) - assert detected.connectors[0].name == "knowledge" - assert detected.connectors[0].sessionCount == 3 - assert detected.items[0].itemType.root == "AGENTS_MD" - assert import_result.import_id == "import-1" - assert recorded_history.import_id == "import-2" - assert windows_setup.started is True - - await client.close() + async with client: + models = await client.models.list(limit=20, include_hidden=False) + model_page = await client.models.list_page(limit=20, include_hidden=False) + apps = await client.apps.list( + cursor="cursor-1", + force_refetch=True, + limit=10, + thread_id="thr-1", + ) + app_page = await client.apps.list_page( + cursor="cursor-1", + force_refetch=True, + limit=10, + thread_id="thr-1", + ) + skills = await client.skills.list( + cwds=["/repo"], + force_reload=True, + ) + skills_result = await client.skills.list_page( + cwds=["/repo"], + force_reload=True, + ) + skill_config = await client.skills.write_config( + path="/repo/.codex/skills/skill-creator/SKILL.md", + enabled=False, + ) + account = await client.account.read(refresh_token=True) + api_key_login = await client.account.login_api_key(api_key="sk-test") + chatgpt_login = await client.account.login_chatgpt() + chatgpt_tokens_login = await client.account.login_chatgpt_tokens( + access_token="access-token", + chatgpt_account_id="acct-1", + chatgpt_plan_type=protocol.PlanType("enterprise"), + ) + canceled_login = await client.account.cancel_login(login_id="login-1") + logout_result = await client.account.logout() + rate_limits = await client.account.read_rate_limits() + config = await client.config.read(cwd="/repo", include_layers=True) + write_result = await client.config.write_value( + key_path="model", + value="gpt-5.4", + merge_strategy="replace", + expected_version="v1", + ) + batch_result = await client.config.batch_write( + edits=[ + protocol.ConfigEdit( + keyPath="apps.demo.enabled", + mergeStrategy="upsert", + value=True, + ) + ], + file_path="/home/user/.codex/config.toml", + ) + requirements = await client.config.read_requirements() + reload_result = await client.config.reload_mcp_servers() + oauth_result = await client.mcp_servers.oauth_login( + client_registration=protocol.McpServerOauthClientRegistration("dcr"), + name="github", + scopes=["repo"], + thread_id="thr-1", + timeout_seconds=30, + ) + mcp_status = await client.mcp_servers.list(cursor="cursor-2", limit=5) + mcp_status_page = await client.mcp_servers.list_page(cursor="cursor-2", limit=5) + mcp_status_alias = await client.mcp_servers.list_status(cursor="cursor-2", limit=5) + mcp_status_page_alias = await client.mcp_servers.list_status_page( + cursor="cursor-2", + limit=5, + ) + feedback = await client.feedback.upload( + classification="bug", + include_logs=True, + extra_log_files=["/tmp/app.log"], + reason="Needs follow-up", + thread_id="thr-1", + ) + command = await client.command.execute( + command=["git", "status"], + cwd="/repo", + sandbox_policy=protocol.WorkspaceWriteSandboxPolicy( + type="workspaceWrite", + networkAccess=True, + ), + timeout_ms=5000, + ) + created_dir = await client.fs.create_directory( + path="/repo/.codex/skills/generated", + recursive=True, + ) + wrote_file = await client.fs.write_file( + path="/repo/.codex/skills/generated/SKILL.md", + data="# Generated Skill\n", + ) + environment = await client.environment.info(environment_id="environment-1") + generated_skill = await client.skills.write_skill( + name="generated", + directory="/repo/.codex/skills/generated", + instructions="# Generated Skill\n", + reload_cwds=["/repo"], + ) + detected = await client.external_agent_config.detect( + cwds=["/repo"], + include_home=True, + max_session_age_days=30, + max_sessions=100, + migration_source="claude", + ) + import_result = await client.external_agent_config.import_items( + migration_items=[ + protocol.ExternalAgentConfigMigrationItem( + itemType="AGENTS_MD", + description="Import CLAUDE.md", + cwd="/repo", + ) + ], + migration_source="claude", + provider_id="provider-1", + source="pytest", + ) + recorded_history = await client.external_agent_config.record_history( + item_type_results=[ + protocol.ExternalAgentConfigImportHistoryRecordTypeResultParams( + itemType="AGENTS_MD", + successes=[], + failures=[], + ) + ], + provider_id="provider-1", + ) + windows_setup = await client.windows_sandbox.setup_start(mode="elevated", cwd="C:/repo") + + assert models[0].display_name == "GPT-5.4" + assert models[0].additional_speed_tiers == ["flex"] + assert models[0].model_specialty == "coding" + assert models[0].multi_agent_version == protocol.MultiAgentVersion("v2") + assert models[0].upgrade_info is not None + assert models[0].upgrade_info.retirement_at == 1800000000 + assert model_page.data[0].display_name == "GPT-5.4" + assert model_page.data[0].additional_speed_tiers == ["flex"] + assert apps[0].id == "demo-app" + assert app_page.data[0].id == "demo-app" + assert skills[0].cwd == "/repo" + assert skills[0].errors[0].message == "missing dependency" + assert skills[0].errors[0].path == "/repo/.codex/skills/broken/SKILL.md" + assert skills[0].skills[0].dependencies is not None + assert skills[0].skills[0].dependencies.tools[0].value == "git" + assert skills[0].skills[0].interface is not None + assert skills[0].skills[0].interface.display_name == "Skill Creator" + assert skills[0].skills[0].plugin_id == "plugin-1" + assert skills[0].skills[0].short_description == "Create or update skills" + assert skills_result.data[0].cwd == "/repo" + assert skill_config.effective_enabled is False + assert account.account is not None + assert account.account.type == "chatgpt" + assert api_key_login.type == "apiKey" + assert chatgpt_login.login_id == "login-1" + assert chatgpt_tokens_login.type == "chatgptAuthTokens" + assert canceled_login.status == "canceled" + assert logout_result == EmptyResult() + assert rate_limits.rate_limits.limitId == "codex" + assert rate_limits.account_id == "acct-1" + assert rate_limits.rate_limit_reset_credits is not None + assert rate_limits.rate_limit_reset_credits.availableCount == 2 + assert rate_limits.rate_limit_upsell == {"banner_type": "usage", "dismissed_at": None} + assert config.config.model == "gpt-5.4" + assert write_result.version == "v2" + assert batch_result.version == "v3" + assert requirements.requirements is not None + assert ( + requirements.requirements.additional_developer_instructions + == "Follow managed policy." + ) + assert requirements.requirements.allow_browser_and_computer_use is True + assert requirements.requirements.auto_review is not None + assert requirements.requirements.auto_review.ignoreRules == ["safe-command"] + assert requirements.requirements.chatgpt_base_url == "https://chatgpt.example.com" + assert requirements.requirements.cli_auth_credentials_store is not None + assert requirements.requirements.cli_auth_credentials_store.root == "file" + assert requirements.requirements.in_app_browser is not None + assert ( + requirements.requirements.in_app_browser.allowExternalBrowserSettingsImport is False + ) + assert requirements.requirements.allowed_sandbox_modes is not None + assert [mode.root for mode in requirements.requirements.allowed_sandbox_modes] == [ + "read-only", + "workspace-write", + ] + assert requirements.requirements.allowed_approvals_reviewers is not None + assert [ + reviewer.root for reviewer in requirements.requirements.allowed_approvals_reviewers + ] == ["user"] + assert requirements.requirements.allowed_web_search_modes is not None + assert [mode.root for mode in requirements.requirements.allowed_web_search_modes] == [ + "disabled", + "live", + ] + assert requirements.requirements.enforce_residency is not None + assert requirements.requirements.enforce_residency.root == "us" + assert requirements.requirements.feature_requirements == { + "personality": {"required": True} + } + assert requirements.requirements.network is not None + assert requirements.requirements.network.enabled is True + assert requirements.requirements.network.allowedDomains == ["api.openai.com"] + assert requirements.requirements.network.deniedDomains == ["example.invalid"] + assert requirements.requirements.network.managedAllowedDomainsOnly is True + assert reload_result == EmptyResult() + assert oauth_result.authorization_url == "https://example.com/oauth" + assert mcp_status[0].name == "github" + assert isinstance(mcp_status[0].auth_status, protocol.McpAuthStatus) + assert mcp_status[0].auth_status.root == "oAuth" + assert mcp_status[0].plugin_id == "plugin-1" + assert mcp_status[0].runtime_status == protocol.McpServerConnectionStatus("connected") + assert isinstance(mcp_status[0].tools["repo_status"], protocol.Tool) + assert mcp_status[0].tools["repo_status"].field_meta == {"origin": "pytest"} + assert mcp_status[0].tools["repo_status"].inputSchema == { + "type": "object", + "properties": {}, + } + assert mcp_status[0].tools["repo_status"].outputSchema == {"type": "object"} + assert isinstance(mcp_status[0].resources[0], protocol.Resource) + assert mcp_status[0].resources[0].field_meta == {"origin": "pytest"} + assert mcp_status[0].resources[0].mimeType == "text/markdown" + assert mcp_status[0].resources[0].uri == "file:///repo/README.md" + assert isinstance(mcp_status[0].resource_templates[0], protocol.ResourceTemplate) + assert mcp_status[0].resource_templates[0].uriTemplate == "file:///repo/{path}" + assert mcp_status_page.data[0].name == "github" + assert mcp_status_alias[0].name == "github" + assert mcp_status_page_alias.data[0].name == "github" + assert feedback.thread_id == "thr-feedback" + assert command.exit_code == 0 + assert isinstance(created_dir, protocol.FsCreateDirectoryResponse) + assert isinstance(wrote_file, protocol.FsWriteFileResponse) + assert environment.cwd == protocol.PathUri("file:///repo") + assert environment.shell == protocol.EnvironmentShellInfo(name="zsh", path="/bin/zsh") + assert generated_skill == protocol.SkillUserInput( + type=protocol.SkillUserInputType("skill"), + name="generated", + path="/repo/.codex/skills/generated/SKILL.md", + ) + assert detected.connectors[0].name == "knowledge" + assert detected.connectors[0].sessionCount == 3 + assert detected.items[0].itemType.root == "AGENTS_MD" + assert import_result.import_id == "import-1" + assert recorded_history.import_id == "import-2" + assert windows_setup.started is True asyncio.run(scenario()) diff --git a/tests/test_app_server_session.py b/tests/test_app_server_session.py index c204c96..ab724d6 100644 --- a/tests/test_app_server_session.py +++ b/tests/test_app_server_session.py @@ -1,11 +1,12 @@ from __future__ import annotations import asyncio +import gc from collections.abc import Mapping from typing import Any import pytest -from pydantic import BaseModel +from pydantic import BaseModel, RootModel from codex.app_server._session import _AsyncSession, _jsonrpc_error_from_exception from codex.app_server.errors import AppServerClosedError, AppServerProtocolError, AppServerRpcError @@ -201,6 +202,28 @@ async def scenario() -> None: asyncio.run(scenario()) +def test_async_session_start_preserves_error_and_notes_cleanup_failure() -> None: + async def scenario() -> None: + class FailingCloseTransport(_FakeTransport): + async def close(self) -> None: + await super().close() + raise RuntimeError("cleanup failed") + + transport = FailingCloseTransport() + transport._fail_send_methods["initialized"] = RuntimeError("notify failed") + session = _AsyncSession(transport) + with pytest.raises(RuntimeError, match="notify failed") as exc_info: + await session.start() + + assert exc_info.value.__notes__ == [ + "Cleanup after start failure also failed: RuntimeError('cleanup failed')" + ] + assert transport.closed + await session.close() + + asyncio.run(asyncio.wait_for(scenario(), timeout=1)) + + def test_async_session_subscription_close_discards_buffered_notifications() -> None: async def scenario() -> None: session = _AsyncSession(_FakeTransport(), AppServerInitializeOptions(strict_protocol=False)) @@ -229,6 +252,7 @@ async def scenario() -> None: with pytest.raises(AppServerProtocolError, match="Unsupported app-server message"): await session.close() + await session.close() asyncio.run(scenario()) @@ -275,3 +299,251 @@ class _WrongRequest(BaseModel): await session.close() asyncio.run(scenario()) + + +def test_async_session_request_serialization_failure_leaves_no_pending_request() -> None: + async def scenario() -> None: + transport = _FakeTransport() + session = _AsyncSession(transport) + await session.start() + + with pytest.raises(TypeError, match="Request params must serialize to an object"): + await session.request("invalid", RootModel[list[int]]([1])) + + assert session._pending == {} + assert all(message.get("method") != "invalid" for message in transport.sent) + assert await session.request("initialize") == {"userAgent": "test-client"} + await session.close() + + asyncio.run(asyncio.wait_for(scenario(), timeout=1)) + + +@pytest.mark.parametrize("failure", ["send", "cancel_send", "cancel_wait"]) +def test_async_session_request_abandonment_cleans_up_and_ignores_late_response( + failure: str, +) -> None: + async def scenario() -> None: + transport = _FakeTransport() + session = _AsyncSession(transport) + await session.start() + send_entered = asyncio.Event() + original_send = transport.send + + async def interrupted_send(message: JsonObject) -> None: + await original_send(message) + if message.get("method") == "abandoned": + send_entered.set() + if failure == "send": + raise RuntimeError("send failed") + if failure == "cancel_send": + await asyncio.Event().wait() + + transport.send = interrupted_send # type: ignore[method-assign] + request = asyncio.create_task(session.request("abandoned")) + await send_entered.wait() + request_id = transport.sent[-1]["id"] + if failure == "send": + with pytest.raises(RuntimeError, match="send failed"): + await request + else: + future = session._pending[request_id] + request.cancel() + with pytest.raises(asyncio.CancelledError): + await request + assert future.cancelled() + + assert session._pending == {} + transport.push({"id": request_id, "error": {"code": -1, "message": "late"}}) + assert await session.request("initialize") == {"userAgent": "test-client"} + await session.close() + + asyncio.run(asyncio.wait_for(scenario(), timeout=1)) + + +@pytest.mark.parametrize("response_error", [False, True]) +def test_async_session_request_cancellation_observes_racing_response(response_error: bool) -> None: + async def scenario() -> None: + transport = _FakeTransport() + session = _AsyncSession(transport) + await session.start() + loop = asyncio.get_running_loop() + unhandled: list[dict[str, Any]] = [] + loop.set_exception_handler(lambda loop, context: unhandled.append(context)) + + async def respond_and_cancel(message: JsonObject) -> None: + transport.sent.append(message) + response: JsonObject = {"id": message["id"]} + if response_error: + response["error"] = {"code": -1, "message": "raced error"} + else: + response["result"] = {"ok": True} + session._handle_response(response) + task = asyncio.current_task() + assert task is not None + task.cancel() + await asyncio.sleep(0) + + transport.send = respond_and_cancel # type: ignore[method-assign] + request = asyncio.create_task(session.request("race")) + with pytest.raises(asyncio.CancelledError): + await request + assert session._pending == {} + del request + gc.collect() + assert unhandled == [] + await session.close() + + asyncio.run(asyncio.wait_for(scenario(), timeout=1)) + + +def test_async_session_close_fails_and_removes_pending_request() -> None: + async def scenario() -> None: + session = _AsyncSession(_FakeTransport()) + await session.start() + request = asyncio.create_task(session.request("pending")) + await asyncio.sleep(0) + assert len(session._pending) == 1 + + await session.close() + + with pytest.raises(AppServerClosedError): + await request + assert session._pending == {} + + asyncio.run(asyncio.wait_for(scenario(), timeout=1)) + + +def test_async_session_close_does_not_replay_reader_failure_reported_by_request() -> None: + async def scenario() -> None: + transport = _FakeTransport() + session = _AsyncSession(transport) + await session.start() + request = asyncio.create_task(session.request("pending")) + await asyncio.sleep(0) + transport.push({"unexpected": "message"}) + + with pytest.raises(AppServerProtocolError, match="Unsupported app-server message"): + await request + assert session._pending == {} + await session.close() + assert transport.closed + + asyncio.run(asyncio.wait_for(scenario(), timeout=1)) + + +@pytest.mark.parametrize("blocked_stage", ["reader", "transport"]) +@pytest.mark.parametrize("cancel_first_waiter", [False, True]) +def test_async_session_close_joins_cleanup_and_survives_waiter_cancellation( + blocked_stage: str, cancel_first_waiter: bool +) -> None: + async def scenario() -> None: + cleanup_entered = asyncio.Event() + cleanup_release = asyncio.Event() + + class BlockingTransport(_FakeTransport): + close_calls = 0 + + async def receive(self) -> JsonObject | None: + try: + return await super().receive() + except asyncio.CancelledError: + if blocked_stage == "reader": + cleanup_entered.set() + await cleanup_release.wait() + raise + + async def close(self) -> None: + self.close_calls += 1 + if blocked_stage == "transport": + cleanup_entered.set() + await cleanup_release.wait() + await super().close() + + transport = BlockingTransport() + session = _AsyncSession(transport) + await session.start() + subscription = session.subscribe_notifications() + first_close = asyncio.create_task(session.close()) + await cleanup_entered.wait() + second_close = asyncio.create_task(session.close()) + await asyncio.sleep(0) + assert not first_close.done() + assert not second_close.done() + + with pytest.raises(AppServerClosedError, match="closed"): + await session.start() + with pytest.raises(AppServerClosedError, match="closed"): + await session.request("too-late") + with pytest.raises(AppServerClosedError, match="closed"): + await session.notify("too-late") + + if cancel_first_waiter: + first_close.cancel() + with pytest.raises(asyncio.CancelledError): + await first_close + assert not second_close.done() + + cleanup_release.set() + await second_close + if not cancel_first_waiter: + await first_close + await session.close() + + assert transport.close_calls == 1 + assert transport.closed + assert session._reader_task is None + assert session._notification_sinks == [] + assert session._pending == {} + with pytest.raises(StopAsyncIteration): + await subscription.next() + + asyncio.run(asyncio.wait_for(scenario(), timeout=1)) + + +@pytest.mark.parametrize("cancel_waiter", [False, True]) +def test_async_session_close_observes_cleanup_failure_and_does_not_replay_it( + cancel_waiter: bool, +) -> None: + async def scenario() -> None: + cleanup_entered = asyncio.Event() + cleanup_release = asyncio.Event() + cleanup_finished = asyncio.Event() + unhandled: list[dict[str, Any]] = [] + asyncio.get_running_loop().set_exception_handler( + lambda loop, context: unhandled.append(context) + ) + + class FailingCloseTransport(_FakeTransport): + async def close(self) -> None: + cleanup_entered.set() + await cleanup_release.wait() + await super().close() + raise RuntimeError("close failed") + + transport = FailingCloseTransport() + session = _AsyncSession(transport) + await session.start() + first_close = asyncio.create_task(session.close()) + await cleanup_entered.wait() + assert session._close_task is not None + session._close_task.add_done_callback(lambda task: cleanup_finished.set()) + + if cancel_waiter: + first_close.cancel() + with pytest.raises(asyncio.CancelledError): + await first_close + cleanup_release.set() + if not cancel_waiter: + with pytest.raises(RuntimeError, match="close failed"): + await first_close + await cleanup_finished.wait() + await session.close() + + assert transport.closed + assert session._reader_task is None + del first_close, session + await asyncio.sleep(0) + gc.collect() + assert unhandled == [] + + asyncio.run(asyncio.wait_for(scenario(), timeout=1)) diff --git a/tests/test_app_server_transports.py b/tests/test_app_server_transports.py index ca588b3..f8e0c7b 100644 --- a/tests/test_app_server_transports.py +++ b/tests/test_app_server_transports.py @@ -371,23 +371,27 @@ async def scenario() -> None: asyncio.run(scenario()) -def test_stdio_transport_close_terminates_and_waits() -> None: +def test_stdio_transport_close_terminates_and_drains_stderr() -> None: async def scenario() -> None: + stderr = _FakeStreamReader([b"stderr line\n", b"\xff\xfe\n", b"last line\n"]) process = _FakeProcess( stdin=_FakeStreamWriter(), stdout=_FakeStreamReader([]), - stderr=_FakeStreamReader([b"stderr line\n"]), + stderr=stderr, ) transport = AsyncStdioTransport() transport._process = process - transport._stderr_task = asyncio.create_task(transport._drain_stderr(process.stderr)) + drain_task = asyncio.create_task(transport._drain_stderr(stderr)) + transport._stderr_task = drain_task await transport.close() assert process.stdin is not None and process.stdin.closed is True assert process.terminated is True assert process.wait_calls >= 1 - assert transport._stderr_lines == ["stderr line"] + assert stderr._chunks == [] + assert drain_task.done() + assert transport._stderr_task is None asyncio.run(scenario()) diff --git a/tests/test_config_overrides.py b/tests/test_config_overrides.py index 60f6fb5..e09a384 100644 --- a/tests/test_config_overrides.py +++ b/tests/test_config_overrides.py @@ -3,16 +3,10 @@ from pathlib import Path import pytest -from pydantic import BaseModel from codex import _binary from codex._binary import bundled_app_server_path, resolve_target_triple from codex.errors import CodexExecError -from codex.output_schema_file import create_output_schema_file - - -class AnswerSchema(BaseModel): - answer: str def test_resolve_target_triple() -> None: @@ -43,30 +37,3 @@ def test_bundled_app_server_path_resolves_when_binary_exists( def test_bundled_app_server_path_raises_when_missing() -> None: with pytest.raises(CodexExecError, match="Bundled codex app-server binary not found"): bundled_app_server_path("missing-target") - - -def test_output_schema_file_lifecycle() -> None: - schema = {"type": "object", "properties": {"answer": {"type": "string"}}} - output_schema = create_output_schema_file(schema) - - assert output_schema.schema_path is not None - schema_path = Path(output_schema.schema_path) - assert schema_path.exists() - output_schema.cleanup() - assert not schema_path.exists() - - -def test_output_schema_file_accepts_pydantic_model_class() -> None: - output_schema = create_output_schema_file(AnswerSchema) - - assert output_schema.schema_path is not None - schema_path = Path(output_schema.schema_path) - assert schema_path.exists() - assert schema_path.read_text(encoding="utf-8") - output_schema.cleanup() - assert not schema_path.exists() - - -def test_output_schema_requires_json_object_or_pydantic_model_class() -> None: - with pytest.raises(ValueError, match="JSON object or a Pydantic model class"): - create_output_schema_file(["not", "an", "object"]) diff --git a/tests/test_dynamic_tools.py b/tests/test_dynamic_tools.py index 36cca72..9cd385d 100644 --- a/tests/test_dynamic_tools.py +++ b/tests/test_dynamic_tools.py @@ -5,7 +5,7 @@ from typing import Annotated import pytest -from pydantic import Field +from pydantic import BaseModel, Field, field_serializer from codex.dynamic_tools import ( _DynamicToolRuntime, @@ -15,6 +15,14 @@ from codex.protocol import types as protocol +class _Ticket(BaseModel): + id: int + + @field_serializer("id") + def serialize_id(self, value: int) -> str: + raise AssertionError("Callback arguments must not be serialized") + + def test_dynamic_tool_derives_schema_from_typed_parameters() -> None: @dynamic_tool def lookup_ticket( @@ -109,6 +117,33 @@ async def preview_image(url: str) -> list[dict[str, str]]: assert response.contentItems[0].root.imageUrl == "https://example.test/image.png" +def test_dynamic_tool_runtime_preserves_validated_nested_arguments() -> None: + @dynamic_tool(description="Sum ticket identifiers.") + def sum_tickets(ticket: _Ticket, related: list[_Ticket], extra: int = 1) -> int: + return ticket.id + sum(item.id for item in related) + extra + + runtime = _DynamicToolRuntime(lambda method, handler, request_model: None) + runtime.activate("thr-1", resolve_dynamic_tools([sum_tickets])) + request = protocol.ItemToolCallRequest.model_validate( + { + "id": "req-1", + "method": "item/tool/call", + "params": { + "callId": "call-1", + "threadId": "thr-1", + "turnId": "turn-1", + "tool": "sum_tickets", + "arguments": {"ticket": {"id": "2"}, "related": [{"id": "3"}]}, + }, + } + ) + + response = asyncio.run(runtime.dispatch(request)) + + assert response.success is True + assert json.loads(response.contentItems[0].root.text) == 6 + + def test_dynamic_tool_rejects_invalid_signatures() -> None: @dynamic_tool(description="Missing annotation.") def missing_annotation(ticket_id) -> str: # type: ignore[no-untyped-def] diff --git a/tests/test_protocol_generation_scripts.py b/tests/test_protocol_generation_scripts.py index fd70115..b11dce5 100644 --- a/tests/test_protocol_generation_scripts.py +++ b/tests/test_protocol_generation_scripts.py @@ -2,6 +2,8 @@ import builtins import importlib.util +import stat +import subprocess import sys from pathlib import Path from types import ModuleType @@ -110,6 +112,79 @@ def fake_run_stage(name: str, command: list[str]) -> None: } +@pytest.mark.parametrize("failure", ["extra", "postprocess", None]) +@pytest.mark.parametrize("destination", ["existing", "absent", "symlink"]) +def test_generate_protocol_types_publishes_only_complete_output( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + failure: str | None, + destination: str, +) -> None: + module = _load_script_module("generate_protocol_types", "scripts/generate_protocol_types.py") + postprocessor = _load_script_module( + "postprocess_protocol_types", "scripts/postprocess_protocol_types.py" + ) + output = tmp_path / "types.py" + original = b"# original contract\n" + target = tmp_path / "linked.py" + if destination == "symlink": + if sys.platform == "win32": + pytest.skip("Creating symlinks requires privileges on Windows") + target.write_bytes(original) + output.symlink_to(target) + elif destination == "existing": + output.write_bytes(original) + output.chmod(0o644) + original_paths = set(tmp_path.iterdir()) + + def fake_run_stage(name: str, command: list[str]) -> None: + if command[0] == "fake-codex": + schema_dir = Path(command[command.index("--out") + 1]) + (schema_dir / "codex_app_server_protocol.schemas.json").write_text("{}") + (schema_dir / "v2").mkdir() + (schema_dir / "v2" / "ExtraResponse.json").write_text("{}") + elif command[0] == "uvx": + generated_path = Path(command[command.index("--output") + 1]) + extra = Path(command[command.index("--input") + 1]).stem == "ExtraResponse" + generated_path.write_text("class Extra: pass\n" if extra else "class Primary: pass\n") + if extra and failure == "extra": + raise subprocess.CalledProcessError(1, command) + else: + assert command[1] == "scripts/postprocess_protocol_types.py" + postprocessor.postprocess_file(Path(command[2])) + if failure == "postprocess": + raise subprocess.CalledProcessError(1, command) + + monkeypatch.setattr(module, "run_stage", fake_run_stage) + monkeypatch.setattr( + sys, + "argv", + ["generate_protocol_types.py", "--codex-bin", "fake-codex", "--output", str(output)], + ) + + if failure is not None: + with pytest.raises(subprocess.CalledProcessError): + module.main() + if destination == "absent": + assert not output.exists() + else: + assert output.read_bytes() == original + assert output.is_symlink() == (destination == "symlink") + else: + assert module.main() == 0 + text = output.read_text() + assert "class Primary: pass" in text + assert "class Extra: pass" in text + assert "from __future__ import annotations" in text + assert not output.is_symlink() + if sys.platform != "win32": + assert stat.S_IMODE(output.stat().st_mode) == 0o600 + + if destination == "symlink": + assert target.read_bytes() == original + assert set(tmp_path.iterdir()) == original_paths | ({output} if failure is None else set()) + + def test_generate_protocol_types_inserts_extra_response_models_before_rebuilds( tmp_path: Path, ) -> None: @@ -219,26 +294,6 @@ def reject_codex_import(name: str, *args: object, **kwargs: object) -> object: assert "class EventMsg(RootModel):" in processed -def test_postprocess_schema_titles_does_not_import_codex_package( - monkeypatch: pytest.MonkeyPatch, -) -> None: - real_import = builtins.__import__ - - def reject_codex_import(name: str, *args: object, **kwargs: object) -> object: - if name == "codex" or name.startswith("codex."): - raise AssertionError(f"unexpected package import: {name}") - return real_import(name, *args, **kwargs) - - monkeypatch.setattr(builtins, "__import__", reject_codex_import) - - module = _load_script_module( - "postprocess_schema_titles_no_package_import", - "scripts/postprocess_schema_titles.py", - ) - - assert module.camelize("turn-completed") == "TurnCompleted" - - def test_postprocess_types_applies_explicit_pipeline_passes() -> None: module = _load_script_module( "postprocess_protocol_types", @@ -340,59 +395,3 @@ class ServerNotification(RootModel[FooNotification | BarNotification]): ' Field(title="ServerNotification"),\n' " ]" ) in processed - - -def test_postprocess_schema_titles_applies_schema_normalization_passes() -> None: - module = _load_script_module( - "postprocess_schema_titles", - "scripts/postprocess_schema_titles.py", - ) - - schema = { - "$defs": { - "ServerNotification": { - "oneOf": [ - { - "properties": { - "method": {"const": "turn-completed"}, - "duration": {"type": "string", "description": "elapsed"}, - "output_tokens": {"type": "number"}, - "nullableValue": {"oneOf": [{"type": "string"}, {"type": "null"}]}, - }, - "required": ["duration", "nullableValue"], - } - ] - }, - "RequestId": {"type": "number"}, - "ExecCommandEndEvent": {"properties": {"exit_code": {"type": "number"}}}, - } - } - - changed, added = module.add_titles(schema) - relaxed, relaxed_count = module.relax_required_for_nullables(schema) - request_id_fixed = module.enforce_request_id_integer(schema) - exit_code_fixed = module.enforce_exec_exit_code_integer(schema) - integers_coerced = module.enforce_integer_fields(schema) - durations_patched = module.enforce_duration_union(schema) - - hoisted = schema["$defs"]["ServerNotification_TurnCompleted"] - - assert module.camelize("turn-completed") == "TurnCompleted" - assert changed is True - assert added == 1 - assert schema["$defs"]["ServerNotification"]["oneOf"] == [ - {"$ref": "#/$defs/ServerNotification_TurnCompleted"} - ] - assert hoisted["title"] == "ServerNotification_TurnCompleted" - assert relaxed is True - assert relaxed_count == 1 - assert hoisted["required"] == ["duration"] - assert request_id_fixed is True - assert schema["$defs"]["RequestId"]["type"] == ["string", "integer"] - assert exit_code_fixed is True - assert schema["$defs"]["ExecCommandEndEvent"]["properties"]["exit_code"]["type"] == "integer" - assert integers_coerced == 1 - assert hoisted["properties"]["output_tokens"]["type"] == "integer" - assert durations_patched == 1 - assert hoisted["properties"]["duration"]["oneOf"][0]["type"] == "string" - assert hoisted["properties"]["duration"]["oneOf"][1]["properties"]["secs"]["type"] == "integer" diff --git a/tests/test_release_workflows.py b/tests/test_release_workflows.py index 78820b9..b081a4e 100644 --- a/tests/test_release_workflows.py +++ b/tests/test_release_workflows.py @@ -1,7 +1,12 @@ from __future__ import annotations +import os +import subprocess +import textwrap from pathlib import Path +import pytest + PINNED_CODEX_BINARY_RELEASE_TAG = "rust-v0.153.4" @@ -48,6 +53,96 @@ def test_release_workflow_rejects_pypi_oversized_files_before_publish() -> None: ) +@pytest.mark.parametrize("version_changed", [False, True]) +def test_next_release_uses_selected_commit_after_branch_advances( + tmp_path: Path, version_changed: bool +) -> None: + workflow = Path(".github/workflows/release-published.yml").read_text() + + def step(name: str) -> str: + return workflow.split(f" - name: {name}\n", 1)[1].split("\n - name:", 1)[0] + + def run_step(name: str, env: dict[str, str]) -> None: + script = textwrap.dedent(step(name).split(" run: |\n", 1)[1]) + script = script.replace("${{ steps.resolve.outputs.version }}", "1.2.3") + script = script.replace("${{ github.event.repository.default_branch }}", "main") + subprocess.run( + ["bash", "-e", "-o", "pipefail", "-c", script], + cwd=repo, + env=env, + check=True, + capture_output=True, + text=True, + ) + + def git(*args: str, cwd: Path) -> str: + return subprocess.run( + ["git", *args], cwd=cwd, check=True, capture_output=True, text=True + ).stdout.strip() + + origin = tmp_path / "origin.git" + repo = tmp_path / "repo" + git("init", "--bare", str(origin), cwd=tmp_path) + git("init", "--initial-branch=main", str(repo), cwd=tmp_path) + git("config", "user.name", "Release test", cwd=repo) + git("config", "user.email", "release@example.invalid", cwd=repo) + git("config", "commit.gpgsign", "false", cwd=repo) + git("config", "tag.gpgsign", "false", cwd=repo) + for relative in ("crates/codex_native/Cargo.toml", "codex/__init__.py"): + path = repo / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("1.2.2" if version_changed else "1.2.3") + git("add", ".", cwd=repo) + git("commit", "-m", "Initial version", cwd=repo) + original_sha = git("rev-parse", "HEAD", cwd=repo) + git("remote", "add", "origin", str(origin), cwd=repo) + git("tag", "next", cwd=repo) + git("push", "origin", "main", "next", cwd=repo) + + if version_changed: + (repo / "crates/codex_native/Cargo.toml").write_text("1.2.3") + (repo / "codex/__init__.py").write_text("1.2.3") + output = tmp_path / "output" + env = dict(os.environ, GITHUB_OUTPUT=str(output), GITHUB_ENV=str(tmp_path / "env")) + run_step("Commit version bump", env) + run_step("Resolve build SHA", env) + build_sha = output.read_text().removeprefix("sha=").strip() + assert (build_sha != original_sha) == version_changed + + # Another push after source selection must not change the tag's source. + other = tmp_path / "other" + git("clone", "--branch", "main", str(origin), str(other), cwd=tmp_path) + git( + "-c", + "user.name=Other", + "-c", + "user.email=other@example.invalid", + "-c", + "commit.gpgsign=false", + "commit", + "--allow-empty", + "-m", + "Branch advances", + cwd=other, + ) + git("push", "origin", "main", cwd=other) + env.update(PLACEHOLDER_TAG="next", VERSION="1.2.3", BUILD_SHA=build_sha) + run_step("Create new semver tag and repoint release", env) + + assert git("rev-parse", "refs/tags/v1.2.3^{commit}", cwd=origin) == build_sha + assert git("rev-parse", "refs/heads/main", cwd=origin) != build_sha + assert workflow.count("ref: ${{ needs.prepare.outputs.build_sha }}") == 2 + assert "BUILD_SHA: ${{ steps.build_sha.outputs.sha }}" in step( + "Create new semver tag and repoint release" + ) + assert workflow.index("Resolve build SHA") < workflow.index( + "Create new semver tag and repoint release" + ) + # Explicit releases retain their existing tags. + for name in ("Create new semver tag and repoint release", "Update GitHub Release to new tag"): + assert "if: ${{ steps.resolve.outputs.mode == 'bump' }}" in step(name) + + def test_autoreview_workflow_fetches_codex_binary_before_action() -> None: workflow = Path(".github/workflows/codex-autoreview.yml").read_text()