Files
codex/sdk/python/tests/test_app_server_goal_controls.py
2026-06-09 16:32:04 -07:00

318 lines
12 KiB
Python

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())