Files
codex/sdk/python/tests/test_turn_subscriptions.py
Ahmed Ibrahim 1a4096e273 Add untrusted external messages to the Python SDK (#44086)
## Why

Applications need to deliver content from other agents, tools, or services with tool-level authority, without treating it as user input or granting authorization.

## What changed

- Export `ExternalMessage` for sync and async `run(...)` and `turn(...)`, accepting text or structured content with a tool name and optional namespace. Send it through `toolOutput` and require CLI 0.151.0 or newer.
- Support starting a turn or joining an active regular turn while preserving external content as function output in history. Keep external messages separate from user-input lists and `steer(...)`.
- Give turn handles independent subscriptions, replaying completed items and latest usage to joining handles. Release consumed transient events and clean up subscriptions on closure, failure, or cancellation.
- Document the authority boundary and add sync and async examples.

## Testing

Add coverage for wire representations, input validation, runtime compatibility, tool authority across resume, active-turn joins, and tool-output truncation. Add subscription tests for replay, concurrent consumers, early completion, cancellation, and cleanup.

GitOrigin-RevId: 6106327085fd9c4bd11b71e20b3d8e74738b8bb5
2026-09-09 06:54:18 +00:00

306 lines
11 KiB
Python

import asyncio
import gc
import threading
from concurrent.futures import ThreadPoolExecutor
from itertools import chain
import pytest
from openai_codex import AsyncCodex
from openai_codex._run import _collect_turn_result
from openai_codex.api import AsyncTurnHandle, TurnHandle
from openai_codex.async_client import AsyncCodexClient
from openai_codex.client import CodexClient
from openai_codex.errors import TransportClosedError
def turn_events(client, *, status="completed"):
scope = {"threadId": "thread-1", "turnId": "turn-1"}
usage = {
"inputTokens": 2,
"cachedInputTokens": 0,
"outputTokens": 3,
"reasoningOutputTokens": 0,
"totalTokens": 5,
}
return [
client._coerce_notification(
"item/completed",
{
**scope,
"completedAtMs": 1,
"item": {
"id": "message",
"type": "agentMessage",
"text": "done",
"phase": "final_answer",
},
},
),
client._coerce_notification(
"thread/tokenUsage/updated", {**scope, "tokenUsage": {"last": usage, "total": usage}}
),
client._coerce_notification(
"turn/completed",
{
"threadId": "thread-1",
"turn": {
"id": "turn-1",
"items": [],
"status": status,
"error": {"message": "model failed"} if status == "failed" else None,
},
},
),
]
def test_late_join_replays_items_already_consumed_by_original_handle():
client = CodexClient()
original = TurnHandle(client, "thread-1", "turn-1")
events = turn_events(client)
client._router.route_notification(events[0])
stream = original.stream()
first = next(stream)
client._router.route_notification(events[1])
usage = next(stream)
joined = TurnHandle(client, "thread-1", "turn-1")
for event in events[2:]:
client._router.route_notification(event)
first_result = _collect_turn_result(chain([first, usage], stream), turn_id="turn-1")
assert joined.run() == first_result
assert first_result.final_response == "done"
assert first_result.usage.last.total_tokens == 5
assert client._router._turn_states == {}
def test_consumed_deltas_are_released_while_turn_is_active():
client = CodexClient()
subscription = client._subscribe_turn_notifications("turn-1")
state = client._router._turn_states["turn-1"]
for index in range(1000):
event = client._coerce_notification(
"item/agentMessage/delta",
{
"threadId": "thread-1",
"turnId": "turn-1",
"itemId": "message",
"delta": str(index),
},
)
client._router.route_notification(event)
assert subscription.next() == event
assert state.events == {}
assert state.completed_items == {}
assert not state.completed
subscription.close()
def test_slow_subscriber_keeps_unread_deltas_until_it_consumes_them():
client = CodexClient()
fast = client._subscribe_turn_notifications("turn-1")
slow = client._subscribe_turn_notifications("turn-1")
event = client._coerce_notification(
"item/agentMessage/delta",
{"threadId": "thread-1", "turnId": "turn-1", "itemId": "message", "delta": "hello"},
)
client._router.route_notification(event)
assert fast.next() == event
assert slow.next() == event
assert client._router._turn_states["turn-1"].events == {}
fast.close()
slow.close()
def test_completion_before_turn_start_response_is_replayed(monkeypatch):
client = CodexClient()
def request_raw(method, params):
assert method == "turn/start"
for event in turn_events(client):
client._router.route_notification(event)
return {"turn": {"id": "turn-1", "status": "inProgress", "items": []}}
monkeypatch.setattr(client, "_request_raw", request_raw)
started = client.turn_start("thread-1", "hello")
handle = TurnHandle(client, "thread-1", started.turn.id)
assert handle.run().final_response == "done"
assert client._router._turn_states == {}
def test_pending_join_preserves_result_after_original_handle_finishes():
client = CodexClient()
original = TurnHandle(client, "thread-1", "turn-1")
with client._router.pending_turn("thread-1"):
for event in turn_events(client):
client._router.route_notification(event)
result = original.run()
client._router.prepare_turn("turn-1", "thread-1")
joined = TurnHandle(client, "thread-1", "turn-1")
assert joined.run() == result
assert client._router._turn_states == {}
def test_failed_start_releases_early_completed_state():
client = CodexClient()
with pytest.raises(ValueError, match="request failed"):
with client._router.pending_turn("thread-1"):
for event in turn_events(client):
client._router.route_notification(event)
raise ValueError("request failed")
assert client._router._turn_states == {}
assert client._router._pending_turn_requests == {}
def test_closing_one_stream_leaves_other_subscriber_intact():
client = CodexClient()
original = TurnHandle(client, "thread-1", "turn-1")
joined = TurnHandle(client, "thread-1", "turn-1")
events = turn_events(client)
client._router.route_notification(events[0])
stream = original.stream()
next(stream)
stream.close()
for event in events[1:]:
client._router.route_notification(event)
assert joined.run().final_response == "done"
assert client._router._turn_states == {}
@pytest.mark.parametrize("failure", ["transport", "model"])
def test_both_handles_observe_failure_and_release_state(failure):
client = CodexClient()
handles = [TurnHandle(client, "thread-1", "turn-1") for _ in range(2)]
if failure == "transport":
client._router.fail_all(TransportClosedError("transport failed"))
else:
for event in turn_events(client, status="failed"):
client._router.route_notification(event)
for handle in handles:
with pytest.raises((TransportClosedError, RuntimeError), match=f"{failure} failed"):
handle.run()
assert client._router._turn_states == {}
def test_abandoned_handle_releases_completed_history():
client = CodexClient()
handle = TurnHandle(client, "thread-1", "turn-1")
for event in turn_events(client):
client._router.route_notification(event)
with client._router._lock:
del handle
gc.collect()
assert client._router._turn_states == {}
def test_cancelled_async_consumer_leaves_other_handle_intact():
async def scenario():
codex = AsyncCodex()
codex._initialized = True
client = codex._client._sync
original = AsyncTurnHandle(codex, "thread-1", "turn-1")
joined = AsyncTurnHandle(codex, "thread-1", "turn-1")
task = asyncio.create_task(original.run())
await asyncio.sleep(0) # Let the stream enter its wait before cancelling it.
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
for event in turn_events(client):
client._router.route_notification(event)
assert (await asyncio.wait_for(joined.run(), timeout=2)).final_response == "done"
assert client._router._turn_states == {}
asyncio.run(scenario())
def test_cancelled_turn_start_releases_result_after_response(monkeypatch):
async def scenario():
client = AsyncCodexClient()
entered = threading.Event()
respond = threading.Event()
released = threading.Event()
subscribe = client._sync._subscribe_turn_notifications
def request_raw(method, params):
entered.set()
assert respond.wait(timeout=2)
for event in turn_events(client._sync):
client._sync._router.route_notification(event)
return {"turn": {"id": "turn-1", "items": [], "status": "completed"}}
def observe_release(turn_id):
subscription = subscribe(turn_id)
close = subscription.close
def close_and_signal():
close()
released.set()
subscription.close = close_and_signal
return subscription
monkeypatch.setattr(client._sync, "_request_raw", request_raw)
monkeypatch.setattr(client._sync, "_subscribe_turn_notifications", observe_release)
task = asyncio.create_task(client.turn_start("thread-1", "hello"))
try:
assert await asyncio.to_thread(entered.wait, 2)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
finally:
respond.set()
assert await asyncio.to_thread(released.wait, 2)
assert client._sync._router._turn_states == {}
asyncio.run(scenario())
def test_cancelled_queued_start_does_not_send_a_request(monkeypatch):
with ThreadPoolExecutor(max_workers=1) as executor:
monkeypatch.setattr("openai_codex.async_client._TURN_START_EXECUTOR", executor)
release_worker = threading.Event()
worker_started = threading.Event()
request_sent = threading.Event()
def occupy_worker():
worker_started.set()
assert release_worker.wait(timeout=5)
occupied = executor.submit(occupy_worker)
assert worker_started.wait(timeout=5)
client = AsyncCodexClient()
monkeypatch.setattr(client._sync, "turn_start", lambda *args: request_sent.set())
async def scenario():
task = asyncio.create_task(client.turn_start("thread-1", "hello"))
await asyncio.sleep(0)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
try:
asyncio.run(scenario())
finally:
release_worker.set()
occupied.result(timeout=5)
executor.submit(lambda: None).result(timeout=5)
assert not request_sent.is_set()
def test_low_level_start_keeps_implicit_registration_and_explicit_unregister(monkeypatch):
client = CodexClient()
def request_raw(method, params):
for event in turn_events(client):
client._router.route_notification(event)
return {"turn": {"id": "turn-1", "status": "completed", "items": []}}
monkeypatch.setattr(client, "_request_raw", request_raw)
started = client.turn_start("thread-1", "hello")
assert client.next_turn_notification(started.turn.id) == turn_events(client)[0]
client.unregister_turn_notifications(started.turn.id)
with pytest.raises(RuntimeError, match="not registered"):
client.next_turn_notification(started.turn.id)
assert client._router._turn_states == {}