fix(sdk): make cleanup reliable and preserve typed tool arguments

.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

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()

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):