Add Python SDK goal turns

This commit is contained in:
Ahmed Ibrahim
2026-06-07 16:09:40 -07:00
parent e093d81982
commit 851a85b369
50 changed files with 2909 additions and 136 deletions
+86 -1
View File
@@ -3,7 +3,10 @@ from __future__ import annotations
from pathlib import Path
from openai_codex.client import CodexClient, _params_dict
from openai_codex.generated.notification_registry import notification_turn_id
from openai_codex.generated.notification_registry import (
notification_thread_id,
notification_turn_id,
)
from openai_codex.generated.v2_all import (
AgentMessageDeltaNotification,
ApprovalsReviewer,
@@ -11,6 +14,8 @@ from openai_codex.generated.v2_all import (
ThreadResumeResponse,
ThreadTokenUsageUpdatedNotification,
TurnCompletedNotification,
TurnStartParams,
TurnStartResponse,
WarningNotification,
)
from openai_codex.models import Notification, UnknownNotification
@@ -26,11 +31,73 @@ def test_generated_params_models_are_snake_case_and_dump_by_alias() -> None:
assert dumped == {"searchTerm": "needle", "limit": 5}
def test_turn_start_sends_goal_only_when_enabled() -> None:
"""The low-level request should preserve false-by-omission wire behavior."""
client = CodexClient()
requests: list[dict[str, object]] = []
def fake_request(method, params, *, response_model):
requests.append({"method": method, "params": params, "response_model": response_model})
return TurnStartResponse.model_validate(
{"turn": {"id": f"turn-{len(requests)}", "items": [], "status": "inProgress"}}
)
client.request = fake_request # type: ignore[method-assign]
client.turn_start("thread-1", "ordinary")
client.turn_start("thread-1", "goal", goal=True)
client.turn_start(
"thread-1",
"typed goal",
TurnStartParams(threadId="thread-1", input=[], goal=True),
)
assert requests == [
{
"method": "turn/start",
"params": {
"threadId": "thread-1",
"input": [{"type": "text", "text": "ordinary"}],
},
"response_model": TurnStartResponse,
},
{
"method": "turn/start",
"params": {
"threadId": "thread-1",
"input": [{"type": "text", "text": "goal"}],
"goal": True,
},
"response_model": TurnStartResponse,
},
{
"method": "turn/start",
"params": {
"threadId": "thread-1",
"input": [{"type": "text", "text": "typed goal"}],
"goal": True,
},
"response_model": TurnStartResponse,
},
]
def test_generated_v2_bundle_has_single_shared_plan_type_definition() -> None:
source = (ROOT / "src" / "openai_codex" / "generated" / "v2_all.py").read_text()
assert source.count("class PlanType(") == 1
def test_generated_turn_start_params_include_goal_mode() -> None:
assert (
TurnStartParams(
threadId="thread-1",
input=[{"type": "text", "text": "goal"}],
goal=True,
).goal
is True
)
def test_thread_resume_response_accepts_auto_review_reviewer() -> None:
"""Generated response models should keep accepting the auto review enum value."""
response = ThreadResumeResponse.model_validate(
@@ -142,6 +209,24 @@ def test_generated_notification_turn_id_handles_known_payload_shapes() -> None:
] == ["turn-1", "turn-2", None]
def test_generated_notification_thread_id_handles_known_payload_shapes() -> None:
"""Generated routing metadata should expose thread ids without field guessing."""
direct = AgentMessageDeltaNotification.model_validate(
{
"delta": "hello",
"itemId": "item-1",
"threadId": "thread-1",
"turnId": "turn-1",
}
)
unscoped = WarningNotification(message="heads up")
assert [notification_thread_id(direct), notification_thread_id(unscoped)] == [
"thread-1",
None,
]
def test_turn_notification_router_demuxes_registered_turns() -> None:
"""The router should deliver out-of-order turn events to the matching queues."""
client = CodexClient()
+598
View File
@@ -0,0 +1,598 @@
import asyncio
import threading
from dataclasses import dataclass
import pytest
from openai_codex.api import AsyncTurnHandle, TurnHandle
from openai_codex.client import CodexClient
from openai_codex.errors import InvalidRequestError
from openai_codex.generated.notification_registry import notification_turn_id
from openai_codex.generated.v2_all import (
TurnCompletedNotification,
TurnInterruptResponse,
TurnStatus,
TurnSteerResponse,
)
def _route(client: CodexClient, method: str, params: dict[str, object]) -> None:
client._router.route_notification(client._coerce_notification(method, params))
def _turn(turn_id: str, status: str, *, started_at: int, completed_at: int | None = None):
turn: dict[str, object] = {
"id": turn_id,
"items": [],
"startedAt": started_at,
"status": status,
}
if completed_at is not None:
turn["completedAt"] = completed_at
return turn
def _goal(thread_id: str, status: str) -> dict[str, object]:
return {
"createdAt": 1,
"objective": "Improve benchmark coverage",
"status": status,
"threadId": thread_id,
"timeUsedSeconds": 0,
"tokensUsed": 0,
"updatedAt": 1,
}
def test_goal_turn_collects_continuations_as_one_logical_turn() -> None:
client = CodexClient()
state = client.register_goal_operation("thread-1")
state.bind("turn-1")
handle = TurnHandle(client, "thread-1", "turn-1", _goal=state)
_route(
client,
"thread/goal/updated",
{"threadId": "thread-1", "turnId": "turn-1", "goal": _goal("thread-1", "active")},
)
_route(
client,
"turn/started",
{"threadId": "thread-1", "turn": _turn("turn-1", "inProgress", started_at=10)},
)
_route(
client,
"item/completed",
{
"threadId": "thread-1",
"turnId": "turn-1",
"item": {"id": "message-1", "type": "agentMessage", "text": "working"},
},
)
_route(
client,
"turn/completed",
{
"threadId": "thread-1",
"turn": _turn("turn-1", "completed", started_at=10, completed_at=12),
},
)
_route(
client,
"turn/started",
{"threadId": "thread-1", "turn": _turn("turn-2", "inProgress", started_at=13)},
)
_route(
client,
"item/completed",
{
"threadId": "thread-1",
"turnId": "turn-2",
"item": {
"id": "message-2",
"type": "agentMessage",
"text": "done",
"phase": "final_answer",
},
},
)
_route(
client,
"thread/goal/updated",
{"threadId": "thread-1", "turnId": "turn-2", "goal": _goal("thread-1", "complete")},
)
_route(
client,
"thread/tokenUsage/updated",
{
"threadId": "thread-1",
"turnId": "turn-2",
"tokenUsage": {
"last": {
"cachedInputTokens": 1,
"inputTokens": 2,
"outputTokens": 3,
"reasoningOutputTokens": 4,
"totalTokens": 9,
},
"total": {
"cachedInputTokens": 10,
"inputTokens": 20,
"outputTokens": 30,
"reasoningOutputTokens": 40,
"totalTokens": 90,
},
},
},
)
_route(
client,
"turn/completed",
{
"threadId": "thread-1",
"turn": _turn("turn-2", "completed", started_at=13, completed_at=15),
},
)
result = handle.run()
assert {
"id": result.id,
"status": result.status,
"started_at": result.started_at,
"completed_at": result.completed_at,
"duration_ms": result.duration_ms,
"final_response": result.final_response,
"item_count": len(result.items),
"total_tokens": result.usage.total.total_tokens if result.usage is not None else None,
"registered_goals": client._router._goal_operations,
} == {
"id": "turn-1",
"status": TurnStatus.completed,
"started_at": 10,
"completed_at": 15,
"duration_ms": 5000,
"final_response": "done",
"item_count": 2,
"total_tokens": 90,
"registered_goals": {},
}
def test_goal_stream_hides_continuation_boundaries_and_rewrites_ids() -> None:
client = CodexClient()
state = client.register_goal_operation("thread-1")
state.bind("turn-1")
handle = TurnHandle(client, "thread-1", "turn-1", _goal=state)
for method, params in [
(
"thread/goal/updated",
{"threadId": "thread-1", "turnId": "turn-1", "goal": _goal("thread-1", "active")},
),
(
"turn/started",
{"threadId": "thread-1", "turn": _turn("turn-1", "inProgress", started_at=10)},
),
(
"turn/completed",
{
"threadId": "thread-1",
"turn": _turn("turn-1", "completed", started_at=10, completed_at=11),
},
),
(
"turn/started",
{"threadId": "thread-1", "turn": _turn("turn-2", "inProgress", started_at=12)},
),
(
"thread/goal/updated",
{"threadId": "thread-1", "turnId": "turn-2", "goal": _goal("thread-1", "complete")},
),
(
"turn/completed",
{
"threadId": "thread-1",
"turn": _turn("turn-2", "completed", started_at=12, completed_at=13),
},
),
]:
_route(client, method, params)
events = list(handle.stream())
lifecycle = [event for event in events if event.method in {"turn/started", "turn/completed"}]
assert {
"lifecycle": [event.method for event in lifecycle],
"turn_ids": [
turn_id
for event in events
if (turn_id := notification_turn_id(event.payload)) is not None
],
"completed": [
event.payload.turn.status
for event in events
if isinstance(event.payload, TurnCompletedNotification)
],
} == {
"lifecycle": ["turn/started", "turn/completed"],
"turn_ids": ["turn-1", "turn-1"],
"completed": [TurnStatus.completed],
}
def test_goal_router_releases_and_wakes_operations_on_transport_failure() -> None:
client = CodexClient()
state = client.register_goal_operation("thread-1")
failure = RuntimeError("transport closed")
client._router.fail_all(failure)
with pytest.raises(RuntimeError, match="transport closed"):
state.next_notification()
assert client._router._goal_operations == {}
def test_closing_unstarted_goal_stream_releases_route_and_controls() -> None:
client = CodexClient()
state = client.register_goal_operation("thread-1")
state.bind("turn-1")
handle = TurnHandle(client, "thread-1", "turn-1", _goal=state)
stream = handle.stream()
stream.close()
assert {
"active_turn": state.active_turn(),
"registered_goals": client._router._goal_operations,
} == {"active_turn": None, "registered_goals": {}}
def test_failed_goal_start_propagates_turn_error() -> None:
client = CodexClient()
state = client.register_goal_operation("thread-1")
state.bind("turn-1")
handle = TurnHandle(client, "thread-1", "turn-1", _goal=state)
failed_turn = _turn("turn-1", "failed", started_at=10, completed_at=11)
failed_turn["error"] = {"message": "failed to start goal"}
_route(
client,
"turn/started",
{"threadId": "thread-1", "turn": _turn("turn-1", "inProgress", started_at=10)},
)
_route(
client,
"turn/completed",
{"threadId": "thread-1", "turn": failed_turn},
)
with pytest.raises(RuntimeError, match="failed to start goal"):
handle.run()
assert client._router._goal_operations == {}
class _SyncControlClient:
def __init__(self, *, fail_first: bool = True) -> None:
self.steer_ids: list[str] = []
self.interrupt_ids: list[str] = []
self.paused = False
self.fail_first = fail_first
def turn_steer(self, _thread_id, turn_id, _input):
self.steer_ids.append(turn_id)
if self.fail_first and len(self.steer_ids) == 1:
raise InvalidRequestError(
-32600,
"expected active turn id `turn-1` but found `turn-2`",
)
return TurnSteerResponse(turn_id=turn_id)
def pause_goal(self, _thread_id):
self.paused = True
def turn_interrupt(self, _thread_id, turn_id):
self.interrupt_ids.append(turn_id)
if self.fail_first and len(self.interrupt_ids) == 1:
raise InvalidRequestError(
-32600,
"expected active turn id turn-1 but found turn-2",
)
return TurnInterruptResponse()
def test_goal_controls_hide_rollover_turn_ids() -> None:
state = CodexClient().register_goal_operation("thread-1")
state.bind("turn-1")
client = _SyncControlClient()
handle = TurnHandle(client, "thread-1", "turn-1", _goal=state) # type: ignore[arg-type]
steered = handle.steer("keep going")
interrupted = handle.interrupt()
assert {
"steer_ids": client.steer_ids,
"public_steer_id": steered.turn_id,
"paused": client.paused,
"interrupt_ids": client.interrupt_ids,
"interrupt": interrupted.model_dump(mode="json"),
} == {
"steer_ids": ["turn-1", "turn-2"],
"public_steer_id": "turn-1",
"paused": True,
"interrupt_ids": ["turn-1", "turn-2"],
"interrupt": {},
}
def test_finished_goal_controls_match_inactive_turn_errors() -> None:
state = CodexClient().register_goal_operation("thread-1")
state.bind("turn-1")
state.finish()
client = _SyncControlClient(fail_first=False)
handle = TurnHandle(client, "thread-1", "turn-1", _goal=state) # type: ignore[arg-type]
with pytest.raises(InvalidRequestError, match="no active turn to steer"):
handle.steer("keep going")
with pytest.raises(InvalidRequestError, match="no active turn to interrupt"):
handle.interrupt()
assert client.paused is False
def test_goal_steer_waits_for_rollover_turn() -> None:
client = CodexClient()
state = client.register_goal_operation("thread-1")
state.bind("turn-1")
control = _SyncControlClient(fail_first=False)
handle = TurnHandle(control, "thread-1", "turn-1", _goal=state) # type: ignore[arg-type]
results: list[TurnSteerResponse] = []
_route(
client,
"thread/goal/updated",
{"threadId": "thread-1", "turnId": "turn-1", "goal": _goal("thread-1", "active")},
)
_route(
client,
"turn/started",
{"threadId": "thread-1", "turn": _turn("turn-1", "inProgress", started_at=10)},
)
_route(
client,
"turn/completed",
{
"threadId": "thread-1",
"turn": _turn("turn-1", "completed", started_at=10, completed_at=11),
},
)
steering = threading.Thread(target=lambda: results.append(handle.steer("continue")))
steering.start()
_route(
client,
"turn/started",
{"threadId": "thread-1", "turn": _turn("turn-2", "inProgress", started_at=12)},
)
steering.join(timeout=5)
assert {
"alive": steering.is_alive(),
"steer_ids": control.steer_ids,
"public_turn_id": results[0].turn_id,
} == {
"alive": False,
"steer_ids": ["turn-2"],
"public_turn_id": "turn-1",
}
def test_goal_interrupt_succeeds_during_rollover() -> None:
client = CodexClient()
state = client.register_goal_operation("thread-1")
state.bind("turn-1")
control = _SyncControlClient()
handle = TurnHandle(control, "thread-1", "turn-1", _goal=state) # type: ignore[arg-type]
_route(
client,
"thread/goal/updated",
{"threadId": "thread-1", "turnId": "turn-1", "goal": _goal("thread-1", "active")},
)
_route(
client,
"turn/started",
{"threadId": "thread-1", "turn": _turn("turn-1", "inProgress", started_at=10)},
)
_route(
client,
"turn/completed",
{
"threadId": "thread-1",
"turn": _turn("turn-1", "completed", started_at=10, completed_at=11),
},
)
response = handle.interrupt()
assert {
"paused": control.paused,
"interrupt_ids": control.interrupt_ids,
"response": response.model_dump(mode="json"),
"interrupted": state.interrupted,
} == {
"paused": True,
"interrupt_ids": [],
"response": {},
"interrupted": True,
}
class _AsyncGoalClient:
def __init__(self, state) -> None:
self.state = state
self.unregistered = False
async def next_goal_notification(self, state):
return await asyncio.to_thread(state.next_notification)
def unregister_goal_operation(self, _state) -> None:
self.unregistered = True
@dataclass(slots=True)
class _AsyncCodexStub:
_client: object
async def _ensure_initialized(self) -> None:
return None
def test_async_goal_stream_matches_sync_logical_lifecycle() -> None:
async def scenario() -> None:
client = CodexClient()
state = client.register_goal_operation("thread-1")
state.bind("turn-1")
async_client = _AsyncGoalClient(state)
codex = _AsyncCodexStub(async_client)
handle = AsyncTurnHandle(codex, "thread-1", "turn-1", _goal=state) # type: ignore[arg-type]
for method, params in [
(
"thread/goal/updated",
{
"threadId": "thread-1",
"turnId": "turn-1",
"goal": _goal("thread-1", "active"),
},
),
(
"turn/started",
{"threadId": "thread-1", "turn": _turn("turn-1", "inProgress", started_at=10)},
),
(
"thread/goal/updated",
{
"threadId": "thread-1",
"turnId": "turn-1",
"goal": _goal("thread-1", "complete"),
},
),
(
"turn/completed",
{
"threadId": "thread-1",
"turn": _turn("turn-1", "completed", started_at=10, completed_at=11),
},
),
]:
_route(client, method, params)
events = [event async for event in handle.stream()]
assert {
"methods": [event.method for event in events],
"ids": [notification_turn_id(event.payload) for event in events],
"unregistered": async_client.unregistered,
} == {
"methods": ["turn/started", "turn/completed"],
"ids": ["turn-1", "turn-1"],
"unregistered": True,
}
asyncio.run(scenario())
def test_cancelling_async_goal_stream_wakes_notification_reader() -> None:
async def scenario() -> None:
state = CodexClient().register_goal_operation("thread-1")
state.bind("turn-1")
async_client = _AsyncGoalClient(state)
codex = _AsyncCodexStub(async_client)
handle = AsyncTurnHandle(codex, "thread-1", "turn-1", _goal=state) # type: ignore[arg-type]
stream = handle.stream()
reading = asyncio.create_task(anext(stream))
await asyncio.sleep(0)
reading.cancel()
with pytest.raises(asyncio.CancelledError):
await reading
assert {"finished": state.is_finished(), "unregistered": async_client.unregistered} == {
"finished": True,
"unregistered": True,
}
asyncio.run(scenario())
class _AsyncControlClient:
def __init__(self) -> None:
self.steer_ids: list[str] = []
self.interrupt_ids: list[str] = []
self.paused = False
async def turn_steer(self, _thread_id, turn_id, _input):
self.steer_ids.append(turn_id)
if len(self.steer_ids) == 1:
raise InvalidRequestError(
-32600,
"expected active turn id `turn-1` but found `turn-2`",
)
return TurnSteerResponse(turn_id=turn_id)
async def pause_goal(self, _thread_id):
self.paused = True
async def turn_interrupt(self, _thread_id, turn_id):
self.interrupt_ids.append(turn_id)
if len(self.interrupt_ids) == 1:
raise InvalidRequestError(
-32600,
"expected active turn id turn-1 but found turn-2",
)
return TurnInterruptResponse()
def test_async_goal_controls_hide_rollover_turn_ids() -> None:
async def scenario() -> None:
state = CodexClient().register_goal_operation("thread-1")
state.bind("turn-1")
client = _AsyncControlClient()
codex = _AsyncCodexStub(client) # type: ignore[arg-type]
handle = AsyncTurnHandle(codex, "thread-1", "turn-1", _goal=state) # type: ignore[arg-type]
steered = await handle.steer("keep going")
interrupted = await handle.interrupt()
assert {
"steer_ids": client.steer_ids,
"public_steer_id": steered.turn_id,
"paused": client.paused,
"interrupt_ids": client.interrupt_ids,
"interrupt": interrupted.model_dump(mode="json"),
} == {
"steer_ids": ["turn-1", "turn-2"],
"public_steer_id": "turn-1",
"paused": True,
"interrupt_ids": ["turn-1", "turn-2"],
"interrupt": {},
}
asyncio.run(scenario())
def test_finished_async_goal_controls_match_inactive_turn_errors() -> None:
async def scenario() -> None:
state = CodexClient().register_goal_operation("thread-1")
state.bind("turn-1")
state.finish()
client = _AsyncControlClient()
codex = _AsyncCodexStub(client)
handle = AsyncTurnHandle(codex, "thread-1", "turn-1", _goal=state) # type: ignore[arg-type]
with pytest.raises(InvalidRequestError, match="no active turn to steer"):
await handle.steer("keep going")
with pytest.raises(InvalidRequestError, match="no active turn to interrupt"):
await handle.interrupt()
assert client.paused is False
asyncio.run(scenario())
@@ -185,6 +185,19 @@ def test_turn_input_methods_accept_string_shortcut() -> None:
)
def test_turn_result_producers_expose_goal_mode() -> None:
"""Goal mode belongs to APIs that produce a turn result or handle."""
funcs = [Thread.run, Thread.turn, AsyncThread.run, AsyncThread.turn]
assert {
fn: (
inspect.signature(fn).parameters["goal"].annotation,
inspect.signature(fn).parameters["goal"].default,
)
for fn in funcs
} == dict.fromkeys(funcs, ("bool", False))
def test_root_exports_approval_mode() -> None:
"""The root package should expose the high-level approval mode enum."""
assert [(mode.name, mode.value) for mode in ApprovalMode] == [
@@ -370,6 +383,7 @@ def test_generated_public_signatures_are_snake_case_and_typed() -> None:
"thread_source",
],
Thread.turn: [
"goal",
"approval_mode",
"cwd",
"effort",
@@ -381,6 +395,7 @@ def test_generated_public_signatures_are_snake_case_and_typed() -> None:
"summary",
],
Thread.run: [
"goal",
"approval_mode",
"cwd",
"effort",
@@ -445,6 +460,7 @@ def test_generated_public_signatures_are_snake_case_and_typed() -> None:
"thread_source",
],
AsyncThread.turn: [
"goal",
"approval_mode",
"cwd",
"effort",
@@ -456,6 +472,7 @@ def test_generated_public_signatures_are_snake_case_and_typed() -> None:
"summary",
],
AsyncThread.run: [
"goal",
"approval_mode",
"cwd",
"effort",