mirror of
https://github.com/openai/codex.git
synced 2026-11-12 21:20:34 +00:00
[5/6] Add Python goal control integration coverage
This commit is contained in:
Ahmed Ibrahim
@@ -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())
|
||||
Reference in New Issue
Block a user