File size: 33,781 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 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 | import json
import os
import re
import subprocess
import threading
import uuid
from _thread import LockType
from collections import deque
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Iterator, TypeVar
from pydantic import BaseModel
from ._goal import _GoalOperationState
from ._initialize_metadata import _split_user_agent
from ._message_router import MessageRouter, _TurnSubscription
from ._runtime_requirements import CheckoutCapabilities, require_runtime_version
from ._version import __version__ as SDK_VERSION
from .errors import CodexError, InvalidRequestError, TransportClosedError
from .generated.notification_registry import NOTIFICATION_MODELS
from .generated.v2_all import (
AccountLoginCompletedNotification,
AgentMessageDeltaNotification,
CancelLoginAccountResponse,
ChatgptDeviceCodeLoginAccountResponse,
ChatgptLoginAccountResponse,
GetAccountParams as V2GetAccountParams,
GetAccountResponse,
IdleThreadStatus,
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,
JsonValue,
Notification,
UnknownNotification,
)
from .retry import retry_on_overload
ModelT = TypeVar("ModelT", bound=BaseModel)
ApprovalHandler = Callable[[str, JsonObject | None], JsonObject]
RUNTIME_PKG_NAME = "openai-codex-cli-bin"
_GOAL_START_TIMEOUT_S = 30.0
@dataclass(slots=True)
class _ThreadStartLock:
lock: LockType = field(default_factory=threading.Lock)
users: int = 0
def _active_turn_id_from_error(exc: InvalidRequestError) -> str | None:
match = re.search(r" but found `?([^`]+)`?$", exc.message)
return match.group(1) if match is not None else None
def _params_dict(
params: (
V2ThreadStartParams
| V2ThreadResumeParams
| V2ThreadListParams
| V2ThreadForkParams
| V2TurnStartParams
| V2GetAccountParams
| V2LoginAccountParams
| JsonObject
| None
),
) -> JsonObject:
if params is None:
return {}
if hasattr(params, "model_dump"):
dumped = params.model_dump(
by_alias=True,
exclude_none=True,
mode="json",
)
if not isinstance(dumped, dict):
raise TypeError("Expected model_dump() to return dict")
return dumped
if isinstance(params, dict):
return params
raise TypeError(f"Expected generated params model or dict, got {type(params).__name__}")
def _installed_codex_path() -> Path:
try:
from codex_cli_bin import bundled_codex_path
except ImportError as exc:
raise FileNotFoundError(
"Unable to locate the pinned Codex runtime. Install the published SDK build "
f"with its {RUNTIME_PKG_NAME} dependency, or set CodexConfig.codex_bin "
"explicitly."
) from exc
return bundled_codex_path()
def _installed_codex_path_dirs() -> tuple[Path, ...]:
try:
from codex_cli_bin import bundled_path_dir
except (ImportError, AttributeError):
return ()
path_dir = bundled_path_dir()
return (path_dir,) if path_dir is not None else ()
def _prepend_path_dirs(env: dict[str, str], path_dirs: tuple[Path, ...]) -> None:
if not path_dirs:
return
path_key = _path_env_key(env)
if os.name == "nt":
for key in list(env):
if key.upper() == "PATH" and key != path_key:
env.pop(key)
path_sep = os.pathsep
existing_path = env.get(path_key, "")
path_dir_values = [str(path_dir) for path_dir in path_dirs]
existing_entries = [
entry for entry in existing_path.split(path_sep) if entry and entry not in path_dir_values
]
env[path_key] = path_sep.join([*path_dir_values, *existing_entries])
def _path_env_key(env: dict[str, str]) -> str:
if os.name != "nt":
return "PATH"
matching_keys = [key for key in env if key.upper() == "PATH"]
if "Path" in matching_keys:
return "Path"
return matching_keys[-1] if matching_keys else "PATH"
@dataclass(frozen=True)
class CodexBinResolverOps:
installed_codex_path: Callable[[], Path]
path_exists: Callable[[Path], bool]
def _default_codex_bin_resolver_ops() -> CodexBinResolverOps:
return CodexBinResolverOps(
installed_codex_path=_installed_codex_path,
path_exists=lambda path: path.exists(),
)
def resolve_codex_bin(config: "CodexConfig", ops: CodexBinResolverOps) -> Path:
if config.codex_bin is not None:
codex_bin = Path(config.codex_bin)
if not ops.path_exists(codex_bin):
raise FileNotFoundError(
f"Codex binary not found at {codex_bin}. Set CodexConfig.codex_bin "
"to a valid binary path."
)
return codex_bin
return ops.installed_codex_path()
def _resolve_codex_bin(config: "CodexConfig") -> Path:
return resolve_codex_bin(config, _default_codex_bin_resolver_ops())
@dataclass(slots=True)
class CodexConfig:
"""Configuration for launching and identifying the local Codex runtime.
Most callers can use ``Codex()`` without configuration. Set ``codex_bin``
only when intentionally using a specific local Codex executable.
"""
codex_bin: str | None = None
launch_args_override: tuple[str, ...] | None = None
config_overrides: tuple[str, ...] = ()
cwd: str | None = None
env: dict[str, str] | None = None
client_name: str = "codex_python_sdk"
client_title: str = "Codex Python SDK"
client_version: str = SDK_VERSION
experimental_api: bool = True
class CodexClient:
"""Synchronous typed JSON-RPC client for `codex app-server` over stdio."""
def __init__(
self,
config: CodexConfig | None = None,
approval_handler: ApprovalHandler | None = None,
) -> None:
self.config = config or CodexConfig()
self._approval_handler = approval_handler or self._default_approval_handler
self._proc: subprocess.Popen[str] | None = None
self._lock = threading.Lock()
self._thread_start_locks_guard = threading.Lock()
self._thread_start_locks: dict[str, _ThreadStartLock] = {}
self._router = MessageRouter()
self._stderr_lines: deque[str] = deque(maxlen=400)
self._stderr_thread: threading.Thread | None = None
self._reader_thread: threading.Thread | None = None
self._runtime_version: str | None = None
self._checkout_capabilities: CheckoutCapabilities | None = None
def __enter__(self) -> "CodexClient":
self.start()
return self
def __exit__(self, _exc_type, _exc, _tb) -> None:
self.close()
def start(self) -> None:
if self._proc is not None:
return
path_dirs: tuple[Path, ...] = ()
if self.config.launch_args_override is not None:
args = list(self.config.launch_args_override)
else:
codex_bin = _resolve_codex_bin(self.config)
if self.config.codex_bin is None:
path_dirs = _installed_codex_path_dirs()
args = [str(codex_bin)]
for kv in self.config.config_overrides:
args.extend(["--config", kv])
args.extend(["app-server", "--listen", "stdio://"])
env = os.environ.copy()
if self.config.env:
env.update(self.config.env)
_prepend_path_dirs(env, path_dirs)
if self.config.launch_args_override is None:
self._checkout_capabilities = CheckoutCapabilities(
command=tuple(args[:-2]), cwd=self.config.cwd, env=env.copy()
)
self._proc = subprocess.Popen(
args,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
encoding="utf-8",
cwd=self.config.cwd,
env=env,
bufsize=1,
)
self._start_stderr_drain_thread()
self._start_reader_thread()
def close(self) -> None:
self._runtime_version = None
self._checkout_capabilities = None
if self._proc is None:
return
proc = self._proc
self._proc = None
if proc.stdin:
proc.stdin.close()
try:
proc.terminate()
proc.wait(timeout=2)
except Exception:
proc.kill()
if self._stderr_thread and self._stderr_thread.is_alive():
self._stderr_thread.join(timeout=0.5)
if self._reader_thread and self._reader_thread.is_alive():
self._reader_thread.join(timeout=0.5)
def initialize(self) -> InitializeResponse:
self._runtime_version = None
result = self.request(
"initialize",
{
"clientInfo": {
"name": self.config.client_name,
"title": self.config.client_title,
"version": self.config.client_version,
},
"capabilities": {
"experimentalApi": self.config.experimental_api,
},
},
response_model=InitializeResponse,
)
version = result.serverInfo.version if result.serverInfo is not None else None
if not version or not version.strip():
_, version = _split_user_agent(result.userAgent or "")
self._runtime_version = version.split()[0] if version and version.strip() else None
self.notify("initialized", None)
return result
def request(
self,
method: str,
params: JsonObject | None,
*,
response_model: type[ModelT],
) -> ModelT:
runtime_fields = {
"turn/start": ("toolOutput", "turnTrigger", "serviceTierForTurn"),
"thread/resume": ("excludeTurns",),
"thread/fork": ("excludeTurns",),
}
supplied_fields = [
field
for field in runtime_fields.get(method, ())
if (params or {}).get(field) is not None
]
if supplied_fields:
try:
if self._runtime_version == "0.0.0":
if self._checkout_capabilities is None:
raise ValueError(
"Cannot verify an unversioned CLI with a custom launch command"
)
supported = self._checkout_capabilities.fields[method]
if missing := set(supplied_fields) - supported:
raise ValueError(
f"The checkout does not support {', '.join(sorted(missing))}"
)
else:
require_runtime_version(self._runtime_version)
except ValueError as exc:
raise CodexError(
f"{method} with {', '.join(supplied_fields)}: {exc}. "
"Configure CodexConfig.codex_bin with a supported CLI."
) from exc
result = self._request_raw(method, params)
if not isinstance(result, dict):
raise CodexError(f"{method} response must be a JSON object")
return response_model.model_validate(result)
def _request_raw(self, method: str, params: JsonObject | None = None) -> JsonValue:
"""Send a JSON-RPC request and wait for the reader thread to route its response."""
request_id = str(uuid.uuid4())
waiter = self._router.create_response_waiter(request_id)
try:
message: JsonObject = {"id": request_id, "method": method}
if params is not None:
message["params"] = params
self._write_message(message)
except BaseException:
self._router.discard_response_waiter(request_id)
raise
item = waiter.get()
if isinstance(item, BaseException):
raise item
return item
def notify(self, method: str, params: JsonObject | None = None) -> None:
"""Send a JSON-RPC notification without waiting for a response."""
message: JsonObject = {"method": method}
if params is not None:
message["params"] = params
self._write_message(message)
def next_notification(self) -> Notification:
"""Return the next notification that is not scoped to an active turn."""
return self._router.next_global_notification()
def register_login_notifications(self, login_id: str) -> None:
"""Start routing notifications for one interactive login attempt."""
self._router.register_login(login_id)
def unregister_login_notifications(self, login_id: str) -> None:
"""Stop routing notifications for one interactive login attempt."""
self._router.unregister_login(login_id)
def next_login_notification(self, login_id: str) -> Notification:
"""Return the next routed notification for the requested login id."""
return self._router.next_login_notification(login_id)
def _subscribe_turn_notifications(self, turn_id: str) -> _TurnSubscription:
return self._router.subscribe_turn(turn_id)
def register_turn_notifications(self, turn_id: str) -> None:
"""Start routing notifications for one turn into its dedicated queue."""
self._router.register_turn(turn_id)
def unregister_turn_notifications(self, turn_id: str) -> None:
"""Stop routing notifications for one turn into its dedicated queue."""
self._router.unregister_turn(turn_id)
def next_turn_notification(self, turn_id: str) -> Notification:
"""Return the next routed notification for the requested turn id."""
return self._router.next_turn_notification(turn_id)
def register_goal_operation(self, thread_id: str) -> _GoalOperationState:
"""Register a private thread-scoped route for a logical goal turn."""
return self._router.register_goal(thread_id)
def reserve_goal_operation(self, thread_id: str) -> _GoalOperationState:
"""Reserve a private thread route before replacing its stored goal."""
return self._router.reserve_goal(thread_id)
def unregister_goal_operation(self, state: _GoalOperationState) -> None:
"""Release routing state for one logical goal turn."""
self._router.unregister_goal(state)
def next_goal_notification(self, state: _GoalOperationState) -> Notification:
"""Wait for the next notification in a logical goal turn."""
return state.next_notification()
def account_login_start(
self,
params: V2LoginAccountParams | JsonObject,
) -> LoginAccountResponse:
response = self.request(
"account/login/start",
_params_dict(params),
response_model=LoginAccountResponse,
)
response_root = response.root
if isinstance(
response_root,
ChatgptLoginAccountResponse | ChatgptDeviceCodeLoginAccountResponse,
):
self.register_login_notifications(response_root.login_id)
return response
def account_login_cancel(self, login_id: str) -> CancelLoginAccountResponse:
return self.request(
"account/login/cancel",
{"loginId": login_id},
response_model=CancelLoginAccountResponse,
)
def account_read(
self,
params: V2GetAccountParams | JsonObject | None = None,
) -> GetAccountResponse:
return self.request(
"account/read",
_params_dict(params),
response_model=GetAccountResponse,
)
def account_logout(self) -> LogoutAccountResponse:
return self.request("account/logout", None, response_model=LogoutAccountResponse)
def thread_start(
self, params: V2ThreadStartParams | JsonObject | None = None
) -> ThreadStartResponse:
return self.request(
"thread/start", _params_dict(params), response_model=ThreadStartResponse
)
def thread_resume(
self,
thread_id: str,
params: V2ThreadResumeParams | JsonObject | None = None,
) -> ThreadResumeResponse:
payload = {"threadId": thread_id, **_params_dict(params)}
return self.request("thread/resume", payload, response_model=ThreadResumeResponse)
def thread_list(
self, params: V2ThreadListParams | JsonObject | None = None
) -> ThreadListResponse:
return self.request("thread/list", _params_dict(params), response_model=ThreadListResponse)
def thread_read(self, thread_id: str, include_turns: bool = False) -> ThreadReadResponse:
return self.request(
"thread/read",
{"threadId": thread_id, "includeTurns": include_turns},
response_model=ThreadReadResponse,
)
def thread_fork(
self,
thread_id: str,
params: V2ThreadForkParams | JsonObject | None = None,
) -> ThreadForkResponse:
payload = {"threadId": thread_id, **_params_dict(params)}
return self.request("thread/fork", payload, response_model=ThreadForkResponse)
def thread_archive(self, thread_id: str) -> ThreadArchiveResponse:
return self.request(
"thread/archive",
{"threadId": thread_id},
response_model=ThreadArchiveResponse,
)
def thread_unarchive(self, thread_id: str) -> ThreadUnarchiveResponse:
return self.request(
"thread/unarchive",
{"threadId": thread_id},
response_model=ThreadUnarchiveResponse,
)
def thread_set_name(self, thread_id: str, name: str) -> ThreadSetNameResponse:
return self.request(
"thread/name/set",
{"threadId": thread_id, "name": name},
response_model=ThreadSetNameResponse,
)
def thread_compact(self, thread_id: str) -> ThreadCompactStartResponse:
return self.request(
"thread/compact/start",
{"threadId": thread_id},
response_model=ThreadCompactStartResponse,
)
def thread_goal_clear(self, thread_id: str) -> ThreadGoalClearResponse:
"""Clear the persisted goal for a thread before replacing it."""
return self.request(
"thread/goal/clear",
{"threadId": thread_id},
response_model=ThreadGoalClearResponse,
)
def thread_goal_set(
self,
thread_id: str,
*,
objective: str | None = None,
status: ThreadGoalStatus | None = None,
) -> ThreadGoalSetResponse:
"""Create or update the persisted goal for a thread."""
payload: JsonObject = {"threadId": thread_id}
if objective is not None:
payload["objective"] = objective
if status is not None:
payload["status"] = status.value
return self.request(
"thread/goal/set",
payload,
response_model=ThreadGoalSetResponse,
)
def pause_goal(self, thread_id: str) -> ThreadGoalSetResponse:
"""Pause the active goal used by a logical goal turn."""
return self.thread_goal_set(thread_id, status=ThreadGoalStatus.paused)
def cancel_goal_operation(self, state: _GoalOperationState) -> None:
"""Best-effort cleanup after a logical goal operation is cancelled."""
try:
self.pause_goal(state.thread_id)
except Exception:
pass
self._interrupt_goal_operation(state)
def _interrupt_goal_operation(self, state: _GoalOperationState) -> None:
turn_id = state.turn_for_interrupt()
if turn_id is None:
return
try:
self.turn_interrupt(state.thread_id, turn_id)
except InvalidRequestError as exc:
if not exc.message.startswith("expected active turn id"):
return
next_turn_id = _active_turn_id_from_error(exc) or state.current_turn()
if next_turn_id is None or next_turn_id == turn_id:
return
try:
self.turn_interrupt(state.thread_id, next_turn_id)
except Exception:
pass
except Exception:
pass
def start_goal_operation(
self,
thread_id: str,
objective: str,
) -> tuple[_GoalOperationState, str]:
"""Start a logical goal and wait for its runtime-generated first turn."""
with self._thread_start_lock(thread_id):
return self._start_goal_operation(thread_id, objective)
def _start_goal_operation(
self,
thread_id: str,
objective: str,
) -> tuple[_GoalOperationState, str]:
thread = self.thread_read(thread_id).thread
if not isinstance(thread.status.root, IdleThreadStatus):
raise InvalidRequestError(
-32600,
f"thread must be idle before starting a goal: {thread_id}",
)
if thread.ephemeral or thread.path is None:
raise InvalidRequestError(
-32600,
f"thread must be persisted before starting a goal: {thread_id}",
)
state = self.reserve_goal_operation(thread_id)
activated = False
try:
self.thread_goal_clear(thread_id)
state.activate_turn_routing()
self.thread_goal_set(
thread_id,
objective=objective,
status=ThreadGoalStatus.active,
)
activated = True
turn_id = state.wait_for_start(_GOAL_START_TIMEOUT_S)
if turn_id is None:
raise CodexError(
"timed out waiting for goal turn to start after "
f"{int(_GOAL_START_TIMEOUT_S)} seconds"
)
return state, turn_id
except BaseException as exc:
if activated or not isinstance(exc, InvalidRequestError):
self.cancel_goal_operation(state)
state.finish()
self.unregister_goal_operation(state)
raise
def turn_start(
self,
thread_id: str,
input_items: list[JsonObject] | JsonObject | str,
params: V2TurnStartParams | JsonObject | None = None,
) -> TurnStartResponse:
"""Start a turn and register its notification queue as early as possible."""
return self._start_turn(thread_id, input_items, params, for_handle=False)[0]
def _start_turn(
self,
thread_id: str,
input_items: list[JsonObject] | JsonObject | str,
params: V2TurnStartParams | JsonObject | None,
for_handle: bool,
) -> tuple[TurnStartResponse, _TurnSubscription | None]:
with self._thread_start_lock(thread_id):
if self._router.has_goal(thread_id):
raise InvalidRequestError(
-32600,
f"thread has an active goal operation: {thread_id}",
)
payload = {
**_params_dict(params),
"threadId": thread_id,
"input": self._normalize_input_items(input_items),
}
with self._router.pending_turn(thread_id) as cursors:
started = self.request("turn/start", payload, response_model=TurnStartResponse)
subscription = self._router.prepare_turn(
started.turn.id, thread_id, cursors, for_handle=for_handle
)
return started, subscription
@contextmanager
def _thread_start_lock(self, thread_id: str) -> Iterator[None]:
with self._thread_start_locks_guard:
entry = self._thread_start_locks.get(thread_id)
if entry is None:
entry = _ThreadStartLock()
self._thread_start_locks[thread_id] = entry
entry.users += 1
try:
with entry.lock:
yield
finally:
with self._thread_start_locks_guard:
entry.users -= 1
if entry.users == 0:
self._thread_start_locks.pop(thread_id, None)
def turn_interrupt(self, thread_id: str, turn_id: str) -> TurnInterruptResponse:
return self.request(
"turn/interrupt",
{"threadId": thread_id, "turnId": turn_id},
response_model=TurnInterruptResponse,
)
def turn_steer(
self,
thread_id: str,
expected_turn_id: str,
input_items: list[JsonObject] | JsonObject | str,
) -> TurnSteerResponse:
return self.request(
"turn/steer",
{
"threadId": thread_id,
"expectedTurnId": expected_turn_id,
"input": self._normalize_input_items(input_items),
},
response_model=TurnSteerResponse,
)
def model_list(self, include_hidden: bool = False) -> ModelListResponse:
return self.request(
"model/list",
{"includeHidden": include_hidden},
response_model=ModelListResponse,
)
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:
return retry_on_overload(
lambda: self.request(method, params, response_model=response_model),
max_attempts=max_attempts,
initial_delay_s=initial_delay_s,
max_delay_s=max_delay_s,
)
def wait_for_turn_completed(self, turn_id: str) -> TurnCompletedNotification:
"""Block on the routed turn stream until the matching completion arrives."""
self.register_turn_notifications(turn_id)
try:
while True:
notification = self.next_turn_notification(turn_id)
if (
notification.method == "turn/completed"
and isinstance(notification.payload, TurnCompletedNotification)
and notification.payload.turn.id == turn_id
):
return notification.payload
finally:
self.unregister_turn_notifications(turn_id)
def wait_for_login_completed(
self,
login_id: str,
) -> AccountLoginCompletedNotification:
"""Block until the matching interactive login attempt completes."""
self.register_login_notifications(login_id)
try:
while True:
notification = self.next_login_notification(login_id)
if (
notification.method == "account/login/completed"
and isinstance(notification.payload, AccountLoginCompletedNotification)
and notification.payload.login_id == login_id
):
return notification.payload
finally:
self.unregister_login_notifications(login_id)
def stream_text(
self,
thread_id: str,
text: str,
params: V2TurnStartParams | JsonObject | None = None,
) -> Iterator[AgentMessageDeltaNotification]:
"""Start a text turn and yield only its agent-message delta payloads."""
started = self.turn_start(thread_id, text, params=params)
turn_id = started.turn.id
self.register_turn_notifications(turn_id)
try:
while True:
notification = self.next_turn_notification(turn_id)
if (
notification.method == "item/agentMessage/delta"
and isinstance(notification.payload, AgentMessageDeltaNotification)
and notification.payload.turn_id == turn_id
):
yield notification.payload
continue
if (
notification.method == "turn/completed"
and isinstance(notification.payload, TurnCompletedNotification)
and notification.payload.turn.id == turn_id
):
break
finally:
self.unregister_turn_notifications(turn_id)
def _coerce_notification(self, method: str, params: object) -> Notification:
params_dict = params if isinstance(params, dict) else {}
model = NOTIFICATION_MODELS.get(method)
if model is None:
return Notification(method=method, payload=UnknownNotification(params=params_dict))
try:
payload = model.model_validate(params_dict)
except Exception: # noqa: BLE001
return Notification(method=method, payload=UnknownNotification(params=params_dict))
return Notification(method=method, payload=payload)
def _normalize_input_items(
self,
input_items: list[JsonObject] | JsonObject | str,
) -> list[JsonObject]:
if isinstance(input_items, str):
return [{"type": "text", "text": input_items}]
if isinstance(input_items, dict):
return [input_items]
return input_items
def _default_approval_handler(self, method: str, params: JsonObject | None) -> JsonObject:
"""Accept approval requests when the caller did not provide a handler."""
if method == "item/commandExecution/requestApproval":
return {"decision": "accept"}
if method == "item/fileChange/requestApproval":
return {"decision": "accept"}
return {}
def _start_stderr_drain_thread(self) -> None:
if self._proc is None or self._proc.stderr is None:
return
def _drain() -> None:
stderr = self._proc.stderr
if stderr is None:
return
for line in stderr:
self._stderr_lines.append(line.rstrip("\n"))
self._stderr_thread = threading.Thread(target=_drain, daemon=True)
self._stderr_thread.start()
def _start_reader_thread(self) -> None:
"""Start the sole stdout reader that fans messages into router queues."""
if self._proc is None or self._proc.stdout is None:
return
self._reader_thread = threading.Thread(target=self._reader_loop, daemon=True)
self._reader_thread.start()
def _reader_loop(self) -> None:
"""Continuously classify transport messages into requests, responses, and events."""
try:
while True:
msg = self._read_message()
if "method" in msg and "id" in msg:
response = self._handle_server_request(msg)
self._write_message({"id": msg["id"], "result": response})
continue
if "method" in msg and "id" not in msg:
method = msg["method"]
if isinstance(method, str):
self._router.route_notification(
self._coerce_notification(method, msg.get("params"))
)
continue
self._router.route_response(msg)
except BaseException as exc:
self._router.fail_all(exc)
def _stderr_tail(self, limit: int = 40) -> str:
return "\n".join(list(self._stderr_lines)[-limit:])
def _handle_server_request(self, msg: dict[str, JsonValue]) -> JsonObject:
method = msg["method"]
params = msg.get("params")
if not isinstance(method, str):
return {}
return self._approval_handler(
method,
params if isinstance(params, dict) else None,
)
def _write_message(self, payload: JsonObject) -> None:
if self._proc is None or self._proc.stdin is None:
raise TransportClosedError("Codex process is not running")
with self._lock:
self._proc.stdin.write(json.dumps(payload) + "\n")
self._proc.stdin.flush()
def _read_message(self) -> dict[str, JsonValue]:
if self._proc is None or self._proc.stdout is None:
raise TransportClosedError("Codex process is not running")
line = self._proc.stdout.readline()
if not line:
raise TransportClosedError(
f"Codex process closed stdout. stderr_tail={self._stderr_tail()[:2000]}"
)
try:
message = json.loads(line)
except json.JSONDecodeError as exc:
raise CodexError(f"Invalid JSON-RPC line: {line!r}") from exc
if not isinstance(message, dict):
raise CodexError(f"Invalid JSON-RPC payload: {message!r}")
return message
def default_codex_home() -> str:
return str(Path.home() / ".codex")
|