File size: 16,245 Bytes
5f25733 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 | from __future__ import annotations
import asyncio
import threading
from collections.abc import Iterator
from concurrent.futures import Future, ThreadPoolExecutor
from contextvars import copy_context
from typing import AsyncIterator, Callable, ParamSpec, TypeVar
from pydantic import BaseModel
from ._goal import _GoalOperationState
from ._message_router import _TurnSubscription
from .client import CodexClient, CodexConfig
from .generated.v2_all import (
AccountLoginCompletedNotification,
AgentMessageDeltaNotification,
CancelLoginAccountResponse,
GetAccountParams as V2GetAccountParams,
GetAccountResponse,
LoginAccountParams as V2LoginAccountParams,
LoginAccountResponse,
LogoutAccountResponse,
ModelListResponse,
ThreadArchiveResponse,
ThreadCompactStartResponse,
ThreadForkParams as V2ThreadForkParams,
ThreadForkResponse,
ThreadGoalClearResponse,
ThreadGoalSetResponse,
ThreadGoalStatus,
ThreadListParams as V2ThreadListParams,
ThreadListResponse,
ThreadReadResponse,
ThreadResumeParams as V2ThreadResumeParams,
ThreadResumeResponse,
ThreadSetNameResponse,
ThreadStartParams as V2ThreadStartParams,
ThreadStartResponse,
ThreadUnarchiveResponse,
TurnCompletedNotification,
TurnInterruptResponse,
TurnStartParams as V2TurnStartParams,
TurnStartResponse,
TurnSteerResponse,
)
from .models import InitializeResponse, JsonObject, Notification
ModelT = TypeVar("ModelT", bound=BaseModel)
ParamsT = ParamSpec("ParamsT")
ReturnT = TypeVar("ReturnT")
# Bound workers while allowing cancellation cleanup to outlive an asyncio waiter.
_TURN_START_EXECUTOR = ThreadPoolExecutor(thread_name_prefix="codex-turn-start")
class AsyncCodexClient:
"""Async wrapper around CodexClient using thread offloading."""
def __init__(self, config: CodexConfig | None = None) -> None:
"""Create the wrapped sync client that owns the transport process."""
self._sync = CodexClient(config=config)
async def __aenter__(self) -> "AsyncCodexClient":
"""Start the Codex process when entering an async context."""
await self.start()
return self
async def __aexit__(self, _exc_type, _exc, _tb) -> None:
"""Close the Codex process when leaving an async context."""
await self.close()
async def _call_sync(
self,
fn: Callable[ParamsT, ReturnT],
/,
*args: ParamsT.args,
**kwargs: ParamsT.kwargs,
) -> ReturnT:
"""Run a blocking sync-client operation without blocking the event loop."""
return await asyncio.to_thread(fn, *args, **kwargs)
@staticmethod
def _next_from_iterator(
iterator: Iterator[AgentMessageDeltaNotification],
) -> tuple[bool, AgentMessageDeltaNotification | None]:
"""Convert StopIteration into a value that can cross asyncio.to_thread."""
try:
return True, next(iterator)
except StopIteration:
return False, None
async def start(self) -> None:
"""Start the wrapped sync client in a worker thread."""
await self._call_sync(self._sync.start)
async def close(self) -> None:
"""Close the wrapped sync client in a worker thread."""
await self._call_sync(self._sync.close)
async def initialize(self) -> InitializeResponse:
"""Initialize the Codex session."""
return await self._call_sync(self._sync.initialize)
def _subscribe_turn_notifications(self, turn_id: str) -> _TurnSubscription:
return self._sync._subscribe_turn_notifications(turn_id)
def register_turn_notifications(self, turn_id: str) -> None:
"""Register a turn notification queue on the wrapped sync client."""
self._sync.register_turn_notifications(turn_id)
def register_login_notifications(self, login_id: str) -> None:
"""Register a login notification queue on the wrapped sync client."""
self._sync.register_login_notifications(login_id)
def unregister_login_notifications(self, login_id: str) -> None:
"""Unregister a login notification queue on the wrapped sync client."""
self._sync.unregister_login_notifications(login_id)
def unregister_turn_notifications(self, turn_id: str) -> None:
"""Unregister a turn notification queue on the wrapped sync client."""
self._sync.unregister_turn_notifications(turn_id)
def register_goal_operation(self, thread_id: str) -> _GoalOperationState:
"""Register a logical goal route on the wrapped sync client."""
return self._sync.register_goal_operation(thread_id)
def unregister_goal_operation(self, state: _GoalOperationState) -> None:
"""Release one logical goal route."""
self._sync.unregister_goal_operation(state)
async def request(
self,
method: str,
params: JsonObject | None,
*,
response_model: type[ModelT],
) -> ModelT:
"""Send a typed JSON-RPC request through the wrapped sync client."""
return await self._call_sync(
self._sync.request,
method,
params,
response_model=response_model,
)
async def account_login_start(
self,
params: V2LoginAccountParams | JsonObject,
) -> LoginAccountResponse:
"""Start one account login attempt through the wrapped sync client."""
return await self._call_sync(self._sync.account_login_start, params)
async def account_login_cancel(self, login_id: str) -> CancelLoginAccountResponse:
"""Cancel one active account login attempt through the wrapped sync client."""
return await self._call_sync(self._sync.account_login_cancel, login_id)
async def account_read(
self,
params: V2GetAccountParams | JsonObject | None = None,
) -> GetAccountResponse:
"""Read current account state through the wrapped sync client."""
return await self._call_sync(self._sync.account_read, params)
async def account_logout(self) -> LogoutAccountResponse:
"""Clear the active account session through the wrapped sync client."""
return await self._call_sync(self._sync.account_logout)
async def thread_start(
self, params: V2ThreadStartParams | JsonObject | None = None
) -> ThreadStartResponse:
"""Start a thread using the wrapped sync client."""
return await self._call_sync(self._sync.thread_start, params)
async def thread_resume(
self,
thread_id: str,
params: V2ThreadResumeParams | JsonObject | None = None,
) -> ThreadResumeResponse:
"""Resume a thread using the wrapped sync client."""
return await self._call_sync(self._sync.thread_resume, thread_id, params)
async def thread_list(
self, params: V2ThreadListParams | JsonObject | None = None
) -> ThreadListResponse:
"""List threads using the wrapped sync client."""
return await self._call_sync(self._sync.thread_list, params)
async def thread_read(self, thread_id: str, include_turns: bool = False) -> ThreadReadResponse:
"""Read a thread using the wrapped sync client."""
return await self._call_sync(self._sync.thread_read, thread_id, include_turns)
async def thread_fork(
self,
thread_id: str,
params: V2ThreadForkParams | JsonObject | None = None,
) -> ThreadForkResponse:
"""Fork a thread using the wrapped sync client."""
return await self._call_sync(self._sync.thread_fork, thread_id, params)
async def thread_archive(self, thread_id: str) -> ThreadArchiveResponse:
"""Archive a thread using the wrapped sync client."""
return await self._call_sync(self._sync.thread_archive, thread_id)
async def thread_unarchive(self, thread_id: str) -> ThreadUnarchiveResponse:
"""Unarchive a thread using the wrapped sync client."""
return await self._call_sync(self._sync.thread_unarchive, thread_id)
async def thread_set_name(self, thread_id: str, name: str) -> ThreadSetNameResponse:
"""Rename a thread using the wrapped sync client."""
return await self._call_sync(self._sync.thread_set_name, thread_id, name)
async def thread_compact(self, thread_id: str) -> ThreadCompactStartResponse:
"""Start thread compaction using the wrapped sync client."""
return await self._call_sync(self._sync.thread_compact, thread_id)
async def thread_goal_clear(self, thread_id: str) -> ThreadGoalClearResponse:
"""Clear the persisted goal through the wrapped sync client."""
return await self._call_sync(self._sync.thread_goal_clear, thread_id)
async def thread_goal_set(
self,
thread_id: str,
*,
objective: str | None = None,
status: ThreadGoalStatus | None = None,
) -> ThreadGoalSetResponse:
"""Create or update a persisted goal through the wrapped sync client."""
return await self._call_sync(
self._sync.thread_goal_set,
thread_id,
objective=objective,
status=status,
)
async def pause_goal(self, thread_id: str) -> ThreadGoalSetResponse:
"""Pause the active goal through the wrapped sync client."""
return await self._call_sync(self._sync.pause_goal, thread_id)
async def cancel_goal_operation(self, state: _GoalOperationState) -> None:
"""Stop continuation work after a logical goal operation is cancelled."""
await self._call_sync(self._sync.cancel_goal_operation, state)
async def start_goal_operation(
self,
thread_id: str,
objective: str,
) -> tuple[_GoalOperationState, str]:
"""Start a logical goal through the wrapped sync client."""
operation: Future[tuple[_GoalOperationState, str]] = Future()
def start_operation() -> None:
try:
operation.set_result(self._sync.start_goal_operation(thread_id, objective))
except BaseException as exc:
operation.set_exception(exc)
worker = threading.Thread(
target=start_operation,
name="codex-goal-start",
daemon=True,
)
worker.start()
try:
return await asyncio.shield(asyncio.wrap_future(operation))
except asyncio.CancelledError:
def cleanup_cancelled_start(
completed: Future[tuple[_GoalOperationState, str]],
) -> None:
try:
state, _ = completed.result()
except BaseException:
return
def stop_cancelled_goal() -> None:
try:
self._sync.cancel_goal_operation(state)
finally:
state.finish()
self._sync.unregister_goal_operation(state)
threading.Thread(
target=stop_cancelled_goal,
name="codex-goal-start-cleanup",
daemon=True,
).start()
operation.add_done_callback(cleanup_cancelled_start)
raise
async def turn_start(
self,
thread_id: str,
input_items: list[JsonObject] | JsonObject | str,
params: V2TurnStartParams | JsonObject | None = None,
) -> TurnStartResponse:
"""Start a turn, releasing an unclaimed result if the caller is cancelled."""
return (await self._start_turn(thread_id, input_items, params, for_handle=False))[0]
async def _start_turn(
self,
thread_id: str,
input_items: list[JsonObject] | JsonObject | str,
params: V2TurnStartParams | JsonObject | None,
for_handle: bool,
) -> tuple[TurnStartResponse, _TurnSubscription | None]:
operation = _TURN_START_EXECUTOR.submit(
copy_context().run, self._sync._start_turn, thread_id, input_items, params, for_handle
)
try:
return await asyncio.wrap_future(operation)
except asyncio.CancelledError:
def discard_cancelled_result(
completed: Future[tuple[TurnStartResponse, _TurnSubscription | None]],
) -> None:
try:
_, subscription = completed.result()
except BaseException:
return
if subscription is not None:
subscription.close()
operation.add_done_callback(discard_cancelled_result)
raise
async def turn_interrupt(self, thread_id: str, turn_id: str) -> TurnInterruptResponse:
"""Interrupt a turn using the wrapped sync client."""
return await self._call_sync(self._sync.turn_interrupt, thread_id, turn_id)
async def turn_steer(
self,
thread_id: str,
expected_turn_id: str,
input_items: list[JsonObject] | JsonObject | str,
) -> TurnSteerResponse:
"""Send steering input to a turn using the wrapped sync client."""
return await self._call_sync(
self._sync.turn_steer,
thread_id,
expected_turn_id,
input_items,
)
async def model_list(self, include_hidden: bool = False) -> ModelListResponse:
"""List models using the wrapped sync client."""
return await self._call_sync(self._sync.model_list, include_hidden)
async def request_with_retry_on_overload(
self,
method: str,
params: JsonObject | None,
*,
response_model: type[ModelT],
max_attempts: int = 3,
initial_delay_s: float = 0.25,
max_delay_s: float = 2.0,
) -> ModelT:
"""Send a typed request with the sync client's overload retry policy."""
return await self._call_sync(
self._sync.request_with_retry_on_overload,
method,
params,
response_model=response_model,
max_attempts=max_attempts,
initial_delay_s=initial_delay_s,
max_delay_s=max_delay_s,
)
async def next_notification(self) -> Notification:
"""Wait for the next global notification without blocking the event loop."""
return await self._call_sync(self._sync.next_notification)
async def next_login_notification(self, login_id: str) -> Notification:
"""Wait for the next notification routed to one login attempt."""
return await self._call_sync(self._sync.next_login_notification, login_id)
async def next_turn_notification(self, turn_id: str) -> Notification:
"""Wait for the next notification routed to one turn."""
return await self._call_sync(self._sync.next_turn_notification, turn_id)
async def next_goal_notification(self, state: _GoalOperationState) -> Notification:
"""Wait for the next notification in a logical goal turn."""
return await self._call_sync(self._sync.next_goal_notification, state)
async def wait_for_login_completed(
self,
login_id: str,
) -> AccountLoginCompletedNotification:
"""Wait for the completion notification routed to one login attempt."""
return await self._call_sync(self._sync.wait_for_login_completed, login_id)
async def wait_for_turn_completed(self, turn_id: str) -> TurnCompletedNotification:
"""Wait for the completion notification routed to one turn."""
return await self._call_sync(self._sync.wait_for_turn_completed, turn_id)
async def stream_text(
self,
thread_id: str,
text: str,
params: V2TurnStartParams | JsonObject | None = None,
) -> AsyncIterator[AgentMessageDeltaNotification]:
"""Stream text deltas from one turn without monopolizing the event loop."""
iterator = self._sync.stream_text(thread_id, text, params)
while True:
has_value, chunk = await asyncio.to_thread(
self._next_from_iterator,
iterator,
)
if not has_value:
break
yield chunk
|