Fix goal failure cleanup in Python SDK

This commit is contained in:
Ahmed Ibrahim
2026-06-08 12:13:06 -07:00
parent 079aaa83b1
commit 4557f39923
6 changed files with 66 additions and 22 deletions
+20 -14
View File
@@ -80,9 +80,6 @@ class _GoalOperationState:
elif isinstance(payload, ThreadTokenUsageUpdatedNotification):
self.usage = payload.token_usage
if (
isinstance(payload, TurnCompletedNotification)
and payload.turn.status in {TurnStatus.failed, TurnStatus.interrupted}
) or (
self.current_turn_id is None
and self.completed_turn is not None
and (self.cleared or _terminal_goal_status(self.status))
@@ -236,7 +233,7 @@ class _GoalStreamCursor:
active: bool = False
cleared: bool = False
def process(self, notification: Notification) -> tuple[list[Notification], bool]:
def process(self, notification: Notification) -> tuple[list[Notification], bool, bool]:
logical_turn_id = self.state.logical_turn_id
if logical_turn_id is None:
raise RuntimeError("goal operation has not been bound to a logical turn id")
@@ -245,24 +242,27 @@ class _GoalStreamCursor:
if isinstance(payload, TurnStartedNotification):
self.active = True
if self.started is not None:
return [], False
return [], False, False
self.started = payload.turn
return [_logical_notification(notification, logical_turn_id)], False
return [_logical_notification(notification, logical_turn_id)], False, False
if isinstance(payload, TurnCompletedNotification):
self.active = False
self.last_completed = payload
if payload.turn.status in {TurnStatus.failed, TurnStatus.interrupted}:
self.state.finish()
return [self._completion(notification.method, payload)], True
return (
[self._completion(notification.method, payload)],
True,
(payload.turn.status == TurnStatus.failed),
)
if self.status is None and not self.cleared:
raise RuntimeError(
"the connected Codex runtime did not activate goal mode for this turn"
)
if self.cleared or _terminal_goal_status(self.status):
self.state.finish()
return [self._completion(notification.method, payload)], True
return [], False
return [self._completion(notification.method, payload)], True, False
return [], False, False
events = [_logical_notification(notification, logical_turn_id)]
if isinstance(payload, ThreadGoalUpdatedNotification):
@@ -281,8 +281,8 @@ class _GoalStreamCursor:
):
self.state.finish()
events.append(self._completion("turn/completed", self.last_completed))
return events, True
return events, False
return events, True, False
return events, False, False
def _completion(
self,
@@ -310,6 +310,7 @@ class _GoalNotificationStream(Iterator[Notification]):
state: _GoalOperationState
next_notification: Callable[[], Notification]
unregister: Callable[[], None]
stop_failed_goal: Callable[[], None]
_cursor: _GoalStreamCursor = field(init=False)
_pending: deque[Notification] = field(default_factory=deque)
_closed: bool = False
@@ -325,7 +326,9 @@ class _GoalNotificationStream(Iterator[Notification]):
raise StopIteration
try:
while not self._pending:
events, completed = self._cursor.process(self.next_notification())
events, completed, failed = self._cursor.process(self.next_notification())
if failed:
self.stop_failed_goal()
self._pending.extend(events)
if completed:
self._finish()
@@ -356,6 +359,7 @@ class _AsyncGoalNotificationStream(AsyncIterator[Notification]):
state: _GoalOperationState
next_notification: Callable[[], Awaitable[Notification]]
unregister: Callable[[], None]
stop_failed_goal: Callable[[], Awaitable[None]]
_cursor: _GoalStreamCursor = field(init=False)
_pending: deque[Notification] = field(default_factory=deque)
_closed: bool = False
@@ -371,7 +375,9 @@ class _AsyncGoalNotificationStream(AsyncIterator[Notification]):
raise StopAsyncIteration
try:
while not self._pending:
events, completed = self._cursor.process(await self.next_notification())
events, completed, failed = self._cursor.process(await self.next_notification())
if failed:
await self.stop_failed_goal()
self._pending.extend(events)
if completed:
self._finish()
+2
View File
@@ -841,6 +841,7 @@ class TurnHandle:
self._goal,
lambda: self._client.next_goal_notification(self._goal),
lambda: self._client.unregister_goal_operation(self._goal),
lambda: self._client.stop_failed_goal(self._goal),
)
def ordinary_stream() -> Iterator[Notification]:
@@ -963,6 +964,7 @@ class AsyncTurnHandle:
self._goal,
next_goal_notification,
lambda: self._codex._client.unregister_goal_operation(self._goal),
lambda: self._codex._client.stop_failed_goal(self._goal),
)
async def ordinary_stream() -> AsyncIterator[Notification]:
@@ -227,6 +227,10 @@ class AsyncCodexClient:
"""Pause the active goal through the wrapped sync client."""
return await self._call_sync(self._sync.pause_goal, thread_id)
async def stop_failed_goal(self, state: _GoalOperationState) -> None:
"""Stop continuation work after a physical goal turn fails."""
await self._call_sync(self._sync.stop_failed_goal, state)
async def start_goal_operation(
self,
thread_id: str,
+24
View File
@@ -501,6 +501,30 @@ class CodexClient:
"""Pause the active goal used by a logical goal turn."""
return self.thread_goal_set(thread_id, status=ThreadGoalStatus.paused)
def stop_failed_goal(self, state: _GoalOperationState) -> None:
"""Best-effort cleanup after a physical goal turn fails."""
try:
self.pause_goal(state.thread_id)
except Exception:
pass
turn_id = state.current_turn()
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 = 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,
+1 -1
View File
@@ -247,7 +247,7 @@ class AppServerHarness:
"""Write config.toml that routes model calls to the mock server."""
config_toml = self.codex_home / "config.toml"
requires_openai_auth = "requires_openai_auth = true\n" if self.requires_openai_auth else ""
goals = "[features]\ngoals = true\n\n" if self.enable_goals else ""
goals = f"[features]\ngoals = {str(self.enable_goals).lower()}\n\n"
config_toml.write_text(
f"""
model = "mock-model"
+15 -7
View File
@@ -466,7 +466,7 @@ def test_goal_interrupt_pauses_continuation_and_leaves_thread_usable(tmp_path) -
def test_terminal_goal_failure_stops_continuation_and_releases_routing(tmp_path) -> None:
"""A failed server turn should end the logical operation without another continuation."""
"""A failed server turn should stop rollover work and leave the thread usable."""
with AppServerHarness(tmp_path, enable_goals=True) as harness:
harness.responses.enqueue_sse(
sse(
@@ -476,9 +476,13 @@ def test_terminal_goal_failure_stops_continuation_and_releases_routing(tmp_path)
]
)
)
harness.responses.enqueue_assistant_message(
"Recovered with an ordinary turn.",
response_id="goal-failure-follow-up",
harness.responses.enqueue_sse(
streaming_response(
"goal-failure-rollover",
"msg-goal-failure-rollover",
["automatic ", "rollover"],
),
delay_between_events_s=0.5,
)
with Codex(config=harness.app_server_config()) as codex:
@@ -486,17 +490,21 @@ def test_terminal_goal_failure_stops_continuation_and_releases_routing(tmp_path)
with pytest.raises(RuntimeError, match="goal model failed"):
thread.run_goal("Fail this goal turn")
harness.responses.enqueue_assistant_message(
"Recovered with an ordinary turn.",
response_id="goal-failure-follow-up",
)
follow_up = thread.run("Run after the goal failure")
harness.responses.wait_for_requests(2)
harness.responses.wait_for_requests(3)
requests = harness.responses.requests()
assert {
"follow_up": (follow_up.status, follow_up.final_response),
"request_count": len(requests),
"follow_up_input": requests[1].message_input_texts("user")[-1:],
"follow_up_input": requests[2].message_input_texts("user")[-1:],
} == {
"follow_up": (TurnStatus.completed, "Recovered with an ordinary turn."),
"request_count": 2,
"request_count": 3,
"follow_up_input": ["Run after the goal failure"],
}