diff --git a/sdk/python/tests/app_server_harness.py b/sdk/python/tests/app_server_harness.py index cf0c84739403..86fa731088a1 100644 --- a/sdk/python/tests/app_server_harness.py +++ b/sdk/python/tests/app_server_harness.py @@ -5,7 +5,7 @@ import shutil import threading import time -from dataclasses import dataclass +from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import Any @@ -82,12 +82,26 @@ def header(self, name: str) -> str | None: return self.headers.get(name.lower()) +@dataclass(frozen=True) +class SseGate: + """Hold an SSE response after its first event has been flushed.""" + + ready: threading.Event = field(default_factory=threading.Event) + release: threading.Event = field(default_factory=threading.Event) + + def wait_until_ready(self, timeout_s: float = 5.0) -> None: + """Wait until the response is active and held before its second event.""" + if not self.ready.wait(timeout_s): + raise AssertionError("SSE response did not reach its gate") + + @dataclass(frozen=True) class MockSseResponse: """One queued SSE response served by the mock Responses API.""" body: str delay_between_events_s: float = 0.0 + gate: SseGate | None = None def chunks(self) -> list[bytes]: """Split an SSE body into event chunks while preserving framing.""" @@ -137,12 +151,14 @@ def enqueue_sse( body: str, *, delay_between_events_s: float = 0.0, + gate: SseGate | None = None, ) -> None: - """Queue one SSE body for the next `/v1/responses` request.""" + """Queue SSE, optionally holding after its first event until released.""" self._responses.put( MockSseResponse( body=body, delay_between_events_s=delay_between_events_s, + gate=gate, ) ) @@ -319,9 +335,15 @@ def do_POST(self) -> None: self.send_response(200) self.send_header("content-type", "text/event-stream") self.end_headers() - for chunk in response.chunks(): - self.wfile.write(chunk) - self.wfile.flush() + for index, chunk in enumerate(response.chunks()): + try: + self.wfile.write(chunk) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + return + if index == 0 and response.gate is not None: + response.gate.ready.set() + response.gate.release.wait() if response.delay_between_events_s: time.sleep(response.delay_between_events_s) diff --git a/sdk/python/tests/test_app_server_goal_controls.py b/sdk/python/tests/test_app_server_goal_controls.py new file mode 100644 index 000000000000..80e3485251e4 --- /dev/null +++ b/sdk/python/tests/test_app_server_goal_controls.py @@ -0,0 +1,317 @@ +import asyncio +import time + +import pytest +from app_server_harness import ( + AppServerHarness, + SseGate, + ev_completed_with_usage, + ev_function_call, + ev_response_created, + sse, +) +from app_server_helpers import agent_message_texts_from_items, streaming_response + +from openai_codex import AsyncCodex, Codex +from openai_codex.generated.v2_all import ( + ThreadGoalGetResponse, + ThreadGoalStatus, + TurnStatus, +) + + +def _enqueue_steerable_goal( + harness: AppServerHarness, + prefix: str, + *, + initial_text: str, + continuation_chunks: list[str], + final_text: str, +) -> SseGate: + continuation_gate = SseGate() + harness.responses.enqueue_assistant_message( + initial_text, + response_id=f"{prefix}-initial", + ) + harness.responses.enqueue_sse( + streaming_response( + f"{prefix}-continuation", + f"msg-{prefix}-continuation", + continuation_chunks, + ), + delay_between_events_s=0.15, + gate=continuation_gate, + ) + harness.responses.enqueue_sse( + sse( + [ + ev_response_created(f"{prefix}-complete-tool"), + ev_function_call( + f"call-{prefix}-complete", + "update_goal", + '{"status":"complete"}', + ), + ev_completed_with_usage( + f"{prefix}-complete-tool", + input_tokens=1, + cached_input_tokens=0, + output_tokens=1, + reasoning_output_tokens=0, + total_tokens=2, + ), + ] + ) + ) + harness.responses.enqueue_assistant_message( + final_text, + response_id=f"{prefix}-final", + ) + return continuation_gate + + +def _enqueue_interruptible_goal( + harness: AppServerHarness, + prefix: str, + *, + initial_text: str, + continuation_chunks: list[str], + follow_up_text: str, +) -> SseGate: + continuation_gate = SseGate() + harness.responses.enqueue_assistant_message( + initial_text, + response_id=f"{prefix}-initial", + ) + harness.responses.enqueue_sse( + streaming_response( + f"{prefix}-continuation", + f"msg-{prefix}-continuation", + continuation_chunks, + ), + delay_between_events_s=0.2, + gate=continuation_gate, + ) + harness.responses.enqueue_assistant_message( + follow_up_text, + response_id=f"{prefix}-follow-up", + ) + return continuation_gate + + +def _stale_routed_turn_id(turn) -> None: + goal = turn._goal + assert goal is not None + goal_state = goal.state + deadline = time.monotonic() + 5 + with goal_state._condition: + while goal_state.current_turn_id in {None, turn.id}: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise AssertionError("continuation turn was not routed") + goal_state._condition.wait(remaining) + goal_state.current_turn_id = turn.id + + +def test_goal_steer_targets_an_active_continuation(tmp_path) -> None: + """Steering should reach the active server turn while returning the logical ID.""" + with AppServerHarness(tmp_path) as harness: + continuation_gate = _enqueue_steerable_goal( + harness, + "steer", + initial_text="Initial work complete.", + continuation_chunks=["continuation ", "work"], + final_text="Steered goal complete.", + ) + + with Codex(config=harness.app_server_config()) as codex: + try: + turn = codex.thread_start().start_goal("Start a goal that needs refinement") + continuation_gate.wait_until_ready() + steer = turn.steer("Prioritize the edge-case coverage.") + finally: + continuation_gate.release.set() + result = turn.run() + requests = harness.responses.wait_for_requests(4) + + assert { + "steer_id": steer.turn_id, + "result": (result.id, result.status, result.final_response), + "messages": agent_message_texts_from_items(result.items), + "steering_input": requests[2].message_input_texts("user")[-1:], + } == { + "steer_id": turn.id, + "result": (turn.id, TurnStatus.completed, "Steered goal complete."), + "messages": [ + "Initial work complete.", + "continuation work", + "Steered goal complete.", + ], + "steering_input": ["Prioritize the edge-case coverage."], + } + + +def test_goal_steer_retries_the_server_reported_rollover_turn(tmp_path) -> None: + """Steering should recover when routed state trails the server continuation.""" + with AppServerHarness(tmp_path) as harness: + continuation_gate = _enqueue_steerable_goal( + harness, + "steer-rollover", + initial_text="Initial rollover work complete.", + continuation_chunks=["rollover ", "work"], + final_text="Rollover steering complete.", + ) + + with Codex(config=harness.app_server_config()) as codex: + try: + turn = codex.thread_start().start_goal("Steer while routing catches up") + continuation_gate.wait_until_ready() + _stale_routed_turn_id(turn) + steer = turn.steer("Steer through the rollover window.") + finally: + continuation_gate.release.set() + result = turn.run() + requests = harness.responses.wait_for_requests(4) + + assert { + "steer_id": steer.turn_id, + "result": (result.id, result.status, result.final_response), + "steering_input": requests[2].message_input_texts("user")[-1:], + } == { + "steer_id": turn.id, + "result": (turn.id, TurnStatus.completed, "Rollover steering complete."), + "steering_input": ["Steer through the rollover window."], + } + + +@pytest.mark.parametrize("stale_route", [False, True]) +def test_goal_interrupts_active_and_stale_routed_continuations(tmp_path, stale_route) -> None: + """Interruption should stop active work despite notification rollover lag.""" + with AppServerHarness(tmp_path) as harness: + continuation_gate = _enqueue_interruptible_goal( + harness, + "interrupt", + initial_text="Initial work complete.", + continuation_chunks=["still ", "working"], + follow_up_text="Ordinary follow-up complete.", + ) + + with Codex(config=harness.app_server_config()) as codex: + try: + thread = codex.thread_start() + turn = thread.start_goal("Start interruptible goal work") + continuation_gate.wait_until_ready() + if stale_route: + _stale_routed_turn_id(turn) + interrupt = turn.interrupt() + finally: + continuation_gate.release.set() + interrupted = turn.run() + follow_up = thread.run("Continue with an ordinary turn") + requests = harness.responses.wait_for_requests(3) + persisted = codex._client.request( + "thread/goal/get", + {"threadId": thread.id}, + response_model=ThreadGoalGetResponse, + ).goal + + assert { + "interrupt": interrupt.model_dump(by_alias=True, mode="json"), + "goal_status": persisted.status if persisted else None, + "interrupted": (interrupted.id, interrupted.status), + "follow_up": (follow_up.status, follow_up.final_response), + "request_count": len(requests), + } == { + "interrupt": {}, + "goal_status": ThreadGoalStatus.paused, + "interrupted": (turn.id, TurnStatus.interrupted), + "follow_up": (TurnStatus.completed, "Ordinary follow-up complete."), + "request_count": 3, + } + + +def test_async_goal_steer_retries_the_server_reported_rollover_turn(tmp_path) -> None: + """Async steering should recover when routed state trails the continuation.""" + + async def scenario() -> None: + with AppServerHarness(tmp_path) as harness: + continuation_gate = _enqueue_steerable_goal( + harness, + "async-steer", + initial_text="Async initial work complete.", + continuation_chunks=["async continuation ", "work"], + final_text="Async steered goal complete.", + ) + + async with AsyncCodex(config=harness.app_server_config()) as codex: + try: + thread = await codex.thread_start() + turn = await thread.start_goal("Start an async goal that needs refinement") + await asyncio.to_thread(continuation_gate.wait_until_ready) + await asyncio.to_thread(_stale_routed_turn_id, turn) + steer = await turn.steer("Prioritize async edge-case coverage.") + finally: + continuation_gate.release.set() + result = await turn.run() + requests = await asyncio.to_thread(harness.responses.wait_for_requests, 4) + + assert { + "steer_id": steer.turn_id, + "result": (result.id, result.status, result.final_response), + "messages": agent_message_texts_from_items(result.items), + "steering_input": requests[2].message_input_texts("user")[-1:], + } == { + "steer_id": turn.id, + "result": (turn.id, TurnStatus.completed, "Async steered goal complete."), + "messages": [ + "Async initial work complete.", + "async continuation work", + "Async steered goal complete.", + ], + "steering_input": ["Prioritize async edge-case coverage."], + } + + asyncio.run(scenario()) + + +def test_async_goal_interrupt_retries_the_server_reported_rollover_turn(tmp_path) -> None: + """Async interruption should survive notification lag during rollover.""" + + async def scenario() -> None: + with AppServerHarness(tmp_path) as harness: + continuation_gate = _enqueue_interruptible_goal( + harness, + "async-interrupt", + initial_text="Async initial work complete.", + continuation_chunks=["async still ", "working"], + follow_up_text="Async ordinary follow-up complete.", + ) + + async with AsyncCodex(config=harness.app_server_config()) as codex: + try: + thread = await codex.thread_start() + turn = await thread.start_goal("Start async interruptible goal work") + await asyncio.to_thread(continuation_gate.wait_until_ready) + await asyncio.to_thread(_stale_routed_turn_id, turn) + interrupt = await turn.interrupt() + finally: + continuation_gate.release.set() + interrupted = await turn.run() + follow_up = await thread.run("Continue with an async ordinary turn") + await asyncio.to_thread(harness.responses.wait_for_requests, 3) + requests = harness.responses.requests() + + assert { + "interrupt": interrupt.model_dump(by_alias=True, mode="json"), + "interrupted": (interrupted.id, interrupted.status), + "follow_up": (follow_up.status, follow_up.final_response), + "request_count": len(requests), + "follow_up_input": requests[2].message_input_texts("user")[-1:], + } == { + "interrupt": {}, + "interrupted": (turn.id, TurnStatus.interrupted), + "follow_up": (TurnStatus.completed, "Async ordinary follow-up complete."), + "request_count": 3, + "follow_up_input": ["Continue with an async ordinary turn"], + } + + asyncio.run(scenario())